diff --git a/cmd/llmsproxy/main.go b/cmd/llmsproxy/main.go index 1aa5a04..b971343 100644 --- a/cmd/llmsproxy/main.go +++ b/cmd/llmsproxy/main.go @@ -12,6 +12,7 @@ import ( "net/http" "os" "os/signal" + "strings" "syscall" "time" @@ -34,6 +35,21 @@ func main() { log.Fatalf("[llmsproxy] gateway: %v", err) } + // Ops hygiene: surface the two most common footguns instead of silently + // running with them. + if keys := c.GatewayKeys(); len(keys) == 0 { + log.Printf("[llmsproxy] WARNING: gateway_keys is EMPTY — without a key every request is rejected") + } else { + for _, k := range keys { + if k == "sk-gw-local-0001" || k == "sk-local-0001" { + log.Printf("[llmsproxy] WARNING: gateway key %q looks like the starter/example key — rotate it before exposing the gateway", k) + } + } + } + if l := c.Listen(); strings.HasPrefix(l, "0.0.0.0:") || strings.HasPrefix(l, "::") { + log.Printf("[llmsproxy] WARNING: listen=%s binds ALL interfaces — bind an internal address in production", l) + } + srv := &http.Server{ Addr: c.Listen(), Handler: gw.Handler(), diff --git a/e2e/e2e_test.go b/e2e/e2e_test.go index 31f3dff..5347f5d 100644 --- a/e2e/e2e_test.go +++ b/e2e/e2e_test.go @@ -13,6 +13,7 @@ import ( "os" "os/exec" "path/filepath" + "runtime" "strings" "sync" "testing" @@ -100,7 +101,15 @@ func buildBinary(t *testing.T) string { t.Helper() dir := t.TempDir() bin := filepath.Join(dir, "llmsproxy") + if runtime.GOOS == "windows" { + bin += ".exe" // go build writes the exact name; Windows needs the suffix to exec + } out, err := exec.Command("go", "build", "-tags", "luajit", "-o", bin, "llmsproxy/cmd/llmsproxy").CombinedOutput() + if err != nil && runtime.GOOS == "windows" { + // local Windows dev box may lack LuaJIT: fall back to the bundled + // Lua runtime (the e2e adapters used here are passthrough-only) + out, err = exec.Command("go", "build", "-o", bin, "llmsproxy/cmd/llmsproxy").CombinedOutput() + } if err != nil { t.Fatalf("build: %v\n%s", err, out) } @@ -166,7 +175,8 @@ func (g *gatewayUnderTest) do(method, path string, body string, authed bool) (*h } // writeConfig writes a temp gateway config pointing at the mock upstreams, -// listening on the given address. +// listening on the given address. Every non-image upstream becomes a chat +// source; the first named one ("good") is the highest-priority AUTO slot. func writeConfig(t *testing.T, dir, listenAddr string, upstreams map[string]*mockUpstream) string { t.Helper() var sb strings.Builder @@ -176,7 +186,8 @@ func writeConfig(t *testing.T, dir, listenAddr string, upstreams map[string]*moc sb.WriteString("adapter_dir: " + filepath.Join(dir, "adapters") + "\n") sb.WriteString("runtime_file: " + filepath.Join(dir, "runtime.json") + "\n") sb.WriteString("sources:\n") - order := []string{"good", "image"} + order := []string{"good", "fallback", "image"} + prio := map[string]int{"good": 100, "fallback": 50} for _, name := range order { u, ok := upstreams[name] if !ok { @@ -185,7 +196,7 @@ func writeConfig(t *testing.T, dir, listenAddr string, upstreams map[string]*moc if name == "image" { sb.WriteString(" - name: imagegen\n base_url: " + u.baseURL + "\n adapter: openai\n models:\n - id: flux-1\n kind: image\n priority: 80\n") } else { - sb.WriteString(" - name: " + name + "\n base_url: " + u.baseURL + "\n adapter: openai\n models:\n - id: " + name + "-m\n priority: 100\n") + sb.WriteString(" - name: " + name + "\n base_url: " + u.baseURL + "\n adapter: openai\n models:\n - id: " + name + "-m\n priority: " + fmt.Sprint(prio[name]) + "\n") } } path := filepath.Join(dir, "config.yaml") @@ -197,8 +208,9 @@ func writeConfig(t *testing.T, dir, listenAddr string, upstreams map[string]*moc func TestEndToEnd(t *testing.T) { upstreams := map[string]*mockUpstream{ - "good": newMockUpstream(t), - "image": newMockUpstream(t), + "good": newMockUpstream(t), + "fallback": newMockUpstream(t), + "image": newMockUpstream(t), } dir := t.TempDir() @@ -279,6 +291,35 @@ func TestEndToEnd(t *testing.T) { } } +// TestEndToEndAuto503: with every AUTO slot failing, the gateway must answer +// 503 whose message summarizes each failed tier/source/model instead of a +// bare "no provider available" (P8). +func TestEndToEndAuto503(t *testing.T) { + upstreams := map[string]*mockUpstream{ + "good": newMockUpstream(t), + } + dir := t.TempDir() + l, err := net.Listen("tcp", "127.0.0.1:0") + if err != nil { + t.Fatalf("pick port: %v", err) + } + addr := l.Addr().String() + l.Close() + cfg := writeConfig(t, dir, addr, upstreams) + bin := buildBinary(t) + g := startGateway(t, bin, addr, cfg) + + upstreams["good"].SetFail(true) + resp, body := g.do("POST", "/v1/chat/completions", + `{"model":"AUTO","messages":[{"role":"user","content":"hi"}]}`, true) + if resp == nil || resp.StatusCode != http.StatusServiceUnavailable { + t.Fatalf("status=%v body=%q", statusOf(resp), body) + } + if !strings.Contains(body, "all auto tiers failed") || !strings.Contains(body, "good/good-m") { + t.Fatalf("503 must summarize the failed slot, body=%q", body) + } +} + func statusOf(resp *http.Response) int { if resp == nil { return -1 diff --git a/internal/core/core.go b/internal/core/core.go index 3b4f3a5..114ed12 100644 --- a/internal/core/core.go +++ b/internal/core/core.go @@ -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() } -} \ No newline at end of file +} diff --git a/internal/gateway/chat.go b/internal/gateway/chat.go index 9b69927..6ccceac 100644 --- a/internal/gateway/chat.go +++ b/internal/gateway/chat.go @@ -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, }) -} \ No newline at end of file +} diff --git a/internal/gateway/gateway_test.go b/internal/gateway/gateway_test.go index 5d17fa5..02b9a78 100644 --- a/internal/gateway/gateway_test.go +++ b/internal/gateway/gateway_test.go @@ -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()) } -} \ No newline at end of file +} diff --git a/internal/gateway/keys.go b/internal/gateway/keys.go index c2aebd9..19a2ae5 100644 --- a/internal/gateway/keys.go +++ b/internal/gateway/keys.go @@ -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" { diff --git a/internal/gateway/server.go b/internal/gateway/server.go index a3de631..99837b9 100644 --- a/internal/gateway/server.go +++ b/internal/gateway/server.go @@ -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 diff --git a/internal/gateway/stats.go b/internal/gateway/stats.go index 5aed7ee..5eb329b 100644 --- a/internal/gateway/stats.go +++ b/internal/gateway/stats.go @@ -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 ..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 ..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...), } } \ No newline at end of file diff --git a/internal/gateway/stats_test.go b/internal/gateway/stats_test.go new file mode 100644 index 0000000..84f077d --- /dev/null +++ b/internal/gateway/stats_test.go @@ -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)) + } +} \ No newline at end of file diff --git a/internal/gateway/ui/index.html b/internal/gateway/ui/index.html index 40230d8..8a2c0af 100644 --- a/internal/gateway/ui/index.html +++ b/internal/gateway/ui/index.html @@ -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() {

${t('dashModel')}

${s.sources ? `

${t('dashSrc')}

` : ''} +

${t('dashStatus')} ${t('statusTag')}

${t('dashKey')}

${t('dashRecs')}

@@ -732,6 +745,7 @@ async function paintStats() {
${t('kpiLat')}
${fmtMs(avg)}
${t('kpiMaxLat')} ${fmtMs(tot.latency_max_ms)}
`; 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) { ${fmtTok(r.tokens)} ${fmtMs(fmtLat(r.latency_sum_ms, r.reqs))}${fmtMs(r.latency_max_ms)}`).join('') + '
'; } +function paintStatusTable(rows) { + const el = $('#tb-status'); if (!el) return; + if (!rows.length) { el.innerHTML = `
${t('noUsage')}
`; return; } + el.innerHTML = `
` + + rows.map(r => ` + `).join('') + '
${t('thCode')}${t('thReqs')}${t('thOk')}${t('thErr')}
${esc(r.name)}${fmtN(r.reqs)}${fmtN(r.ok)}${fmtN(r.err)}
'; +} function paintKeyTable(rows, keyNames) { const el = $('#tb-key'); if (!el) return; if (!rows.length) { el.innerHTML = `
${t('noUsage')}
`; 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() {
`; paintSort(); } +function healthTag(it) { + const st = sortStateMap.get(it.id + '|' + (it.src || '*')); + if (!st) return ''; + const bits = []; + if (st.cooling) bits.push(`${esc(t('sortCooling'))}`); + if (st.fail_count > 0) bits.push(`${esc(t('sortFail') + '×' + st.fail_count)}`); + if (st.pref !== 0) bits.push(`${esc(st.pref > 0 ? '+' + st.pref : '' + st.pref)}`); + if (!bits.length) return ''; + return `${bits.join('')}`; +} 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) { ${esc(it.src === '*' ? '+' : s)} ${esc(it.id)}${esc(it.src === '*' ? t('kAnySrc') : it.src)} ${it.meta ? `${esc(quantBadge(it.meta.quota, it.meta.period, it.meta.hours))}` : ''} + ${healthTag(it)} × `; @@ -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 => { diff --git a/internal/provider/provider.go b/internal/provider/provider.go index f785e8f..2a5467f 100644 --- a/internal/provider/provider.go +++ b/internal/provider/provider.go @@ -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] + "..." -} \ No newline at end of file +} diff --git a/internal/provider/provider_test.go b/internal/provider/provider_test.go index d24c548..46d9f44 100644 --- a/internal/provider/provider_test.go +++ b/internal/provider/provider_test.go @@ -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() -} \ No newline at end of file + if err := <-done; err != nil { + t.Fatalf("first chat: %v", err) + } +} diff --git a/internal/provider/registry.go b/internal/provider/registry.go index 3990a34..67a2aaf 100644 --- a/internal/provider/registry.go +++ b/internal/provider/registry.go @@ -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 } \ No newline at end of file diff --git a/internal/scheduler/scheduler.go b/internal/scheduler/scheduler.go index 2dd1edf..4dc4853 100644 --- a/internal/scheduler/scheduler.go +++ b/internal/scheduler/scheduler.go @@ -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 -} \ No newline at end of file +} diff --git a/internal/scheduler/scheduler_test.go b/internal/scheduler/scheduler_test.go new file mode 100644 index 0000000..a53ae74 --- /dev/null +++ b/internal/scheduler/scheduler_test.go @@ -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) + } +} diff --git a/internal/types/types.go b/internal/types/types.go index 1f07c2b..4493f22 100644 --- a/internal/types/types.go +++ b/internal/types/types.go @@ -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 { diff --git a/plan.md b/plan.md new file mode 100644 index 0000000..b30190d --- /dev/null +++ b/plan.md @@ -0,0 +1,228 @@ +# ModelRouter 调度层重构方案（基于线上实测） + +> 调研对象：192.168.2.60 生产实例（`/usr/local/bin/llmsproxy -config /etc/llmsproxy/config.yaml`，systemd 托管） +> 调研时间：2026-08-10，审计文件 `/etc/llmsproxy/runtime.json.audit.jsonl`（3978 行 / 1042 条请求记录） + +--- + +## 一、线上实测问题（证据链） + +### P1（核心，用户实测）编辑 AUTO 优先级不重载退避状态 +- 证据：`/api/status` 实时返回 `qijiar available=false, live_available=true`、`zen available=false, live_available=true`。 +- 根因：`Core.SaveAutoRules`（core.go:232）只写 runtime.json，不重建 Provider；退避状态存在 `Provider.health`（provider.go:24-51），只有 `rebuildRegistry`（core.go:290）会重建。**用户在 WebUI 改完优先级后，被退避锁死的源依旧被跳过。** +- 线上复现：AUTO 链 4 槽 `[qijiar/gpt-5.5→zen/flash-free→frank/gpt-5.6-sol→zen/nemotron]`，实时请求 `POST /v1/chat/completions {"model":"AUTO"}` → `{"error":{"message":"no provider available","type":"upstream_error"}}`，0ms 返回；其中 2 个源探活确认在线。 + +### P2 源级黑名单（粒度错） +- 证据：`zen` 因 `429 FreeUsage` 退避（审计 ×4）→ 其 9 个配置模型全部不可调度；`gpt-5.5` 的 qijiar 一次失败 → 源整体跳过。 +- 根因：健康标志是 `(源)` 级而非 `(源, 模型)` 级（provider.go:24）；一个模型失败黑掉全源。 + +### P3 永久黑名单（无法自愈） +- 证据：`frank` 收到一次 `401 API_KEY_DISABLED` → `markPermanent()`（provider.go:48-51, 281-283）；此后 AUTO 永远跳过它，且有 30 分钟周期的显式探测持续打它产生 502（502 又推高退避）。 +- 根因：`permanent` 无过期、无重置入口；显式请求路径不检查 `Available()`，照打不误。 + +### P4 并发打满 = 排队 60s，而非立即切换 +- 证据/源码：`Provider.Acquire`（provider.go:306-324）满则等待 `QueueTimeout`（默认 60s，config.go:154）；AUTO 逐槽串行，打满的档位拖死整条链。 +- 无 `TryAcquire` 语义，路由层无法区分"忙"与"故障"。 + +### P5 探活与调度状态脱节 +- 探活（GET /models）不污染调度状态是正确设计，但**探活结果没有任何一条通道能恢复调度健康**——线上 qijiar/zen "探活通、调度死" 的矛盾即因此产生。 + +### P6 流式成功不清退避 +- 证据/源码：`ChatStream`（provider.go:406-490）全程无 `reportOK()`；流式源一旦退避，后续成功流也不复位。 + +### P7 两套 AUTO 逻辑并存，配置互相矛盾 +- 线上：config.yaml 里 qijiar priority 95/90/85、deepseek 40/30；runtime AUTO 链却是 `tier1 gpt-5.5 → tier2 zen/flash-free → tier3 frank/gpt-5.6-sol → tier4 zen/nemotron`。 +- `Registry.Resolve("AUTO")`（registry.go:83-121，按"源的最大模型 priority + 健康优先"）与 `autoPlans`（chat.go:630，按 runtime 槽位）行为不一致；deepseek 全部模型不在 AUTO 链中，但运维以为"priority 高"会被 AUTO 选中。 + +### P8 错误码策略粗糙 +- `deepseek` 402 欠费 ×17：走 `ReportStatus` 不触发任何退避也不在 UI 提示（402 不属于 401/403/429/5xx 分支），欠费源被无限重试。 +- 400 Invalid schema ×19（客户端工具定义问题）同样无区分。 +- 链全灭时只返回 `no provider available`（12 条），无分槽错误汇总，无法定位是谁挂了。 + +### P9 其他（顺带） +- 种子 admin key `sk-gw-local-0001` 未更换（README 明确要求换）；`listen: 0.0.0.0:8081` 全端口暴露。 +- AUTO 链只有 4 槽，13+ 配置模型不参与 AUTO，且 webui 编辑链后无任何"健康复位"提示。 +- 审计文件 jsonl 无限增长（当前 547KB）；`AppendAudit` 每事件一次文件 open。 + +--- + +## 二、目标架构（定稿） + +### 2.1 分层（横向三层 + 状态基座横切） + +``` +┌────────────────────────────────────────────────┐ +│ 直连调度 AUTO 调度 │ ← 同级，共享底座 +│ source:model 直派 链快照 + 偏好表 │ 直连失败即返回（不重试） +├────────────────────────────────────────────────┤ +│ 源抽象层（每源） │ +│ · 信号量 = max_concurrent，TryAcquire 非阻塞 │ ← 唯一"忙/闲"判定 +│ · Chat / ChatStream / Image / 事件上报 │ +├────────────────────────────────────────────────┤ +│ adapter 池层（每适配器，启动预热） │ +│ worker = Σ(使用该适配器的源 × max_concurrent) │ +└────────────────────────────────────────────────┘ + ▲ 状态基座（不属任何层） + ModelState 表 (source,model) → {pref, failCount, cooldownUntil} + tier 游标表 tier → atomic next（记录上次分配位置） +``` + +事件流：所有成功/失败/冷却由**源抽象层**上报到状态基座；auto 与直连都只读。直连失败同样写入状态。 + +### 2.2 数据结构 + +```go +// 调度链：保存时整体新建、原子替换；调度期只读快照 +type Chain struct { Tiers []TierNode } // 按 tier 降序 +type TierNode struct { Tier int; Slots []*Slot } +type Slot struct { // 静态配置，链构建时冻结配额 + Model, Source string + Kind string + Quota int64; Period string; Hours int64 + state *ModelState // 指针 → 跨链存活 +} + +// 状态基座（key=(source,model)，配置移除后回收） +type ModelState struct { + pref atomic.Int64 // +1/-5，clamp[-20,+20] + failCount atomic.Int64 + cooldownUntil atomic.Int64 // 惰性指数退避，无定时器 +} +``` + +### 2.3 调度算法（AUTO） + +``` +snap := chain.Snapshot() // O(n) 拷指针，调度期链只读 +for tier := snap.Tiers 降序 { + 初筛 := 剔除 [冷却中 | 配额超限(slot.Quota>0 && windowTokens>=Quota)] 的槽 + if 初筛为空 → 该 tier 顺延，记录原因 + base := cursor[tier].Add(1) // per-tier 原子游标（记录上次位置） + ordered := 稳定排序(初筛, 按 pref 降序) // 负面模型沉底 + for i := base; i < base+len(ordered); i++ { + slot := ordered[i % n] + if slot.state 冷却中 → continue // 硬跳过（复验） + if !src.TryAcquire() → continue // 忙 = 软跳过（不记分） + resp, err := src.Chat(slot.model) + if err != nil { + src.Release() + state.RecordFailure() // failCount++, pref-5, 冷却 5s·2^n(≤30min) + continue // 单请求内不重试已失败槽 + } + state.RecordSuccess() // pref+1, failCount/cooldown 清零 + return resp, slot.model, slot.source + } +} +return 503 { 各 tier 错误汇总 } // 全部档位失败才报错，可定位 +``` + +关键语义： +- **主序 = 游标轮转（均衡），偏好只做同起点排序（自适应）**；负偏好"沉底不跳过"——仅当同 tier 无非负候选时才尝试负面模型，成功 +1 自愈。**无永久黑名单、无定时器**。 +- **忙 ≠ 失败**：不扣分、不计冷却；同 tier 全忙 → 有界等待（≤2s 轮询 TryAcquire，尊重 ctx）再顺延，避免高 QPS 时全线降级。 +- **冷却是唯一硬跳过**：`cooldownUntil = lastFail + min(5s·2^failCount, 30min)`，必然到期 → 自愈。 +- 401/403：failCount 一次打高（不设永久位），可被冷却到期/手动重置恢复。 +- 流式：失败→切换仅限首 chunk 前；首块后固定。首块前失败 -5，正常结束 +1。 +- 生图：不进 AUTO 链，直连 `kind:image` 模型，同一状态机制。 + +### 2.4 生命周期（修 P1） + +| 事件 | 动作 | +|---|---| +| WebUI 保存 AUTO 链 | 构造新 `Chain` → 写锁原子替换 → **链内全部 ModelState 冷却清零（偏好保留）** | +| 增删/编辑源 | `rebuildRegistry` 重建 Provider，ModelState 复用/回收 | +| (source,model) 移出配置 | 回收其 ModelState（偏好一并丢弃） | +| 进行中请求 | 不受 swap 影响（入口 Snapshot 已拷贝引用） | + +### 2.5 直连 + +``` +p := resolve(source, model) // 唯一归属，无 AUTO 逻辑 +p==nil → 400/404；!TryAcquire → 429 busy（快速失败） +成功/失败 → RecordSuccess/RecordFailure（同样写状态基座），返回，不重试 +``` + +--- + +## 三、实施计划（分阶段上线） + +### Phase 0 — 止血热修（改动最小，当天可上） +1. `core.go`：`SaveAutoRules` 成功后调用 `rebuildRegistry()`；给 `Provider` 增加 `ResetHealth()`（清 failCount/permanent/cooldown），rebuild 时顺带重置。 +2. `provider.go`：`ChatStream` 成功（收到 `[DONE]` 或正常结束）调用 `reportOK()`。 +3. 状态页：`/api/status` 增加 `last_backoff`、`permanent` 展示 + admin 可"重置源健康"按钮（调用 ResetHealth）。 +- 验证：改优先级 → AUTO 立即按新链调度，qijiar/zen 场景恢复；回归 `go test`。 + +### Phase 1 — 状态基座落地（provider.go 重构） +1. 引入 `ModelState`（每 (source,model)：pref/failCount/cooldownUntil，原子）。 +2. `Provider` 增加 `TryAcquire(ctx) error`（非阻塞）与 `RecordFailure(model, code)` / `RecordSuccess(model)`；`ReportStatus` 改为按模型记账与 401/403 不再永久化。 +3. 删除 `permanent` 语义（由有界冷却 + 重置通道取代）。 +- 验证：单测（退避按模型隔离、429/401 行为、TryAcquire 满即返）。 + +### Phase 2 — 调度层重写（scheduler.go + chat.go） +1. 新增 `Chain/TierNode/Slot` 与 Snapshot/原子替换；`autoPlans`、`rotateSameTier` 删除，`singleChatAuto`/`streamChatAuto` 按 2.3 伪码重写。 +2. `Registry.Resolve` 移除 AUTO 排序/健康优先职责（只留归属 + pin 解析；`AUTOChain`/`Default` 不再用于调度）。 +3. 503 响应携带分槽错误汇总（哪个 tier 哪个源什么错）。 +4. 同 tier 全忙时 ≤2s 有界等待再降级。 +- 验证：scheduler 单测（tier 分桶、游标均衡、偏好沉底自愈、配额初筛、忙不记分）；e2e 增加"上游 429 → 切同 tier → 再切下 tier"、"编辑优先级后退避清零"用例。 + +### Phase 3 — 收尾 +1. UI：优先级页展示链上模型的冷却/偏好状态；状态页按模型展示健康。 +2. WebUI 保存链时前端提示"健康状态已复位"。 +3. 审计：错误码分类（402 欠费、400 schema 计入独立统计），jsonl 轮转（按天/大小）。 +4. 运维：更换 admin key；`listen` 收敛到内网地址（0.0.0.0:8081 → 127.0.0.1 或内网 IP + 反向代理）。 + +### 回归与上线 +- 构建：`go build -tags luajit`（生产机 Debian amd64 已具备工具链，可直接在 60 上编译）；本地 Windows 需先补 LuaJIT 库才能链接。 +- 测试：`go test -tags luajit ./...` 全绿后按 Phase 0→3 灰度，每个 Phase 观察审计中 `no provider available` 计数与 502 分布。 +- 观察指标：`no provider available` 计数归零；qijiar/zen 由 `available=false` 恢复；`zen 429` 触发时仅 zen 对应模型短冷却，deepseek/qijiar 不受牵连。 + +--- + +## 四、与既有代码的对应改动文件 + +| 文件 | 改动 | +|---|---| +| internal/provider/provider.go | health→ModelState 表；TryAcquire；ReportStatus/ResetHealth；ChatStream reportOK | +| internal/provider/registry.go | 删除 Resolve(AUTO) 排序/健康逻辑，保留归属/pin | +| internal/scheduler/scheduler.go | 新增 Chain/Slot/游标/偏好排序迭代器 | +| internal/gateway/chat.go | 删除 autoPlans/rotateSameTier，重写单次/流式 AUTO；错误汇总 | +| internal/core/core.go | SaveAutoRules 原子换链 + 冷却清零；ModelState 生命周期 | +| internal/gateway/api.go / keys.go | 状态页/优先级页 UI 展示与重置接口 | +| e2e/e2e_test.go 等 | 新场景用例 | + +--- + +## 五、实施状态追踪（每项改动后更新） + +### Phase 0 — 止血热修 + +- [x] **P0-1（2026-08-10）`core.go`**：`SaveAutoRules` 持久化后调用 `rebuildRegistry()`——新建 Provider 即退避归零，**编辑优先级立即生效**（修 P1 主诉）；新增 `Core.ResetHealth()` 供管理端手动清退避。 +- [x] **P0-2（2026-08-10）`provider.go`**：新增 `Provider.ResetHealth()` / `Provider.HealthInfo()`（failCount/cooldownUntil/permanent 只读暴露）；`ChatStream` 流式正常结束（`[DONE]`/EOF、非客户端断开、非读错误）调用 `reportOK()`——流式成功可恢复退避（修 P6）。 +- [x] **P0-3（2026-08-10）`registry.go` + `server.go`**：`SourceStatus` 增加 `fail_count`/`backoff_until`/`permanent` 字段（状态页可分辨"探活通但调度退避"）；新增 `POST /api/status/reset`（admin 专属，写审计）。 +- [x] **P0-4（2026-08-10）验证**：`go build ./internal/...` 通过；`go test ./internal/config/...` 通过（Windows 缺 LuaJIT 库，带 cgo 的包无法本地链接，待生产机验证）。 + +> 说明：P0-1 采用"重建 Provider"实现退避归零（非目标架构的 ModelState 粒度），属临时止血；Phase 1 引入 per-(源,模型) ModelState 后，`ResetHealth` 语义将迁移到模型级。 + +### Phase 1 — 状态基座落地 +- [x] **P1-1（2026-08-10）`provider.go`**：`ModelState` 表落地（pref±1/-5、failCount、cooldownUntil 原子，惰性求值无定时器）；新增 `ErrBusy`、`TryAcquire`（非阻塞，满即返）、`ModelAvailable(model)`、`RecordFailure(model, code)` / `RecordSuccess(model)`、`Provider.Pref(model)`；`ReportStatus` 改按模型记账（401/403 → 打满档冷却 + 双倍偏好惩罚，**不再永久化**；5xx/429 → 指数退避；400/402 不惩罚）；`ResetHealth` 清全源模型状态；`HealthInfo` 源级聚合（permanent 恒 false）；`Chat/ChatStream/Image` 全部按命中模型记账 + 内部 `Acquire` 换 `TryAcquire`（忙即 429，不再排队 60s，修 P4）；流式干净收尾 `RecordSuccess`（修 P6）。 +- [x] **P1-2（2026-08-10）`chat.go`**：直连/生图路径 `errors.Is(err, ErrBusy)` → HTTP 429（`upstreamErrStatus`）；AUTO 槽改用 `ModelAvailable(slot.model)` 模型级冷却跳过（frank 401 后其模型不被 AUTO 反复打，修 P3 的一半——永久黑名单已随 P1-1 移除）。 +- [x] **P1-3（2026-08-10）单测（`provider_test.go`）**：模型级退避隔离（m1 失败 m2 照常）、401 → failCount 打满 capN + 冷却 ~30min + 偏好 -10（非永久，reset 即恢复）、单次失败冷却 ≈5s 且成功复位 +1、TryAcquire 满即返、Chat 忙时快速 ErrBusy、流式成功清理冷却（P6 回归）。 +- [x] **P1-4（2026-08-10）本地验证**：`go build ./internal/...`、`go vet ./internal/...`、`go test ./internal/config/...` 全绿；provider/gateway 测试因 Windows 缺 `-llua` 仅静态检查通过，待生产机 `-tags luajit` 全量跑（同 P0-4 约束）。 +- [ ] 生产机（192.168.2.60）`go test -tags luajit ./internal/provider/...` 回归 — 待 Phase 2/3 完成后一并部署验证。 + +### Phase 2 — 调度层重写 +- [x] **P2-1（2026-08-10）`scheduler.go`**：`Chain/TierNode/Slot/Rule` 落地（tier 降序；槽静态配置，同 tier 保持配置序；per-tier `next atomic.Int64` 游标始 -1，首请求从配置序开始）；`BuildChain` 丢弃 provider 解析失败的槽；`runTier` 按 2.3 语义重写——冷却复验硬跳过、busy 软跳过不记分、**硬失败记账后同档继续下一槽**（单请求不重试已失败槽）、整档遍历完才顺延；`chainDrive` 初筛（配额/冷却）→ 同 tier 按 `Pref` 稳定降序排序 → `NextStart()` 起步轮转 → 全忙/全冷却有界等待（`busyWait` 2s/`busyPoll` 100ms，期间刷新冷却）再降级；`ChainErr` 携带各档 `TierError` + 顺延原因，Error() 输出 `all auto tiers failed: tier N src/model: err; ...`；`ChainChat/ChainChatStream` 返回 (resp, source, model, err)；直连 `Chat/ChatStream/Image` 保持不变；**scheduler 不再 import provider**（`FromRegistry` 移入 gateway 为 `toScheduler`），scheduler 单测不拉 LuaJIT 链接。 +- [x] **P2-2（2026-08-10）`registry.go` + `core.go`**：`Registry.Resolve` 移除 AUTO 排序/健康优先（AUTO→按配置序全量仅用于生图/tool 锚定；未知模型→nil→gateway 404；`AUTOChain`/`Default` 删除）；`Core` 持 `autoChain atomic.Pointer[scheduler.Chain]` + `AutoChain()`；`buildAutoChain`（过滤 image-kind 与已不存在槽，**槽 Source 规范化为所有权源**使汇总/audit/配额窗口键一致）随 `rebuildRegistry` 与 `SaveAutoRules` 重建；`SaveAutoRules` = 持久化 → 原子换链 → **链内每槽 `ResetModelCooldown`（偏好保留，不重建 Provider）**——修 P1 主诉（编辑优先级立即生效）。 +- [x] **P2-3（2026-08-10）`chat.go`**：AUTO 分支改走 `AutoChain()`+`quotaExhausted`（`Quota<=0` 不过滤；`AutoPeriodSeconds`+`stats.WindowTokens` 实时判定）；`singleChatAuto`/`streamChatAuto` 按 `ChainChat`/`ChainChatStream` 重写，总失败在**任何 SSE 字节前**写 JSON；`upstreamErrStatus`：`ChainErr`→503（错误消息即分档汇总）、`ErrBusy`→429、其余→502；失败记录取首个 `TierError` 填 audit source/model；直连无候选→404 `model_not_found`。 +- [x] **P2-4（2026-08-10）单测（`scheduler_test.go`）**：tier 分桶/降序、同 tier 游标轮转交替（s1,s2,s1,s2）、负偏好沉底仍可达（硬失败换槽后由负面槽承接）、busy 跳过不记分、整档全忙有界等待（实测 <1s）后降级、配额耗尽槽不调度、全灭 503 汇总（Tiers+Skipped 文本断言）、流式首 chunk 前失败换槽、游标首帧从 0 起。 +- [x] **P2-5（2026-08-10）本地验证（Windows 已补齐 Lua）**：把 golua 自带 Lua 5.1 头对应的源码编成 `liblua.a` 放入 golua 模块目录（golua 官方 Windows 做法），本机 `go build/vet/test ./...` 全绿——**provider/gateway/lua/e2e 全部首次真正跑通**，并借此揪出三处从未被发现的存量问题：① `RecordFailure(auth)` 只把 failCount 上限用于冷却算式、未落盘计数器（已修：auth 时 `failCount.Store(backoffCapN)`）；② gateway 测试种子 key 未进 runtime store 导致全 401（已修：测试 cfg 补 `GatewayKeys`）；③ e2e failover 用例只有一个 chat 源、AUTO 全灭必然 503（已修：新增第二 chat 源 `fallback`，真故障转移）；e2e `buildBinary` Windows 回退无 tag 构建（bundled Lua，透传适配器行为一致），生产机仍优先 `-tags luajit`。 +- [x] **P2-6（2026-08-10）Phase 2 场景测试**：gateway 新增——同 tier 硬失败顺延承接（`TestChatAutoChainTierFailover`）、全灭 503 含 `a/a-m` 分档汇总（`TestChatAutoChain503Summary`）、配额耗尽槽跳过且不再打上游（`TestChatAutoQuotaSkip`）、`PUT /api/auto` 后冷却立即复位并恢复调度（`TestAutoSaveResetsCooldown`，修 P1 回归）；e2e 新增 `TestEndToEndAuto503`（真实二进制 503 汇总）。全部本地跑通。 +- [ ] **P2-7**：生产机（192.168.2.60）`go test -tags luajit ./...` 最终回归 + 灰度部署（luajit 与 bundled Lua 的适配器行为差异由生产验证兜底）。 + +### Phase 3 — 收尾 +- [x] **P3-1（2026-08-10）UI 链上健康展示**：`GET /api/auto` 增加 `states`（`Core.AutoSlotStates` 遍历当前 Chain 槽 × `Provider.ModelHealthInfo`：pref/failCount/cooldownUntil/cooling）；优先级页每块按 `model|source` 渲染徽标（冷却红/失败橙×N/±偏好蓝，title 说明）；状态页新增"状态码分布"卡片。 +- [x] **P3-2（2026-08-10）保存链健康复位提示**：`saveSort` toast 变更为 `排序已保存并热重载 · 链上冷却已复位`（zh/en）。 +- [x] **P3-3（2026-08-10）审计分类与 jsonl 轮转**：`Stats.byStatus map[int]*Stat`（402 欠费/400 schema 等按状态码独立计数，不触发 provider 退避，`Snapshot.by_status` 有序输出）；audit 文件超 `auditRotateBytes`(64MB, var 可测) 轮转 `rename ..old` 并保留最新 `auditKeepOld`(10) 份——`rotateAuditLocked` 持 mu 在 `Record`/`AppendAudit` 内触发。 +- [x] **P3-4（2026-08-10）运维项（代码部分）**：`main.go` 启动告警——gateway_keys 为空 / 命中种子 key（sk-gw-local-0001 等）提示轮换、listen 绑定 0.0.0.0/:: 提示收敛内网。生产实践（换 admin key、内网绑定）随本次上线执行。 +- [x] **P3-5（2026-08-10）测试**：新增 `stats_test.go`（by_status 断言、轮转→保留上限→Record 路径轮转）；gateway 新增 `TestAutoStatesReportChainHealth`（states 契约：失败后 fail_count>0+cooling，PUT 复位归零）；`go vet ./...` + `go test ./...` 全绿。 +- [x] **P3-6（2026-08-10）WebUI 右键菜单无法关闭（用户实测）**：根因——关闭依赖 `document` 冒泡阶段 once-click 监听，而优先级页块/密钥砖自己的 click 处理 `stopPropagation()` 阻断冒泡 → 菜单永不关闭；修复：改为 **document 捕获阶段**全局 `click`+`mousedown` 关闭（捕获先于一切目标处理器，stopPropagation 无法拦截），优先级页与密钥页共用 `showCtx` 一并修复。 +- [ ] **P3-7**：推送 origin → 生产机 pull → `-tags luajit` 全量回归（含 P2-7）→ 部署 `/usr/local/bin/llmsproxy` + 重启 service → 观察。 \ No newline at end of file