jianf/ModelRouter
1
0
Fork 0
mirror of https://gitcode.com/JianFeeeee/ModelRouter.git synced 2026-09-20 00:48:00 +00:00
Code Issues Packages Projects Releases Wiki Activity

feat: AUTO chain rewrite — silent failover+busy skip+pref round-robin+503 tier summary; chain edits reset slot cooldowns (P0/P1); stats by_status + audit jsonl rotation; UI priority-page health badges & status-code card; ctx-menu capture-phase close (outside-press guard); main.go ops warnings; local bundled-Lua verified tests (3 latent bugs fixed); plan.md

Browse Source
This commit is contained in:
JianFeeeee
2026-08-10 23:58:31 +08:00
parent 397af36fbb
commit 88802f9ef6
17 changed files with 2060 additions and 433 deletions
Download Patch File Download Diff File Expand all files Collapse all files

117
internal/core/core.go
View File

@ -11,6 +11,7 @@ import (
"path/filepath"
"sort"
"strings"
"sync/atomic"
"time"
"llmsproxy/internal/config"
@ -26,6 +27,7 @@ type Core struct {
store *config.Store
scheduler *scheduler.Scheduler
registry *provider.Registry
autoChain atomic.Pointer[scheduler.Chain]
}
// New builds the core from a config file plus runtime overlay.
@ -138,6 +140,11 @@ 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 (immutable after build;
// a rebuilt chain is swapped in atomically). nil before the first build or
// when no auto slots could be resolved.
func (c *Core) AutoChain() *scheduler.Chain { return c.autoChain.Load() }
func (c *Core) DefaultModel() string { return c.cfg.DefaultModel }
func (c *Core) GatewayKeys() []string { return c.cfg.GatewayKeys }
@ -228,9 +235,34 @@ func cleanScopes(entries []config.ModelScope) []config.ModelScope {
return clean
}
// SaveAutoRules persists the AUTO scheduling slots.
// SaveAutoRules persists the AUTO scheduling slots, rebuilds the chain and
// clears the cooldown of every slot in it — preference scores are kept, so a
// reliably good model keeps its edge while an edited chain applies
// immediately. Providers are NOT rebuilt here (their per-model state survives
// the edit, plan 2.4 lifecycle); rebuildRegistry covers source edits.
func (c *Core) SaveAutoRules(entries []config.ModelScope) error {
return c.store.SaveAutoRules(cleanScopes(entries))
if err := c.store.SaveAutoRules(cleanScopes(entries)); err != nil {
return err
}
c.buildAutoChain()
if ch := c.autoChain.Load(); ch != nil {
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)
}
}
}
}
return nil
}
// ResetHealth clears the scheduling backoff state of every provider (admin
// UI action). Unlike SaveAutoRules this does not touch the chain itself.
func (c *Core) ResetHealth() {
for _, p := range c.registry.Providers() {
p.ResetHealth()
}
}
// Registry resolves model -> owning provider.
@ -303,9 +335,88 @@ func (c *Core) rebuildRegistry() error {
} else {
c.registry.Replace(providers)
}
c.buildAutoChain()
return nil
}
// buildAutoChain rebuilds the AUTO chain snapshot from the persisted rules
// against the current providers. Slots whose (model, source) no longer exists
// and image-kind models are dropped; a chain with no slots makes AUTO
// requests answer "no auto slot configured".
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.store.AutoRules()
sr := make([]scheduler.Rule, 0, len(rules))
for _, e := range rules {
r := scheduler.Rule{
Model: e.Model,
Source: e.Source,
Tier: e.Tier,
Quota: e.TokenQuota,
Period: e.Period,
Hours: e.Hours,
}
if r.Source == "" {
// canonicalize to the owning source so summaries/audit/quota
// windows always carry a real source name
if p := c.registry.ProviderForSlot(e.Model, ""); p != nil {
r.Source = p.Name()
}
}
sr = append(sr, r)
}
c.autoChain.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"`
}
// AutoSlotStates returns per-slot health (preference, failure count,
// cooldown) for every slot of the current AUTO chain, mirroring the chain
// order so the priority-page UI can annotate its blocks.
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)
out = append(out, AutoSlotState{
Model: sl.Model,
Source: sl.Source,
Pref: pref,
FailCount: fail,
CooldownUntil: until,
Cooling: until > now,
})
}
}
return out
}
// Reload re-reads the runtime store and rebuilds sources (adapter reload is not
// strictly needed since adapters are loaded into the VM at startup; uploaded
// adapters are placed in the adapter dir and loaded by the web UI).
@ -408,4 +519,4 @@ func (c *Core) Close() {
if c.vm != nil {
c.vm.Stop()
}
}
}

363
internal/gateway/chat.go
View File

@ -3,6 +3,7 @@ package gateway
import (
"context"
"encoding/json"
"errors"
"fmt"
"net/http"
"strings"
@ -17,15 +18,15 @@ import (
// chatRequest mirrors the OpenAI chat completions request the gateway accepts.
type chatRequest struct {
Model string `json:"model"`
Messages []types.ChatMessage `json:"messages"`
Temperature *float64 `json:"temperature,omitempty"`
MaxTokens int `json:"max_tokens,omitempty"`
Stream bool `json:"stream,omitempty"`
Tools []interface{} `json:"tools,omitempty"`
ToolChoice interface{} `json:"tool_choice,omitempty"`
DisableThinking bool `json:"disable_thinking"`
ExtraBody map[string]interface{} `json:"extra_body,omitempty"`
Model string `json:"model"`
Messages []types.ChatMessage `json:"messages"`
Temperature *float64 `json:"temperature,omitempty"`
MaxTokens int `json:"max_tokens,omitempty"`
Stream bool `json:"stream,omitempty"`
Tools []interface{} `json:"tools,omitempty"`
ToolChoice interface{} `json:"tool_choice,omitempty"`
DisableThinking bool `json:"disable_thinking"`
ExtraBody map[string]interface{} `json:"extra_body,omitempty"`
}
// ChatCompletion is the non-streaming OpenAI response object.
@ -60,9 +61,9 @@ type ChatChunk struct {
}
type ChunkChoice struct {
Index int `json:"index"`
Delta RespMessage `json:"delta"`
FinishReason *string `json:"finish_reason"`
Index int `json:"index"`
Delta RespMessage `json:"delta"`
FinishReason *string `json:"finish_reason"`
}
var seq int64
@ -279,9 +280,9 @@ func (g *Gateway) handleChat(w http.ResponseWriter, r *http.Request) {
model = g.core.DefaultModel()
}
if isAuto(model) {
plans := g.autoPlans()
if len(plans) == 0 {
writeError(w, http.StatusServiceUnavailable, "no_provider", "no auto slot available (quota exhausted or none configured)")
chain := g.core.AutoChain()
if chain == nil || len(chain.Tiers) == 0 {
writeError(w, http.StatusServiceUnavailable, "no_provider", "no auto slot configured")
return
}
if msg := g.checkModelScope(r.Context(), "AUTO"); msg != "" {
@ -306,12 +307,21 @@ func (g *Gateway) handleChat(w http.ResponseWriter, r *http.Request) {
Type: "chat",
OK: false,
}
// quotaExhausted reports a slot whose token window has been used up;
// exhausted slots are dropped from scheduling without penalty.
quotaExhausted := func(sl *scheduler.Slot) bool {
if sl.Quota <= 0 {
return false
}
win := AutoPeriodSeconds(sl.Period, sl.Hours)
return g.stats.WindowTokens(sl.Model, sl.Source, win) >= sl.Quota
}
if req.Stream {
rec.Type = "stream"
g.streamChatAuto(w, ctx, plans, inner, rec)
g.streamChatAuto(w, ctx, chain, inner, rec, quotaExhausted)
return
}
g.singleChatAuto(w, ctx, plans, inner, rec)
g.singleChatAuto(w, ctx, chain, inner, rec, quotaExhausted)
return
}
if !isAuto(model) {
@ -322,7 +332,7 @@ func (g *Gateway) handleChat(w http.ResponseWriter, r *http.Request) {
}
cands, effective := g.resolveCands(r.Context(), &req)
if len(cands) == 0 {
writeError(w, http.StatusServiceUnavailable, "no_provider", "no LLM source configured")
writeError(w, http.StatusNotFound, "model_not_found", fmt.Sprintf("model %q is not configured", model))
return
}
if effective == "" {
@ -472,17 +482,43 @@ func estimateTextTokens(parts ...interface{}) int64 {
return int64(n/3 + 1)
}
// toScheduler adapts concrete providers to the scheduler.Provider interface.
// It lives here (not in the scheduler package) so scheduler tests do not pull
// in the provider package and with it the Lua runtime's link requirements.
func toScheduler(cands []*provider.Provider) []scheduler.Provider {
out := make([]scheduler.Provider, len(cands))
for i, p := range cands {
out[i] = p
}
return out
}
// upstreamErrStatus maps a scheduling error to its HTTP status: a failed
// AUTO chain answers 503 with its per-tier summary, a busy source (every
// concurrency slot in use) is a transient capacity condition answered with
// 429 so clients fail fast, while other upstream failures stay 502.
func upstreamErrStatus(err error) int {
var ce *scheduler.ChainErr
if errors.As(err, &ce) {
return http.StatusServiceUnavailable
}
if errors.Is(err, provider.ErrBusy) {
return http.StatusTooManyRequests
}
return http.StatusBadGateway
}
func (g *Gateway) singleChat(w http.ResponseWriter, ctx context.Context, cands []*provider.Provider, req *types.ChatRequest, effective string, rec *Req) {
rec.LatMs = 0
t0 := time.Now()
resp, usedSrc, usedModel, err := g.core.Scheduler().Chat(ctx, scheduler.FromRegistry(cands), req)
resp, usedSrc, usedModel, err := g.core.Scheduler().Chat(ctx, toScheduler(cands), req)
rec.LatMs = time.Since(t0).Milliseconds()
if err != nil {
rec.OK = false
rec.Status = http.StatusBadGateway
rec.Status = upstreamErrStatus(err)
rec.Err = err.Error()
g.writeRec(rec)
writeError(w, http.StatusBadGateway, "upstream_error", err.Error())
writeError(w, rec.Status, "upstream_error", err.Error())
return
}
rec.OK = true
@ -538,12 +574,12 @@ func (g *Gateway) streamChat(w http.ResponseWriter, ctx context.Context, cands [
rec.LatMs = time.Since(t0).Milliseconds()
g.writeRec(rec)
}()
chunks, _, usedModel, err := g.core.Scheduler().ChatStream(ctx, scheduler.FromRegistry(cands), req)
chunks, _, usedModel, err := g.core.Scheduler().ChatStream(ctx, toScheduler(cands), req)
if err != nil {
rec.OK = false
rec.Status = http.StatusBadGateway
rec.Status = upstreamErrStatus(err)
rec.Err = err.Error()
writeError(w, http.StatusBadGateway, "upstream_error", err.Error())
writeError(w, rec.Status, "upstream_error", err.Error())
return
}
if usedModel != "" {
@ -611,146 +647,66 @@ func (g *Gateway) streamChat(w http.ResponseWriter, ctx context.Context, cands [
}
}
// autoPlan is one schedulable AUTO slot: a model id pinned to its provider
// with an optional token quota window. Quota-exhausted slots are skipped.
type autoPlan struct {
p *provider.Provider
model string
tier int
quota int64
win int64
}
// autoPlans builds the schedulable AUTO slots from the persisted rules. A
// slot is schedulable while its model is available and (when quota > 0) the
// tokens used within its reset window are below the quota. Slots are returned
// tiered high→low, and within the same tier the order is rotated round-robin
// so concurrent requests spread evenly across equal-priority sources (still
// with failover to the next slot if one errors).
func (g *Gateway) autoPlans() []autoPlan {
rules := g.core.AutoRules()
if len(rules) == 0 {
return nil
}
plans := make([]autoPlan, 0, len(rules))
for _, e := range rules {
p := g.core.ProviderForSlot(e.Model, e.Source)
if p == nil {
continue
}
if m := p.ModelByID(e.Model); m != nil && m.Kind == "image" {
continue
}
win := AutoPeriodSeconds(e.Period, e.Hours)
if e.TokenQuota > 0 && g.stats.WindowTokens(e.Model, e.Source, win) >= e.TokenQuota {
continue
}
plans = append(plans, autoPlan{p: p, model: e.Model, tier: e.Tier, quota: e.TokenQuota, win: win})
}
return g.rotateSameTier(plans)
}
// rotateSameTier reorders the leading plan of each consecutive same-tier run
// using a global round-robin counter, so requests distribute across
// equal-priority sources while preserving tier ordering and in-tier failover.
func (g *Gateway) rotateSameTier(plans []autoPlan) []autoPlan {
if len(plans) < 2 {
// allow single slot without varying
return plans
}
rot := int(g.autoRR.Add(1))
out := make([]autoPlan, 0, len(plans))
for i := 0; i < len(plans); {
j := i
for j < len(plans) && plans[j].tier == plans[i].tier {
j++
}
run := plans[i:j]
if len(run) > 1 {
off := rot % len(run)
run = append(run[off:], run[:off]...)
}
out = append(out, run...)
i = j
}
return out
}
// singleChatAuto runs a non-streaming AUTO request slot by slot: each slot
// pins its own model; a slot whose provider errors out is skipped. The first
// slot to answer wins; when every slot fails, the recorded error is from the
// last one.
func (g *Gateway) singleChatAuto(w http.ResponseWriter, ctx context.Context, plans []autoPlan, req *types.ChatRequest, rec *Req) {
// singleChatAuto runs a non-streaming AUTO request down the chain (see
// scheduler.ChainChat): tiers descending, per-tier round-robin ordered by
// preference, cooldown as the only hard skip, busy slots skipped without
// penalty and a bounded busy wait. When every tier fails, the response is a
// 503 carrying the per-tier error summary (which source/model failed why).
func (g *Gateway) singleChatAuto(w http.ResponseWriter, ctx context.Context, chain *scheduler.Chain, req *types.ChatRequest, rec *Req, quotaExhausted func(*scheduler.Slot) bool) {
rec.LatMs = 0
t0 := time.Now()
var lastErr error
var lastSrc, lastModel string
for _, pl := range plans {
if !pl.p.Available() {
continue
resp, usedSrc, usedModel, err := g.core.Scheduler().ChainChat(ctx, chain, req, quotaExhausted)
rec.LatMs = time.Since(t0).Milliseconds()
if err != nil {
rec.OK = false
rec.Status = upstreamErrStatus(err)
rec.Err = err.Error()
if ce, ok := err.(*scheduler.ChainErr); ok && len(ce.Tiers) > 0 {
rec.Source = ce.Tiers[0].Source
rec.Model = ce.Tiers[0].Model
}
r := *req
r.Model = pl.model
lastSrc, lastModel = pl.p.Name(), pl.model
resp, usedSrc, usedModel, err := g.core.Scheduler().Chat(ctx, scheduler.FromRegistry([]*provider.Provider{pl.p}), &r)
if err != nil {
lastErr = err
continue
}
rec.LatMs = time.Since(t0).Milliseconds()
rec.OK = true
rec.Status = http.StatusOK
rec.Prompt = int64(resp.TokenUsage.Prompt)
if rec.Prompt == 0 {
rec.Prompt = estimatePromptTokens(&r)
}
rec.Compl = int64(resp.TokenUsage.Completion)
if rec.Compl == 0 {
rec.Compl = estimateTextTokens(resp.Content, resp.ReasoningContent, resp.ToolCalls)
}
rec.Source = usedSrc
rec.Model = usedModel
g.writeRec(rec)
msg := RespMessage{Role: "assistant", Content: resp.Content}
if resp.ReasoningContent != "" {
msg.ReasoningContent = resp.ReasoningContent
}
if len(resp.ToolCalls) > 0 {
msg.ToolCalls = toolCallsWire(resp.ToolCalls)
}
out := ChatCompletion{
ID: newID(),
Object: "chat.completion",
Created: time.Now().Unix(),
Model: usedModel,
Choices: []ChatChoice{{Index: 0, Message: msg, FinishReason: resp.FinishReason}},
}
if resp.TokenUsage.Total > 0 || resp.TokenUsage.Prompt > 0 || resp.TokenUsage.Completion > 0 {
out.Usage = &resp.TokenUsage
}
writeJSON(w, http.StatusOK, out)
writeError(w, rec.Status, "upstream_error", err.Error())
return
}
rec.LatMs = time.Since(t0).Milliseconds()
if lastErr == nil {
lastErr = fmt.Errorf("no provider available")
rec.OK = true
rec.Status = http.StatusOK
rec.Prompt = int64(resp.TokenUsage.Prompt)
if rec.Prompt == 0 {
rec.Prompt = estimatePromptTokens(req)
}
rec.OK = false
rec.Status = http.StatusBadGateway
rec.Err = lastErr.Error()
if rec.Model == "" {
rec.Model = lastModel
}
if rec.Source == "" {
rec.Source = lastSrc
rec.Compl = int64(resp.TokenUsage.Completion)
if rec.Compl == 0 {
rec.Compl = estimateTextTokens(resp.Content, resp.ReasoningContent, resp.ToolCalls)
}
rec.Source = usedSrc
rec.Model = usedModel
g.writeRec(rec)
writeError(w, http.StatusBadGateway, "upstream_error", lastErr.Error())
msg := RespMessage{Role: "assistant", Content: resp.Content}
if resp.ReasoningContent != "" {
msg.ReasoningContent = resp.ReasoningContent
}
if len(resp.ToolCalls) > 0 {
msg.ToolCalls = toolCallsWire(resp.ToolCalls)
}
out := ChatCompletion{
ID: newID(),
Object: "chat.completion",
Created: time.Now().Unix(),
Model: usedModel,
Choices: []ChatChoice{{Index: 0, Message: msg, FinishReason: resp.FinishReason}},
}
if resp.TokenUsage.Total > 0 || resp.TokenUsage.Prompt > 0 || resp.TokenUsage.Completion > 0 {
out.Usage = &resp.TokenUsage
}
writeJSON(w, http.StatusOK, out)
}
// streamChatAuto streams an AUTO request. It stays pinned to the first slot
// whose stream begins; a slot that fails to connect is skipped.
func (g *Gateway) streamChatAuto(w http.ResponseWriter, ctx context.Context, plans []autoPlan, req *types.ChatRequest, rec *Req) {
// streamChatAuto streams an AUTO request down the chain. A slot is abandoned
// only before its first chunk (connect error / non-200 / busy); once a stream
// starts it stays pinned. Total failure writes a JSON 503 (with the per-tier
// summary) before any SSE byte is sent.
func (g *Gateway) streamChatAuto(w http.ResponseWriter, ctx context.Context, chain *scheduler.Chain, req *types.ChatRequest, rec *Req, quotaExhausted func(*scheduler.Slot) bool) {
rec.LatMs = 0
t0 := time.Now()
rec.OK = true
@ -759,6 +715,23 @@ func (g *Gateway) streamChatAuto(w http.ResponseWriter, ctx context.Context, pla
rec.LatMs = time.Since(t0).Milliseconds()
g.writeRec(rec)
}()
chunks, usedSrc, usedModel, err := g.core.Scheduler().ChainChatStream(ctx, chain, req, quotaExhausted)
if err != nil {
rec.OK = false
rec.Status = upstreamErrStatus(err)
rec.Err = err.Error()
if ce, ok := err.(*scheduler.ChainErr); ok && len(ce.Tiers) > 0 {
rec.Source = ce.Tiers[0].Source
rec.Model = ce.Tiers[0].Model
}
writeError(w, rec.Status, "upstream_error", err.Error())
return
}
if usedModel != "" {
rec.Model = usedModel
}
rec.Source = usedSrc
rec.Prompt = estimatePromptTokens(req)
w.Header().Set("Content-Type", "text/event-stream")
w.Header().Set("Cache-Control", "no-cache")
w.Header().Set("Connection", "keep-alive")
@ -780,68 +753,38 @@ func (g *Gateway) streamChatAuto(w http.ResponseWriter, ctx context.Context, pla
return true
}
if !send(ChatChunk{
ID: id, Object: "chat.completion.chunk", Created: created, Model: "auto",
ID: id, Object: "chat.completion.chunk", Created: created, Model: rec.Model,
Choices: []ChunkChoice{{Index: 0, Delta: RespMessage{Role: "assistant"}}},
}) {
return
}
var lastErr error
var lastSrc, lastModel string
for _, pl := range plans {
if !pl.p.Available() {
continue
for ck := range chunks {
chunk := ChatChunk{
ID: id, Object: "chat.completion.chunk", Created: created, Model: rec.Model,
}
r := *req
r.Model = pl.model
lastSrc, lastModel = pl.p.Name(), pl.model
chunks, usedSrc, usedModel, err := g.core.Scheduler().ChatStream(ctx, scheduler.FromRegistry([]*provider.Provider{pl.p}), &r)
if err != nil {
lastErr = err
continue
delta := RespMessage{Role: "assistant", Content: ck.Content}
if ck.ReasoningContent != "" {
delta.ReasoningContent = ck.ReasoningContent
}
if usedModel != "" {
rec.Model = usedModel
rec.Source = usedSrc
if len(ck.ToolCalls) > 0 {
delta.ToolCalls = ck.ToolCalls
}
rec.Prompt = estimatePromptTokens(&r)
for ck := range chunks {
chunk := ChatChunk{
ID: id, Object: "chat.completion.chunk", Created: created, Model: rec.Model,
}
delta := RespMessage{Role: "assistant", Content: ck.Content}
if ck.ReasoningContent != "" {
delta.ReasoningContent = ck.ReasoningContent
}
if len(ck.ToolCalls) > 0 {
delta.ToolCalls = ck.ToolCalls
}
choice := ChunkChoice{Index: 0, Delta: delta}
if ck.Done {
stop := "stop"
choice.FinishReason = &stop
}
chunk.Choices = []ChunkChoice{choice}
rec.Compl += int64(len(ck.Content)+len(ck.ReasoningContent)+len(ck.ToolCalls)) / 3
if !send(chunk) {
return
}
choice := ChunkChoice{Index: 0, Delta: delta}
if ck.Done {
stop := "stop"
choice.FinishReason = &stop
}
chunk.Choices = []ChunkChoice{choice}
rec.Compl += int64(len(ck.Content)+len(ck.ReasoningContent)+len(ck.ToolCalls)) / 3
if !send(chunk) {
return
}
return
}
if lastErr == nil {
lastErr = fmt.Errorf("no provider available")
}
rec.OK = false
rec.Status = http.StatusBadGateway
rec.Err = lastErr.Error()
if rec.Model == "" {
rec.Model = lastModel
}
if rec.Source == "" {
rec.Source = lastSrc
}
errEvent, _ := json.Marshal(map[string]interface{}{"error": map[string]string{"message": lastErr.Error(), "type": "upstream_error"}})
fmt.Fprintf(w, "data: %s\n\n", errEvent)
stop := "stop"
send(ChatChunk{
ID: id, Object: "chat.completion.chunk", Created: created, Model: rec.Model,
Choices: []ChunkChoice{{Index: 0, Delta: RespMessage{}, FinishReason: &stop}},
})
fmt.Fprintf(w, "data: [DONE]\n\n")
if flusher != nil {
flusher.Flush()
@ -889,13 +832,13 @@ func (g *Gateway) handleImage(w http.ResponseWriter, r *http.Request) {
defer done()
rec := &Req{Key: keyID(reqKey(r.Context())), Type: "image", Model: model, Source: firstSource(cands), OK: false}
t0 := time.Now()
resp, usedSrc, err := g.core.Scheduler().Image(r.Context(), scheduler.FromRegistry(cands), &req)
resp, usedSrc, err := g.core.Scheduler().Image(r.Context(), toScheduler(cands), &req)
rec.LatMs = time.Since(t0).Milliseconds()
if err != nil {
rec.Status = http.StatusBadGateway
rec.Status = upstreamErrStatus(err)
rec.Err = err.Error()
g.writeRec(rec)
writeError(w, http.StatusBadGateway, "upstream_error", err.Error())
writeError(w, rec.Status, "upstream_error", err.Error())
return
}
if usedSrc != "" {
@ -909,4 +852,4 @@ func (g *Gateway) handleImage(w http.ResponseWriter, r *http.Request) {
Created: time.Now().Unix(),
Data: resp.ImageData,
})
}
}

189
internal/gateway/gateway_test.go
View File

@ -40,6 +40,7 @@ func newTestGateway(t *testing.T, srcs ...config.Source) *Gateway {
cfg := &config.Config{
AdapterDir: filepath.Join(t.TempDir(), "adapters"),
RuntimeFile: filepath.Join(t.TempDir(), "runtime.json"),
GatewayKeys: []string{"sk-test"},
Sources: srcs,
}
if err := cfg.ApplyDefaults(); err != nil {
@ -69,6 +70,192 @@ func doReq(t *testing.T, g *Gateway, method, path, body string) *httptest.Respon
return rr
}
// upstreamCtrl toggles a mocked upstream's behavior between requests.
type upstreamCtrl struct {
status int // 0 = healthy; else every request fails with that status
hits int // chat call count
}
// upstream returns a mocked OpenAI upstream driven by ctrl.status.
func upstream(t *testing.T, ctrl *upstreamCtrl) *httptest.Server {
t.Helper()
return httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
ctrl.hits++
if ctrl.status != 0 {
w.WriteHeader(ctrl.status)
fmt.Fprint(w, `{"error":"boom"}`)
return
}
w.Header().Set("Content-Type", "application/json")
fmt.Fprintf(w, `{"choices":[{"message":{"content":"pong"},"finish_reason":"stop"}],"usage":{"prompt_tokens":3,"completion_tokens":1,"total_tokens":4}}`)
}))
}
// TestChatAutoChainTierFailover: AUTO chain, first slot hard-fails, the pass
// moves on within the same tier and the request is served by the next slot.
func TestChatAutoChainTierFailover(t *testing.T) {
a, b := &upstreamCtrl{status: 500}, &upstreamCtrl{}
aUp := upstream(t, a)
bUp := upstream(t, b)
defer aUp.Close()
defer bUp.Close()
g := newTestGateway(t,
config.Source{Name: "a", BaseURL: aUp.URL, Adapter: "openai", Models: []config.Model{{ID: "a-m", Priority: 100}}},
config.Source{Name: "b", BaseURL: bUp.URL, Adapter: "openai", Models: []config.Model{{ID: "b-m", Priority: 10}}},
)
rr := doReq(t, g, "POST", "/v1/chat/completions",
`{"model":"AUTO","messages":[{"role":"user","content":"hi"}]}`)
if rr.Code != 200 {
t.Fatalf("status=%d body=%s", rr.Code, rr.Body.String())
}
var cc ChatCompletion
_ = json.Unmarshal(rr.Body.Bytes(), &cc)
if cc.Model != "b-m" {
t.Fatalf("AUTO served %q, want b-m", cc.Model)
}
if a.hits == 0 || b.hits == 0 {
t.Fatalf("hit counts a=%d b=%d, want both > 0", a.hits, b.hits)
}
}
// TestChatAutoChain503Summary: every AUTO slot fails -> 503 whose message
// names each failed tier/source/model.
func TestChatAutoChain503Summary(t *testing.T) {
a, b := &upstreamCtrl{status: 500}, &upstreamCtrl{status: 500}
aUp := upstream(t, a)
bUp := upstream(t, b)
defer aUp.Close()
defer bUp.Close()
g := newTestGateway(t,
config.Source{Name: "a", BaseURL: aUp.URL, Adapter: "openai", Models: []config.Model{{ID: "a-m", Priority: 100}}},
config.Source{Name: "b", BaseURL: bUp.URL, Adapter: "openai", Models: []config.Model{{ID: "b-m", Priority: 10}}},
)
rr := doReq(t, g, "POST", "/v1/chat/completions",
`{"model":"AUTO","messages":[{"role":"user","content":"hi"}]}`)
if rr.Code != http.StatusServiceUnavailable {
t.Fatalf("status=%d body=%s", rr.Code, rr.Body.String())
}
if !strings.Contains(rr.Body.String(), "all auto tiers failed") ||
!strings.Contains(rr.Body.String(), "a/a-m") ||
!strings.Contains(rr.Body.String(), "b/b-m") {
t.Fatalf("503 must summarize every tier, body=%s", rr.Body.String())
}
}
// TestChatAutoQuotaSkip: a slot whose token quota is exhausted is dropped
// from scheduling; with no other slot the chain answers 503 naming the quota.
func TestChatAutoQuotaSkip(t *testing.T) {
ctrl := &upstreamCtrl{}
up := upstream(t, ctrl)
defer up.Close()
g := newTestGateway(t,
config.Source{Name: "a", BaseURL: up.URL, Adapter: "openai", Models: []config.Model{{ID: "a-m", Priority: 100}}},
)
// one slot with an hourly quota of 1 token
rr := doReq(t, g, "PUT", "/api/auto",
`{"rules":[{"model":"a-m","tier":0,"token_quota":1,"period":"hour"}]}`)
if rr.Code != 200 {
t.Fatalf("put auto status=%d body=%s", rr.Code, rr.Body.String())
}
// first request consumes 4 tokens -> quota exhausted
rr = doReq(t, g, "POST", "/v1/chat/completions",
`{"model":"AUTO","messages":[{"role":"user","content":"hi"}]}`)
if rr.Code != 200 {
t.Fatalf("first status=%d body=%s", rr.Code, rr.Body.String())
}
// second request must skip the exhausted slot and fail 503
rr = doReq(t, g, "POST", "/v1/chat/completions",
`{"model":"AUTO","messages":[{"role":"user","content":"hi"}]}`)
if rr.Code != http.StatusServiceUnavailable {
t.Fatalf("quota status=%d body=%s", rr.Code, rr.Body.String())
}
if !strings.Contains(rr.Body.String(), "quota exhausted") {
t.Fatalf("503 must name the quota reason, body=%s", rr.Body.String())
}
if ctrl.hits != 1 {
t.Fatalf("upstream hits = %d, want 1 (exhausted slot must not be called)", ctrl.hits)
}
}
// TestAutoStatesReportChainHealth: GET /api/auto reports per-slot health for
// the priority-page UI; a chain edit resets the failure state to zero.
func TestAutoStatesReportChainHealth(t *testing.T) {
ctrl := &upstreamCtrl{status: 500}
up := upstream(t, ctrl)
defer up.Close()
g := newTestGateway(t,
config.Source{Name: "a", BaseURL: up.URL, Adapter: "openai", Models: []config.Model{{ID: "a-m", Priority: 100}}},
)
doReq(t, g, "POST", "/v1/chat/completions",
`{"model":"AUTO","messages":[{"role":"user","content":"hi"}]}`)
fetch := func() []core.AutoSlotState {
rr := doReq(t, g, "GET", "/api/auto", "")
if rr.Code != 200 {
t.Fatalf("get auto status=%d body=%s", rr.Code, rr.Body.String())
}
var body struct {
Rules []config.ModelScope `json:"rules"`
States []core.AutoSlotState `json:"states"`
}
if err := json.Unmarshal(rr.Body.Bytes(), &body); err != nil {
t.Fatalf("unmarshal: %v", err)
}
return body.States
}
st := fetch()
if len(st) != 1 || st[0].Model != "a-m" || st[0].Source != "a" {
t.Fatalf("want 1 slot a/a-m, got %#v", st)
}
if st[0].FailCount == 0 || !st[0].Cooling {
t.Fatalf("slot must report the failure (fail=%d cooling=%v)", st[0].FailCount, st[0].Cooling)
}
doReq(t, g, "PUT", "/api/auto",
`{"rules":[{"model":"a-m","tier":0}]}`)
st = fetch()
if st[0].FailCount != 0 || st[0].Cooling {
t.Fatalf("edit must reset health, got %#v", st[0])
}
}
// TestAutoSaveResetsCooldown: editing the AUTO chain clears the cooldown of
// its slots, so a fixed upstream is schedulable again without waiting (P1).
func TestAutoSaveResetsCooldown(t *testing.T) {
ctrl := &upstreamCtrl{status: 500}
up := upstream(t, ctrl)
defer up.Close()
g := newTestGateway(t,
config.Source{Name: "a", BaseURL: up.URL, Adapter: "openai", Models: []config.Model{{ID: "a-m", Priority: 100}}},
)
rr := doReq(t, g, "POST", "/v1/chat/completions",
`{"model":"AUTO","messages":[{"role":"user","content":"hi"}]}`)
if rr.Code != http.StatusServiceUnavailable {
t.Fatalf("expect 503 while upstream down, got %d", rr.Code)
}
p := g.core.ProviderForSlot("a-m", "a")
if p == nil || p.ModelAvailable("a-m") {
t.Fatal("a-m must be cooling after the failure")
}
// editing the chain (same rules) must clear the cooldown immediately
rr = doReq(t, g, "PUT", "/api/auto",
`{"rules":[{"model":"a-m","tier":0}]}`)
if rr.Code != 200 {
t.Fatalf("put auto status=%d body=%s", rr.Code, rr.Body.String())
}
if !p.ModelAvailable("a-m") {
t.Fatal("SaveAutoRules must reset the slot cooldown")
}
// healed upstream -> AUTO serves again on the next request
ctrl.status = 0
rr = doReq(t, g, "POST", "/v1/chat/completions",
`{"model":"AUTO","messages":[{"role":"user","content":"hi"}]}`)
if rr.Code != 200 {
t.Fatalf("AUTO after reset status=%d body=%s", rr.Code, rr.Body.String())
}
}
func TestChatSingle(t *testing.T) {
up := mockUpstream()
defer up.Close()
@ -446,4 +633,4 @@ func TestAPIChatInternal(t *testing.T) {
if !strings.Contains(rr.Body.String(), "pong") {
t.Fatalf("api chat body=%s", rr.Body.String())
}
}
}

5
internal/gateway/keys.go
View File

@ -142,7 +142,10 @@ func (g *Gateway) allowedModels(ctx context.Context) []config.ModelScope {
// current rules; PUT /api/auto replaces them (admin only).
func (g *Gateway) handleAutoAPI(w http.ResponseWriter, r *http.Request) {
if r.Method == http.MethodGet {
writeJSON(w, http.StatusOK, map[string]interface{}{"rules": g.core.AutoRules()})
writeJSON(w, http.StatusOK, map[string]interface{}{
"rules": g.core.AutoRules(),
"states": g.core.AutoSlotStates(),
})
return
}
if reqRole(r.Context()) != "admin" {

30
internal/gateway/server.go
View File

@ -15,7 +15,6 @@ import (
"net/url"
"strings"
"sync"
"sync/atomic"
"time"
"llmsproxy/internal/config"
@ -27,12 +26,11 @@ var uiFS embed.FS
// Gateway is the HTTP handler for the OpenAI-compatible endpoint + web UI.
type Gateway struct {
core *core.Core
ui http.Handler
stats *Stats
probeMu sync.Mutex
lastProbe time.Time
autoRR atomic.Uint64
core *core.Core
ui http.Handler
stats *Stats
probeMu sync.Mutex
lastProbe time.Time
}
func New(c *core.Core, gatewayKeys []string) (*Gateway, error) {
@ -131,6 +129,8 @@ func (g *Gateway) routes(w http.ResponseWriter, r *http.Request) {
g.handleChat(w, r)
case r.URL.Path == "/api/status":
g.handleStatusAPI(w, r)
case r.URL.Path == "/api/status/reset":
g.handleResetHealth(w, r)
case r.URL.Path == "/api/stats" || strings.HasPrefix(r.URL.Path, "/api/stats/"):
g.handleStatsAPI(w, r)
case r.URL.Path == "/api/keys" || strings.HasPrefix(r.URL.Path, "/api/keys/"):
@ -389,6 +389,22 @@ func (g *Gateway) ensureProbe(ctx context.Context) {
}
}
// handleResetHealth (admin) clears the per-source backoff state so a fixed
// upstream or an edited AUTO priority chain becomes schedulable immediately.
func (g *Gateway) handleResetHealth(w http.ResponseWriter, r *http.Request) {
if r.Method != http.MethodPost {
writeError(w, http.StatusMethodNotAllowed, "method_not_allowed", "use POST")
return
}
if reqRole(r.Context()) != "admin" {
writeError(w, http.StatusForbidden, "forbidden", "admin role required")
return
}
g.core.ResetHealth()
g.stats.AppendAudit("config", map[string]interface{}{"action": "reset_health", "key": keyID(reqKey(r.Context()))})
writeJSON(w, http.StatusOK, map[string]interface{}{"ok": true})
}
func (g *Gateway) handleStatusAPI(w http.ResponseWriter, r *http.Request) {
g.ensureProbe(r.Context())
host := r.Host

94
internal/gateway/stats.go
View File

@ -3,7 +3,11 @@ package gateway
import (
"bufio"
"encoding/json"
"fmt"
"os"
"path/filepath"
"sort"
"strconv"
"sync"
"time"
)
@ -56,6 +60,7 @@ type Stats struct {
bySrc map[string]*Stat
byKeyModel map[string]map[string]*Stat
byKeySrc map[string]map[string]*Stat
byStatus map[int]*Stat // per http status code aggregates (incl. 402/400)
recs []Req
maxRecs int
auditPath string
@ -64,6 +69,15 @@ type Stats struct {
const hourSec = 3600
// auditRotateBytes rotates the audit file once it grows past this size (the
// file is renamed to <path>.<unix>.old and a fresh one is started); pruning
// keeps at most auditKeepOld rotated files. Both are vars so tests can shrink
// the threshold.
var (
auditRotateBytes int64 = 64 << 20
auditKeepOld = 10
)
func NewStats(maxRecords int) *Stats {
if maxRecords <= 0 {
maxRecords = 3000
@ -74,6 +88,7 @@ func NewStats(maxRecords int) *Stats {
bySrc: map[string]*Stat{},
byKeyModel: map[string]map[string]*Stat{},
byKeySrc: map[string]map[string]*Stat{},
byStatus: map[int]*Stat{},
modelHour: map[string]map[int64]int64{},
maxRecs: maxRecords,
}
@ -98,6 +113,10 @@ func inc(m map[string]*Stat, name string, r Req) {
a = &Stat{}
m[name] = a
}
incStatus(a, name, r)
}
func incStatus(a *Stat, name string, r Req) {
a.Reqs++
if r.OK {
a.OK++
@ -160,6 +179,15 @@ func (s *Stats) Record(r Req) {
s.byKeySrc[r.Key] = ks
}
inc(ks, r.Source, r)
if r.Status > 0 {
name := strconv.Itoa(r.Status)
a := s.byStatus[r.Status]
if a == nil {
a = &Stat{}
s.byStatus[r.Status] = a
}
incStatus(a, name, r)
}
// window bucket for quota enforcement (per source-model pair, per unix hour)
tok := r.Prompt + r.Compl
if tok > 0 && r.Model != "" {
@ -187,28 +215,32 @@ func (s *Stats) Record(r Req) {
s.recs = s.recs[len(s.recs)-s.maxRecs:]
}
if s.auditPath != "" {
if f, err := os.OpenFile(s.auditPath, os.O_CREATE|os.O_APPEND|os.O_WRONLY, 0644); err == nil {
if b, err := json.Marshal(r); err == nil {
_, _ = f.Write(append(b, '\n'))
}
_ = f.Close()
s.rotateAuditLocked()
appendAuditLine(s.auditPath, r)
}
}
// rotateAuditLocked renames the audit file to <path>.<unix>.old once it
// exceeds auditRotateBytes and prunes old files beyond auditKeepOld, keeping
// the newest ones. Caller must hold s.mu.
func (s *Stats) rotateAuditLocked() {
if s.auditPath == "" || auditRotateBytes <= 0 {
return
}
if fi, err := os.Stat(s.auditPath); err == nil && fi.Size() < auditRotateBytes {
return
}
ts := time.Now().Unix()
if os.Rename(s.auditPath, fmt.Sprintf("%s.%d.old", s.auditPath, ts)) == nil {
old, _ := filepath.Glob(s.auditPath + ".*.old")
sort.Sort(sort.Reverse(sort.StringSlice(old)))
for i := auditKeepOld; i < len(old); i++ {
_ = os.Remove(old[i])
}
}
}
// AppendAudit writes a generic event line (access log entry, login event,
// config change, …) to the same audit file without touching the aggregates.
func (s *Stats) AppendAudit(obj string, data map[string]interface{}) {
s.mu.Lock()
path := s.auditPath
s.mu.Unlock()
if path == "" {
return
}
row := map[string]interface{}{"obj": obj, "time": time.Now().UnixMilli()}
for k, v := range data {
row[k] = v
}
func appendAuditLine(path string, row interface{}) {
f, err := os.OpenFile(path, os.O_CREATE|os.O_APPEND|os.O_WRONLY, 0644)
if err != nil {
return
@ -219,6 +251,22 @@ func (s *Stats) AppendAudit(obj string, data map[string]interface{}) {
}
}
// AppendAudit writes a generic event line (access log entry, login event,
// config change, …) to the same audit file without touching the aggregates.
func (s *Stats) AppendAudit(obj string, data map[string]interface{}) {
row := map[string]interface{}{"obj": obj, "time": time.Now().UnixMilli()}
for k, v := range data {
row[k] = v
}
s.mu.Lock()
defer s.mu.Unlock()
if s.auditPath == "" {
return
}
s.rotateAuditLocked()
appendAuditLine(s.auditPath, row)
}
// ModelTokens returns the tokens consumed per model for one gateway key id
// (used for per-model token quota enforcement).
func (s *Stats) ModelTokens(key string) map[string]int64 {
@ -376,12 +424,22 @@ func (s *Stats) Snapshot(limit int, key string) map[string]interface{} {
total.LatMax = a.LatMax
}
}
bs := make([]agrRow, 0, len(s.byStatus))
for code := range s.byStatus {
bs = append(bs, agrRow{Name: strconv.Itoa(code), Stat: *s.byStatus[code]})
}
sort.Slice(bs, func(i, j int) bool {
ci, _ := strconv.Atoi(bs[i].Name)
cj, _ := strconv.Atoi(bs[j].Name)
return ci < cj
})
return map[string]interface{}{
"active": s.active,
"total": total,
"by_key": rows(byKey),
"by_model": rows(byModel),
"by_source": rows(bySrc),
"by_status": bs,
"records": append([]Req(nil), recs...),
}
}

88
internal/gateway/stats_test.go Normal file
View File

@ -0,0 +1,88 @@
package gateway
import (
"fmt"
"os"
"path/filepath"
"testing"
)
func TestStatsByStatus(t *testing.T) {
s := NewStats(100)
s.Record(Req{Key: "k", Model: "m", Source: "s", Status: 200, OK: true})
s.Record(Req{Key: "k", Model: "m", Source: "s", Status: 402, OK: false})
s.Record(Req{Key: "k", Model: "m", Source: "s", Status: 400, OK: false})
snap := s.Snapshot(0, "")
bs, ok := snap["by_status"].([]agrRow)
if !ok {
t.Fatalf("by_status missing: %#v", snap["by_status"])
}
if len(bs) != 3 {
t.Fatalf("want 3 status buckets, got %d: %#v", len(bs), bs)
}
if bs[0].Name != "200" || bs[0].OK != 1 || bs[0].Err != 0 {
t.Fatalf("bucket 200 wrong: %#v", bs[0])
}
if bs[1].Name != "400" || bs[1].Err != 1 {
t.Fatalf("bucket 400 wrong: %#v", bs[1])
}
if bs[2].Name != "402" || bs[2].Err != 1 {
t.Fatalf("bucket 402 wrong: %#v", bs[2])
}
}
func TestAuditRotation(t *testing.T) {
dir := t.TempDir()
path := filepath.Join(dir, "audit.jsonl")
s := NewStats(10)
s.LoadAudit(path)
oldRotate, oldKeep := auditRotateBytes, auditKeepOld
auditRotateBytes, auditKeepOld = 64, 10
defer func() { auditRotateBytes, auditKeepOld = oldRotate, oldKeep }()
oldFiles := func() []string {
matches, _ := filepath.Glob(path + ".*.old")
return matches
}
for i := 0; i < 3; i++ {
s.AppendAudit("ev", map[string]interface{}{"i": i})
}
if got := len(oldFiles()); got != 1 {
t.Fatalf("want 1 rotated file after first overflow, got %d", got)
}
if b, err := os.ReadFile(path); err != nil || len(b) == 0 {
t.Fatalf("active audit file must continue appending: %v %d bytes", err, len(b))
}
// seed 12 fake old files; the next rotation must prune back to keep=10
for i := 1; i <= 12; i++ {
name := fmt.Sprintf("%s.%010d.old", path, i)
_ = os.WriteFile(name, []byte("x\n"), 0644)
}
s.AppendAudit("ev", map[string]interface{}{"i": 98})
s.AppendAudit("ev", map[string]interface{}{"i": 99})
if got := len(oldFiles()); got != auditKeepOld {
t.Fatalf("want keeper %d old files, got %d", auditKeepOld, got)
}
}
func TestAuditRotationRecords(t *testing.T) {
dir := t.TempDir()
path := filepath.Join(dir, "audit.jsonl")
s := NewStats(10)
s.LoadAudit(path)
oldRotate := auditRotateBytes
auditRotateBytes = 64
defer func() { auditRotateBytes = oldRotate }()
for i := 0; i < 5; i++ {
s.Record(Req{Key: "k", Model: "m", Source: "s", Status: 200, OK: true})
}
matches, _ := filepath.Glob(path + ".*.old")
if len(matches) != 1 {
t.Fatalf("Record must rotate too: got %d old files", len(matches))
}
}

55
internal/gateway/ui/index.html
View File

@ -306,6 +306,14 @@ html[data-theme="dark"] .dropzone.dragover, html[data-theme="dark"] .dropzone:ho
.scr-block .scr-tag { flex:0 0 auto; font-family:ui-monospace,Menlo,Consolas,monospace; font-size:10.5px;
padding:2px 7px; border-radius:9px; background:rgba(0,0,0,.24); color:#ffe9a8;
border:1px solid rgba(255,220,130,.35); cursor:pointer; }
.scr-block .scr-htag { flex:0 0 auto; display:flex; gap:4px; align-items:center;
font-family:ui-monospace,Menlo,Consolas,monospace; font-size:10px; font-weight:700; cursor:help; }
.scr-htag .ht-cool { padding:2px 6px; border-radius:9px; background:rgba(255,80,80,.28); color:#ffd9d9;
border:1px solid rgba(255,120,120,.5); }
.scr-htag .ht-fail { padding:2px 6px; border-radius:9px; background:rgba(255,160,60,.2); color:#ffd9a8;
border:1px solid rgba(255,180,90,.42); }
.scr-htag .ht-pref { padding:2px 6px; border-radius:9px; background:rgba(120,180,255,.18); color:#cfe3ff;
border:1px solid rgba(150,190,255,.38); }
.scr-block .scr-grip { flex:0 0 auto; display:flex; flex-direction:column; gap:2px; padding:6px 4px;
margin-left:2px; border-radius:6px; cursor:grab; background:rgba(255,255,255,.18);
box-shadow:inset 0 1px 2px rgba(0,0,0,.18); transition:background .12s; touch-action:none; }
@ -471,9 +479,11 @@ const STR = {
seedWarnTitle:'请更换初始管理员密钥', seedWarnText:'当前登录的是配置文件中的初始密钥，明文写入 config.yaml、存在泄露风险。请在下方创建新的管理员密钥，用新密钥登录后删除此初始密钥。', seedWarnGo:'去更换密钥', seedWarnLater:'稍后', seedWarnDismiss:'本次不再提示',
sortTitle:'拖拽积木配置模型优先级', sortHint:'每行 = 一个优先级档位，行从上到下优先级递减；同一行的模型并排，视为同优先级。按住积木右侧 ⠿ 把手拖动：拖到行内 = 放入该档位或调整同档顺序，拖到行与行之间的缝隙 = 提升或降低到新档位。生图模型不参与排序。', sortDragGrip:'拖拽前须按住把手',
sortSave:'保存排序', sortReset:'重置', sortAdd:'添加档位', sortSaved:'排序已保存并热重载', sortNoChange:'无变更', sortHintSave:'点击保存排序后生效',
sortCooling:'冷却', sortFail:'失败', sortHealthTip:'冷却 / 失败次数 / 偏好分 实时状态', sortHealthReset:'链上冷却已复位',
sortSource:'源', sortPrio:'优先级 %s', sortEmpty:'该源暂无模型',
kpiActive:'活跃请求', kpiReqs:'总请求', kpiOk:'成功率', kpiTokens:'Tokens', kpiLat:'平均延迟', kpiMaxLat:'最大延迟',
dashModel:'模型用量', dashSrc:'源用量与延迟', dashKey:'密钥用量', dashRecs:'请求记录', exportCsv:'导出 CSV', expWeek:'近一周', expMonth:'近一月', expYear:'近一年', expRange:'自定义范围', expStart:'开始日期', expEnd:'结束日期', expDownload:'下载', expKeysCsv:'导出密钥用量',
dashStatus:'状态码分布', thCode:'状态码', statusTag:'状态码分类统计（含 402 欠费 / 400 schema 错误；两者不计入上游退避但单独计数）',
thModel:'模型', thSrc:'源', thKey:'密钥', thReqs:'请求', thOk:'成功', thErr:'失败',
thPrompt:'输入 Tokens', thCompl:'输出 Tokens', thAvgLat:'平均延迟', thMaxLat:'最长延迟',
thTime:'时间', thType:'类型', thStatus:'状态', thLatMs:'延迟',
@ -525,9 +535,11 @@ kMeTitle:'My key', kMeRole:'Role', kMeModels:'Models I can use', kMeHint:'Keys c
seedWarnTitle:'Replace the initial admin key', seedWarnText:'You are logged in with the seed key from config.yaml. It is plaintext in the config file and a security risk. Create a new admin key below, log in with it, then delete this seed key.', seedWarnGo:'Change my key', seedWarnLater:'Later', seedWarnDismiss:'Don\'t ask again',
sortTitle:'Drag blocks to set model priority', sortHint:'Each row = one priority tier, rows go high→low; models on the same row sit side by side and share that priority. Grab the ⠿ handle on the right of a block to drag: drop into a row = join that tier or reorder within it, drop into the gap between rows = move up/down a tier. Image models stay out.', sortDragGrip:'grab the handle to drag',
sortSave:'Save order', sortReset:'Reset', sortAdd:'Add slot', sortSaved:'Order saved & hot-reloaded', sortNoChange:'No changes', sortHintSave:'Click Save for it to take effect',
sortCooling:'cooling', sortFail:'fail', sortHealthTip:'live cooldown / failures / preference score', sortHealthReset:'chain cooldowns reset',
sortSource:'source', sortPrio:'priority %s', sortEmpty:'no models in this source',
kpiActive:'Active requests', kpiReqs:'Requests', kpiOk:'Success rate', kpiTokens:'Tokens', kpiLat:'Avg latency', kpiMaxLat:'Max latency',
dashModel:'Model usage', dashSrc:'Source usage & latency', dashKey:'Key usage', dashRecs:'Request records', exportCsv:'Export CSV', expWeek:'Last week', expMonth:'Last month', expYear:'Last year', expRange:'Custom range', expStart:'Start date', expEnd:'End date', expDownload:'Download',
dashStatus:'Status codes', thCode:'Code', statusTag:'Per-status aggregates — 402 quota / 400 schema errors are counted here but never back off the provider',
thModel:'Model', thSrc:'Source', thKey:'Key', thReqs:'Requests', thOk:'OK', thErr:'Err',
thPrompt:'Prompt Tokens', thCompl:'Completion Tokens', thAvgLat:'Avg latency', thMaxLat:'Max latency',
thTime:'Time', thType:'Type', thStatus:'Status', thLatMs:'Latency',
@ -658,6 +670,7 @@ async function renderStatus() {
<div class="card"><h2>${t('dashModel')}</h2><div id="tb-model"></div></div>
${s.sources ? `<div class="card"><h2>${t('dashSrc')}</h2><div id="tb-src"></div></div>` : ''}
</div>
<div class="card"><h2>${t('dashStatus')} <span class="muted" style="font-weight:400;font-size:12px">${t('statusTag')}</span></h2><div id="tb-status"></div></div>
<div class="card"><h2>${t('dashKey')}<span class="grow"></span><button class="ghost small" onclick="openExportModal()">${t('exportCsv')}</button></h2><div id="tb-key"></div></div>
<div class="card"><h2><span>${t('dashRecs')}</span><span class="grow"></span><button class="ghost small" onclick="openExportModal()">${t('exportCsv')}</button></h2>
<div class="filter-line">
@ -732,6 +745,7 @@ async function paintStats() {
<div class="kpi"><div class="k-lab">${t('kpiLat')}</div><div class="k-val">${fmtMs(avg)}</div><div class="k-sub">${t('kpiMaxLat')} ${fmtMs(tot.latency_max_ms)}</div></div>`;
paintModelTable(st.by_model || []);
paintSrcTable(st.by_source || []);
paintStatusTable(st.by_status || []);
paintKeyTable(st.by_key || [], st.key_names || {});
paintRecords(st.records || [], st.key_names || {});
renderKeySelect((st.by_key || []).map(k => k.name));
@ -762,6 +776,13 @@ function paintSrcTable(rows) {
<td class="num">${fmtTok(r.tokens)}</td>
<td class="num">${fmtMs(fmtLat(r.latency_sum_ms, r.reqs))}</td><td class="num">${fmtMs(r.latency_max_ms)}</td></tr>`).join('') + '</table></div>';
}
function paintStatusTable(rows) {
const el = $('#tb-status'); if (!el) return;
if (!rows.length) { el.innerHTML = `<div class="muted">${t('noUsage')}</div>`; return; }
el.innerHTML = `<div class="tbl-wrap"><table><tr><th>${t('thCode')}</th><th class="num">${t('thReqs')}</th><th class="num">${t('thOk')}</th><th class="num">${t('thErr')}</th></tr>` +
rows.map(r => `<tr><td><b class="${+r.name >= 400 ? 'errc' : 'okc'}">${esc(r.name)}</b></td>
<td class="num">${fmtN(r.reqs)}</td><td class="num okc">${fmtN(r.ok)}</td><td class="num errc">${fmtN(r.err)}</td></tr>`).join('') + '</table></div>';
}
function paintKeyTable(rows, keyNames) {
const el = $('#tb-key'); if (!el) return;
if (!rows.length) { el.innerHTML = `<div class="muted">${t('noUsage')}</div>`; return; }
@ -1139,10 +1160,15 @@ function srcColor(name) {
}
function srcShort(name) { return (name || '?').slice(0, 2).toUpperCase(); }
const sortState = { lanes: [], origin: null, drag: null };
let sortStateMap = new Map();
async function renderSort() {
const j = await api('/api/sources');
let autoR = [];
try { autoR = (await api('/api/auto')).rules || []; } catch (e) {}
try {
const a = await api('/api/auto');
autoR = a.rules || [];
sortStateMap = new Map((a.states || []).map(st => [st.model + '|' + (st.source || '*'), st]));
} catch (e) {}
const byModel = new Map();
const byPair = new Map();
const sourceRows = new Map();
@ -1194,6 +1220,16 @@ async function renderSort() {
</div>`;
paintSort();
}
function healthTag(it) {
const st = sortStateMap.get(it.id + '|' + (it.src || '*'));
if (!st) return '';
const bits = [];
if (st.cooling) bits.push(`<span class="ht-cool">${esc(t('sortCooling'))}</span>`);
if (st.fail_count > 0) bits.push(`<span class="ht-fail">${esc(t('sortFail') + '×' + st.fail_count)}</span>`);
if (st.pref !== 0) bits.push(`<span class="ht-pref">${esc(st.pref > 0 ? '+' + st.pref : '' + st.pref)}</span>`);
if (!bits.length) return '';
return `<span class="scr-htag" title="${escAttr(t('sortHealthTip'))}">${bits.join('')}</span>`;
}
function scrBlockHtml(it, isFirst, li, ji, extraClass) {
const c = srcColor(it.src);
const s = srcShort(it.src);
@ -1206,6 +1242,7 @@ function scrBlockHtml(it, isFirst, li, ji, extraClass) {
<span class="scr-ico">${esc(it.src === '*' ? '+' : s)}</span>
<span class="scr-name">${esc(it.id)}<em class="scr-srcname">${esc(it.src === '*' ? t('kAnySrc') : it.src)}</em></span>
${it.meta ? `<span class="scr-tag">${esc(quantBadge(it.meta.quota, it.meta.period, it.meta.hours))}</span>` : ''}
${healthTag(it)}
<span class="scr-x" title="${escAttr(t('kDelB2'))}" onclick="event.stopPropagation();scrDelSlot('${li}','${ji}')">×</span>
<span class="scr-grip"><i></i><i></i><i></i></span>
</div>`;
@ -1521,7 +1558,7 @@ async function saveSort() {
try {
await persistAuto();
sortState.origin = JSON.stringify(sortState.lanes);
toast(t('sortSaved'));
toast(t('sortSaved') + ' · ' + t('sortHealthReset'));
} catch (e) { toast(e.message); }
}
@ -1986,10 +2023,20 @@ function showCtx(x, y, items) {
const r = w.getBoundingClientRect();
w.style.left = Math.max(6, Math.min(x, window.innerWidth - r.width - 6)) + 'px';
w.style.top = Math.max(6, Math.min(y, window.innerHeight - r.height - 6)) + 'px';
setTimeout(() => document.addEventListener('click', hideCtx2, { once: true }), 10);
}
function hideCtx2() { hideCtx(); }
function hideCtx() { if (ctxEl) { ctxEl.remove(); ctxEl = null; } }
// Close the ctx menu on any primary click/press OUTSIDE the menu. Both
// listeners run in the CAPTURE phase, so they fire even when the clicked
// element stops propagation (priority blocks and key bricks call
// stopPropagation in their own click handlers, which would otherwise keep the
// menu open forever). Presses INSIDE the menu are left alone: the menu's own
// click handler closes it after running the item action.
document.addEventListener('click', e => {
if (e.button === 0 && !(ctxEl && ctxEl.contains(e.target))) hideCtx();
}, true);
document.addEventListener('mousedown', e => {
if (e.button === 0 && !(ctxEl && ctxEl.contains(e.target))) hideCtx();
}, true);
/* cross-canvas brick dragging */
function bindBrickDrag(b) {
b.addEventListener('dragstart', e => {

324
internal/provider/provider.go
View File

@ -13,6 +13,7 @@ import (
"net/http"
"strings"
"sync"
"sync/atomic"
"time"
"llmsproxy/internal/config"
@ -20,36 +21,109 @@ import (
"llmsproxy/internal/types"
)
// health tracks availability with exponential backoff.
type health struct {
failCount int
unavailableUntil time.Time
permanent bool
// ---- per-(source,model) scheduling state ----
// ModelState is the scheduling state of one (source, model) pair: a soft
// preference score used to order candidates within a priority tier, a
// consecutive-failure counter driving exponential cooldown, and a hard
// cooldown deadline. All fields are atomic and cooldown expiry is evaluated
// lazily (no timers, no goroutines). A model is never permanently blacklisted:
// cooldown always expires and any success resets the state, so a fixed
// upstream recovers on its own.
type ModelState struct {
pref atomic.Int64 // +1 per success / -5 per failure, clamped
failCount atomic.Int64
cooldownUntil atomic.Int64 // unix seconds; 0 = schedulable
}
func (h *health) reset() { h.failCount = 0; h.unavailableUntil = time.Time{}; h.permanent = false }
const (
prefFailStep = 5
prefMin = -20
prefMax = 20
backoffBase = 5 * time.Second
backoffCap = 30 * time.Minute
// failures at/above this count back off at the capped duration
backoffCapN = 10
)
func (h *health) available() bool {
if h.permanent {
return false
func clampPref(a *atomic.Int64, lo, hi int64) {
for {
cur := a.Load()
if cur < lo {
if a.CompareAndSwap(cur, lo) {
return
}
continue
}
if cur > hi {
if a.CompareAndSwap(cur, hi) {
return
}
continue
}
return
}
return time.Now().After(h.unavailableUntil)
}
func (h *health) backoff() {
h.failCount++
cooldown := 5 * time.Second * time.Duration(1<<(h.failCount-1))
if cooldown > 30*time.Minute {
cooldown = 30 * time.Minute
// Available reports whether the model may be scheduled right now (cooldown
// expired or not yet set).
func (s *ModelState) Available() bool {
return s.cooldownUntil.Load() <= time.Now().Unix()
}
// RecordFailure counts one consecutive failure and schedules exponential
// cooldown (5s, 10s, 20s … capped at 30min). auth marks 401/403 credential
// failures: it jumps straight to the capped cooldown and doubles the
// preference penalty, but the model still recovers when the cooldown expires.
func (s *ModelState) RecordFailure(auth bool) {
n := s.failCount.Add(1)
if auth && n < backoffCapN {
n = backoffCapN
s.failCount.Store(n) // persist the cap so FailCount() reports it too
}
h.unavailableUntil = time.Now().Add(cooldown)
var cd time.Duration
if n >= backoffCapN {
cd = backoffCap
} else {
cd = backoffBase * time.Duration(1<<(n-1))
if cd > backoffCap {
cd = backoffCap
}
}
s.cooldownUntil.Store(time.Now().Add(cd).Unix())
pen := int64(prefFailStep)
if auth {
pen *= 2
}
s.pref.Add(-pen)
clampPref(&s.pref, prefMin, prefMax)
}
func (h *health) markPermanent() {
h.permanent = true
h.unavailableUntil = time.Time{}
// RecordSuccess resets the failure counter and cooldown and bumps the
// preference score by one.
func (s *ModelState) RecordSuccess() {
s.failCount.Store(0)
s.cooldownUntil.Store(0)
s.pref.Add(1)
clampPref(&s.pref, prefMin, prefMax)
}
func (s *ModelState) reset() {
s.failCount.Store(0)
s.cooldownUntil.Store(0)
s.pref.Store(0)
}
// Pref is the current preference score (higher = preferred).
func (s *ModelState) Pref() int64 { return s.pref.Load() }
// FailCount is the number of consecutive failures.
func (s *ModelState) FailCount() int64 { return s.failCount.Load() }
// CooldownUntil is the unix timestamp until which the model is cooled; 0 when
// schedulable.
func (s *ModelState) CooldownUntil() int64 { return s.cooldownUntil.Load() }
// Provider is a single configured upstream LLM source.
type Provider struct {
cfg config.Source
@ -59,11 +133,11 @@ type Provider struct {
mu sync.Mutex
sem chan struct{}
health health
states map[string]*ModelState // key = model id
lastProbe struct {
ok bool
err string
at int64
ok bool
err string
at int64
}
}
@ -74,17 +148,21 @@ func New(cfg config.Source, vm *lua.VM) *Provider {
adapter: cfg.Adapter,
client: &http.Client{Timeout: cfg.Timeout},
sem: make(chan struct{}, cfg.MaxConcurrent),
states: map[string]*ModelState{},
}
if cfg.MaxConcurrent <= 0 {
p.sem = nil
}
for _, m := range cfg.Models {
p.states[m.ID] = &ModelState{}
}
return p
}
func (p *Provider) Name() string { return p.cfg.Name }
func (p *Provider) Adapter() string { return p.cfg.Adapter }
func (p *Provider) MaxConcurrent() int { return p.cfg.MaxConcurrent }
func (p *Provider) Config() *config.Source { return &p.cfg }
func (p *Provider) Name() string { return p.cfg.Name }
func (p *Provider) Adapter() string { return p.cfg.Adapter }
func (p *Provider) MaxConcurrent() int { return p.cfg.MaxConcurrent }
func (p *Provider) Config() *config.Source { return &p.cfg }
// Models returns the model ids exposed by this source.
func (p *Provider) Models() []string {
@ -174,10 +252,59 @@ func (p *Provider) ImageURL() string {
// ---- availability ----
// ErrBusy is returned when every concurrency slot of a source is in use. It
// is a soft signal: schedulers skip a busy candidate without recording any
// failure (busy is not a failure) and gateways map it to HTTP 429. It aliases
// types.ErrBusy so the scheduler layer (which must not depend on this package)
// can detect busy via the shared sentinel.
var ErrBusy = types.ErrBusy
// Available reports whether the source is schedulable at source level: at
// least one of its models is not cooling down. Per-model scheduling decisions
// must use ModelAvailable instead.
func (p *Provider) Available() bool {
p.mu.Lock()
defer p.mu.Unlock()
return p.health.available()
for _, s := range p.states {
if s.Available() {
return true
}
}
return false
}
// ModelAvailable reports whether the exact model is schedulable right now
// (its cooldown expired). An unknown model id is treated as available.
func (p *Provider) ModelAvailable(model string) bool {
p.mu.Lock()
defer p.mu.Unlock()
if s, ok := p.states[model]; ok {
return s.Available()
}
return true
}
// Pref returns the adaptive preference score of a model (higher = preferred
// within a priority tier). Used by the AUTO chain to order same-tier slots.
func (p *Provider) Pref(model string) int64 {
return p.state(model).Pref()
}
// ResetModelCooldown clears the cooldown and failure counter of a single
// model while preserving its preference score. Called after AUTO-chain edits
// so edited slots become schedulable immediately (a stored preference for a
// reliably good model is kept).
func (p *Provider) ResetModelCooldown(model string) {
s := p.state(model)
s.failCount.Store(0)
s.cooldownUntil.Store(0)
}
// ModelHealthInfo exposes the per-model scheduling state for the web UI.
// An unknown model id reports zeros.
func (p *Provider) ModelHealthInfo(model string) (pref, failCount, cooldownUntil int64) {
s := p.state(model)
return s.Pref(), s.FailCount(), s.CooldownUntil()
}
// Probe performs a lightweight reachability + auth check against the source.
@ -274,35 +401,106 @@ func (p *Provider) LastProbe() (bool, string, int64) {
return p.lastProbe.ok, p.lastProbe.err, p.lastProbe.at
}
// ReportStatus records an upstream HTTP status for backoff decisions.
func (p *Provider) ReportStatus(code int) {
// state returns the ModelState for a model id, creating it on first use so
// dynamically requested models are still tracked. The registry keeps states
// alive across provider rebuilds only for configured models; a lazily created
// state simply lives for the provider's lifetime.
func (p *Provider) state(model string) *ModelState {
p.mu.Lock()
defer p.mu.Unlock()
s, ok := p.states[model]
if !ok {
s = &ModelState{}
p.states[model] = s
}
return s
}
// RecordFailure records a failed downstream attempt on model: consecutive
// failure count +1 and exponential cooldown (5s·2^n, capped at 30min).
// code 401/403 is treated as a credential problem: the cooldown jumps to the
// cap and the preference penalty doubles, but the model still recovers when
// the cooldown expires (no permanent blacklist). code 0 = transport failure.
func (p *Provider) RecordFailure(model string, code int) {
p.state(model).RecordFailure(code == 401 || code == 403)
}
// RecordSuccess resets the model's failure counter / cooldown and bumps its
// preference by one.
func (p *Provider) RecordSuccess(model string) {
p.state(model).RecordSuccess()
}
// ReportStatus records an upstream HTTP status for the given model and drives
// the (source, model) backoff state. 401/403 → capped self-healing cooldown
// with doubled penalty; 429 and 5xx → normal exponential backoff. Other codes
// (400 client schema errors, 402 billing errors) are not penalized here —
// they surface via the status page / audit instead.
func (p *Provider) ReportStatus(model string, code int) {
if code == 401 || code == 403 {
p.health.markPermanent()
p.RecordFailure(model, code)
return
}
if code >= 500 || code == 429 {
p.health.backoff()
p.RecordFailure(model, code)
}
}
func (p *Provider) reportError() {
// ResetHealth resets the scheduling state of every model of this source
// (cooldown and preference to zero), so the source becomes fully schedulable
// again. Called after AUTO-chain edits and from the admin UI.
func (p *Provider) ResetHealth() {
p.mu.Lock()
p.health.backoff()
p.mu.Unlock()
defer p.mu.Unlock()
for _, s := range p.states {
s.reset()
}
}
func (p *Provider) reportOK() {
// HealthInfo exposes the source-level backoff state for the status page: the
// highest failure count and the latest cooldown deadline across all models of
// this source. permanent is always false — the permanent-blacklist semantics
// were removed; every cooldown expires on its own.
func (p *Provider) HealthInfo() (failCount int, until time.Time, permanent bool) {
p.mu.Lock()
p.health.reset()
p.mu.Unlock()
defer p.mu.Unlock()
now := time.Now().Unix()
for _, s := range p.states {
if n := int(s.FailCount()); n > failCount {
failCount = n
}
if t := s.CooldownUntil(); t > now && t > until.Unix() {
until = time.Unix(t, 0)
}
}
return failCount, until, false
}
// ---- concurrency limiting ----
// TryAcquire takes one concurrency slot without blocking: it returns nil when
// a slot is free and ErrBusy when the source is at capacity. A nil semaphore
// (MaxConcurrent <= 0) means unlimited and always succeeds. TryAcquire is the
// single busy/idle signal for schedulers; a busy source is skipped, never
// penalized.
func (p *Provider) TryAcquire(ctx context.Context) error {
if p.sem == nil {
return nil
}
select {
case p.sem <- struct{}{}:
return nil
case <-ctx.Done():
return ctx.Err()
default:
return ErrBusy
}
}
// Acquire waits for a free concurrency slot (bounded by cfg.QueueTimeout),
// or context cancel. The HTTP call itself is not truncated.
// or context cancel. The HTTP call itself is not truncated. Direct requests
// historically queued here; the scheduler now prefers TryAcquire so a full
// source fails fast instead of blocking the whole chain.
func (p *Provider) Acquire(ctx context.Context) error {
if p.sem == nil {
return nil
@ -367,11 +565,14 @@ func (p *Provider) buildHeaders(body, url string) (http.Header, error) {
// ---- chat ----
// Chat performs a non-streaming round trip and returns the unified response.
// It fails fast with ErrBusy when the source is at capacity; success/failure
// is recorded against the resolved (source, model) scheduling state.
func (p *Provider) Chat(ctx context.Context, req *types.ChatRequest) (*types.UnifiedResponse, error) {
if err := p.Acquire(ctx); err != nil {
if err := p.TryAcquire(ctx); err != nil {
return nil, err
}
defer p.Release()
model := p.ModelFor(req.Model)
body, err := marshalTransform(p.vm, p.adapter, "transform_request", req)
if err != nil {
@ -383,11 +584,11 @@ func (p *Provider) Chat(ctx context.Context, req *types.ChatRequest) (*types.Uni
}
raw, status, err := p.do(ctx, p.URL(), body, hdrs)
if err != nil {
p.reportError()
p.RecordFailure(model, 0)
return nil, err
}
if status != 200 {
p.ReportStatus(status)
p.ReportStatus(model, status)
return nil, fmt.Errorf("api error %d: %s", status, truncate(raw, 500))
}
unified, err := p.vm.Transform(p.adapter, "transform_response", raw)
@ -398,15 +599,21 @@ func (p *Provider) Chat(ctx context.Context, req *types.ChatRequest) (*types.Uni
if err := json.Unmarshal([]byte(unified), &out); err != nil {
return nil, fmt.Errorf("unmarshal unified response: %w (body: %s)", err, unified)
}
p.reportOK()
p.RecordSuccess(model)
return &out, nil
}
// ChatStream performs a streaming round trip, emitting unified chunks.
// ChatStream performs a streaming round trip, emitting unified chunks. It
// fails fast with ErrBusy when the source is at capacity. Only a failure
// before the first chunk (connect error or non-200 status) is recorded
// against the (source, model) state; afterwards the stream is pinned. A clean
// end ([DONE] or EOF without read errors, and no client disconnect) counts as
// success and resets the cooldown.
func (p *Provider) ChatStream(ctx context.Context, req *types.ChatRequest) (<-chan types.UnifiedChunk, error) {
if err := p.Acquire(ctx); err != nil {
if err := p.TryAcquire(ctx); err != nil {
return nil, err
}
model := p.ModelFor(req.Model)
req.Stream = true
body, err := marshalTransform(p.vm, p.adapter, "transform_request", req)
if err != nil {
@ -432,14 +639,14 @@ func (p *Provider) ChatStream(ctx context.Context, req *types.ChatRequest) (<-ch
ch := make(chan types.UnifiedChunk, 64)
sel := <-rc
if sel.err != nil {
p.reportError()
p.RecordFailure(model, 0)
p.Release()
return nil, sel.err
}
if sel.resp.StatusCode != 200 {
raw, _ := io.ReadAll(sel.resp.Body)
sel.resp.Body.Close()
p.ReportStatus(sel.resp.StatusCode)
p.ReportStatus(model, sel.resp.StatusCode)
p.Release()
return nil, fmt.Errorf("api error %d: %s", sel.resp.StatusCode, truncate(string(raw), 500))
}
@ -485,16 +692,25 @@ func (p *Provider) ChatStream(ctx context.Context, req *types.ChatRequest) (<-ch
return
}
}
// The stream ended cleanly ([DONE] seen or EOF without an upstream
// read error): record success so a previously cooled model can be
// retried. A client disconnect or mid-stream read error is neither
// success nor failure for scheduling purposes.
if ctx.Err() == nil && scanner.Err() == nil {
p.RecordSuccess(model)
}
}()
return ch, nil
}
// Image generates images via /v1/images/generations.
// Image generates images via /v1/images/generations. Same scheduling-state
// accounting as Chat: fail fast on busy, record per (source, model).
func (p *Provider) Image(ctx context.Context, req *types.ImageGenRequest) (*types.UnifiedResponse, error) {
if err := p.Acquire(ctx); err != nil {
if err := p.TryAcquire(ctx); err != nil {
return nil, err
}
defer p.Release()
model := p.ModelFor(req.Model)
b, _ := json.Marshal(req)
transformed, err := p.vm.Transform(p.adapter+"_image", "transform_request", string(b))
@ -508,11 +724,11 @@ func (p *Provider) Image(ctx context.Context, req *types.ImageGenRequest) (*type
}
raw, status, err := p.do(ctx, p.ImageURL(), transformed, hdrs)
if err != nil {
p.reportError()
p.RecordFailure(model, 0)
return nil, err
}
if status != 200 {
p.ReportStatus(status)
p.ReportStatus(model, status)
return nil, fmt.Errorf("image api error %d: %s", status, truncate(raw, 500))
}
var out types.UnifiedResponse
@ -520,7 +736,7 @@ func (p *Provider) Image(ctx context.Context, req *types.ImageGenRequest) (*type
unified, terr := p.vm.Transform(p.adapter+"_image", "transform_response", raw)
if terr == nil && unified != raw {
if err := json.Unmarshal([]byte(unified), &out); err == nil {
p.reportOK()
p.RecordSuccess(model)
return &out, nil
}
}
@ -529,7 +745,7 @@ func (p *Provider) Image(ctx context.Context, req *types.ImageGenRequest) (*type
return nil, fmt.Errorf("unmarshal image response: %w", err)
}
out.ImageData = img.Data
p.reportOK()
p.RecordSuccess(model)
return &out, nil
}
@ -593,4 +809,4 @@ func truncate(s string, n int) string {
return s
}
return s[:n] + "..."
}
}

205
internal/provider/provider_test.go
View File

@ -3,11 +3,12 @@ package provider
import (
"context"
"encoding/json"
"errors"
"fmt"
"net/http"
"net/http/httptest"
"path/filepath"
"sync"
"sync/atomic"
"testing"
"time"
@ -90,6 +91,49 @@ func TestProviderChatStream(t *testing.T) {
}
}
// TestStreamSuccessClearsBackoff guards P6: a clean streaming end must reset
// a previously cooled (source, model) pair.
func TestStreamSuccessClearsBackoff(t *testing.T) {
var fail atomic.Bool
fail.Store(true)
up := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
if fail.Load() {
w.WriteHeader(500)
return
}
w.Header().Set("Content-Type", "text/event-stream")
fmt.Fprint(w, "data: {\"choices\":[{\"delta\":{\"content\":\"a\"}}]}\n\n")
fmt.Fprint(w, "data: [DONE]\n\n")
}))
defer up.Close()
p := newTestProvider(t, src("mock", up.URL, "openai", "m"))
if _, err := p.Chat(context.Background(), &types.ChatRequest{
Model: "m",
Messages: []types.ChatMessage{{Role: "user", Content: types.StringContent("x")}},
}); err == nil {
t.Fatal("expected first chat to fail")
}
if p.ModelAvailable("m") {
t.Fatal("m must be cooling after the failed chat")
}
fail.Store(false)
ch, err := p.ChatStream(context.Background(), &types.ChatRequest{
Model: "m",
Messages: []types.ChatMessage{{Role: "user", Content: types.StringContent("x")}},
})
if err != nil {
t.Fatalf("stream: %v", err)
}
for range ch {
}
if !p.ModelAvailable("m") {
t.Fatal("clean stream must clear the cooldown")
}
if st := p.state("m"); st.FailCount() != 0 {
t.Fatalf("fail count after clean stream = %d", st.FailCount())
}
}
func TestProviderImage(t *testing.T) {
up := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
fmt.Fprint(w, `{"created":123,"data":[{"b64_json":"QUJD"}]}`)
@ -105,12 +149,71 @@ func TestProviderImage(t *testing.T) {
}
}
func TestProviderBackoff(t *testing.T) {
func TestProviderBackoffPerModel(t *testing.T) {
up := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
w.WriteHeader(500)
fmt.Fprint(w, "boom")
}))
defer up.Close()
// two models on one source: a failure on m1 must not blacklist m2
p := newTestProvider(t, src("mock", up.URL, "openai", "m1", "m2"))
_, err := p.Chat(context.Background(), &types.ChatRequest{
Model: "m1",
Messages: []types.ChatMessage{{Role: "user", Content: types.StringContent("x")}},
})
if err == nil {
t.Fatal("expected error")
}
if p.ModelAvailable("m1") {
t.Fatal("expected m1 to be cooling down")
}
if !p.ModelAvailable("m2") {
t.Fatal("m2 must stay schedulable (per-model isolation)")
}
if !p.Available() {
t.Fatal("source must stay available while any model is schedulable")
}
if st := p.state("m1"); st.FailCount() != 1 {
t.Fatalf("fail count = %d, want 1", st.FailCount())
}
}
func TestProviderAuthFailureSelfHeals(t *testing.T) {
up := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
w.WriteHeader(401)
fmt.Fprint(w, `{"error":"API_KEY_DISABLED"}`)
}))
defer up.Close()
p := newTestProvider(t, src("mock2", up.URL, "openai", "m2"))
_, err := p.Chat(context.Background(), &types.ChatRequest{
Model: "m2",
Messages: []types.ChatMessage{{Role: "user", Content: types.StringContent("x")}},
})
if err == nil {
t.Fatal("expected error")
}
st := p.state("m2")
if st.FailCount() != backoffCapN {
t.Fatalf("auth failure must jump to capped count, got %d", st.FailCount())
}
if until := st.CooldownUntil(); until <= time.Now().Add(25*time.Minute).Unix() {
t.Fatalf("auth failure must cool near the cap (until=%d)", until)
}
if st.Pref() != -2*int64(prefFailStep) {
t.Fatalf("auth failure pref penalty must be doubled, got %d", st.Pref())
}
// not permanent: the reset channel and a later success both restore it
st.reset()
if !p.ModelAvailable("m2") {
t.Fatal("reset must restore schedulability")
}
}
func TestModelStateCooldownAndRecovery(t *testing.T) {
up := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
w.WriteHeader(500)
}))
defer up.Close()
p := newTestProvider(t, src("mock", up.URL, "openai", "m"))
_, err := p.Chat(context.Background(), &types.ChatRequest{
Model: "m",
@ -119,55 +222,79 @@ func TestProviderBackoff(t *testing.T) {
if err == nil {
t.Fatal("expected error")
}
if p.Available() {
t.Fatal("expected provider to be in backoff")
st := p.state("m")
if st.FailCount() != 1 {
t.Fatalf("fail count = %d", st.FailCount())
}
// 401 -> permanent
up2 := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
w.WriteHeader(401)
}))
defer up2.Close()
p2 := newTestProvider(t, src("mock2", up2.URL, "openai", "m2"))
p2.Chat(context.Background(), &types.ChatRequest{Model: "m2", Messages: []types.ChatMessage{{Role: "user", Content: types.StringContent("x")}}})
if p2.Available() {
t.Fatal("expected permanent unavailability on 401")
// one failure -> 5s cooldown from now
until := st.CooldownUntil()
want := time.Now().Add(backoffBase).Unix()
if until < want-2 || until > want+2 {
t.Fatalf("cooldown = %d, want ~%d", until, want)
}
// success resets everything and bumps the preference
st.RecordSuccess()
if !p.ModelAvailable("m") {
t.Fatal("success must clear cooldown")
}
if st.FailCount() != 0 {
t.Fatalf("fail count after success = %d", st.FailCount())
}
if st.Pref() != 1-int64(prefFailStep) {
t.Fatalf("pref after one failure (-5) then success (+1) = %d, want %d", st.Pref(), 1-int64(prefFailStep))
}
}
func TestProviderConcurrencyCap(t *testing.T) {
func TestTryAcquire(t *testing.T) {
p := newTestProvider(t, src("mock", "http://127.0.0.1:1", "openai", "m"))
p.cfg.MaxConcurrent = 1
p.sem = make(chan struct{}, 1)
if err := p.TryAcquire(context.Background()); err != nil {
t.Fatalf("first acquire: %v", err)
}
if err := p.TryAcquire(context.Background()); !errors.Is(err, ErrBusy) {
t.Fatalf("second acquire = %v, want ErrBusy", err)
}
p.Release()
if err := p.TryAcquire(context.Background()); err != nil {
t.Fatalf("acquire after release: %v", err)
}
p.Release()
}
func TestChatBusyFailsFast(t *testing.T) {
release := make(chan struct{})
started := make(chan struct{}, 100)
started := make(chan struct{}, 10)
up := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
started <- struct{}{}
<-release
fmt.Fprint(w, `{"choices":[{"message":{"content":"ok"}}]}`)
}))
defer up.Close()
// cap 2
p := newTestProvider(t, src("mock", up.URL, "openai", "m"))
p.cfg.MaxConcurrent = 2
p.sem = make(chan struct{}, 2)
p.cfg.MaxConcurrent = 1
p.sem = make(chan struct{}, 1)
var wg sync.WaitGroup
for i := 0; i < 6; i++ {
wg.Add(1)
go func() {
defer wg.Done()
p.Chat(context.Background(), &types.ChatRequest{Model: "m", Messages: []types.ChatMessage{{Role: "user", Content: types.StringContent("x")}}})
}()
}
// wait until 2 requests started
deadline := time.Now().Add(2 * time.Second)
for len(started) < 2 {
if time.Now().After(deadline) {
t.Fatal("timeout waiting for first two")
}
time.Sleep(5 * time.Millisecond)
}
time.Sleep(100 * time.Millisecond)
if len(started) > 2 {
t.Fatalf("more than 2 concurrent: %d", len(started))
done := make(chan error, 1)
go func() {
_, err := p.Chat(context.Background(), &types.ChatRequest{
Model: "m",
Messages: []types.ChatMessage{{Role: "user", Content: types.StringContent("x")}},
})
done <- err
}()
<-started // first request holds the only slot
// second request must fail fast with ErrBusy instead of queueing
_, err2 := p.Chat(context.Background(), &types.ChatRequest{
Model: "m",
Messages: []types.ChatMessage{{Role: "user", Content: types.StringContent("x")}},
})
if !errors.Is(err2, ErrBusy) {
t.Fatalf("second chat err = %v, want ErrBusy", err2)
}
close(release)
wg.Wait()
}
if err := <-done; err != nil {
t.Fatalf("first chat: %v", err)
}
}

111
internal/provider/registry.go
View File

@ -75,63 +75,39 @@ func (r *Registry) ModelList() []string {
return out
}
// Resolve returns the ordered candidate providers to try for a request,
// honoring explicit model selection or AUTO (priority order, healthy first).
// Resolve returns the provider (or providers) serving a requested model,
// owning no AUTO scheduling logic anymore: AUTO chat scheduling is driven by
// the scheduler chain built from the runtime rules (see core/SaveAutoRules
// and scheduler.Chain).
//
// model "" or "AUTO" -> all sources sorted by (priority desc, healthy first).
// Otherwise the owning provider, if healthy; else its source anyway.
// model "" or "AUTO" -> every provider in configured order. Used only by the
// image path (which then filters to image-capable sources) and tool-call
// anchoring; chat AUTO requests go through the chain instead.
// Otherwise the owning provider; "source-model"/"source:model"/"source/model"
// pinning resolves first; an unknown model resolves to nil (gateway answers
// 404) instead of silently falling back to the AUTO chain.
func (r *Registry) Resolve(model string) []*Provider {
r.mu.RLock()
defer r.mu.RUnlock()
model = strings.TrimSpace(model)
if model == "" || strings.EqualFold(model, "AUTO") {
// priority chain across all models
type cand struct {
prov *Provider
priority int
}
var cands []cand
seen := map[string]bool{}
for _, p := range r.providers {
prio := -1
for _, m := range p.cfg.Models {
if m.Priority > prio {
prio = m.Priority
}
}
if prio < 0 {
prio = 0
}
cands = append(cands, cand{p, prio})
seen[p.Name()] = true
}
sort.SliceStable(cands, func(i, j int) bool {
if cands[i].priority != cands[j].priority {
return cands[i].priority > cands[j].priority
}
// healthy preferred at same priority
return cands[i].prov.Available() && !cands[j].prov.Available()
})
out := make([]*Provider, 0, len(cands))
for _, c := range cands {
out = append(out, c.prov)
}
out := make([]*Provider, len(r.providers))
copy(out, r.providers)
return out
}
// explicit model
if p, ok := r.byModel[strings.ToLower(model)]; ok {
// switch to the owning source but pin the model via request
return []*Provider{p}
}
// "source-model" / "source:model" / "source/model" pinning — disambiguates
// duplicate model ids across sources.
if p := r.ResolvePinned(model); p != nil {
return []*Provider{p}
}
// unknown model -> fall back to default/AUTO chain
return r.AUTOChain()
// explicit model
if p, ok := r.byModel[strings.ToLower(model)]; ok {
// switch to the owning source but pin the model via request
return []*Provider{p}
}
return nil
}
// EffectiveModel strips a "source-model" / "source:model" / "source/model"
@ -174,11 +150,6 @@ func (r *Registry) ResolvePinned(model string) *Provider {
return nil
}
// AUTOChain returns the priority-sorted providers for AUTO.
func (r *Registry) AUTOChain() []*Provider {
return r.Resolve("AUTO")
}
// ProviderForModel returns the provider owning the model id (nil if unknown).
func (r *Registry) ProviderForModel(model string) *Provider {
r.mu.RLock()
@ -220,15 +191,6 @@ func (r *Registry) ProviderForSlot(model, source string) *Provider {
return nil
}
// Default returns the highest-priority available provider.
func (r *Registry) Default() *Provider {
chain := r.AUTOChain()
if len(chain) == 0 {
return nil
}
return chain[0]
}
// ModelStatus is a web-UI friendly snapshot per source.
type SourceStatus struct {
Name string `json:"name"`
@ -241,6 +203,9 @@ type SourceStatus struct {
LiveAvailable bool `json:"live_available"`
LastError string `json:"last_error,omitempty"`
LastChecked int64 `json:"last_checked,omitempty"`
FailCount int `json:"fail_count,omitempty"`
BackoffUntil int64 `json:"backoff_until,omitempty"`
Permanent bool `json:"permanent,omitempty"`
}
// ProbeAll runs a live reachability check for every provider (in parallel).
@ -267,19 +232,27 @@ func (r *Registry) Status() []SourceStatus {
out := make([]SourceStatus, 0, len(r.providers))
for _, p := range r.providers {
live, lastErr, lastAt := p.LastProbe()
s := SourceStatus{
Name: p.Name(),
Adapter: p.Adapter(),
BaseURL: p.Config().BaseURL,
Models: p.Models(),
Available: p.Available(),
Healthy: p.Available(),
MaxConcurrent: p.MaxConcurrent(),
LiveAvailable: live,
LastError: lastErr,
LastChecked: lastAt,
}
out = append(out, s)
fails, until, perm := p.HealthInfo()
backoffUntil := int64(0)
if !until.IsZero() {
backoffUntil = until.Unix()
}
s := SourceStatus{
Name: p.Name(),
Adapter: p.Adapter(),
BaseURL: p.Config().BaseURL,
Models: p.Models(),
Available: p.Available(),
Healthy: p.Available(),
MaxConcurrent: p.MaxConcurrent(),
LiveAvailable: live,
LastError: lastErr,
LastChecked: lastAt,
FailCount: fails,
BackoffUntil: backoffUntil,
Permanent: perm,
}
out = append(out, s)
}
return out
}

326
internal/scheduler/scheduler.go
View File

@ -1,16 +1,29 @@
// Package scheduler implements request scheduling across providers: per-source
// concurrency caps (acquire with wait = queuing), AUTO model fallback chains,
// and exponential backoff via provider health.
// Package scheduler implements request scheduling across providers: direct
// fallback scheduling over candidate lists, and the AUTO chain (tiers with
// per-tier round-robin cursors, preference ordering, token-quota windows and
// per-(source,model) cooldown awareness) per the target architecture in
// plan.md.
package scheduler
import (
"context"
"errors"
"fmt"
"sort"
"strings"
"sync/atomic"
"time"
"llmsproxy/internal/provider"
"llmsproxy/internal/types"
)
// busyWait is how long a fully-busy tier is polled for a free slot before the
// request falls through to the next tier (bounded wait, plan 2.3).
var busyWait = 2 * time.Second
// busyPoll is the polling interval while waiting for a busy tier.
var busyPoll = 100 * time.Millisecond
// Scheduler drives one chat tool call across the candidate provider chain.
type Scheduler struct {
// MaxRetries how many fallback providers to try before failing.
@ -27,29 +40,297 @@ func New(maxRetries int) *Scheduler {
// Provider is the minimal interface the scheduler needs to schedule over.
type Provider interface {
Name() string
Available() bool
ModelFor(reqModel string) string
ModelAvailable(model string) bool
Pref(model string) int64
Chat(ctx context.Context, req *types.ChatRequest) (*types.UnifiedResponse, error)
ChatStream(ctx context.Context, req *types.ChatRequest) (<-chan types.UnifiedChunk, error)
Image(ctx context.Context, req *types.ImageGenRequest) (*types.UnifiedResponse, error)
}
// FromRegistry converts *provider.Provider slices to the scheduler interface.
func FromRegistry(ps []*provider.Provider) []Provider {
out := make([]Provider, len(ps))
for i, p := range ps {
out[i] = p
}
return out
// ---- AUTO chain ----
// Rule is one persisted AUTO chain slot (mirror of config.ModelScope).
type Rule struct {
Model string
Source string
Tier int
Quota int64
Period string
Hours int64
}
// Slot is one schedulable chain position: a model pinned to its provider,
// with an optional token-quota window. Slots are immutable after build.
type Slot struct {
Model string
Source string
Quota int64
Period string
Hours int64
Prov Provider
}
// TierNode is one priority tier. Slots keep their configured order (the
// stable base for preference ordering). next is the round-robin cursor: it
// holds the last used slot index (-1 = none yet), so the very first request
// starts at the configured order and later ones rotate.
type TierNode struct {
Tier int
Slots []*Slot
next atomic.Int64
}
// NextStart advances the tier cursor and returns the start index for the next
// scheduling run (first run: index 0).
func (tn *TierNode) NextStart() int64 {
return tn.next.Add(1)
}
// Chain is the immutable AUTO scheduling plan. A rebuilt chain is swapped in
// atomically; per-tier cursors live inside the chain and are shared across
// requests (rotation state resets when the chain is rebuilt, e.g. after
// editing the rules — acceptable, the swap also resets cooldowns).
type Chain struct {
Tiers []*TierNode // descending tier order
}
// TierErrors is the per-tier failure summary carried by ChainErr. Errors
// (TierError or skipped-tier reasons) are collected in tier order.
type TierError struct {
Tier int
Source string
Model string
Err error
}
// ChainErr is returned by chain scheduling when every AUTO tier failed. Its
// message summarizes each failed tier (which source/model and why) so a 503
// names the culprits instead of the bare "no provider available".
type ChainErr struct {
Tiers []TierError
Skipped []string // whole-tier reasons (cooling / quota / all busy)
}
func (e *ChainErr) Error() string {
var b strings.Builder
b.WriteString("all auto tiers failed: ")
first := true
for _, t := range e.Tiers {
if !first {
b.WriteString("; ")
}
first = false
fmt.Fprintf(&b, "tier %d %s/%s: %v", t.Tier, t.Source, t.Model, t.Err)
}
for _, s := range e.Skipped {
if !first {
b.WriteString("; ")
}
first = false
b.WriteString(s)
}
return b.String()
}
// BuildChain groups rules into descending tiers and resolves each slot's
// provider via prov. Rules whose provider resolves to nil are dropped (the
// source no longer serves the model). Slot order within a tier follows the
// configured rule order.
func BuildChain(rules []Rule, prov func(model, source string) Provider) *Chain {
byTier := map[int][]*Slot{}
var tiers []int
for _, r := range rules {
p := prov(r.Model, r.Source)
if p == nil {
continue
}
if _, ok := byTier[r.Tier]; !ok {
tiers = append(tiers, r.Tier)
}
byTier[r.Tier] = append(byTier[r.Tier], &Slot{
Model: r.Model,
Source: r.Source,
Quota: r.Quota,
Period: r.Period,
Hours: r.Hours,
Prov: p,
})
}
sort.Slice(tiers, func(i, j int) bool { return tiers[i] > tiers[j] })
ch := &Chain{}
for _, t := range tiers {
tn := &TierNode{Tier: t, Slots: byTier[t]}
tn.next.Store(-1)
ch.Tiers = append(ch.Tiers, tn)
}
return ch
}
// tierResult is the outcome of one scheduling run over one tier.
type tierResult struct {
resp *types.UnifiedResponse
chunks <-chan types.UnifiedChunk
src string
model string
hard []TierError // hard failures seen in this pass (nil = none)
}
// runTier executes one tier pass starting at the round-robin base index.
// Cooldown is the only hard skip (re-verified per slot); a busy slot is
// skipped without any penalty; a hard failure is recorded and the pass moves
// on to the next slot (plan 2.3: "单请求内不重试已失败槽" — the failed slot is
// not retried, the others still are). hard == nil and no success means every
// candidate was merely busy/cooling, so the caller may wait a bounded time.
func runTier(ctx context.Context, tn *TierNode, cands []*Slot, base int64, req *types.ChatRequest, stream bool) tierResult {
n := len(cands)
var hard []TierError
for i := 0; i < n; i++ {
sl := cands[(int(base)+i)%n]
if !sl.Prov.ModelAvailable(sl.Model) {
continue
}
r := *req
r.Model = sl.Model
if stream {
chunks, err := sl.Prov.ChatStream(ctx, &r)
if err == nil {
return tierResult{chunks: chunks, src: sl.Source, model: sl.Model}
}
if ctx.Err() != nil {
return tierResult{}
}
if errors.Is(err, types.ErrBusy) {
continue
}
hard = append(hard, TierError{Tier: tn.Tier, Source: sl.Source, Model: sl.Model, Err: err})
continue
}
resp, err := sl.Prov.Chat(ctx, &r)
if err == nil {
return tierResult{resp: resp, src: sl.Source, model: sl.Model}
}
if ctx.Err() != nil {
return tierResult{}
}
if errors.Is(err, types.ErrBusy) {
continue
}
hard = append(hard, TierError{Tier: tn.Tier, Source: sl.Source, Model: sl.Model, Err: err})
}
return tierResult{hard: hard}
}
// chainDrive runs a request down the chain (plan 2.3): tiers descending,
// per-tier round-robin starting at the tier cursor, same-tier runs ordered by
// preference (negative prefs sink but stay reachable). Quota-exhausted and
// cooling slots are filtered up front; a fully busy tier is polled for a
// bounded time before falling through. Failures are summarized in *ChainErr
// for the caller to map to HTTP 503.
func (s *Scheduler) chainDrive(ctx context.Context, chain *Chain, req *types.ChatRequest, exhausted func(*Slot) bool, stream bool) (*types.UnifiedResponse, <-chan types.UnifiedChunk, string, string, error) {
if chain == nil || len(chain.Tiers) == 0 {
return nil, nil, "", "", fmt.Errorf("no auto slot configured")
}
var ce ChainErr
for _, tn := range chain.Tiers {
// initial filter: quota-exhausted and cooling slots are dropped
var cands []*Slot
for _, sl := range tn.Slots {
if exhausted != nil && exhausted(sl) {
continue
}
if !sl.Prov.ModelAvailable(sl.Model) {
continue
}
cands = append(cands, sl)
}
if len(cands) == 0 {
ce.Skipped = append(ce.Skipped, fmt.Sprintf("tier %d: no schedulable slot (cooling or quota exhausted)", tn.Tier))
continue
}
// preference orders a same-tier run; stable so equal prefs keep order
sort.SliceStable(cands, func(i, j int) bool {
return cands[i].Prov.Pref(cands[i].Model) > cands[j].Prov.Pref(cands[j].Model)
})
base := tn.NextStart()
res := runTier(ctx, tn, cands, base, req, stream)
if res.resp != nil || res.chunks != nil {
return res.resp, res.chunks, res.src, res.model, nil
}
if ctx.Err() != nil {
return nil, nil, "", "", ctx.Err()
}
if len(res.hard) > 0 {
ce.Tiers = append(ce.Tiers, res.hard...)
continue // hard failures: fall through to the next tier, no waiting
}
// every candidate was busy or cooling: bounded poll before downgrading
deadline := time.Now().Add(busyWait)
for {
select {
case <-ctx.Done():
return nil, nil, "", "", ctx.Err()
case <-time.After(busyPoll):
}
done := time.Now().After(deadline)
if done {
ce.Skipped = append(ce.Skipped, fmt.Sprintf("tier %d: no free slot within %v", tn.Tier, busyWait))
break
}
// refresh candidates: cooldowns may have expired meanwhile
var again []*Slot
for _, sl := range cands {
if sl.Prov.ModelAvailable(sl.Model) {
again = append(again, sl)
}
}
if len(again) == 0 {
ce.Skipped = append(ce.Skipped, fmt.Sprintf("tier %d: no free slot within %v", tn.Tier, busyWait))
break
}
res = runTier(ctx, tn, again, base, req, stream)
if res.resp != nil || res.chunks != nil {
return res.resp, res.chunks, res.src, res.model, nil
}
if ctx.Err() != nil {
return nil, nil, "", "", ctx.Err()
}
if len(res.hard) > 0 {
ce.Tiers = append(ce.Tiers, res.hard...)
break // hard failure while waiting: stop waiting, fall through
}
}
}
if len(ce.Tiers) == 0 && len(ce.Skipped) == 0 {
return nil, nil, "", "", fmt.Errorf("no auto slot configured")
}
return nil, nil, "", "", &ce
}
// ChainChat runs a non-streaming AUTO request down the chain. exhausted, when
// non-nil, decides slot token-quota exhaustion. Returns the response, the
// serving source and the exact model id used; on total failure a *ChainErr
// summarizing every tier.
func (s *Scheduler) ChainChat(ctx context.Context, chain *Chain, req *types.ChatRequest, exhausted func(*Slot) bool) (*types.UnifiedResponse, string, string, error) {
resp, _, src, model, err := s.chainDrive(ctx, chain, req, exhausted, false)
return resp, src, model, err
}
// ChainChatStream runs a streaming AUTO request down the chain. A slot is
// abandoned only on connect failures / busy (before its first chunk); after a
// stream starts it is pinned. Same return contract as ChainChat.
func (s *Scheduler) ChainChatStream(ctx context.Context, chain *Chain, req *types.ChatRequest, exhausted func(*Slot) bool) (<-chan types.UnifiedChunk, string, string, error) {
_, chunks, src, model, err := s.chainDrive(ctx, chain, req, exhausted, true)
return chunks, src, model, err
}
// ---- direct scheduling ----
// Chat runs a chat request across cands, falling back on failure. Each
// candidate receives a request pinned to its own model (ModelFor), so an AUTO
// chain fallback switches the model id per provider instead of reusing the
// first candidate's model name.
//
// On success it returns the response together with the name of the provider
// and the exact model id that actually served the request (used for stats).
// candidate receives a request pinned to its own model (ModelFor), so a
// fallback switches the model id per provider instead of reusing the first
// candidate's model name. On success it returns the response together with
// the name of the provider and the exact model id that served the request.
func (s *Scheduler) Chat(ctx context.Context, cands []Provider, req *types.ChatRequest) (*types.UnifiedResponse, string, string, error) {
attempts := s.MaxRetries + 1
var lastErr error
@ -74,9 +355,10 @@ func (s *Scheduler) Chat(ctx context.Context, cands []Provider, req *types.ChatR
return nil, "", "", lastErr
}
// ChatStream runs a streaming chat across cands, falling back early on connect
// errors. The request model is pinned per candidate like Chat. On success it
// returns the chunk channel plus the serving provider name and model id.
// ChatStream runs a streaming chat across cands, falling back early on
// connect errors. The request model is pinned per candidate like Chat. On
// success it returns the chunk channel plus the serving provider name and
// model id.
func (s *Scheduler) ChatStream(ctx context.Context, cands []Provider, req *types.ChatRequest) (<-chan types.UnifiedChunk, string, string, error) {
attempts := s.MaxRetries + 1
var lastErr error
@ -113,4 +395,4 @@ func (s *Scheduler) Image(ctx context.Context, cands []Provider, req *types.Imag
lastErr = fmt.Errorf("no provider available")
}
return nil, "", lastErr
}
}

283
internal/scheduler/scheduler_test.go Normal file
View File

@ -0,0 +1,283 @@
package scheduler
import (
"context"
"errors"
"fmt"
"strings"
"sync/atomic"
"testing"
"time"
"llmsproxy/internal/types"
)
// fakeProvider is an in-memory Provider used to exercise chain scheduling
// deterministically without a Lua runtime.
type fakeProvider struct {
name string
model string
pref atomic.Int64
available atomic.Bool
busy atomic.Bool
fail atomic.Bool
chatHits atomic.Int64
}
func fakeProv(name, model string) *fakeProvider {
f := &fakeProvider{name: name, model: model}
f.available.Store(true)
return f
}
func (f *fakeProvider) Name() string { return f.name }
func (f *fakeProvider) ModelFor(reqModel string) string { return f.model }
func (f *fakeProvider) ModelAvailable(model string) bool { return f.available.Load() }
func (f *fakeProvider) Pref(model string) int64 { return f.pref.Load() }
func (f *fakeProvider) Chat(ctx context.Context, req *types.ChatRequest) (*types.UnifiedResponse, error) {
f.chatHits.Add(1)
if f.busy.Load() {
return nil, types.ErrBusy
}
if f.fail.Load() {
return nil, fmt.Errorf("upstream error")
}
return &types.UnifiedResponse{Content: f.name, FinishReason: "stop", TokenUsage: types.TokenUsage{Prompt: 1, Completion: 1, Total: 2}}, nil
}
func (f *fakeProvider) ChatStream(ctx context.Context, req *types.ChatRequest) (<-chan types.UnifiedChunk, error) {
if f.busy.Load() {
return nil, types.ErrBusy
}
if f.fail.Load() {
return nil, fmt.Errorf("upstream error")
}
ch := make(chan types.UnifiedChunk, 2)
ch <- types.UnifiedChunk{Content: f.name}
ch <- types.UnifiedChunk{Done: true}
close(ch)
return ch, nil
}
func (f *fakeProvider) Image(ctx context.Context, req *types.ImageGenRequest) (*types.UnifiedResponse, error) {
return nil, errors.New("no image")
}
// lookup resolves (model, source) -> provider for chain builders in tests.
type lookup func(m, s string) Provider
func bySource(ps ...*fakeProvider) lookup {
return func(m, s string) Provider {
for _, p := range ps {
if p.name == s {
return p
}
}
return nil
}
}
func chatReq() *types.ChatRequest {
return &types.ChatRequest{Messages: []types.ChatMessage{{Role: "user", Content: types.StringContent("hi")}}}
}
func TestBuildChainTierOrdering(t *testing.T) {
a, b, c, d := fakeProv("s1", "a"), fakeProv("s2", "b"), fakeProv("s3", "c"), fakeProv("s4", "d")
ch := BuildChain([]Rule{
{Tier: 1, Model: "a", Source: "s1"},
{Tier: 3, Model: "c", Source: "s3"},
{Tier: 2, Model: "b", Source: "s2"},
{Tier: 1, Model: "d", Source: "s4"},
{Tier: 9, Model: "gone", Source: "missing"},
}, bySource(a, b, c, d))
if len(ch.Tiers) != 3 {
t.Fatalf("tiers = %d, want 3", len(ch.Tiers))
}
if ch.Tiers[0].Tier != 3 || ch.Tiers[1].Tier != 2 || ch.Tiers[2].Tier != 1 {
t.Fatalf("tier order = %d,%d,%d, want 3,2,1", ch.Tiers[0].Tier, ch.Tiers[1].Tier, ch.Tiers[2].Tier)
}
// same-tier slots keep configured order, unresolvable rule is dropped
if len(ch.Tiers[2].Slots) != 2 || ch.Tiers[2].Slots[0].Model != "a" || ch.Tiers[2].Slots[1].Model != "d" {
t.Fatalf("tier 1 slots = %+v", ch.Tiers[2].Slots)
}
}
func TestChainRoundRobin(t *testing.T) {
a, b := fakeProv("s1", "a"), fakeProv("s2", "b")
ch := BuildChain([]Rule{
{Tier: 0, Model: "a", Source: "s1"},
{Tier: 0, Model: "b", Source: "s2"},
}, bySource(a, b))
s := New(0)
var got []string
for i := 0; i < 4; i++ {
_, src, _, err := s.ChainChat(context.Background(), ch, chatReq(), nil)
if err != nil {
t.Fatalf("iter %d: %v", i, err)
}
got = append(got, src)
}
want := []string{"s1", "s2", "s1", "s2"}
for i := range want {
if got[i] != want[i] {
t.Fatalf("rr order = %v, want %v", got, want)
}
}
}
func TestChainPreferenceSinksButStaysReachable(t *testing.T) {
neg := fakeProv("neg", "n")
good := fakeProv("good", "g")
neg.pref.Store(-5)
good.fail.Store(true)
ch := BuildChain([]Rule{
{Tier: 0, Model: "n", Source: "neg"},
{Tier: 0, Model: "g", Source: "good"},
}, bySource(neg, good))
s := New(0)
resp, src, model, err := s.ChainChat(context.Background(), ch, chatReq(), nil)
if err != nil {
t.Fatalf("chain: %v", err)
}
// higher pref (good) is tried first and hard-fails; the pass moves on to
// the negative-pref slot, which sinks but stays reachable: with the real
// provider its success would RecordSuccess (+1 pref, self-heal)
if src != "neg" || model != "n" || resp.Content != "neg" {
t.Fatalf("served src=%q model=%q content=%q", src, model, resp.Content)
}
if good.chatHits.Load() == 0 {
t.Fatal("higher-pref slot must be tried first")
}
}
func TestChainBusySkipsWithoutPenalty(t *testing.T) {
a, b := fakeProv("s1", "a"), fakeProv("s2", "b")
a.busy.Store(true)
ch := BuildChain([]Rule{
{Tier: 0, Model: "a", Source: "s1"},
{Tier: 0, Model: "b", Source: "s2"},
}, bySource(a, b))
s := New(0)
_, src, _, err := s.ChainChat(context.Background(), ch, chatReq(), nil)
if err != nil {
t.Fatalf("chain: %v", err)
}
if src != "s2" {
t.Fatalf("src = %q, want s2", src)
}
if a.chatHits.Load() == 0 {
t.Fatal("busy slot must have been attempted")
}
}
func TestChainAllBusyBoundedWaitThenNextTier(t *testing.T) {
oldWait, oldPoll := busyWait, busyPoll
busyWait, busyPoll = 60*time.Millisecond, 10*time.Millisecond
t.Cleanup(func() { busyWait, busyPoll = oldWait, oldPoll })
a, b, c := fakeProv("s1", "a"), fakeProv("s2", "b"), fakeProv("s3", "c")
a.busy.Store(true)
b.busy.Store(true)
ch := BuildChain([]Rule{
{Tier: 5, Model: "a", Source: "s1"},
{Tier: 5, Model: "b", Source: "s2"},
{Tier: 4, Model: "c", Source: "s3"},
}, bySource(a, b, c))
s := New(0)
t0 := time.Now()
_, src, _, err := s.ChainChat(context.Background(), ch, chatReq(), nil)
el := time.Since(t0)
if err != nil {
t.Fatalf("chain: %v", err)
}
if src != "s3" {
t.Fatalf("src = %q, want s3 (downgrade after bounded wait)", src)
}
if el > time.Second {
t.Fatalf("busy wait not bounded: %v", el)
}
}
func TestChainQuotaExhausted(t *testing.T) {
a, b := fakeProv("s1", "a"), fakeProv("s2", "b")
ch := BuildChain([]Rule{
{Tier: 0, Model: "a", Source: "s1", Quota: 100, Period: "hour"},
{Tier: 0, Model: "b", Source: "s2"},
}, bySource(a, b))
s := New(0)
exhausted := func(sl *Slot) bool { return sl.Source == "s1" && sl.Quota > 0 }
_, src, _, err := s.ChainChat(context.Background(), ch, chatReq(), exhausted)
if err != nil {
t.Fatalf("chain: %v", err)
}
if src != "s2" {
t.Fatalf("src = %q, want s2 (quota slot dropped)", src)
}
if a.chatHits.Load() != 0 {
t.Fatal("quota-exhausted slot must not be called")
}
}
func TestChainErrSummary(t *testing.T) {
a, b, c := fakeProv("s1", "a"), fakeProv("s2", "b"), fakeProv("s3", "c")
a.fail.Store(true)
b.fail.Store(true)
c.available.Store(false) // whole tier 1 cooling
ch := BuildChain([]Rule{
{Tier: 0, Model: "a", Source: "s1"},
{Tier: 0, Model: "b", Source: "s2"},
{Tier: 1, Model: "c", Source: "s3"},
}, bySource(a, b, c))
s := New(0)
_, _, _, err := s.ChainChat(context.Background(), ch, chatReq(), nil)
var ce *ChainErr
if !errors.As(err, &ce) {
t.Fatalf("err = %v, want *ChainErr", err)
}
if len(ce.Tiers) != 2 || ce.Tiers[0].Source != "s1" || ce.Tiers[1].Source != "s2" {
t.Fatalf("tiers = %+v", ce.Tiers)
}
if len(ce.Skipped) != 1 {
t.Fatalf("skipped = %+v", ce.Skipped)
}
msg := ce.Error()
if !strings.Contains(msg, "s1") || !strings.Contains(msg, "s2") || !strings.Contains(msg, "tier 1: no schedulable slot") {
t.Fatalf("summary = %q", msg)
}
}
func TestChainStreamFallsBackBeforeFirstChunk(t *testing.T) {
a, b := fakeProv("s1", "a"), fakeProv("s2", "b")
a.fail.Store(true)
ch := BuildChain([]Rule{
{Tier: 0, Model: "a", Source: "s1"},
{Tier: 0, Model: "b", Source: "s2"},
}, bySource(a, b))
s := New(0)
chunks, src, model, err := s.ChainChatStream(context.Background(), ch, chatReq(), nil)
if err != nil {
t.Fatalf("chain stream: %v", err)
}
if src != "s2" || model != "b" {
t.Fatalf("src=%q model=%q", src, model)
}
var text string
for ck := range chunks {
text += ck.Content
}
if text != "s2" {
t.Fatalf("text = %q", text)
}
}
func TestChainResetCooldownAfterSwap(t *testing.T) {
a := fakeProv("s1", "a")
ch := BuildChain([]Rule{{Tier: 0, Model: "a", Source: "s1"}}, bySource(a))
// a freshly built chain must schedule from index 0 (cursor starts at -1)
if base := ch.Tiers[0].NextStart(); base != 0 {
t.Fatalf("first start = %d, want 0", base)
}
}

8
internal/types/types.go
View File

@ -4,9 +4,17 @@ package types
import (
"encoding/json"
"errors"
"time"
)
// ErrBusy is the soft "source at capacity" sentinel shared by the provider
// layer (returns it) and the scheduler layer (reacts to it): busy is not a
// failure, so no cooldown/preference penalty is recorded, and gateways map
// it to HTTP 429. Defined here so the scheduler does not depend on the
// provider package (which pulls in the Lua runtime).
var ErrBusy = errors.New("provider busy")
// ---- OpenAI wire request (gateway input) ----
type ChatRequest struct {