Files
ModelRouter/internal/provider/provider.go
JianFeeeee 21cd0429a2 feat(provider): half-cooldown probe so recovered slots return immediately
A slot that failed ten times landed in a 30-minute cooldown and was a hard
skip for the whole window: ModelAvailable() said no, the scheduler dropped the
candidate, and the ONLY way back was a success that could never happen. In
practice an upstream that blipped for a minute — or whose quota reset in
seconds — stayed unusable for half an hour.

Cooldown is now a window with a probe gate instead of a wall:

  * backoffCap 30min -> 5min. 401/403 no longer implicitly "jump to the cap"
    but get an explicit authCooldown (10min); 429 keeps its 30s fixed window.
  * ModelState records cooldownFrom alongside cooldownUntil, so the window has
    a measurable midpoint, plus a single-token probe permit (probeInFlight).
  * TryProbe() hands out that permit only in the SECOND half of the window: a
    freshly failed slot stays completely silent, and past the midpoint exactly
    one request may go through as a probe.
  * ProbeSlots() = clamp(max_concurrent/10, 1, 2), so a source configured for
    max_concurrent=10 lets exactly one request probe.
  * A probe is a REAL request: success runs RecordSuccess(), which clears the
    cooldown and failCount on the spot and returns the slot to full rotation.
    Failure reopens the window, pushing the next probe to the NEW midpoint
    instead of retrying immediately.

Quota exhaustion is modelled separately (RecordQuotaExhausted): running out of
allowance is not the upstream being broken, so it must not feed the exponential
ladder or inflate failCount. It cools until the allowance window can plausibly
have reset (capped at 30min so a manual top-up is noticed) and is recognised
from the adapter-normalized error text via ReportStatusReason.

Scheduler side: collectCands() splits a tier into normal + probe candidates
with probes strictly LAST, so probe traffic is what a tier falls back to, never
what it prefers, and the round-robin cursor rotates over the normal head only.
releaseProbes() returns every permit on all exits (success, hard failure, ctx
cancel, tier fallthrough); the 90-line inlined busy-poll loop is extracted into
pollBusyTier() to keep those exits auditable. Direct Chat/ChatStream/Image take
the same gate.

Also fixes a real blacklist bug found on the way: a slot whose pref sank to
prefMin was refused by ModelAvailable forever, long after every cooldown had
expired. Such a slot is now probeable, and one success lifts it off the floor.

Verified live against a controllable upstream (tier1 flaky, tier2 healthy):
failure ladder 5s/9s/20s/40s/79s; after healing the upstream, 10 AUTO requests
across the silent half produced ZERO upstream calls; past the midpoint a single
AUTO request produced exactly ONE upstream call, cleared the remaining 39s of
cooldown, and the next three requests were served by the recovered slot.
2026-08-30 08:04:50 +08:00

1466 lines
46 KiB
Go

// Package provider binds a configured source + Lua adapter and performs the
// HTTP call / stream / image generation against the upstream LLM, with
// per-source concurrency limiting and availability backoff.
package provider
import (
"bufio"
"bytes"
"context"
"encoding/json"
"fmt"
"io"
"log"
"net"
"net/http"
"strconv"
"strings"
"sync"
"sync/atomic"
"time"
"llmsproxy/internal/config"
"llmsproxy/internal/lua"
"llmsproxy/internal/types"
)
// ---- 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
// cooldownFrom is the unix second the current cooldown window opened. It
// is the reference point for the half-cooldown probe gate: probing is
// allowed only in the second half of [cooldownFrom, cooldownUntil), so a
// freshly failed slot stays fully silent for a while instead of being
// hammered immediately.
cooldownFrom atomic.Int64
// probeInFlight is the single-token probe permit. At most one probe
// request may be in flight per (source, model) while cooling; the token is
// released by ProbeDone regardless of outcome.
probeInFlight atomic.Bool
}
const (
prefFailStep = 5
prefMin = -20
prefMax = 20
backoffBase = 5 * time.Second
// backoffCap bounds the exponential schedule. It used to be 30min, which
// in practice meant "a slot whose upstream blipped is unusable for half an
// hour even though its quota recovered in seconds". Combined with the
// half-cooldown probe gate below a much shorter cap is safe: a still-dead
// upstream costs one probe request per cap/2, while a recovered one is back
// in full rotation on its first probe success.
backoffCap = 5 * time.Minute
// failures at/above this count back off at the capped duration
backoffCapN = 10
)
// authCooldown is the cooldown applied to 401/403 credential failures. A bad
// key does not fix itself in seconds, so it cools longer than a transport
// blip — but it is still a self-healing cooldown (a rotated key is picked up by
// the next probe), never a permanent blacklist.
const authCooldown = 10 * time.Minute
// maxQuotaCooldown bounds a quota-exhausted cooldown. Even when the reported
// reset window is huge (monthly quotas), the slot is re-probed at least this
// often so a manually topped-up allowance is noticed.
const maxQuotaCooldown = 30 * time.Minute
// quotaDefaultCooldown is used when the source gives no hint about its quota
// window length. It is long enough not to hammer an exhausted allowance and
// short enough that the half-cooldown probe retries within minutes.
const quotaDefaultCooldown = 10 * time.Minute
// rateLimitCooldown is the fixed cooldown applied on HTTP 429. A rate-limit
// rejection means "too fast", not "broken": unlike transport/5xx failures it
// must NOT escalate exponentially to backoffCap, or a quota-rich source gets
// locked out for up to 30 minutes while its allowance is untouched. A short
// fixed window lets the slot re-enter the rotation quickly.
const rateLimitCooldown = 30 * time.Second
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
}
}
// 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 backoffCap). auth marks 401/403
// credential failures: they get the longer fixed authCooldown and a doubled
// preference penalty, but the model still recovers when the cooldown expires.
func (s *ModelState) RecordFailure(auth bool) {
n := s.failCount.Add(1)
var cd time.Duration
switch {
case auth:
// Persist the cap so FailCount() reports "this is a hard problem" and a
// following non-auth failure does not restart the ladder from the bottom.
if n < backoffCapN {
n = backoffCapN
s.failCount.Store(n)
}
cd = authCooldown
case n >= backoffCapN:
cd = backoffCap
default:
cd = backoffBase * time.Duration(1<<(n-1))
if cd > backoffCap {
cd = backoffCap
}
}
s.openCooldown(cd)
pen := int64(prefFailStep)
if auth {
pen *= 2
}
s.pref.Add(-pen)
clampPref(&s.pref, prefMin, prefMax)
}
// openCooldown starts a fresh cooldown window of length cd. Recording both
// ends of the window (not just the deadline) is what makes the half-cooldown
// probe gate possible, and resetting cooldownFrom on every new failure means a
// failed probe pushes the next probe to the middle of the NEW window instead of
// retrying immediately.
func (s *ModelState) openCooldown(cd time.Duration) {
now := time.Now()
s.cooldownFrom.Store(now.Unix())
s.cooldownUntil.Store(now.Add(cd).Unix())
}
// RecordRateLimit records an HTTP 429: a short fixed cooldown instead of the
// exponential schedule, so a merely-throttled source recovers in seconds and
// its remaining quota stays usable. The preference penalty still applies so
// other slots are preferred while cooling.
func (s *ModelState) RecordRateLimit() {
s.failCount.Add(1)
s.openCooldown(rateLimitCooldown)
s.pref.Add(-prefFailStep)
clampPref(&s.pref, prefMin, prefMax)
}
// RecordQuotaExhausted marks the model as out of allowance until its quota
// window resets. Unlike a failure this is not "the upstream is broken", so it
// must not feed the exponential ladder: the slot is unusable until the window
// boundary and immediately usable again after it. resetIn is the time left in
// the current quota window; non-positive values fall back to the short
// rate-limit window (the caller could not determine a period).
func (s *ModelState) RecordQuotaExhausted(resetIn time.Duration) {
if resetIn <= 0 {
resetIn = rateLimitCooldown
}
if resetIn > maxQuotaCooldown {
resetIn = maxQuotaCooldown
}
s.openCooldown(resetIn)
// No failCount bump: an exhausted quota is not an error streak, and
// letting it inflate failCount would make the NEXT real failure jump
// straight to the capped backoff.
s.pref.Add(-prefFailStep)
clampPref(&s.pref, prefMin, prefMax)
}
// RecordSuccess resets the failure counter and cooldown and bumps the
// preference score by one. A single success is enough to lift a slot off the
// prefMin floor (prefMin+1 > prefMin), which is what lets one successful probe
// return a written-off slot to normal rotation.
func (s *ModelState) RecordSuccess() {
s.failCount.Store(0)
s.cooldownUntil.Store(0)
s.cooldownFrom.Store(0)
s.pref.Add(1)
clampPref(&s.pref, prefMin, prefMax)
}
func (s *ModelState) reset() {
s.failCount.Store(0)
s.cooldownUntil.Store(0)
s.cooldownFrom.Store(0)
s.probeInFlight.Store(false)
s.pref.Store(0)
}
// TryProbe attempts to claim the single probe permit for a model that normal
// scheduling refuses. It returns true only when a probe is both needed and due:
//
// - cooling: allowed once the window passed its midpoint (the first half stays
// completely silent).
// - cooldown expired but the preference score sank to prefMin: ModelAvailable
// still refuses the slot, so a probe is the ONLY way back. Without this a
// slot that failed prefMin/prefFailStep times in a row would be blacklisted
// forever despite every cooldown having long expired.
//
// A healthy, non-cooling model returns false: the caller should schedule it
// normally. On true the caller MUST call ProbeDone exactly once.
func (s *ModelState) TryProbe() bool {
now := time.Now().Unix()
until := s.cooldownUntil.Load()
if until <= now {
if s.pref.Load() > prefMin {
return false // schedulable normally, no probe needed
}
// preference floor: probe to earn the way out
return s.probeInFlight.CompareAndSwap(false, true)
}
from := s.cooldownFrom.Load()
if from <= 0 || from >= until {
return false // no measurable window (legacy/degenerate state)
}
if now < from+(until-from)/2 {
return false // first half of the window: stay silent
}
return s.probeInFlight.CompareAndSwap(false, true)
}
// ProbeDone releases the probe permit claimed by TryProbe.
func (s *ModelState) ProbeDone() { s.probeInFlight.Store(false) }
// Probing reports whether a probe request currently holds the permit.
func (s *ModelState) Probing() bool { return s.probeInFlight.Load() }
// CooldownFrom is the unix timestamp the current cooldown window opened; 0
// when the model is not cooling.
func (s *ModelState) CooldownFrom() int64 { return s.cooldownFrom.Load() }
// 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
vm *lua.VM
adapter string
// client bounds a whole non-streaming request (dial+read body). stream
// is used for SSE: no overall timeout (a long stream must not be cut),
// only the transport's ResponseHeaderTimeout bounds time-to-first-byte.
client *http.Client
stream *http.Client
mu sync.Mutex
sem chan struct{}
states map[string]*ModelState // key = model id
// proactive rate limiting: nextOK is the earliest unix-nano time the next
// request may leave; throttle() spaces requests 60s/RPM apart. Zero when
// cfg.RPM <= 0 (unlimited).
rateMu sync.Mutex
nextOK int64
rpmGap time.Duration
lastProbe struct {
ok bool
err string
at int64
}
}
func New(cfg config.Source, vm *lua.VM) *Provider {
if cfg.Timeout <= 0 {
cfg.Timeout = config.DefaultSourceTimeout
}
if cfg.MaxConcurrent <= 0 {
cfg.MaxConcurrent = config.DefaultSourceConcurrency
}
// Shared transport: ResponseHeaderTimeout bounds how long we wait for the
// first response byte (applies to both paths); the stream client has no
// client-level Timeout so the SSE body can run past the header timeout.
tr := &http.Transport{
Proxy: http.ProxyFromEnvironment,
DialContext: (&net.Dialer{Timeout: 30 * time.Second, KeepAlive: 30 * time.Second}).DialContext,
ForceAttemptHTTP2: true,
MaxIdleConns: 100,
IdleConnTimeout: 90 * time.Second,
TLSHandshakeTimeout: 10 * time.Second,
ExpectContinueTimeout: 1 * time.Second,
ResponseHeaderTimeout: cfg.Timeout,
}
p := &Provider{
cfg: cfg,
vm: vm,
adapter: cfg.Adapter,
client: &http.Client{Timeout: cfg.Timeout, Transport: tr},
stream: &http.Client{Transport: tr},
sem: make(chan struct{}, cfg.MaxConcurrent),
states: map[string]*ModelState{},
}
if cfg.MaxConcurrent <= 0 {
p.sem = nil
}
if cfg.RPM > 0 {
p.rpmGap = time.Minute / time.Duration(cfg.RPM)
}
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 }
// Models returns the model ids exposed by this source.
func (p *Provider) Models() []string {
out := make([]string, 0, len(p.cfg.Models))
for _, m := range p.cfg.Models {
out = append(out, m.ID)
}
return out
}
// ModelByID returns the model definition if owned by this source.
func (p *Provider) ModelByID(id string) *config.Model {
for i := range p.cfg.Models {
if p.cfg.Models[i].ID == id {
return &p.cfg.Models[i]
}
}
return nil
}
// ModelIDFold returns the configured model id matching id case-insensitively
// ("" when no model matches). AUTO-chain slot models are normalized through
// this so cooldown state, quota windows and the upstream model id all refer
// to the exact configured spelling.
func (p *Provider) ModelIDFold(id string) string {
for _, m := range p.cfg.Models {
if strings.EqualFold(m.ID, id) {
return m.ID
}
}
return ""
}
// ModelFor resolves the model name this provider should send upstream.
// If the requested model is not owned by this provider (e.g. an AUTO chain
// fallback), it returns this provider's highest-priority chat model instead.
func (p *Provider) ModelFor(reqModel string) string {
if reqModel == "" || isAutoID(reqModel) {
return p.bestChatModel()
}
if p.ModelByID(reqModel) != nil {
return reqModel
}
return p.bestChatModel()
}
// bestChatModel returns the highest-priority chat-kind model of this source.
func (p *Provider) bestChatModel() string {
bestID, bestPrio := "", -1
for _, m := range p.cfg.Models {
if m.Kind != "" && m.Kind != "chat" {
continue
}
if m.Priority > bestPrio {
bestPrio = m.Priority
bestID = m.ID
}
}
if bestID == "" && len(p.cfg.Models) > 0 {
bestID = p.cfg.Models[0].ID
}
return bestID
}
// bestImageModel returns the highest-priority image-kind model of this source
// (fallback: first configured model). Used for AUTO image generation so a
// mixed source never sends a chat model to /v1/images/generations.
func (p *Provider) bestImageModel() string {
bestID, bestPrio := "", -1
for _, m := range p.cfg.Models {
if m.Kind != "" && m.Kind != "image" {
continue
}
if m.Priority > bestPrio {
bestPrio = m.Priority
bestID = m.ID
}
}
if bestID == "" && len(p.cfg.Models) > 0 {
bestID = p.cfg.Models[0].ID
}
return bestID
}
// IsAutoID reports whether s is an AUTO routing placeholder.
func isAutoID(s string) bool {
s = strings.TrimSpace(s)
return s == "" || strings.EqualFold(s, "AUTO")
}
// Endpoint resolves the upstream chat path.
func (p *Provider) Endpoint() string {
if p.cfg.Endpoint != "" {
return p.cfg.Endpoint
}
if ep := p.vm.Endpoint(p.adapter); ep != "" {
return ep
}
return "/chat/completions"
}
// ImageEndpoint resolves the upstream image-generation path.
func (p *Provider) ImageEndpoint() string {
if p.cfg.ImageEndpoint != "" {
return p.cfg.ImageEndpoint
}
if ep := p.vm.Endpoint(p.adapter + "_image"); ep != "" {
return ep
}
return "/v1/images/generations"
}
func (p *Provider) URL() string {
return strings.TrimRight(p.cfg.BaseURL, "/") + p.Endpoint()
}
func (p *Provider) ImageURL() string {
return strings.TrimRight(p.cfg.BaseURL, "/") + p.ImageEndpoint()
}
// ---- 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()
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() && s.Pref() > prefMin
}
return true
}
// ModelSchedulable is the scheduling gate that replaces a bare ModelAvailable
// check. It reports whether the model may take a request right now and whether
// doing so is a cooldown PROBE rather than normal traffic:
//
// - (true, false): fully available, schedule normally.
// - (true, true): cooling, but past the half-cooldown mark and this caller
// claimed the probe permit. The caller MUST call ProbeDone(model) once the
// attempt finished, and SHOULD prefer non-probe candidates first.
// - (false, false): cooling and not probeable (first half of the window, the
// permit is taken, or the slot is persistently bad).
//
// Probing is what stops a recovered upstream from sitting out a whole cooldown:
// the probe is a real request, so success immediately clears the cooldown via
// RecordSuccess and the slot is back in full rotation.
func (p *Provider) ModelSchedulable(model string) (ok bool, isProbe bool) {
p.mu.Lock()
s, known := p.states[model]
p.mu.Unlock()
if !known {
return true, false
}
if s.Available() && s.Pref() > prefMin {
return true, false
}
if s.TryProbe() {
return true, true
}
return false, false
}
// ProbeDone releases a probe permit claimed through ModelSchedulable.
func (p *Provider) ProbeDone(model string) { p.state(model).ProbeDone() }
// ProbeSlots is how many concurrent probe requests this source tolerates while
// a model cools: one per 10 configured concurrency slots, at least 1 and never
// more than 2. A source configured for max_concurrent=10 therefore lets exactly
// one request through to probe, which is the intended "small-scale probe"
// behaviour; a large source still cannot flood a cooling upstream.
func (p *Provider) ProbeSlots() int {
n := p.cfg.MaxConcurrent / 10
if n < 1 {
return 1
}
if n > 2 {
return 2
}
return n
}
// 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()
}
// ModelProbeInfo exposes the probe-gate state for the web UI: when the current
// cooldown window opened, when probing becomes allowed (its midpoint), and
// whether a probe is in flight right now. All zero when not cooling.
func (p *Provider) ModelProbeInfo(model string) (from, probeAfter int64, probing bool) {
s := p.state(model)
until := s.CooldownUntil()
from = s.CooldownFrom()
if until <= 0 || from <= 0 || from >= until {
return 0, 0, s.Probing()
}
return from, from + (until-from)/2, s.Probing()
}
// Probe performs a lightweight reachability + auth check against the source.
// It first tries GET <base>/models (fast, ~1s for OpenAI-compatible upstreams)
// and only falls back to a 1-token chat call when that endpoint is
// unavailable. It does NOT touch the health/backoff state so probing never
// disables a source.
func (p *Provider) Probe(ctx context.Context) (bool, string) {
ok, msg := p.probeModels(ctx)
if !ok && msg == "" {
ok, msg = p.probeChat(ctx)
}
p.mu.Lock()
p.lastProbe.ok = ok
p.lastProbe.err = msg
p.lastProbe.at = time.Now().Unix()
p.mu.Unlock()
return ok, msg
}
// probeModels GETs <base>/models. Returns (true,…) when reachable, (false,
// errortext) on an auth/permanent failure, and (false,"") when the endpoint
// simply isn't available so the caller can fall back to a chat probe.
func (p *Provider) probeModels(ctx context.Context) (bool, string) {
u := strings.TrimRight(p.cfg.BaseURL, "/") + "/models"
req, err := http.NewRequestWithContext(ctx, http.MethodGet, u, nil)
if err != nil {
return false, ""
}
if hdrs, herr := p.buildHeaders("{}", u); herr == nil {
req.Header = hdrs
}
resp, err := p.client.Do(req)
if err != nil {
return false, ""
}
defer resp.Body.Close()
raw, _ := io.ReadAll(resp.Body)
switch {
case resp.StatusCode == 200:
return true, ""
case resp.StatusCode == 404 || resp.StatusCode == 405:
return false, ""
default:
return false, p.apiErrReason(resp.StatusCode, string(raw))
}
}
// probeChat sends a minimal single-token chat request to the chat endpoint.
func (p *Provider) probeChat(ctx context.Context) (bool, string) {
ok := false
msg := ""
model := p.bestChatModel()
if pm := p.ModelByID(model); pm != nil && pm.Kind == "image" {
model = ""
}
if model == "" {
if ms := p.Models(); len(ms) > 0 {
model = ms[0]
}
}
if model == "" {
msg = "no chat model configured"
} else {
probe := map[string]interface{}{
"model": model,
"messages": []map[string]interface{}{{"role": "user", "content": "hi"}},
"max_tokens": 1,
}
body, err := json.Marshal(probe)
if err == nil {
var hdr http.Header
if hdrs, herr := p.buildHeaders(string(body), p.URL()); herr == nil {
hdr = hdrs
}
raw, status, derr := p.do(ctx, p.URL(), string(body), hdr)
if derr != nil {
msg = derr.Error()
} else if status == 200 {
ok = true
} else {
msg = p.apiErrReason(status, raw)
}
} else {
msg = err.Error()
}
}
return ok, msg
}
func (p *Provider) LastProbe() (bool, string, int64) {
p.mu.Lock()
defer p.mu.Unlock()
return p.lastProbe.ok, p.lastProbe.err, p.lastProbe.at
}
// 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 → long self-healing authCooldown
// with doubled penalty; 5xx → normal exponential backoff; 429 → short fixed
// cooldown (rateLimitCooldown): a rate rejection means "too fast", not
// "broken", so the slot must re-enter rotation quickly instead of escalating
// to a long lockout that wastes remaining quota.
//
// Prefer ReportStatusReason when the upstream error text is available: a
// quota-exhausted rejection is cooled to its window boundary instead of being
// treated as a failure streak.
func (p *Provider) ReportStatus(model string, code int) {
p.ReportStatusReason(model, code, "")
}
// ReportStatusReason is ReportStatus plus the upstream error text, which lets a
// quota/allowance rejection be distinguished from a generic rate limit. Quota
// exhaustion is not an error streak: it is cooled until the allowance window
// can plausibly have reset, and never inflates the exponential ladder.
func (p *Provider) ReportStatusReason(model string, code int, reason string) {
if code == 401 || code == 403 {
p.RecordFailure(model, code)
return
}
if code == 429 || code == 402 {
if quotaExhaustedReason(reason) {
p.state(model).RecordQuotaExhausted(p.quotaResetIn(model))
return
}
p.state(model).RecordRateLimit()
return
}
if code >= 500 {
p.RecordFailure(model, code)
}
}
// quotaResetIn estimates how long the model's allowance window still has to
// run. Upstreams do not report this, so it is taken from the source's optional
// `quota_hours` meta hint when present (aligned to the epoch, matching the
// gateway's own quota-window convention) and otherwise falls back to
// quotaDefaultCooldown. Either way the value is capped by maxQuotaCooldown, so
// a wrong hint can only delay recovery to that bound — and the half-cooldown
// probe fires at half of it.
func (p *Provider) quotaResetIn(model string) time.Duration {
hours := p.quotaHours()
if hours <= 0 {
return quotaDefaultCooldown
}
win := int64(hours * float64(time.Hour/time.Second))
if win <= 0 {
return quotaDefaultCooldown
}
elapsed := time.Now().Unix() % win
return time.Duration(win-elapsed) * time.Second
}
// quotaHours reads the optional per-source `quota_hours` meta hint describing
// how long its allowance window is. Absent or unparseable = 0 (unknown).
func (p *Provider) quotaHours() float64 {
v, ok := p.cfg.Meta["quota_hours"]
if !ok {
return 0
}
switch n := v.(type) {
case float64:
return n
case int:
return float64(n)
case int64:
return float64(n)
case string:
f, err := strconv.ParseFloat(strings.TrimSpace(n), 64)
if err != nil {
return 0
}
return f
}
return 0
}
// quotaKeywords are the upstream phrases that mean "you are out of allowance"
// rather than "you are going too fast". Matching is case-insensitive substring
// matching against the adapter-normalized error reason.
var quotaKeywords = []string{
"insufficient_quota",
"insufficient quota",
"quota exceeded",
"quota_exceeded",
"exceeded your current quota",
"out of credit",
"insufficient balance",
"insufficient_user_quota",
"billing_hard_limit_reached",
"credit balance is too low",
"余额不足",
"额度不足",
"额度已用完",
}
func quotaExhaustedReason(reason string) bool {
if reason == "" {
return false
}
low := strings.ToLower(reason)
for _, k := range quotaKeywords {
if strings.Contains(low, k) {
return true
}
}
return false
}
// 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()
defer p.mu.Unlock()
for _, s := range p.states {
s.reset()
}
}
// 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()
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. 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
}
var qCtx context.Context
var cancel context.CancelFunc
if p.cfg.QueueTimeout > 0 {
qCtx, cancel = context.WithTimeout(ctx, p.cfg.QueueTimeout)
} else {
qCtx, cancel = context.WithCancel(ctx)
}
defer cancel()
select {
case p.sem <- struct{}{}:
return nil
case <-qCtx.Done():
return qCtx.Err()
}
}
func (p *Provider) Release() {
if p.sem == nil {
return
}
<-p.sem
}
// ---- proactive rate limiting ----
// Throttle blocks until this source's rate-limit window allows the next
// request, spacing outbound requests at least rpmGap apart (60s/RPM). It is a
// no-op when no rpm cap is configured. The wait happens AFTER the concurrency
// slot is taken, so a queued request still counts against max_concurrent —
// bounded by the caller's ctx (queue timeout / client disconnect).
func (p *Provider) Throttle(ctx context.Context) error {
if p.rpmGap <= 0 {
return nil
}
p.rateMu.Lock()
now := time.Now().UnixNano()
wait := p.nextOK - now
if wait > 0 {
p.nextOK += int64(p.rpmGap) // reserve the next window for the follower
p.rateMu.Unlock()
t := time.NewTimer(time.Duration(wait))
defer t.Stop()
select {
case <-t.C:
return nil
case <-ctx.Done():
return ctx.Err()
}
}
p.nextOK = now + int64(p.rpmGap)
p.rateMu.Unlock()
return nil
}
// ---- request construction ----
func (p *Provider) buildHeaders(body, url string) (http.Header, error) {
meta := map[string]interface{}{
"url": url,
"method": http.MethodPost,
"body": body,
"api_key": p.cfg.APIKey,
"timestamp": types.Now(),
"source": map[string]interface{}{
"name": p.cfg.Name,
"meta": p.cfg.Meta,
},
}
hdrs, err := p.vm.BuildHeaders(p.adapter, meta)
if err != nil {
return nil, err
}
h := http.Header{}
h.Set("Content-Type", "application/json")
for k, v := range p.cfg.Headers {
h.Set(k, v)
}
for k, v := range hdrs {
if _, ok := p.cfg.Headers[k]; !ok {
h.Set(k, v)
}
}
if h.Get("Authorization") == "" && p.cfg.APIKey != "" {
h.Set("Authorization", "Bearer "+p.cfg.APIKey)
}
return h, nil
}
// ---- 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.TryAcquire(ctx); err != nil {
return nil, err
}
defer p.Release()
if err := p.Throttle(ctx); err != nil {
return nil, err
}
model := p.ModelFor(req.Model)
body, err := marshalTransform(p.vm, p.adapter, "transform_request", req)
if err != nil {
return nil, err
}
hdrs, err := p.buildHeaders(body, p.URL())
if err != nil {
return nil, err
}
raw, status, err := p.do(ctx, p.URL(), body, hdrs)
if err != nil {
// a client disconnect or cancelled context is neither a success nor
// a failure for scheduling purposes — only upstream errors count
if ctx.Err() == nil {
p.RecordFailure(model, 0)
}
return nil, err
}
if status != 200 {
reason := p.apiErrReason(status, raw)
p.ReportStatusReason(model, status, reason)
return nil, fmt.Errorf("%s", reason)
}
unified, err := p.vm.Transform(p.adapter, "transform_response", raw)
if err != nil {
// adapter produced unusable output: a real failure the slot must
// back off from, otherwise a broken adapter source is retried at
// full latency forever
if ctx.Err() == nil {
p.RecordFailure(model, 0)
}
return nil, err
}
var out types.UnifiedResponse
if err := json.Unmarshal([]byte(unified), &out); err != nil {
if ctx.Err() == nil {
p.RecordFailure(model, 0)
}
return nil, fmt.Errorf("unmarshal unified response: %w (body: %s)", err, unified)
}
p.RecordSuccess(model)
return &out, nil
}
// 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.TryAcquire(ctx); err != nil {
return nil, err
}
if err := p.Throttle(ctx); err != nil {
p.Release()
return nil, err
}
model := p.ModelFor(req.Model)
req.Stream = true
body, err := marshalTransform(p.vm, p.adapter, "transform_request", req)
if err != nil {
p.Release()
return nil, err
}
hdrs, err := p.buildHeaders(body, p.URL())
if err != nil {
p.Release()
return nil, err
}
type respOrErr struct {
resp *http.Response
err error
}
rc := make(chan respOrErr, 1)
go func() {
resp, err := p.doRawStream(ctx, p.URL(), body, hdrs)
rc <- respOrErr{resp, err}
}()
inner := make(chan types.UnifiedChunk, 64)
sel := <-rc
if sel.err != nil {
// client disconnect/cancel before the first byte: not a scheduling
// failure (a healthy source must not be cooled by client cancellations)
if ctx.Err() == nil {
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()
reason := p.apiErrReason(sel.resp.StatusCode, string(raw))
p.ReportStatusReason(model, sel.resp.StatusCode, reason)
p.Release()
return nil, fmt.Errorf("%s", reason)
}
go func() {
defer p.Release()
defer close(inner)
defer sel.resp.Body.Close()
scanner := bufio.NewScanner(sel.resp.Body)
scanner.Buffer(make([]byte, 0, 64*1024), 1024*1024)
var chunks int
var realChunks int
var doneSeen bool
var doneSent bool
// toolIdx maps a per-stream tool fragment's upstream index (Anthropic
// content-block ordinal, which is sparse when thinking/text blocks
// precede tool_use) to a dense 0-based OpenAI tool_call ordinal so
// clients that accumulate fragments by index reassemble multi-tool
// responses correctly.
toolIdx := newToolIndexRemapper()
for scanner.Scan() {
line := strings.TrimSpace(scanner.Text())
if line == "" || !strings.HasPrefix(line, "data:") {
continue
}
data := strings.TrimSpace(strings.TrimPrefix(line, "data:"))
if data == "" {
continue
}
if data == "[DONE]" {
doneSeen = true
// adapters that already emitted their terminating done chunk
// (with the real finish reason) must not get a second,
// reason-less done from [DONE] — it would override the true
// finish_reason downstream.
if !doneSent {
select {
case inner <- types.UnifiedChunk{Done: true}:
case <-ctx.Done():
}
}
continue
}
unified, err := p.vm.Transform(p.adapter, "transform_stream_chunk", data)
if err != nil || unified == "" {
continue
}
if unified == data {
unified = standardSSEChunk(data)
if unified == "" {
continue
}
}
var ck types.UnifiedChunk
if err := json.Unmarshal([]byte(unified), &ck); err != nil {
continue
}
ck.ToolCalls = toolIdx.remap(ck.ToolCalls)
if ck.Done {
doneSent = true
}
if !errorOnlyChunk(ck) {
realChunks++
}
chunks++
select {
case inner <- ck:
case <-ctx.Done():
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. A 200 that produced
// zero chunks and no [DONE] is an empty stream, i.e. a failure
// before the first chunk — record it so the slot can fall back.
// A stream whose only payload was error-only done chunks (e.g.
// finish_reason:"network_error") is likewise a failure, not a success.
if ctx.Err() == nil && scanner.Err() == nil {
if realChunks > 0 || (doneSeen && chunks == 0) {
p.RecordSuccess(model)
} else {
p.RecordFailure(model, 0)
}
}
}()
// Hold back the first chunk to validate the stream actually carries
// content: some upstreams answer HTTP 200 with a degenerate stream whose
// only payload is an error finish reason (zen free pool sends
// finish_reason:"network_error" with empty delta). Failing the candidate
// here — before any byte reaches the gateway — lets the scheduler fall
// through to the next source instead of serving the client an empty reply.
type heldChunk struct {
ck types.UnifiedChunk
ok bool
}
var first heldChunk
select {
case ck, ok := <-inner:
first = heldChunk{ck, ok}
case <-ctx.Done():
return nil, ctx.Err()
}
if !first.ok {
return nil, fmt.Errorf("provider %s: empty stream", p.Name())
}
if errorOnlyChunk(first.ck) {
go func() {
for range inner {
}
}()
return nil, fmt.Errorf("provider %s: upstream returned %q stream",
p.Name(), first.ck.FinishReason)
}
out := make(chan types.UnifiedChunk, 64)
go func() {
defer close(out)
out <- first.ck
for ck := range inner {
out <- ck
}
}()
return out, nil
}
// errorOnlyChunk reports whether ck carries nothing but an upstream error
// signal: a done chunk with a non-standard finish reason and zero content,
// tool calls, reasoning text or usage. Standard OpenAI finish reasons are
// never classified as errors, so legitimate instant-empty completions
// (finish_reason:"stop", no output) still reach the client.
// toolIndexRemapper maps per-stream tool_call fragment indices from a sparse
// upstream ordinal (e.g. Anthropic content_block index, which skips over
// thinking/text blocks) to a dense 0-based OpenAI tool_call ordinal.
type toolIndexRemapper struct {
seen map[int]int // upstream index -> dense ordinal
next int // next dense ordinal to assign
}
func newToolIndexRemapper() toolIndexRemapper {
return toolIndexRemapper{seen: make(map[int]int)}
}
func (t *toolIndexRemapper) remap(raw json.RawMessage) json.RawMessage {
if len(raw) == 0 {
return raw
}
// Decode into generic maps so every upstream field is preserved verbatim;
// only the index value is rewritten.
var frags []map[string]json.RawMessage
if err := json.Unmarshal(raw, &frags); err != nil || len(frags) == 0 {
return raw
}
changed := false
for i := range frags {
orig := 0
if v, ok := frags[i]["index"]; ok {
if err := json.Unmarshal(v, &orig); err != nil {
continue
}
}
mapped, ok := t.seen[orig]
if !ok {
mapped = t.next
t.next++
t.seen[orig] = mapped
}
if mapped != orig {
b, err := json.Marshal(mapped)
if err != nil {
continue
}
frags[i]["index"] = b
changed = true
}
}
if !changed {
return raw
}
out, err := json.Marshal(frags)
if err != nil {
return raw
}
return out
}
func errorOnlyChunk(ck types.UnifiedChunk) bool {
if !ck.Done || ck.FinishReason == "" {
return false
}
switch ck.FinishReason {
case "stop", "length", "tool_calls", "function_call", "content_filter":
return false
}
return ck.Content == "" && len(ck.ToolCalls) == 0 &&
ck.ReasoningContent == "" && ck.Usage == nil
}
// apiErrReason builds the client-facing reason for a non-200 upstream
// response: the adapter's optional transform_error hook wins (per-source
// protocol knowledge lives in Lua), otherwise clients get a uniform
// "unknown error" while the raw body stays in the server log for debugging.
func (p *Provider) apiErrReason(status int, raw string) string {
if reason, ok, err := p.vm.TransformError(p.adapter, status, raw); err == nil && ok {
if trimmed := strings.TrimSpace(reason); trimmed != "" {
return fmt.Sprintf("api error %d: %s", status, types.OneLine(trimmed, 200))
}
}
log.Printf("[provider] unhandled upstream error body (adapter %q lacks transform_error): status=%d body=%.300s",
p.adapter, status, types.OneLine(raw, 300))
return fmt.Sprintf("api error %d: unknown error", status)
}
// 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.TryAcquire(ctx); err != nil {
return nil, err
}
defer p.Release()
// AUTO/unknown model: resolve to this source's image model, never to the
// best chat model (a mixed source would otherwise send a chat id to
// /v1/images/generations). An explicitly pinned model is honored as-is.
if isAutoID(req.Model) || p.ModelByID(req.Model) == nil {
r := *req
r.Model = p.bestImageModel()
req = &r
}
model := req.Model
b, _ := json.Marshal(req)
transformed, err := p.vm.Transform(p.adapter+"_image", "transform_request", string(b))
if err != nil {
// fall back to passthrough adapter (openai-style)
transformed = string(b)
}
hdrs, err := p.buildHeaders(transformed, p.ImageURL())
if err != nil {
return nil, err
}
raw, status, err := p.do(ctx, p.ImageURL(), transformed, hdrs)
if err != nil {
// client disconnect/cancel: not a scheduling failure
if ctx.Err() == nil {
p.RecordFailure(model, 0)
}
return nil, err
}
if status != 200 {
reason := p.apiErrReason(status, raw)
p.ReportStatusReason(model, status, reason)
return nil, fmt.Errorf("%s", reason)
}
var out types.UnifiedResponse
// try adapter transform_response; if missing, parse standard openai image format
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 {
out.Model = model // actual model served (AUTO resolved to bestImageModel)
p.RecordSuccess(model)
return &out, nil
}
}
var img types.ImageGenResponse
if err := json.Unmarshal([]byte(raw), &img); err != nil {
if ctx.Err() == nil {
p.RecordFailure(model, 0)
}
return nil, fmt.Errorf("unmarshal image response: %w", err)
}
out.ImageData = img.Data
out.Model = model // actual model served (AUTO resolved to bestImageModel)
p.RecordSuccess(model)
return &out, nil
}
// ---- http helpers ----
func (p *Provider) do(ctx context.Context, url, body string, hdr http.Header) (string, int, error) {
resp, err := p.doRaw(ctx, url, body, hdr)
if err != nil {
return "", 0, err
}
defer resp.Body.Close()
raw, _ := io.ReadAll(resp.Body)
return string(raw), resp.StatusCode, nil
}
func (p *Provider) doRaw(ctx context.Context, url, body string, hdr http.Header) (*http.Response, error) {
httpReq, err := http.NewRequestWithContext(ctx, http.MethodPost, url, bytes.NewReader([]byte(body)))
if err != nil {
return nil, err
}
httpReq.Header = hdr
return p.client.Do(httpReq)
}
// doRawStream is the streaming variant of doRaw: it uses the no-overall-timeout
// stream client so a long SSE body is not cut by client.Timeout. The transport
// still enforces ResponseHeaderTimeout on time-to-first-byte.
func (p *Provider) doRawStream(ctx context.Context, url, body string, hdr http.Header) (*http.Response, error) {
httpReq, err := http.NewRequestWithContext(ctx, http.MethodPost, url, bytes.NewReader([]byte(body)))
if err != nil {
return nil, err
}
httpReq.Header = hdr
return p.stream.Do(httpReq)
}
func marshalTransform(vm *lua.VM, adapter, fn string, v interface{}) (string, error) {
b, err := json.Marshal(v)
if err != nil {
return "", err
}
out, err := vm.Transform(adapter, fn, string(b))
if err != nil {
return "", err
}
return out, nil
}
func standardSSEChunk(data string) string {
var raw struct {
Choices []struct {
Delta struct {
Content string `json:"content"`
} `json:"delta"`
FinishReason *string `json:"finish_reason"`
} `json:"choices"`
UpstreamUsage struct {
PromptTokens int `json:"prompt_tokens"`
CompletionTokens int `json:"completion_tokens"`
TotalTokens int `json:"total_tokens"`
Prompt int `json:"prompt"`
Completion int `json:"completion"`
Total int `json:"total"`
PromptCacheHit int `json:"prompt_cache_hit_tokens"`
PromptCacheMiss int `json:"prompt_cache_miss_tokens"`
PromptTokensDetails *struct {
CachedTokens int `json:"cached_tokens"`
} `json:"prompt_tokens_details"`
} `json:"usage"`
}
if err := json.Unmarshal([]byte(data), &raw); err != nil {
return ""
}
if len(raw.Choices) == 0 && raw.UpstreamUsage.Total == 0 && raw.UpstreamUsage.TotalTokens == 0 {
return ""
}
var usage *types.TokenUsage
pu := raw.UpstreamUsage
if pu.Total > 0 || pu.TotalTokens > 0 {
usage = &types.TokenUsage{
Prompt: pickFirst(pu.PromptTokens, pu.Prompt),
Completion: pickFirst(pu.CompletionTokens, pu.Completion),
Total: pickFirst(pu.TotalTokens, pu.Total),
PromptCacheHit: pu.PromptCacheHit,
PromptCacheMiss: pu.PromptCacheMiss,
}
if pu.PromptTokensDetails != nil && pu.PromptTokensDetails.CachedTokens > 0 {
usage.PromptTokensDetails = &types.PromptTokensDetails{
CachedTokens: pu.PromptTokensDetails.CachedTokens,
}
}
}
finish := ""
done := false
if len(raw.Choices) > 0 && raw.Choices[0].FinishReason != nil {
// empty-string finish reasons (sensenova sends "" on every chunk)
// are not a finish signal
if fr := *raw.Choices[0].FinishReason; fr != "" {
finish, done = fr, true
}
}
out, _ := json.Marshal(types.UnifiedChunk{
Content: func() string {
if len(raw.Choices) > 0 {
return raw.Choices[0].Delta.Content
}
return ""
}(),
Done: done,
FinishReason: finish,
Usage: usage,
})
return string(out)
}
// pickFirst returns a if non-zero, else b (for usage keys that may appear in
// either standard *_tokens or legacy short form).
func pickFirst(a, b int) int {
if a != 0 {
return a
}
return b
}