feat(llm): LLM 配置轮询(故障转移),某个配置没额度/失效时自动切下一个
默认关闭。开启后「未指定模型」的 agent 在当前配置不可用(402 余额不足 /
401 key 失效 / 429 限流 / 5xx)时自动切到下一个配置;agent 绑定或任务 pin
了配置的默认不轮询(绑定即独占,可由 llm_pool_bind_fallback 打开兜底)。
链序 = 激活配置 + 其余按 priority DESC,排除 pool_exclude 的。只给
llm_profiles 加两列,不引入关联表,所以 agents/tasks/conversations 对
llm_profile_id 的四处引用一行未改。
核心约束:只在本次调用尚未 yield 出任何事件时切换成员。已经吐出
text_delta / tool_use_start 后中途断流一律原样上抛(交给 agentcore 的
resume),否则重发同一 request 会重复输出、污染会话历史。SDK 的 doStream
保证 4xx/5xx 在读 body 前返回,所以目标场景都落在这个安全窗口内。
context.Canceled / DeadlineExceeded 绝不触发切换——那是用户取消或任务
超时,误判会白烧备用 key 并污染 a6bd180 的终止原因诊断。400 也不切
(请求本身有问题,换模型结果一样)。
熔断:402/401/403/404 一次即熔断,429/5xx/网络错误连续 3 次;冷却
60s→5min→30min 退避,到期半开试探,成功清零。Registry 是进程级的,独立
于 provider 缓存,重建链路不丢状态,且全进程共享(A 任务撞见的 402,B 任务
立刻知道)。状态旁路落库,重启只恢复尚未到期的冷却窗口。
上下文窗口装不下当前请求的成员直接跳过(必然 400),不把主力模型的窗口
压到全链最小值。全员熔断时仍探链首,不让引擎停摆。
Pool 包在 llmrec 外面、每个成员自带 recorder,失败调用会以 status=error
+ 正确的 profile 名进 llm_records;切换/熔断走 log.Printf,自动出现在
系统日志页。
顺带修两个既有问题:
- DecomposeGoals 原本自己 NewProvider(),绕过限速器/录制/轮询,且每个
任务新建一个实例——改为由 server 传入共享 provider
- saveLLMConfig 会把 priority/pool_exclude 重置为零值
前端:LLM 页新增轮询卡片(总开关/兜底开关/实际链路顺序/健康徽章/冷却
倒计时/立即恢复),配置编辑区加优先级与「不参与轮询」。
只有一个配置时链长度为 1,poolChain 返回 nil 走裸 provider,行为与开关
打开前完全一致。设计见 docs/LLM轮询设计.md。
This commit is contained in:
Autumn-27 committed
1 parent
89c5e9b26f
commit
0c653a2135
19 files changed
+1750
-50
No files matched your search
@@ -0,0 +1,149 @@
|
||||
// Package llmpool implements LLM failover ("轮询"): a Provider decorator that
|
||||
// walks an ordered chain of LLM profiles and moves to the next one when the
|
||||
// current one can't serve the request — out of credit, revoked key, rate-limited
|
||||
// past the SDK's own retries, or down.
|
||||
//
|
||||
// The chain is built by the server from llm_profiles (active profile first, then
|
||||
// by priority), so the whole process shares ONE chain and ONE circuit-breaker
|
||||
// registry: when a task discovers a profile is out of credit, every other task
|
||||
// skips it immediately.
|
||||
package llmpool
|
||||
|
||||
import (
|
||||
"sync"
|
||||
"time"
|
||||
)
|
||||
|
||||
// backoff is the cooling-off ladder, indexed by trip count: the 1st trip cools
|
||||
// for a minute, the 2nd for five, everything after that for half an hour. A
|
||||
// profile that is merely rate-limited recovers quickly; one that keeps failing
|
||||
// stops being probed every round.
|
||||
var backoff = []time.Duration{time.Minute, 5 * time.Minute, 30 * time.Minute}
|
||||
|
||||
// State is one profile's circuit-breaker state.
|
||||
type State struct {
|
||||
Fails int // consecutive failures; reset by a success
|
||||
Trips int // total trips, indexes the backoff ladder
|
||||
OpenUntil time.Time // zero / past = closed (usable)
|
||||
LastError string
|
||||
LastAt time.Time
|
||||
}
|
||||
|
||||
// Open reports whether the breaker is currently open (profile should be skipped).
|
||||
func (s State) Open() bool { return time.Now().Before(s.OpenUntil) }
|
||||
|
||||
// softTripAfter is how many consecutive TRANSIENT failures (429 / 5xx / network)
|
||||
// trip the breaker. Deterministic failures (no credit, bad key) trip on the first.
|
||||
const softTripAfter = 3
|
||||
|
||||
// Registry holds the circuit-breaker state of every profile, keyed by profile id.
|
||||
// It is process-wide and outlives any single Pool instance, so rebuilding the pool
|
||||
// (saving an unrelated profile, toggling a setting) never clears what we've learned
|
||||
// about which backends are broken.
|
||||
type Registry struct {
|
||||
mu sync.Mutex
|
||||
m map[int64]*State
|
||||
|
||||
// persist / forget mirror state to the DB so a cooling-off window survives a
|
||||
// restart. Both may be nil (no DB); both are called off the hot path.
|
||||
persist func(id int64, st State)
|
||||
forget func(id int64)
|
||||
}
|
||||
|
||||
// NewRegistry builds an empty registry. persist/forget may be nil.
|
||||
func NewRegistry(persist func(id int64, st State), forget func(id int64)) *Registry {
|
||||
return &Registry{m: map[int64]*State{}, persist: persist, forget: forget}
|
||||
}
|
||||
|
||||
// Restore seeds state loaded from the DB at startup (bypasses persistence).
|
||||
func (r *Registry) Restore(id int64, st State) {
|
||||
r.mu.Lock()
|
||||
defer r.mu.Unlock()
|
||||
cp := st
|
||||
r.m[id] = &cp
|
||||
}
|
||||
|
||||
// Get returns a copy of one profile's state.
|
||||
func (r *Registry) Get(id int64) State {
|
||||
r.mu.Lock()
|
||||
defer r.mu.Unlock()
|
||||
if st := r.m[id]; st != nil {
|
||||
return *st
|
||||
}
|
||||
return State{}
|
||||
}
|
||||
|
||||
// IsOpen reports whether a profile is in its cooling-off window.
|
||||
func (r *Registry) IsOpen(id int64) bool { return r.Get(id).Open() }
|
||||
|
||||
// Pass records a successful call: clears the failure counters so a profile that
|
||||
// recovers is fully trusted again (and drops the persisted row).
|
||||
func (r *Registry) Pass(id int64) {
|
||||
r.mu.Lock()
|
||||
st := r.m[id]
|
||||
if st == nil || (st.Fails == 0 && st.Trips == 0 && st.OpenUntil.IsZero()) {
|
||||
r.mu.Unlock()
|
||||
return // already clean — nothing to write
|
||||
}
|
||||
delete(r.m, id)
|
||||
forget := r.forget
|
||||
r.mu.Unlock()
|
||||
if forget != nil {
|
||||
forget(id)
|
||||
}
|
||||
}
|
||||
|
||||
// Trip records a failed call. hard=true marks a deterministic failure (no credit,
|
||||
// invalid key, missing model) which opens the breaker immediately; hard=false is a
|
||||
// transient one (429 / 5xx / network) that needs softTripAfter in a row. Returns
|
||||
// true when this call is what opened the breaker, so the caller can log it once.
|
||||
func (r *Registry) Trip(id int64, errMsg string, hard bool) (tripped bool) {
|
||||
r.mu.Lock()
|
||||
st := r.m[id]
|
||||
if st == nil {
|
||||
st = &State{}
|
||||
r.m[id] = st
|
||||
}
|
||||
st.Fails++
|
||||
st.LastError = errMsg
|
||||
st.LastAt = time.Now()
|
||||
if hard || st.Fails >= softTripAfter {
|
||||
st.Trips++
|
||||
d := backoff[len(backoff)-1]
|
||||
if st.Trips-1 < len(backoff) {
|
||||
d = backoff[st.Trips-1]
|
||||
}
|
||||
st.OpenUntil = time.Now().Add(d)
|
||||
st.Fails = 0 // counted into this trip; start fresh for the half-open probe
|
||||
tripped = true
|
||||
}
|
||||
snap := *st
|
||||
persist := r.persist
|
||||
r.mu.Unlock()
|
||||
if persist != nil {
|
||||
persist(id, snap)
|
||||
}
|
||||
return tripped
|
||||
}
|
||||
|
||||
// Reset clears one profile's state — the UI's "立即恢复" action.
|
||||
func (r *Registry) Reset(id int64) {
|
||||
r.mu.Lock()
|
||||
delete(r.m, id)
|
||||
forget := r.forget
|
||||
r.mu.Unlock()
|
||||
if forget != nil {
|
||||
forget(id)
|
||||
}
|
||||
}
|
||||
|
||||
// Snapshot returns a copy of all tracked state, for the status API.
|
||||
func (r *Registry) Snapshot() map[int64]State {
|
||||
r.mu.Lock()
|
||||
defer r.mu.Unlock()
|
||||
out := make(map[int64]State, len(r.m))
|
||||
for id, st := range r.m {
|
||||
out[id] = *st
|
||||
}
|
||||
return out
|
||||
}
|
||||
+261
@@ -0,0 +1,261 @@
|
||||
package llmpool
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"fmt"
|
||||
"iter"
|
||||
"log"
|
||||
"regexp"
|
||||
"strconv"
|
||||
"strings"
|
||||
"sync/atomic"
|
||||
|
||||
"github.com/Autumn-27/norma/llm"
|
||||
)
|
||||
|
||||
// Member is one LLM profile in the chain, already built into a provider (wrapped
|
||||
// with the recorder by the caller, so a failed attempt is still recorded under
|
||||
// its own profile name).
|
||||
type Member struct {
|
||||
ID int64 // llm_profiles.id
|
||||
Name string // profile name, for logs / UI
|
||||
Model string
|
||||
Format string // "anthropic" | "openai"
|
||||
Priority int // the profile's configured priority (display only)
|
||||
Active bool // is_default (display only)
|
||||
// Rank is the ordering key the caller assigned: higher goes first, and members
|
||||
// sharing a Rank take turns leading (load-spreading across duplicate keys).
|
||||
// The caller encodes "active profile heads the chain" as a Rank above every
|
||||
// user-settable priority, so this type needs no policy of its own.
|
||||
Rank int
|
||||
// WindowTokens is the profile's context window in tokens. A member whose
|
||||
// window can't hold the request is skipped rather than made to fail on it.
|
||||
WindowTokens int
|
||||
Prov llm.Provider
|
||||
}
|
||||
|
||||
// RankActive is the Rank the caller gives the chain head (the active profile, or
|
||||
// an explicitly bound one) so it always outranks any configured priority.
|
||||
const RankActive = int(^uint(0)>>1) - 1
|
||||
|
||||
// ErrExhausted is returned when every member of the chain failed.
|
||||
var ErrExhausted = errors.New("LLM 轮询：所有配置均不可用")
|
||||
|
||||
// Pool is an llm.Provider that fails over across an ordered chain of members.
|
||||
// It is safe for concurrent use: members are immutable after construction and
|
||||
// all mutable state lives in the shared Registry.
|
||||
type Pool struct {
|
||||
members []*Member // in chain order (active first, then priority DESC)
|
||||
health *Registry
|
||||
rr atomic.Uint64 // rotates the starting point within an equal-priority group
|
||||
}
|
||||
|
||||
// New builds a Pool over members (already in chain order). Returns nil when the
|
||||
// chain is empty. A single-member chain is still a valid Pool — it just behaves
|
||||
// exactly like the bare provider.
|
||||
func New(members []*Member, health *Registry) *Pool {
|
||||
if len(members) == 0 {
|
||||
return nil
|
||||
}
|
||||
if health == nil {
|
||||
health = NewRegistry(nil, nil)
|
||||
}
|
||||
return &Pool{members: members, health: health}
|
||||
}
|
||||
|
||||
// Members returns the chain in order (read-only).
|
||||
func (p *Pool) Members() []*Member { return p.members }
|
||||
|
||||
// Head returns the first member of the chain.
|
||||
func (p *Pool) Head() *Member { return p.members[0] }
|
||||
|
||||
// Stream implements llm.Provider with failover.
|
||||
//
|
||||
// The one hard rule: a member may only be abandoned BEFORE it has yielded any
|
||||
// event. Once text or a tool_use has reached the caller, re-sending the same
|
||||
// request to another model would duplicate output and corrupt the conversation
|
||||
// history — so a mid-stream failure is surfaced as-is and left to the agent
|
||||
// harness's resume logic. Fortunately the failures this exists for (402 no
|
||||
// credit, 401 bad key, 429, 5xx) all surface during request establishment,
|
||||
// before the body is read, so they always land in the safe window.
|
||||
func (p *Pool) Stream(ctx context.Context, req llm.CompletionRequest) iter.Seq2[llm.StreamEvent, error] {
|
||||
order := p.order(req)
|
||||
return func(yield func(llm.StreamEvent, error) bool) {
|
||||
var lastErr error
|
||||
for i, m := range order {
|
||||
emitted := false
|
||||
var failed error
|
||||
for ev, err := range m.Prov.Stream(ctx, req) {
|
||||
if err != nil && !emitted && shouldFailover(ctx, err) {
|
||||
failed = err
|
||||
break // safe window: nothing reached the caller yet
|
||||
}
|
||||
emitted = true
|
||||
if !yield(ev, err) {
|
||||
return // caller stopped consuming (cancel / early exit)
|
||||
}
|
||||
if err != nil {
|
||||
return // terminal error already handed to the caller
|
||||
}
|
||||
}
|
||||
if failed == nil {
|
||||
p.health.Pass(m.ID) // completed (or failed in a non-failover way)
|
||||
return
|
||||
}
|
||||
lastErr = failed
|
||||
hard := isHardFailure(failed)
|
||||
if p.health.Trip(m.ID, trimErr(failed), hard) {
|
||||
log.Printf("[llmpool] 配置 %q(%s) 已熔断：%s", m.Name, m.Model, trimErr(failed))
|
||||
}
|
||||
if i+1 < len(order) {
|
||||
n := order[i+1]
|
||||
log.Printf("[llmpool] LLM 故障转移：%q(%s) → %q(%s)，原因：%s",
|
||||
m.Name, m.Model, n.Name, n.Model, trimErr(failed))
|
||||
}
|
||||
}
|
||||
if lastErr == nil {
|
||||
lastErr = ErrExhausted
|
||||
}
|
||||
log.Printf("[llmpool] 轮询链已耗尽(%d 个配置全部失败)，最后错误：%s", len(order), trimErr(lastErr))
|
||||
yield(llm.StreamEvent{}, fmt.Errorf("%w：%v", ErrExhausted, lastErr))
|
||||
}
|
||||
}
|
||||
|
||||
// order picks the members to try, in order: skip those in a cooling-off window
|
||||
// and those whose context window can't hold this request, then rotate within each
|
||||
// equal-priority group so same-priority profiles share the load. Never returns an
|
||||
// empty slice — if everything is filtered out, the head of the chain is tried
|
||||
// anyway, since stalling the engine is worse than one more failed request.
|
||||
func (p *Pool) order(req llm.CompletionRequest) []*Member {
|
||||
est := estimateTokens(req)
|
||||
var open []*Member
|
||||
for _, m := range p.members {
|
||||
if p.health.IsOpen(m.ID) {
|
||||
continue
|
||||
}
|
||||
if m.WindowTokens > 0 && est > m.WindowTokens {
|
||||
continue // would 400 on length — not a useful failover target
|
||||
}
|
||||
open = append(open, m)
|
||||
}
|
||||
if len(open) == 0 {
|
||||
return p.members[:1] // last resort: probe the head rather than stall
|
||||
}
|
||||
return rotateGroups(open, p.rr.Add(1)-1)
|
||||
}
|
||||
|
||||
// rotateGroups rotates each run of equal-ranked members by n, so profiles sharing
|
||||
// a priority take turns going first (free load-spreading across duplicate keys).
|
||||
// The chain head has its own Rank and always stays at the front.
|
||||
func rotateGroups(in []*Member, n uint64) []*Member {
|
||||
out := make([]*Member, 0, len(in))
|
||||
for i := 0; i < len(in); {
|
||||
j := i + 1
|
||||
for j < len(in) && in[j].Rank == in[i].Rank {
|
||||
j++
|
||||
}
|
||||
g := in[i:j]
|
||||
if len(g) > 1 {
|
||||
off := int(n % uint64(len(g)))
|
||||
for k := range g {
|
||||
out = append(out, g[(k+off)%len(g)])
|
||||
}
|
||||
} else {
|
||||
out = append(out, g...)
|
||||
}
|
||||
i = j
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
// statusRe pulls the HTTP status out of the SDK's error text, which is formatted
|
||||
// as "<prefix>: status <code>: <body>" (norma/llm/retry.go). The SDK exposes no
|
||||
// typed error, so the string is what we have.
|
||||
var statusRe = regexp.MustCompile(`status (\d{3})`)
|
||||
|
||||
// statusOf returns the HTTP status carried by err, or 0 if it isn't one.
|
||||
func statusOf(err error) int {
|
||||
m := statusRe.FindStringSubmatch(err.Error())
|
||||
if m == nil {
|
||||
return 0
|
||||
}
|
||||
code, _ := strconv.Atoi(m[1])
|
||||
return code
|
||||
}
|
||||
|
||||
// shouldFailover reports whether err justifies trying the next profile.
|
||||
//
|
||||
// Never fails over on context cancellation — that's the user stopping a task or a
|
||||
// task-level timeout, and burning a backup key on it would both waste credit and
|
||||
// pollute the run's termination diagnosis. Never on 400 either: a malformed or
|
||||
// over-long request fails identically everywhere.
|
||||
func shouldFailover(ctx context.Context, err error) bool {
|
||||
if err == nil {
|
||||
return false
|
||||
}
|
||||
if ctx.Err() != nil || errors.Is(err, context.Canceled) || errors.Is(err, context.DeadlineExceeded) {
|
||||
return false
|
||||
}
|
||||
switch code := statusOf(err); {
|
||||
case code == 0:
|
||||
return true // no status → transport-level failure (reset / DNS / timeout)
|
||||
case code == 400:
|
||||
return false // bad or over-long request: identical everywhere
|
||||
case code == 401, code == 402, code == 403, code == 404, code == 408, code == 429:
|
||||
return true
|
||||
case code >= 500:
|
||||
return true
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
// isHardFailure reports whether the failure is deterministic (the profile will
|
||||
// keep failing until a human fixes it) rather than transient. Hard failures open
|
||||
// the breaker on the first occurrence.
|
||||
func isHardFailure(err error) bool {
|
||||
switch statusOf(err) {
|
||||
case 401, 402, 403, 404:
|
||||
return true
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
// trimErr shortens an error for logs/UI — provider bodies can be long.
|
||||
func trimErr(err error) string {
|
||||
s := strings.TrimSpace(strings.ReplaceAll(err.Error(), "\n", " "))
|
||||
if len(s) > 300 {
|
||||
s = s[:300] + "…"
|
||||
}
|
||||
return s
|
||||
}
|
||||
|
||||
// estimateTokens roughly sizes a request so members whose context window clearly
|
||||
// can't hold it are skipped. Deliberately crude (~3.5 chars/token) and biased to
|
||||
// over-estimate slightly; it only needs to tell "fits" from "nowhere near".
|
||||
func estimateTokens(req llm.CompletionRequest) int {
|
||||
n := 0
|
||||
for _, s := range req.System {
|
||||
n += len(s)
|
||||
}
|
||||
for _, m := range req.Messages {
|
||||
n += blocksLen(m.Content)
|
||||
}
|
||||
for _, t := range req.Tools {
|
||||
// InputSchema is a decoded map, so its serialized size isn't available
|
||||
// cheaply — charge a flat ~120 chars per top-level property instead.
|
||||
n += len(t.Name) + len(t.Description) + len(t.InputSchema)*120
|
||||
}
|
||||
return n * 2 / 7 // ≈ len/3.5
|
||||
}
|
||||
|
||||
func blocksLen(bs []llm.ContentBlock) int {
|
||||
n := 0
|
||||
for _, b := range bs {
|
||||
n += len(b.Text) + len(b.Thinking) + len(b.Input) + len(b.Name)
|
||||
if len(b.Content) > 0 {
|
||||
n += blocksLen(b.Content)
|
||||
}
|
||||
}
|
||||
return n
|
||||
}
|
||||
@@ -0,0 +1,359 @@
|
||||
package llmpool
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"fmt"
|
||||
"iter"
|
||||
"strings"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/Autumn-27/norma/llm"
|
||||
)
|
||||
|
||||
// fakeProv is a scripted provider: script[i] is what the i-th call yields —
|
||||
// some events, then optionally an error.
|
||||
type fakeProv struct {
|
||||
name string
|
||||
calls int
|
||||
events [][]llm.StreamEvent // events emitted before the error, per call
|
||||
errs []error // error to end each call with (nil = clean finish)
|
||||
}
|
||||
|
||||
// at returns script entry n, repeating the last one once the script runs out, so
|
||||
// a provider defined as "always succeeds" / "always 402" keeps behaving that way
|
||||
// across repeated calls.
|
||||
func at[T any](s []T, n int) (T, bool) {
|
||||
var zero T
|
||||
if len(s) == 0 {
|
||||
return zero, false
|
||||
}
|
||||
if n >= len(s) {
|
||||
n = len(s) - 1
|
||||
}
|
||||
return s[n], true
|
||||
}
|
||||
|
||||
func (f *fakeProv) Stream(ctx context.Context, req llm.CompletionRequest) iter.Seq2[llm.StreamEvent, error] {
|
||||
n := f.calls
|
||||
f.calls++
|
||||
return func(yield func(llm.StreamEvent, error) bool) {
|
||||
if evs, ok := at(f.events, n); ok {
|
||||
for _, e := range evs {
|
||||
if !yield(e, nil) {
|
||||
return
|
||||
}
|
||||
}
|
||||
}
|
||||
err, _ := at(f.errs, n)
|
||||
if err != nil {
|
||||
yield(llm.StreamEvent{}, err)
|
||||
return
|
||||
}
|
||||
yield(llm.StreamEvent{Type: llm.SEMessageStop}, nil)
|
||||
}
|
||||
}
|
||||
|
||||
// ok builds a provider that always succeeds with one text delta.
|
||||
func okProv(name, text string) *fakeProv {
|
||||
return &fakeProv{name: name, events: [][]llm.StreamEvent{{{Type: llm.SETextDelta, Text: text}}}}
|
||||
}
|
||||
|
||||
// failProv builds a provider that fails immediately (before any event) with status.
|
||||
func failProv(name string, status int) *fakeProv {
|
||||
return &fakeProv{name: name, errs: []error{fmt.Errorf("anthropic: status %d: nope", status)}}
|
||||
}
|
||||
|
||||
func member(id int64, name string, rank int, p llm.Provider) *Member {
|
||||
return &Member{ID: id, Name: name, Model: "m" + name, Rank: rank, Prov: p}
|
||||
}
|
||||
|
||||
// drain consumes a stream, returning the concatenated text and terminal error.
|
||||
func drain(seq iter.Seq2[llm.StreamEvent, error]) (string, error) {
|
||||
var sb strings.Builder
|
||||
for ev, err := range seq {
|
||||
if err != nil {
|
||||
return sb.String(), err
|
||||
}
|
||||
if ev.Type == llm.SETextDelta {
|
||||
sb.WriteString(ev.Text)
|
||||
}
|
||||
}
|
||||
return sb.String(), nil
|
||||
}
|
||||
|
||||
func TestFailoverOnNoCredit(t *testing.T) {
|
||||
a, b := failProv("a", 402), okProv("b", "hello")
|
||||
p := New([]*Member{member(1, "a", 10, a), member(2, "b", 5, b)}, NewRegistry(nil, nil))
|
||||
|
||||
got, err := drain(p.Stream(context.Background(), llm.CompletionRequest{}))
|
||||
if err != nil {
|
||||
t.Fatalf("expected failover to succeed, got %v", err)
|
||||
}
|
||||
if got != "hello" {
|
||||
t.Fatalf("text = %q, want %q", got, "hello")
|
||||
}
|
||||
if a.calls != 1 || b.calls != 1 {
|
||||
t.Fatalf("calls: a=%d b=%d, want 1/1", a.calls, b.calls)
|
||||
}
|
||||
}
|
||||
|
||||
// A 402 is deterministic: one failure must open the breaker, so the NEXT request
|
||||
// skips that profile entirely instead of paying for another round-trip.
|
||||
func TestHardFailureTripsBreakerImmediately(t *testing.T) {
|
||||
a, b := failProv("a", 402), okProv("b", "x")
|
||||
reg := NewRegistry(nil, nil)
|
||||
p := New([]*Member{member(1, "a", 10, a), member(2, "b", 5, b)}, reg)
|
||||
|
||||
_, _ = drain(p.Stream(context.Background(), llm.CompletionRequest{}))
|
||||
if !reg.IsOpen(1) {
|
||||
t.Fatal("402 should have tripped the breaker on the first failure")
|
||||
}
|
||||
_, _ = drain(p.Stream(context.Background(), llm.CompletionRequest{}))
|
||||
if a.calls != 1 {
|
||||
t.Fatalf("tripped profile was called again: calls=%d, want 1", a.calls)
|
||||
}
|
||||
if b.calls != 2 {
|
||||
t.Fatalf("fallback calls=%d, want 2", b.calls)
|
||||
}
|
||||
}
|
||||
|
||||
// 429 is transient: the SDK already retried, but we shouldn't write a profile off
|
||||
// until it fails repeatedly.
|
||||
func TestSoftFailureNeedsRepeats(t *testing.T) {
|
||||
reg := NewRegistry(nil, nil)
|
||||
for i := 1; i < softTripAfter; i++ {
|
||||
if reg.Trip(1, "429", false) {
|
||||
t.Fatalf("tripped after %d transient failures, want %d", i, softTripAfter)
|
||||
}
|
||||
}
|
||||
if !reg.Trip(1, "429", false) {
|
||||
t.Fatalf("should trip on failure #%d", softTripAfter)
|
||||
}
|
||||
if !reg.IsOpen(1) {
|
||||
t.Fatal("breaker should be open")
|
||||
}
|
||||
}
|
||||
|
||||
// A success must fully clear the counters, so an intermittent profile never
|
||||
// accumulates its way to a trip.
|
||||
func TestPassResetsCounters(t *testing.T) {
|
||||
reg := NewRegistry(nil, nil)
|
||||
reg.Trip(1, "429", false)
|
||||
reg.Trip(1, "429", false)
|
||||
reg.Pass(1)
|
||||
if got := reg.Get(1).Fails; got != 0 {
|
||||
t.Fatalf("fails=%d after Pass, want 0", got)
|
||||
}
|
||||
if reg.Trip(1, "429", false) {
|
||||
t.Fatal("tripped immediately after a success — counters were not reset")
|
||||
}
|
||||
}
|
||||
|
||||
func TestBackoffLadderGrows(t *testing.T) {
|
||||
reg := NewRegistry(nil, nil)
|
||||
var prev time.Duration
|
||||
for i := range 4 {
|
||||
reg.Reset(1)
|
||||
st := State{Trips: i}
|
||||
reg.Restore(1, st)
|
||||
reg.Trip(1, "402", true)
|
||||
d := time.Until(reg.Get(1).OpenUntil)
|
||||
// Non-decreasing, with a second of slack: the ladder plateaus at its last
|
||||
// rung, and each Trip stamps its own time.Now().
|
||||
if i > 0 && d < prev-time.Second {
|
||||
t.Fatalf("trip #%d cools for %v, shorter than the previous %v", i+1, d, prev)
|
||||
}
|
||||
prev = d
|
||||
}
|
||||
if prev < 25*time.Minute {
|
||||
t.Fatalf("ladder tops out at %v, want ~30m", prev)
|
||||
}
|
||||
}
|
||||
|
||||
// The safety rule: once output has reached the caller, a mid-stream failure must
|
||||
// NOT be retried on another model — that would duplicate the assistant turn.
|
||||
func TestNoFailoverAfterEmit(t *testing.T) {
|
||||
a := &fakeProv{
|
||||
name: "a",
|
||||
events: [][]llm.StreamEvent{{{Type: llm.SETextDelta, Text: "partial"}}},
|
||||
errs: []error{fmt.Errorf("anthropic: status 500: mid-stream drop")},
|
||||
}
|
||||
b := okProv("b", "full")
|
||||
p := New([]*Member{member(1, "a", 10, a), member(2, "b", 5, b)}, NewRegistry(nil, nil))
|
||||
|
||||
got, err := drain(p.Stream(context.Background(), llm.CompletionRequest{}))
|
||||
if err == nil {
|
||||
t.Fatal("mid-stream error should surface, not be swallowed by a failover")
|
||||
}
|
||||
if got != "partial" {
|
||||
t.Fatalf("text = %q, want the partial output %q", got, "partial")
|
||||
}
|
||||
if b.calls != 0 {
|
||||
t.Fatalf("fell over to the backup after emitting output (calls=%d) — would duplicate the turn", b.calls)
|
||||
}
|
||||
}
|
||||
|
||||
// Cancelling a task must not burn a backup key, and must not be diagnosed as an
|
||||
// LLM fault.
|
||||
func TestNoFailoverOnCancel(t *testing.T) {
|
||||
ctx, cancel := context.WithCancel(context.Background())
|
||||
cancel()
|
||||
a := &fakeProv{name: "a", errs: []error{context.Canceled}}
|
||||
b := okProv("b", "x")
|
||||
reg := NewRegistry(nil, nil)
|
||||
p := New([]*Member{member(1, "a", 10, a), member(2, "b", 5, b)}, reg)
|
||||
|
||||
_, err := drain(p.Stream(ctx, llm.CompletionRequest{}))
|
||||
if !errors.Is(err, context.Canceled) {
|
||||
t.Fatalf("err = %v, want context.Canceled", err)
|
||||
}
|
||||
if b.calls != 0 {
|
||||
t.Fatalf("failed over on cancellation (backup calls=%d)", b.calls)
|
||||
}
|
||||
if reg.Get(1).Fails != 0 {
|
||||
t.Fatal("cancellation counted as a profile failure")
|
||||
}
|
||||
}
|
||||
|
||||
func TestShouldFailoverByStatus(t *testing.T) {
|
||||
cases := []struct {
|
||||
status int
|
||||
want bool
|
||||
}{
|
||||
{400, false}, // bad/over-long request fails identically everywhere
|
||||
{401, true}, {402, true}, {403, true}, {404, true},
|
||||
{408, true}, {429, true}, {500, true}, {503, true},
|
||||
{200, false},
|
||||
}
|
||||
for _, c := range cases {
|
||||
err := fmt.Errorf("openai: status %d: body", c.status)
|
||||
if got := shouldFailover(context.Background(), err); got != c.want {
|
||||
t.Errorf("status %d: shouldFailover=%v, want %v", c.status, got, c.want)
|
||||
}
|
||||
}
|
||||
// A transport error carries no status and must fail over.
|
||||
if !shouldFailover(context.Background(), errors.New("dial tcp: connection reset by peer")) {
|
||||
t.Error("network error should fail over")
|
||||
}
|
||||
}
|
||||
|
||||
func TestHardVsSoftClassification(t *testing.T) {
|
||||
for _, s := range []int{401, 402, 403, 404} {
|
||||
if !isHardFailure(fmt.Errorf("status %d: x", s)) {
|
||||
t.Errorf("status %d should be a hard failure", s)
|
||||
}
|
||||
}
|
||||
for _, s := range []int{408, 429, 500, 502, 503} {
|
||||
if isHardFailure(fmt.Errorf("status %d: x", s)) {
|
||||
t.Errorf("status %d should be transient, not hard", s)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// A member whose context window can't hold the request is a guaranteed 400 —
|
||||
// skip it rather than spend a round-trip proving it.
|
||||
func TestSkipsMembersTooSmallForRequest(t *testing.T) {
|
||||
small, big := okProv("small", "s"), okProv("big", "b")
|
||||
ms := member(1, "small", 10, small)
|
||||
ms.WindowTokens = 1000
|
||||
mb := member(2, "big", 5, big)
|
||||
mb.WindowTokens = 1_000_000
|
||||
p := New([]*Member{ms, mb}, NewRegistry(nil, nil))
|
||||
|
||||
req := llm.CompletionRequest{Messages: []llm.Message{llm.UserText(strings.Repeat("x", 100_000))}}
|
||||
got, err := drain(p.Stream(context.Background(), req))
|
||||
if err != nil {
|
||||
t.Fatalf("unexpected error: %v", err)
|
||||
}
|
||||
if got != "b" {
|
||||
t.Fatalf("text=%q, want the large-window member to serve it", got)
|
||||
}
|
||||
if small.calls != 0 {
|
||||
t.Fatalf("sent a request to a member that can't hold it (calls=%d)", small.calls)
|
||||
}
|
||||
}
|
||||
|
||||
// Everything tripped: probing the head beats stalling the engine outright.
|
||||
func TestAllTrippedStillProbesHead(t *testing.T) {
|
||||
a, b := okProv("a", "a"), okProv("b", "b")
|
||||
reg := NewRegistry(nil, nil)
|
||||
reg.Trip(1, "402", true)
|
||||
reg.Trip(2, "402", true)
|
||||
p := New([]*Member{member(1, "a", 10, a), member(2, "b", 5, b)}, reg)
|
||||
|
||||
got, err := drain(p.Stream(context.Background(), llm.CompletionRequest{}))
|
||||
if err != nil {
|
||||
t.Fatalf("unexpected error: %v", err)
|
||||
}
|
||||
if got != "a" || a.calls != 1 {
|
||||
t.Fatalf("text=%q a.calls=%d — want the head probed as a last resort", got, a.calls)
|
||||
}
|
||||
}
|
||||
|
||||
func TestExhaustedChainReportsClearly(t *testing.T) {
|
||||
a, b := failProv("a", 402), failProv("b", 402)
|
||||
p := New([]*Member{member(1, "a", 10, a), member(2, "b", 5, b)}, NewRegistry(nil, nil))
|
||||
|
||||
_, err := drain(p.Stream(context.Background(), llm.CompletionRequest{}))
|
||||
if !errors.Is(err, ErrExhausted) {
|
||||
t.Fatalf("err = %v, want it to wrap ErrExhausted", err)
|
||||
}
|
||||
}
|
||||
|
||||
// Equal-rank members take turns leading, so duplicate keys share the load.
|
||||
func TestEqualRankRotates(t *testing.T) {
|
||||
a, b := okProv("a", "a"), okProv("b", "b")
|
||||
head := member(1, "head", RankActive, failProv("head", 402))
|
||||
p := New([]*Member{head, member(2, "a", 5, a), member(3, "b", 5, b)}, NewRegistry(nil, nil))
|
||||
|
||||
first := map[string]int{}
|
||||
for range 4 {
|
||||
txt, _ := drain(p.Stream(context.Background(), llm.CompletionRequest{}))
|
||||
first[txt]++
|
||||
}
|
||||
if first["a"] == 0 || first["b"] == 0 {
|
||||
t.Fatalf("same-rank members did not rotate: %v", first)
|
||||
}
|
||||
}
|
||||
|
||||
// The head keeps its own rank and never joins a rotation group.
|
||||
func TestHeadAlwaysFirst(t *testing.T) {
|
||||
head := okProv("head", "H")
|
||||
p := New([]*Member{
|
||||
member(1, "head", RankActive, head),
|
||||
member(2, "a", 5, okProv("a", "a")),
|
||||
member(3, "b", 5, okProv("b", "b")),
|
||||
}, NewRegistry(nil, nil))
|
||||
for range 5 {
|
||||
if txt, _ := drain(p.Stream(context.Background(), llm.CompletionRequest{})); txt != "H" {
|
||||
t.Fatalf("head was not tried first, got %q", txt)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// A single-member chain must behave exactly like the bare provider.
|
||||
func TestSingleMemberPassthrough(t *testing.T) {
|
||||
a := okProv("a", "solo")
|
||||
p := New([]*Member{member(1, "a", RankActive, a)}, NewRegistry(nil, nil))
|
||||
got, err := drain(p.Stream(context.Background(), llm.CompletionRequest{}))
|
||||
if err != nil || got != "solo" {
|
||||
t.Fatalf("got %q / %v, want solo / nil", got, err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestNewEmptyChainIsNil(t *testing.T) {
|
||||
if New(nil, nil) != nil {
|
||||
t.Fatal("empty chain should yield a nil pool so callers use the bare provider")
|
||||
}
|
||||
}
|
||||
|
||||
// An expired cooling-off window must not keep a profile out of the chain.
|
||||
func TestExpiredWindowIsClosed(t *testing.T) {
|
||||
reg := NewRegistry(nil, nil)
|
||||
reg.Restore(1, State{Trips: 1, OpenUntil: time.Now().Add(-time.Second)})
|
||||
if reg.IsOpen(1) {
|
||||
t.Fatal("expired window should read as closed")
|
||||
}
|
||||
}
|
||||
Reference in new issue
Block a user