feat(llm): add profile failover and pool health

This commit is contained in:
pphua committed 2026-08-24 19:28:18 +08:00
1 parent d4b21f0bdb
commit 18926c129f
15 files changed
+2507 -761

No files matched your search

+27 -22
View File
@@ -8,8 +8,14 @@
package agent
import (
"bytes"
"context"
"fmt"
"io"
"net/http"
"net/url"
"os"
"regexp"
"strings"
"time"
@@ -36,9 +42,14 @@ type Config struct {
// ContextWindowK is the model's context window in K tokens (user-configured),
// used to size compaction thresholds. 0 = default; see CompactionWindow.
ContextWindowK int
// ReasoningEffort selects the thinking mode: "" = 默认(不发送思考参数);
// "off" = 关闭(thinking.type=disabled); "low"/"medium"/"high"/"max" = 开启并设强度.
// NewProvider derives the provider-specific request fields from it.
// ThinkingType 独立控制思考「开关」字段(thinking.type):
// "" = 不发送(默认,兼容不支持该字段的模型); "disabled" = 显式关闭;
// "enabled" = 开启. 与 ReasoningEffort 完全解耦——有些接口没有 thinking 字段、
// 只靠强度参数就能激活思考,故两者可各自单独设置.
ThinkingType string
// ReasoningEffort 独立控制思考「强度」字段:
// "" = 不发送(默认); "low"/"medium"/"high"/"xhigh"/"max" = 对应强度.
// OpenAI 映射为顶层 reasoning_effort;Anthropic 映射为 output_config.effort.
ReasoningEffort string
}
@@ -168,32 +179,27 @@ func (c Config) Provider() string {
// limiter lives on the single provider instance — so planner + all workers +
// main agent (which share this provider) are bounded by one shared rate limit.
func (c Config) NewProvider() (llm.Provider, error) {
client, err := quotaAwareHTTPClient(c.Proxy)
if err != nil {
return nil, err
}
lc := llm.Config{
Format: c.Format,
BaseURL: c.BaseURL,
APIKey: c.APIKey,
Model: c.Model,
Proxy: c.Proxy,
}
// Derive the provider thinking params from the single UI selector:
// "" → omit both (provider default); "off" → disable; else → enable + effort.
switch c.ReasoningEffort {
case "":
// leave both empty → omitted from the request
case "off":
lc.ThinkingType = "disabled"
default:
lc.ThinkingType = "enabled"
lc.ReasoningEffort = c.ReasoningEffort
Format: c.Format,
BaseURL: c.BaseURL,
APIKey: c.APIKey,
Model: c.Model,
HTTPClient: client,
}
// 思考开关与强度两个字段各自透传(空 = 该字段不发送)。二者解耦:
// 可只发 thinking.type、只发 effort、都发、或都不发。
lc.ThinkingType = c.ThinkingType
lc.ReasoningEffort = c.ReasoningEffort
if c.RatePerSecond > 0 || c.RatePerMinute > 0 {
lc.RateLimit = &llm.RateLimit{PerSecond: c.RatePerSecond, PerMinute: c.RatePerMinute}
}
return llm.NewProvider(lc)
}
<<<<<<< Updated upstream
=======
// IsQuotaExhaustedMessage deliberately recognizes only explicit balance,
// billing, credit, or quota-exhaustion signals. Generic 429/rate-limit text,
// authentication failures, network errors, and server failures are excluded.
@@ -281,7 +287,6 @@ func quotaAwareHTTPClient(proxy string) (*http.Client, error) {
return &http.Client{Transport: quotaAwareTransport{base: transport}}, nil
}
>>>>>>> Stashed changes
// TestConnection makes a minimal real completion to verify the provider/model/
// endpoint/key actually work. Returns the round-trip latency.
func TestConnection(ctx context.Context, c Config) (time.Duration, error) {
+131
View File
@@ -0,0 +1,131 @@
package agent
import (
"context"
"io"
"net/http"
"net/http/httptest"
"strings"
"sync/atomic"
"testing"
"github.com/Autumn-27/norma/llm"
)
type roundTripperFunc func(*http.Request) (*http.Response, error)
func (f roundTripperFunc) RoundTrip(req *http.Request) (*http.Response, error) { return f(req) }
func TestIsQuotaExhaustedMessage(t *testing.T) {
t.Parallel()
positive := []string{
`{"error":{"code":"insufficient_quota"}}`,
`RESOURCE_EXHAUSTED`,
`You exceeded your current quota, please check your plan and billing details.`,
`billing_not_active`,
`credit balance is too low`,
`账户余额不足,请充值`,
}
for _, message := range positive {
if !IsQuotaExhaustedMessage(message) {
t.Errorf("expected quota classification for %q", message)
}
}
negative := []string{
`status 429: rate limit exceeded`,
`RESOURCE_EXHAUSTED: rate limit exceeded`,
`too many requests per minute`,
`quota exceeded for quota metric GenerateRequestsPerMinutePerProjectPerBaseModel`,
`RESOURCE_EXHAUSTED: TPM quota exceeded`,
`tokens per minute quota exceeded`,
`rate_limit_exceeded: requests per second`,
`status 401: invalid api key`,
`status 401: insufficient_quota`,
`HTTP 403: billing_hard_limit_reached`,
`status 500: internal server error`,
`status=503: insufficient_quota`,
`context length exceeded`,
}
for _, message := range negative {
if IsQuotaExhaustedMessage(message) {
t.Errorf("unexpected quota classification for %q", message)
}
}
}
func TestQuotaAwareTransportOnlyNormalizesExplicitQuota429(t *testing.T) {
t.Parallel()
tests := []struct {
name string
status int
body string
wantStatus int
}{
{name: "quota", status: http.StatusTooManyRequests, body: `{"code":"insufficient_quota"}`, wantStatus: http.StatusPaymentRequired},
{name: "ordinary rate limit", status: http.StatusTooManyRequests, body: `{"message":"rate limit exceeded"}`, wantStatus: http.StatusTooManyRequests},
{name: "server error", status: http.StatusInternalServerError, body: `insufficient_quota`, wantStatus: http.StatusInternalServerError},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
transport := quotaAwareTransport{base: roundTripperFunc(func(*http.Request) (*http.Response, error) {
return &http.Response{
StatusCode: tt.status,
Status: http.StatusText(tt.status),
Body: io.NopCloser(strings.NewReader(tt.body)),
Header: make(http.Header),
}, nil
})}
req, err := http.NewRequestWithContext(context.Background(), http.MethodPost, "https://example.invalid", nil)
if err != nil {
t.Fatal(err)
}
resp, err := transport.RoundTrip(req)
if err != nil {
t.Fatal(err)
}
defer resp.Body.Close()
if resp.StatusCode != tt.wantStatus {
t.Fatalf("status=%d, want %d", resp.StatusCode, tt.wantStatus)
}
gotBody, err := io.ReadAll(resp.Body)
if err != nil {
t.Fatal(err)
}
if string(gotBody) != tt.body {
t.Fatalf("body=%q, want %q", gotBody, tt.body)
}
})
}
}
func TestProviderDoesNotRetryExplicitQuota429(t *testing.T) {
t.Parallel()
var requests atomic.Int32
upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
requests.Add(1)
w.Header().Set("content-type", "application/json")
w.WriteHeader(http.StatusTooManyRequests)
_, _ = io.WriteString(w, `{"error":{"code":"insufficient_quota","message":"You exceeded your current quota"}}`)
}))
defer upstream.Close()
provider, err := ConfigFrom("openai", "test-model", upstream.URL, "test-key", "").NewProvider()
if err != nil {
t.Fatal(err)
}
var streamErr error
for _, err := range provider.Stream(context.Background(), llm.CompletionRequest{
Messages: []llm.Message{llm.UserText("ping")},
}) {
if err != nil {
streamErr = err
}
}
if streamErr == nil || !IsQuotaExhaustedMessage(streamErr.Error()) {
t.Fatalf("expected explicit quota error, got %v", streamErr)
}
if got := requests.Load(); got != 1 {
t.Fatalf("explicit quota request retried %d times, want exactly one request", got)
}
}
+335 -21
View File
@@ -1,8 +1,11 @@
package db
import (
"context"
"database/sql"
"encoding/json"
"errors"
"fmt"
"time"
)
@@ -22,15 +25,36 @@ type LLMProfile struct {
// ContextWindowK is the model's context window in K tokens, used to size
// compaction thresholds. 0 = use a 200K default; capped at 1000 (1M).
ContextWindowK int `json:"context_window_k"`
// ReasoningEffort is the thinking-mode selector. "" = 默认(不发送思考参数);
// "off" = 显式关闭(thinking.type=disabled); "low"/"medium"/"high"/"max" =
// 开启思考并设强度. Mapped to the provider request in agent.Config.NewProvider.
// ThinkingType 独立控制思考「开关」(thinking.type):"" = 不发送(默认);
// "disabled" = 显式关闭; "enabled" = 开启. 与 ReasoningEffort 解耦.
ThinkingType string `json:"thinking_type"`
// ReasoningEffort 独立控制思考「强度」:"" = 不发送(默认);
// "low"/"medium"/"high"/"xhigh"/"max" = 对应强度. 见 agent.Config.NewProvider.
ReasoningEffort string `json:"reasoning_effort"`
IsDefault bool `json:"is_default"`
// Priority orders the failover chain: higher goes first. The ACTIVE profile
// (IsDefault) always heads the chain regardless of this value.
Priority int `json:"priority"`
// PoolExclude=true keeps this profile out of the failover chain — it stays
// usable when an agent/task binds it explicitly, it just never gets picked up
// as a fallback target.
PoolExclude bool `json:"pool_exclude"`
}
// profileCols is the read column list (hint variant, no api key) shared by the
// list query; profileColsKey is the same with api_key for the single-row loads.
const profileCols = `id,name,format,COALESCE(base_url,''),COALESCE(proxy,''),model,COALESCE(api_key_hint,''),rate_per_second,rate_per_minute,context_window_k,COALESCE(reasoning_effort,''),is_default,priority,pool_exclude,COALESCE(thinking_type,'')`
const profileColsKey = `id,name,format,COALESCE(base_url,''),COALESCE(proxy,''),model,COALESCE(api_key,''),rate_per_second,rate_per_minute,context_window_k,COALESCE(reasoning_effort,''),is_default,priority,pool_exclude,COALESCE(thinking_type,'')`
// scanProfile reads one row in the profileCols / profileColsKey column order. The
// 7th column lands in APIKeyHint or APIKey depending on which list the caller used.
func scanProfile(sc interface{ Scan(...any) error }, into *string, p *LLMProfile) error {
return sc.Scan(&p.ID, &p.Name, &p.Format, &p.BaseURL, &p.Proxy, &p.Model, into,
&p.RatePerSecond, &p.RatePerMinute, &p.ContextWindowK, &p.ReasoningEffort, &p.IsDefault, &p.Priority, &p.PoolExclude, &p.ThinkingType)
}
func (d *DB) ListProfiles() ([]*LLMProfile, error) {
rows, err := d.Query(`SELECT id,name,format,COALESCE(base_url,''),COALESCE(proxy,''),model,COALESCE(api_key_hint,''),rate_per_second,rate_per_minute,context_window_k,COALESCE(reasoning_effort,''),is_default FROM llm_profiles ORDER BY id`)
rows, err := d.Query(`SELECT ` + profileCols + ` FROM llm_profiles ORDER BY id`)
if err != nil {
return nil, err
}
@@ -38,7 +62,7 @@ func (d *DB) ListProfiles() ([]*LLMProfile, error) {
var out []*LLMProfile
for rows.Next() {
var p LLMProfile
if err := rows.Scan(&p.ID, &p.Name, &p.Format, &p.BaseURL, &p.Proxy, &p.Model, &p.APIKeyHint, &p.RatePerSecond, &p.RatePerMinute, &p.ContextWindowK, &p.ReasoningEffort, &p.IsDefault); err != nil {
if err := scanProfile(rows, &p.APIKeyHint, &p); err != nil {
return nil, err
}
out = append(out, &p)
@@ -49,8 +73,7 @@ func (d *DB) ListProfiles() ([]*LLMProfile, error) {
// ActiveProfile returns the default (active) profile with its api key, or nil.
func (d *DB) ActiveProfile() (*LLMProfile, error) {
var p LLMProfile
err := d.QueryRow(`SELECT id,name,format,COALESCE(base_url,''),COALESCE(proxy,''),model,COALESCE(api_key,''),rate_per_second,rate_per_minute,context_window_k,COALESCE(reasoning_effort,''),is_default FROM llm_profiles WHERE is_default LIMIT 1`).
Scan(&p.ID, &p.Name, &p.Format, &p.BaseURL, &p.Proxy, &p.Model, &p.APIKey, &p.RatePerSecond, &p.RatePerMinute, &p.ContextWindowK, &p.ReasoningEffort, &p.IsDefault)
err := scanProfile(d.QueryRow(`SELECT `+profileColsKey+` FROM llm_profiles WHERE is_default LIMIT 1`), &p.APIKey, &p)
if err == sql.ErrNoRows {
return nil, nil
}
@@ -61,14 +84,37 @@ func (d *DB) ActiveProfile() (*LLMProfile, error) {
// Used to run a task on a specific (non-default) LLM profile.
func (d *DB) ProfileByID(id int64) (*LLMProfile, error) {
var p LLMProfile
err := d.QueryRow(`SELECT id,name,format,COALESCE(base_url,''),COALESCE(proxy,''),model,COALESCE(api_key,''),rate_per_second,rate_per_minute,context_window_k,COALESCE(reasoning_effort,''),is_default FROM llm_profiles WHERE id=$1`, id).
Scan(&p.ID, &p.Name, &p.Format, &p.BaseURL, &p.Proxy, &p.Model, &p.APIKey, &p.RatePerSecond, &p.RatePerMinute, &p.ContextWindowK, &p.ReasoningEffort, &p.IsDefault)
err := scanProfile(d.QueryRow(`SELECT `+profileColsKey+` FROM llm_profiles WHERE id=$1`, id), &p.APIKey, &p)
if err == sql.ErrNoRows {
return nil, nil
}
return &p, err
}
// PoolProfiles returns the failover chain in run order, api keys included: the
// active profile first, then every other keyed profile that isn't excluded, by
// priority DESC (id ASC to stay stable). Profiles without an api key can't serve
// a request, so they never enter the chain. The ordering IS the policy — callers
// walk the slice front to back.
func (d *DB) PoolProfiles() ([]*LLMProfile, error) {
rows, err := d.Query(`SELECT ` + profileColsKey + ` FROM llm_profiles
WHERE COALESCE(api_key,'') <> '' AND (is_default OR NOT pool_exclude)
ORDER BY is_default DESC, priority DESC, id ASC`)
if err != nil {
return nil, err
}
defer rows.Close()
var out []*LLMProfile
for rows.Next() {
var p LLMProfile
if err := scanProfile(rows, &p.APIKey, &p); err != nil {
return nil, err
}
out = append(out, &p)
}
return out, rows.Err()
}
// SaveProfile inserts (id==0) or updates a profile. Empty apiKey on update keeps existing.
func (d *DB) SaveProfile(p *LLMProfile) (int64, error) {
hint := p.APIKeyHint
@@ -77,24 +123,251 @@ func (d *DB) SaveProfile(p *LLMProfile) (int64, error) {
}
if p.ID == 0 {
var id int64
err := d.QueryRow(`INSERT INTO llm_profiles(name,format,base_url,proxy,model,api_key,api_key_hint,rate_per_second,rate_per_minute,context_window_k,reasoning_effort)
VALUES ($1,$2,NULLIF($3,''),NULLIF($4,''),$5,NULLIF($6,''),NULLIF($7,''),$8,$9,$10,$11) RETURNING id`,
p.Name, p.Format, p.BaseURL, p.Proxy, p.Model, p.APIKey, hint, p.RatePerSecond, p.RatePerMinute, p.ContextWindowK, p.ReasoningEffort).Scan(&id)
err := d.QueryRow(`INSERT INTO llm_profiles(name,format,base_url,proxy,model,api_key,api_key_hint,rate_per_second,rate_per_minute,context_window_k,reasoning_effort,priority,pool_exclude,thinking_type)
VALUES ($1,$2,NULLIF($3,''),NULLIF($4,''),$5,NULLIF($6,''),NULLIF($7,''),$8,$9,$10,$11,$12,$13,$14) RETURNING id`,
p.Name, p.Format, p.BaseURL, p.Proxy, p.Model, p.APIKey, hint, p.RatePerSecond, p.RatePerMinute, p.ContextWindowK, p.ReasoningEffort, p.Priority, p.PoolExclude, p.ThinkingType).Scan(&id)
return id, err
}
if p.APIKey == "" {
_, err := d.Exec(`UPDATE llm_profiles SET name=$1,format=$2,base_url=NULLIF($3,''),proxy=NULLIF($4,''),model=$5,rate_per_second=$6,rate_per_minute=$7,context_window_k=$8,reasoning_effort=$9 WHERE id=$10`,
p.Name, p.Format, p.BaseURL, p.Proxy, p.Model, p.RatePerSecond, p.RatePerMinute, p.ContextWindowK, p.ReasoningEffort, p.ID)
_, err := d.Exec(`UPDATE llm_profiles SET name=$1,format=$2,base_url=NULLIF($3,''),proxy=NULLIF($4,''),model=$5,rate_per_second=$6,rate_per_minute=$7,context_window_k=$8,reasoning_effort=$9,priority=$10,pool_exclude=$11,thinking_type=$12 WHERE id=$13`,
p.Name, p.Format, p.BaseURL, p.Proxy, p.Model, p.RatePerSecond, p.RatePerMinute, p.ContextWindowK, p.ReasoningEffort, p.Priority, p.PoolExclude, p.ThinkingType, p.ID)
return p.ID, err
}
_, err := d.Exec(`UPDATE llm_profiles SET name=$1,format=$2,base_url=NULLIF($3,''),proxy=NULLIF($4,''),model=$5,api_key=$6,api_key_hint=$7,rate_per_second=$8,rate_per_minute=$9,context_window_k=$10,reasoning_effort=$11 WHERE id=$12`,
p.Name, p.Format, p.BaseURL, p.Proxy, p.Model, p.APIKey, hint, p.RatePerSecond, p.RatePerMinute, p.ContextWindowK, p.ReasoningEffort, p.ID)
_, err := d.Exec(`UPDATE llm_profiles SET name=$1,format=$2,base_url=NULLIF($3,''),proxy=NULLIF($4,''),model=$5,api_key=$6,api_key_hint=$7,rate_per_second=$8,rate_per_minute=$9,context_window_k=$10,reasoning_effort=$11,priority=$12,pool_exclude=$13,thinking_type=$14 WHERE id=$15`,
p.Name, p.Format, p.BaseURL, p.Proxy, p.Model, p.APIKey, hint, p.RatePerSecond, p.RatePerMinute, p.ContextWindowK, p.ReasoningEffort, p.Priority, p.PoolExclude, p.ThinkingType, p.ID)
return p.ID, err
}
var (
ErrActiveLLMProfileDelete = errors.New("cannot delete the active LLM profile; activate another profile first")
ErrLLMProfileReferencesChanged = errors.New("LLM profile references changed while deleting; retry the request")
ErrLLMProfileNotFound = errors.New("LLM profile not found")
)
func (d *DB) DeleteProfile(id int64) error {
_, err := d.Exec(`DELETE FROM llm_profiles WHERE id=$1`, id)
return err
return d.DeleteProfileContext(context.Background(), id)
}
// DeleteProfileContext removes a non-default profile while preserving the task
// failover cursor. Reference changes are retried because a task, agent, or
// conversation may start pointing at the profile between the initial scan and
// the profile row lock. A bound prevents a continuously changing workload from
// keeping an HTTP request alive forever.
func (d *DB) DeleteProfileContext(ctx context.Context, id int64) error {
const (
maxAttempts = 8
maxDuration = 15 * time.Second
)
ctx, cancel := context.WithTimeout(ctx, maxDuration)
defer cancel()
for attempt := 0; attempt < maxAttempts; attempt++ {
if err := ctx.Err(); err != nil {
return err
}
retry, err := d.deleteProfile(ctx, id)
if err != nil || !retry {
return err
}
}
return fmt.Errorf("%w: profile %d", ErrLLMProfileReferencesChanged, id)
}
func (d *DB) deleteProfile(ctx context.Context, id int64) (bool, error) {
tx, err := d.BeginTx(ctx, nil)
if err != nil {
return false, err
}
defer tx.Rollback()
// Reference setters lock their task/agent/conversation row before an
// llm_profiles row can be locked by a foreign-key check. Keep deletion in the
// same order: the old profile -> child order deadlocked with a setter's child
// -> profile order. Sorting also gives concurrent profile deletions a stable
// order when their chains overlap multiple tasks.
lockedTasks, err := lockProfileReferenceRows(ctx, tx, `SELECT t.id
FROM tasks t
WHERE t.llm_profile_id=$1
OR t.active_llm_profile_id=$1
OR EXISTS (
SELECT 1 FROM task_llm_profiles x
WHERE x.task_id=t.id AND x.profile_id=$1
)
ORDER BY t.id
FOR UPDATE OF t`, id)
if err != nil {
return false, err
}
lockedAgents, err := lockProfileReferenceRows(ctx, tx, `SELECT id FROM agents
WHERE llm_profile_id=$1
ORDER BY id
FOR UPDATE`, id)
if err != nil {
return false, err
}
lockedConversations, err := lockProfileReferenceRows(ctx, tx, `SELECT id FROM conversations
WHERE llm_profile_id=$1
ORDER BY id
FOR UPDATE`, id)
if err != nil {
return false, err
}
var isDefault bool
if err := tx.QueryRowContext(ctx, `SELECT is_default FROM llm_profiles WHERE id=$1 FOR UPDATE`, id).Scan(&isDefault); err != nil {
if err == sql.ErrNoRows {
return false, ErrLLMProfileNotFound
}
return false, err
}
if isDefault {
return false, ErrActiveLLMProfileDelete
}
type affectedTask struct {
id int64
position int
wasActive bool
}
// A task may have committed a new reference after the first statement took
// its snapshot but before this transaction acquired the profile lock. The
// profile lock now prevents further references; retry if that committed task
// was not part of the task-first lock set. Never acquire a new task lock while
// holding the profile lock, because that would recreate the inversion.
rows, err := tx.QueryContext(ctx, `SELECT ref_kind, ref_id FROM (
SELECT 'task'::text AS ref_kind, t.id AS ref_id
FROM tasks t
WHERE t.llm_profile_id=$1
OR t.active_llm_profile_id=$1
OR EXISTS (
SELECT 1 FROM task_llm_profiles x
WHERE x.task_id=t.id AND x.profile_id=$1
)
UNION ALL
SELECT 'agent', a.id FROM agents a WHERE a.llm_profile_id=$1
UNION ALL
SELECT 'conversation', c.id FROM conversations c WHERE c.llm_profile_id=$1
) refs
ORDER BY ref_kind, ref_id`, id)
if err != nil {
return false, err
}
for rows.Next() {
var (
kind string
rowID int64
)
if err := rows.Scan(&kind, &rowID); err != nil {
rows.Close()
return false, err
}
locked := false
switch kind {
case "task":
_, locked = lockedTasks[rowID]
case "agent":
_, locked = lockedAgents[rowID]
case "conversation":
_, locked = lockedConversations[rowID]
}
if !locked {
rows.Close()
return true, nil
}
}
if err := rows.Err(); err != nil {
rows.Close()
return false, err
}
if err := rows.Close(); err != nil {
return false, err
}
rows, err = tx.QueryContext(ctx, `SELECT x.task_id, x.position, COALESCE(t.active_llm_profile_id=$1, false)
FROM task_llm_profiles x
JOIN tasks t ON t.id=x.task_id
WHERE x.profile_id=$1
ORDER BY x.task_id`, id)
if err != nil {
return false, err
}
var affected []affectedTask
for rows.Next() {
var task affectedTask
if err := rows.Scan(&task.id, &task.position, &task.wasActive); err != nil {
rows.Close()
return false, err
}
affected = append(affected, task)
}
if err := rows.Err(); err != nil {
rows.Close()
return false, err
}
if err := rows.Close(); err != nil {
return false, err
}
if _, err := tx.ExecContext(ctx, `DELETE FROM llm_profiles WHERE id=$1`, id); err != nil {
return false, err
}
// Only a successor after the deleted cursor is eligible. If there is none,
// clear the explicit chain so the task falls back to its Agent/global provider;
// profiles before a manually selected cursor must never be revived.
for _, task := range affected {
if !task.wasActive {
if _, err := tx.ExecContext(ctx, `UPDATE tasks SET llm_chain_revision=llm_chain_revision+1 WHERE id=$1`, task.id); err != nil {
return false, err
}
continue
}
var next int64
err := tx.QueryRowContext(ctx, `SELECT profile_id FROM task_llm_profiles
WHERE task_id=$1 AND position>$2 AND status='ready'
ORDER BY position
LIMIT 1`, task.id, task.position).Scan(&next)
if err != nil && err != sql.ErrNoRows {
return false, err
}
if err == sql.ErrNoRows {
if _, err := tx.ExecContext(ctx, `DELETE FROM task_llm_profiles WHERE task_id=$1`, task.id); err != nil {
return false, err
}
if _, err := tx.ExecContext(ctx, `UPDATE tasks
SET active_llm_profile_id=NULL, llm_profile_id=NULL, llm_chain_revision=llm_chain_revision+1
WHERE id=$1`, task.id); err != nil {
return false, err
}
continue
}
if _, err := tx.ExecContext(ctx, `UPDATE tasks
SET active_llm_profile_id=$2, llm_profile_id=$2, llm_chain_revision=llm_chain_revision+1
WHERE id=$1`, task.id, next); err != nil {
return false, err
}
}
if err := tx.Commit(); err != nil {
return false, err
}
return false, nil
}
func lockProfileReferenceRows(ctx context.Context, tx *sql.Tx, query string, profileID int64) (map[int64]struct{}, error) {
rows, err := tx.QueryContext(ctx, query, profileID)
if err != nil {
return nil, err
}
defer rows.Close()
locked := make(map[int64]struct{})
for rows.Next() {
var id int64
if err := rows.Scan(&id); err != nil {
return nil, err
}
locked[id] = struct{}{}
}
if err := rows.Err(); err != nil {
return nil, err
}
return locked, nil
}
// SetActiveProfile makes one profile the global default (single-default invariant).
@@ -107,9 +380,13 @@ func (d *DB) SetActiveProfile(id int64) error {
if _, err := tx.Exec(`UPDATE llm_profiles SET is_default=false WHERE is_default`); err != nil {
return err
}
if _, err := tx.Exec(`UPDATE llm_profiles SET is_default=true WHERE id=$1`, id); err != nil {
res, err := tx.Exec(`UPDATE llm_profiles SET is_default=true WHERE id=$1`, id)
if err != nil {
return err
}
if n, _ := res.RowsAffected(); n != 1 {
return ErrLLMProfileNotFound
}
return tx.Commit()
}
@@ -272,8 +549,45 @@ func (d *DB) SetAgentMaxTurns(key string, maxTurns int) error {
// the binding (id == nil) so the agent follows the task/conversation pin, else the
// global active profile. Precedence at runtime: agent binding → task/conv pin → active.
func (d *DB) SetAgentLLMProfile(key string, id *int64) error {
_, err := d.Exec(`UPDATE agents SET llm_profile_id=$1 WHERE key=$2`, id, key)
return err
tx, err := d.Begin()
if err != nil {
return err
}
defer tx.Rollback()
// DeleteProfile locks reference rows before the profile row. Keep the same
// order here so a concurrent rebind cannot form a child/profile deadlock.
var agentID int64
if err := tx.QueryRow(`SELECT id FROM agents WHERE key=$1 FOR UPDATE`, key).Scan(&agentID); err != nil {
if err == sql.ErrNoRows {
// Preserve the previous UPDATE semantics: an unknown key is a no-op.
return tx.Commit()
}
return err
}
if err := lockLLMProfileForReference(tx, id); err != nil {
return err
}
if _, err := tx.Exec(`UPDATE agents SET llm_profile_id=$1 WHERE id=$2`, id, agentID); err != nil {
return err
}
return tx.Commit()
}
// lockLLMProfileForReference makes the profile side of the shared lock-order
// protocol explicit. Callers must already own the referencing child row.
func lockLLMProfileForReference(tx *sql.Tx, profileID *int64) error {
if profileID == nil {
return nil
}
var lockedID int64
if err := tx.QueryRow(`SELECT id FROM llm_profiles WHERE id=$1 FOR KEY SHARE`, *profileID).Scan(&lockedID); err != nil {
if err == sql.ErrNoRows {
return ErrLLMProfileNotFound
}
return err
}
return nil
}
// SetAgentWebSearch toggles whether an agent uses network search (still gated by
+62
View File
@@ -1,10 +1,69 @@
package db
import (
"context"
"encoding/json"
"errors"
"testing"
)
// TestPoolProfilesOrder pins the failover chain query: keyless profiles can't
// serve a request and excluded ones aren't fallback targets, so neither belongs
// in the chain; the rest come back by priority, highest first.
// Deliberately does NOT touch is_default — flipping the active profile would be a
// side effect on the shared dev database.
func TestPoolProfilesOrder(t *testing.T) {
d, err := Open(testDSN(t))
if err != nil {
t.Skipf("postgres unavailable (%v) — skipping", err)
}
defer d.Close()
mk := func(name string, priority int, exclude bool, key string) int64 {
id, err := d.SaveProfile(&LLMProfile{
Name: name, Format: "openai", Model: "m", APIKey: key,
Priority: priority, PoolExclude: exclude,
})
if err != nil {
t.Fatal(err)
}
t.Cleanup(func() { d.Exec(`DELETE FROM llm_profiles WHERE id=$1`, id) })
return id
}
lo := mk("t-pool-lo", 1, false, "k1")
hi := mk("t-pool-hi", 9, false, "k2")
mk("t-pool-excluded", 99, true, "k3") // excluded despite the top priority
mk("t-pool-nokey", 50, false, "") // no key → cannot serve anything
chain, err := d.PoolProfiles()
if err != nil {
t.Fatal(err)
}
var got []int64
for _, p := range chain {
switch p.Name {
case "t-pool-lo", "t-pool-hi":
got = append(got, p.ID)
case "t-pool-excluded":
t.Fatal("pool_exclude profile entered the failover chain")
case "t-pool-nokey":
t.Fatal("keyless profile entered the failover chain")
}
}
if len(got) != 2 || got[0] != hi || got[1] != lo {
t.Fatalf("chain order = %v, want [hi=%d lo=%d]", got, hi, lo)
}
}
func TestDeleteProfileContextHonorsCancellation(t *testing.T) {
ctx, cancel := context.WithCancel(context.Background())
cancel()
var d DB
if err := d.DeleteProfileContext(ctx, 1); !errors.Is(err, context.Canceled) {
t.Fatalf("DeleteProfileContext error=%v, want context cancellation", err)
}
}
func TestConfigStores(t *testing.T) {
d, err := Open(testDSN(t))
if err != nil {
@@ -21,6 +80,9 @@ func TestConfigStores(t *testing.T) {
if err := d.SetActiveProfile(pid); err != nil {
t.Fatal(err)
}
if err := d.DeleteProfile(pid); !errors.Is(err, ErrActiveLLMProfileDelete) {
t.Fatalf("deleting active profile error=%v, want %v", err, ErrActiveLLMProfileDelete)
}
act, err := d.ActiveProfile()
if err != nil || act == nil || act.APIKey != "secret123" {
t.Fatalf("active profile/key: %+v err=%v", act, err)
+57
View File
@@ -0,0 +1,57 @@
package db
import "time"
// LLM failover health: one row per profile tracking its circuit-breaker state.
// The authoritative copy lives in memory (llmpool.Registry); this table only
// survives a restart so a cooling-off window isn't silently reset by one.
// LLMHealth is one profile's circuit-breaker state.
type LLMHealth struct {
ProfileID int64 `json:"profile_id"`
Fails int `json:"fails"` // consecutive failures; cleared on success
Trips int `json:"trips"` // total trips, drives the backoff ladder
OpenUntil *time.Time `json:"open_until"` // nil/past = closed (healthy)
LastError string `json:"last_error"`
LastAt time.Time `json:"last_at"`
}
// LoadLLMHealth returns the profiles still in an UNEXPIRED cooling-off window.
// Expired rows are deliberately skipped: after a restart a profile that has
// finished cooling should be treated as healthy again and re-probed on its next
// call, not resurrected as broken.
func (d *DB) LoadLLMHealth() ([]LLMHealth, error) {
rows, err := d.Query(`SELECT profile_id,fails,trips,open_until,COALESCE(last_error,''),last_at
FROM llm_profile_health WHERE open_until IS NOT NULL AND open_until > now()`)
if err != nil {
return nil, err
}
defer rows.Close()
var out []LLMHealth
for rows.Next() {
var h LLMHealth
if err := rows.Scan(&h.ProfileID, &h.Fails, &h.Trips, &h.OpenUntil, &h.LastError, &h.LastAt); err != nil {
return nil, err
}
out = append(out, h)
}
return out, rows.Err()
}
// SaveLLMHealth upserts one profile's circuit-breaker state.
func (d *DB) SaveLLMHealth(h LLMHealth) error {
_, err := d.Exec(`
INSERT INTO llm_profile_health(profile_id,fails,trips,open_until,last_error,last_at)
VALUES ($1,$2,$3,$4,$5,now())
ON CONFLICT (profile_id) DO UPDATE SET
fails=EXCLUDED.fails, trips=EXCLUDED.trips, open_until=EXCLUDED.open_until,
last_error=EXCLUDED.last_error, last_at=now()`,
h.ProfileID, h.Fails, h.Trips, h.OpenUntil, h.LastError)
return err
}
// ClearLLMHealth drops one profile's state (manual "recover now" from the UI).
func (d *DB) ClearLLMHealth(profileID int64) error {
_, err := d.Exec(`DELETE FROM llm_profile_health WHERE profile_id=$1`, profileID)
return err
}
+149
View File
@@ -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
View File
@@ -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
}
+359
View File
@@ -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")
}
}
+136 -39
View File
@@ -11,25 +11,55 @@ import (
"strings"
"time"
"github.com/Autumn-27/artex/db"
"github.com/Autumn-27/norma/llm"
"github.com/Autumn-27/norma/transcript"
"github.com/Autumn-27/artex/db"
)
type taskIDContextKey struct{}
// WithTaskID attaches the owning task registry id to an LLM call. Session ids
// are based on exploration ids, which are not interchangeable with task ids.
func WithTaskID(ctx context.Context, taskID string) context.Context {
taskID = strings.TrimSpace(taskID)
if taskID == "" {
return ctx
}
return context.WithValue(ctx, taskIDContextKey{}, taskID)
}
// TaskIDFrom returns the explicit task registry id attached by the task runtime.
func TaskIDFrom(ctx context.Context) string {
if ctx == nil {
return ""
}
taskID, _ := ctx.Value(taskIDContextKey{}).(string)
return strings.TrimSpace(taskID)
}
// Recorder wraps an llm.Provider and records every completion call.
type Recorder struct {
inner llm.Provider
pg *db.DB
model string // model name (from config, not in CompletionRequest)
prof string // LLM profile name (from llm_profiles)
enabled func() bool // reports whether recording is currently on; nil = always record
model string // model name (from config, not in CompletionRequest)
prof string // LLM profile name (from llm_profiles)
// thinkingType / reasoningEffort 是配置级思考参数(思考开关 / 思考强度)。它们在
// norma 的 buildBody() 里从 provider 配置注入真正的 HTTP body,不出现在
// CompletionRequest 上,故 Recorder 需在此单独带一份,序列化时写进录制。
thinkingType string
reasoningEffort string
enabled func() bool // reports whether recording is currently on; nil = always record
}
// Wrap returns a Provider that records calls to pg when enabled() reports true.
// profName is the LLM profile name (e.g. "default"); may be empty. A nil enabled
// predicate records unconditionally (backward-compatible default).
func Wrap(inner llm.Provider, pg *db.DB, model, profName string, enabled func() bool) *Recorder {
return &Recorder{inner: inner, pg: pg, model: model, prof: profName, enabled: enabled}
// profName is the LLM profile name (e.g. "default"); may be empty. thinkingType /
// reasoningEffort are the config-level thinking params actually sent to the API
// (empty = not sent). A nil enabled predicate records unconditionally.
func Wrap(inner llm.Provider, pg *db.DB, model, profName, thinkingType, reasoningEffort string, enabled func() bool) *Recorder {
return &Recorder{
inner: inner, pg: pg, model: model, prof: profName,
thinkingType: thinkingType, reasoningEffort: reasoningEffort, enabled: enabled,
}
}
// parseSession extracts task id and worker role from session strings like
@@ -63,64 +93,123 @@ func parseSession(s string) (taskID, worker string) {
// harness via transcript.WithSessionID; reading it per-call is race-free even when
// planner and multiple workers share one Recorder instance.
func (r *Recorder) Stream(ctx context.Context, req llm.CompletionRequest) iter.Seq2[llm.StreamEvent, error] {
// Recording off → passthrough with zero overhead: no request serialization,
// no response accumulation, no DB insert.
if r.enabled != nil && !r.enabled() {
return r.inner.Stream(ctx, req)
}
// Body recording (the heavy debug trace: full request/response) is gated by the
// llm_record setting. Lightweight usage metering always runs — it powers token
// stats and must be complete even for interrupted/failed runs, so it is NOT
// gated. Only body serialization + response accumulation are skipped when off.
recordBodies := r.enabled == nil || r.enabled()
start := time.Now()
session := transcript.SessionIDFrom(ctx)
taskID, worker := parseSession(session)
parsedID, worker := parseSession(session)
expID := db.ParseExpID(parsedID)
taskID := TaskIDFrom(ctx)
if taskID == "" {
// Backward compatibility for non-task callers. For task calls, the task
// runtime always supplies the registry id explicitly.
taskID = parsedID
}
// Serialize the request for storage.
reqBody := serializeRequest(req)
// Serialize the request only when storing bodies (this is the expensive part).
reqBody := ""
if recordBodies {
reqBody = r.serializeRequest(req)
}
return func(yield func(llm.StreamEvent, error) bool) {
var (
textBuf strings.Builder
textBuf strings.Builder
thinkingBuf strings.Builder
usage llm.Usage
stopReason string
streamErr error
usage llm.Usage
stopReason string
streamErr error
)
finish := func(err error) {
status := "ok"
if err != nil {
status = "error"
}
// Lightweight metering row — always written.
r.recordUsage(taskID, expID, worker, usage, int(time.Since(start).Milliseconds()), status)
// Heavy trace row — only when body recording is on.
if recordBodies {
r.record(req, session, taskID, worker, reqBody, start, textBuf.String(), thinkingBuf.String(), usage, stopReason, err)
}
}
for ev, err := range r.inner.Stream(ctx, req) {
if err != nil {
streamErr = err
// Record the failed call, then propagate.
r.record(req, session, taskID, worker, reqBody, start, textBuf.String(), thinkingBuf.String(), usage, "", err)
finish(streamErr)
if !yield(ev, err) {
return
}
return
}
// Accumulate response content.
// Always track usage (cheap); accumulate text/thinking only for bodies.
// Anthropic (and the other providers) split token usage across events:
// message_start carries ONLY the input side (input + cache), message_delta
// ONLY the output. They must be FOLDED with Add — overwriting on delta
// would zero out the input already counted at start (mirrors the SDK's own
// llm.Accumulator; see norma/llm/accumulate.go).
switch ev.Type {
case llm.SETextDelta:
textBuf.WriteString(ev.Text)
if recordBodies {
textBuf.WriteString(ev.Text)
}
case llm.SEThinkingDelta:
thinkingBuf.WriteString(ev.Text)
if recordBodies {
thinkingBuf.WriteString(ev.Text)
}
case llm.SEMessageStart:
usage.Add(ev.Usage)
case llm.SEMessageDelta:
if ev.StopReason != "" {
stopReason = ev.StopReason
}
usage = ev.Usage
case llm.SEMessageStart:
if ev.Usage.InputTokens > 0 {
usage = ev.Usage
}
usage.Add(ev.Usage)
}
if !yield(ev, err) {
return
}
}
// Stream completed normally — record it.
r.record(req, session, taskID, worker, reqBody, start, textBuf.String(), thinkingBuf.String(), usage, stopReason, streamErr)
// Stream completed normally.
finish(streamErr)
}
}
// record persists one LLM call to PostgreSQL (async, best-effort).
// recordUsage appends one lightweight metering row to llm_usage (no bodies). Skips
// zero-token calls with no model, which carry nothing worth metering.
func (r *Recorder) recordUsage(taskID string, expID int64, worker string, usage llm.Usage, latencyMs int, status string) {
if r.pg == nil {
return
}
if r.model == "" && usage.InputTokens == 0 && usage.OutputTokens == 0 &&
usage.CacheReadTokens == 0 && usage.CacheWriteTokens == 0 {
return
}
err := r.pg.InsertLLMUsage(&db.LLMUsage{
TaskID: taskID,
ExplorationID: expID,
Worker: worker,
Model: r.model,
ProfileName: r.prof,
LatencyMs: latencyMs,
InputTokens: usage.InputTokens,
OutputTokens: usage.OutputTokens,
CacheRead: usage.CacheReadTokens,
CacheWrite: usage.CacheWriteTokens,
Status: status,
})
if err != nil {
log.Printf("[llmusage] insert: %v", err)
}
}
// record persists one LLM call to PostgreSQL before the provider stream returns.
// Keeping the write inside the owning task operation means task deletion can
// drain calls and then remove records without a late async insert recreating one.
func (r *Recorder) record(req llm.CompletionRequest, session, taskID, worker, reqBody string, start time.Time, text, thinking string, usage llm.Usage, stopReason string, streamErr error) {
latency := int(time.Since(start).Milliseconds())
status := "ok"
@@ -169,20 +258,28 @@ func (r *Recorder) record(req llm.CompletionRequest, session, taskID, worker, re
ResponseBody: string(respBody),
}
go func() {
if err := r.pg.InsertLLMRecord(rec); err != nil {
log.Printf("[llmrec] insert: %v", err)
}
}()
if err := r.pg.InsertLLMRecord(rec); err != nil {
log.Printf("[llmrec] insert: %v", err)
}
}
// serializeRequest builds a JSON representation of the completion request.
func serializeRequest(req llm.CompletionRequest) string {
func (r *Recorder) serializeRequest(req llm.CompletionRequest) string {
m := map[string]any{
"system": req.System,
"messages": req.Messages,
"max_tokens": req.MaxTokens,
}
// 记录本次调用实际发出的思考参数。type 采用「有效值」:每请求覆盖 req.Thinking
// 优先于配置级 thinkingType(与 norma buildBody 的判定一致,如 compaction 摘要会
// 强制 disabled);effort 无每请求覆盖,直接取配置值。两者皆空则不写 thinking 字段。
effType := r.thinkingType
if req.Thinking != "" {
effType = req.Thinking
}
if effType != "" || r.reasoningEffort != "" {
m["thinking"] = map[string]string{"type": effType, "effort": r.reasoningEffort}
}
if len(req.Tools) > 0 {
// Store tool names only (full schemas are huge).
names := make([]string, len(req.Tools))
+29
View File
@@ -0,0 +1,29 @@
package llmrec
import (
"context"
"testing"
)
func TestTaskIDContextUsesExplicitRegistryID(t *testing.T) {
ctx := WithTaskID(context.Background(), " 42 ")
if got := TaskIDFrom(ctx); got != "42" {
t.Fatalf("TaskIDFrom()=%q want 42", got)
}
if got := TaskIDFrom(WithTaskID(ctx, " ")); got != "42" {
t.Fatalf("blank task id should preserve parent context, got %q", got)
}
if got := TaskIDFrom(nil); got != "" {
t.Fatalf("nil context returned %q", got)
}
}
func TestParseSessionFallbackIsExplorationScoped(t *testing.T) {
taskID, worker := parseSession("exp12-worker-i99")
if taskID != "12" || worker != "worker" {
t.Fatalf("parseSession()=(%q,%q)", taskID, worker)
}
if taskID, worker := parseSession("not-a-task"); taskID != "" || worker != "" {
t.Fatalf("unexpected non-session parse: (%q,%q)", taskID, worker)
}
}
+43
View File
@@ -0,0 +1,43 @@
package server
import (
"sync"
"testing"
)
func TestEngineActiveLLMCallsConcurrent(t *testing.T) {
e := NewEngine(nil)
const taskID = "42"
const calls = 64
var started sync.WaitGroup
var release sync.WaitGroup
var finished sync.WaitGroup
started.Add(calls)
finished.Add(calls)
release.Add(1)
for range calls {
go func() {
defer finished.Done()
e.BeginLLMCall(taskID)
started.Done()
release.Wait()
e.EndLLMCall(taskID)
}()
}
started.Wait()
if got := e.ActiveLLMCalls(taskID); got != calls {
t.Fatalf("active calls = %d, want %d", got, calls)
}
release.Done()
finished.Wait()
if got := e.ActiveLLMCalls(taskID); got != 0 {
t.Fatalf("active calls after completion = %d, want 0", got)
}
// A defensive extra end must never expose a negative UI count.
e.EndLLMCall(taskID)
if got := e.ActiveLLMCalls(taskID); got != 0 {
t.Fatalf("active calls after extra end = %d, want 0", got)
}
}
+239
View File
@@ -0,0 +1,239 @@
package server
import (
"log"
"time"
"github.com/Autumn-27/artex/agent"
"github.com/Autumn-27/artex/db"
"github.com/Autumn-27/artex/llmpool"
"github.com/Autumn-27/norma/llm"
)
// LLM 轮询(故障转移)的服务端接线。设计见 docs/LLM轮询设计.md:
// - 全局激活配置这条路径(agent 未绑定、任务未 pin)才轮询;
// - 绑定/pin 的路径默认独占该配置,失败即失败(可由 llm_pool_bind_fallback 打开兜底);
// - 链序 = 激活配置 → 其余按 priority DESC,排除 pool_exclude 的;
// - 熔断状态进程级共享(s.llmHealth),重建 pool 不清空。
// newLLMHealthRegistry builds the process-wide circuit-breaker registry, mirroring
// state into PG so a cooling-off window survives a restart. Writes are async and
// best-effort — the in-memory copy is authoritative.
func newLLMHealthRegistry(pg *db.DB) *llmpool.Registry {
if pg == nil {
return llmpool.NewRegistry(nil, nil)
}
persist := func(id int64, st llmpool.State) {
h := db.LLMHealth{ProfileID: id, Fails: st.Fails, Trips: st.Trips, LastError: st.LastError}
if !st.OpenUntil.IsZero() {
t := st.OpenUntil
h.OpenUntil = &t
}
go func() {
if err := pg.SaveLLMHealth(h); err != nil {
log.Printf("[llmpool] 熔断状态落库失败: %v", err)
}
}()
}
forget := func(id int64) {
go func() { _ = pg.ClearLLMHealth(id) }()
}
reg := llmpool.NewRegistry(persist, forget)
// Restore only windows that haven't expired (LoadLLMHealth filters), so a
// profile that finished cooling while we were down comes back healthy.
if rows, err := pg.LoadLLMHealth(); err == nil {
for _, h := range rows {
st := llmpool.State{Fails: h.Fails, Trips: h.Trips, LastError: h.LastError, LastAt: h.LastAt}
if h.OpenUntil != nil {
st.OpenUntil = *h.OpenUntil
}
reg.Restore(h.ProfileID, st)
log.Printf("[llmpool] 恢复熔断状态: 配置 #%d 冷却至 %s", h.ProfileID, st.OpenUntil.Format(time.RFC3339))
}
}
return reg
}
// poolMember builds one chain member from a profile, reusing the per-profile
// provider cache so every agent pointing at the same profile shares one provider
// (hence one rate limiter). Returns nil when the profile can't be built.
func (s *Server) poolMember(p *db.LLMProfile, rank int) *llmpool.Member {
prov, cfg, ok := s.providerForProfile(p.ID)
if !ok {
return nil
}
return &llmpool.Member{
ID: p.ID, Name: p.Name, Model: p.Model, Format: p.Format,
Priority: p.Priority, Active: p.IsDefault, Rank: rank,
WindowTokens: cfg.CompactionWindow(), Prov: prov,
}
}
// poolChain reads the failover chain from the DB. headID/headProv/headCfg describe
// the member that must lead the chain — the globally active profile on the normal
// path, or an explicitly bound one when bound-profile fallback is on. The head is
// passed in (rather than resolved through the cache) so applyLLM's already-built
// provider is reused instead of duplicated.
//
// Returns nil when failover is off or the chain has fewer than two usable members
// — callers then use the bare provider, which is byte-for-byte the old behavior.
func (s *Server) poolChain(headID int64, headProv llm.Provider, headCfg agent.Config) *llmpool.Pool {
if s.m == nil || s.m.pg == nil || !s.m.LLMPoolEnabled() {
return nil
}
profs, err := s.m.pg.PoolProfiles()
if err != nil {
log.Printf("[llmpool] 读取轮询链失败: %v", err)
return nil
}
var head *db.LLMProfile
for _, p := range profs {
if p.ID == headID {
head = p
break
}
}
if head == nil { // head excluded from the chain (bound + pool_exclude) — still leads
if p, err := s.m.pg.ProfileByID(headID); err == nil && p != nil {
head = p
} else {
return nil
}
}
members := []*llmpool.Member{{
ID: head.ID, Name: head.Name, Model: head.Model, Format: head.Format,
Priority: head.Priority, Active: head.IsDefault, Rank: llmpool.RankActive,
WindowTokens: headCfg.CompactionWindow(), Prov: headProv,
}}
for _, p := range profs {
if p.ID == headID {
continue
}
// The globally active profile outranks the others when it isn't the head
// (bound-fallback chains), so it stays the first fallback tried.
rank := p.Priority
if p.IsDefault {
rank = llmpool.RankActive - 1
}
if m := s.poolMember(p, rank); m != nil {
members = append(members, m)
}
}
if len(members) < 2 {
return nil // nothing to fail over to
}
return llmpool.New(members, s.llmHealth)
}
// poolForActive wraps the globally-active provider in the failover chain. Returns
// prov unchanged when failover is off or there's nothing to fall back to.
func (s *Server) poolForActive(activeID int64, prov llm.Provider, cfg agent.Config) llm.Provider {
pool := s.poolChain(activeID, prov, cfg)
if pool == nil {
return prov
}
names := make([]string, 0, len(pool.Members()))
for _, m := range pool.Members() {
names = append(names, m.Name+"/"+m.Model)
}
log.Printf("[llmpool] LLM 轮询已启用,链路(%d): %v", len(names), names)
return pool
}
// poolForBinding wraps a BOUND profile's provider so it falls back to the chain.
// Only active when both the failover master switch and the bound-fallback switch
// are on; otherwise a bound profile stays exclusive (fails hard), which is the
// default and the documented precedence.
func (s *Server) poolForBinding(id int64, prov llm.Provider, cfg agent.Config) llm.Provider {
if s.m == nil || !s.m.LLMPoolEnabled() || !s.m.LLMPoolBindFallback() {
return prov
}
if pool := s.poolChain(id, prov, cfg); pool != nil {
return pool
}
return prov
}
// LLMPoolMemberStatus is one chain entry as shown in the UI.
type LLMPoolMemberStatus struct {
ProfileID string `json:"profile_id"`
Name string `json:"name"`
Model string `json:"model"`
Format string `json:"format"`
Priority int `json:"priority"`
Active bool `json:"active"` // is the globally-activated profile
Excluded bool `json:"excluded"` // pool_exclude — not a failover target
// Health: state is "ok" | "degraded" (failing but not tripped) | "tripped".
State string `json:"state"`
Fails int `json:"fails"`
Trips int `json:"trips"`
CooldownSecs int `json:"cooldown_secs"` // remaining cooling-off seconds; 0 = none
LastError string `json:"last_error,omitempty"`
LastAt string `json:"last_at,omitempty"`
}
// llmPoolStatus reports the whole picture for the LLM page: whether failover is
// on, the resolved chain order, and each profile's breaker state. Excluded
// profiles are listed too (flagged), so the user can see why one isn't in line.
func (s *Server) llmPoolStatus() map[string]any {
out := map[string]any{
"enabled": false,
"bind_fallback": false,
"chain": []LLMPoolMemberStatus{},
}
if s.m == nil || s.m.pg == nil {
return out
}
out["enabled"] = s.m.LLMPoolEnabled()
out["bind_fallback"] = s.m.LLMPoolBindFallback()
all, err := s.m.pg.ListProfiles()
if err != nil {
return out
}
health := s.llmHealth.Snapshot()
now := time.Now()
// Chain order: active first, then priority DESC / id ASC — the same ordering
// PoolProfiles applies, recomputed here so excluded ones can be shown in place.
chain := make([]LLMPoolMemberStatus, 0, len(all))
for _, p := range all {
st := health[p.ID]
m := LLMPoolMemberStatus{
ProfileID: i64s(p.ID), Name: p.Name, Model: p.Model, Format: p.Format,
Priority: p.Priority, Active: p.IsDefault, Excluded: p.PoolExclude,
State: "ok", Fails: st.Fails, Trips: st.Trips,
LastError: st.LastError,
}
if st.Open() {
m.State = "tripped"
m.CooldownSecs = int(st.OpenUntil.Sub(now).Seconds()) + 1
} else if st.Fails > 0 {
m.State = "degraded"
}
if !st.LastAt.IsZero() {
m.LastAt = st.LastAt.Format(time.RFC3339)
}
chain = append(chain, m)
}
sortPoolStatus(chain)
out["chain"] = chain
return out
}
// sortPoolStatus orders the status rows exactly like the live chain: active first,
// then priority DESC, then id ASC (the incoming slice is already id-ordered, so a
// stable insertion sort on the first two keys is enough).
func sortPoolStatus(in []LLMPoolMemberStatus) {
for i := 1; i < len(in); i++ {
for j := i; j > 0 && poolLess(in[j], in[j-1]); j-- {
in[j], in[j-1] = in[j-1], in[j]
}
}
}
func poolLess(a, b LLMPoolMemberStatus) bool {
if a.Active != b.Active {
return a.Active
}
return a.Priority > b.Priority
}
+41
View File
@@ -0,0 +1,41 @@
package server
import "testing"
// The status list drives the "轮询顺序" strip in the UI, so it has to match the
// order PoolProfiles actually runs: active first, then priority DESC, then id ASC
// (input arrives id-ordered, so equal priorities must keep their relative order).
func TestSortPoolStatusMatchesChainOrder(t *testing.T) {
in := []LLMPoolMemberStatus{
{ProfileID: "1", Name: "low", Priority: 0},
{ProfileID: "2", Name: "active", Priority: 0, Active: true},
{ProfileID: "3", Name: "high", Priority: 10},
{ProfileID: "4", Name: "mid-a", Priority: 5},
{ProfileID: "5", Name: "mid-b", Priority: 5},
}
sortPoolStatus(in)
want := []string{"active", "high", "mid-a", "mid-b", "low"}
for i, w := range want {
if in[i].Name != w {
got := make([]string, len(in))
for j, m := range in {
got[j] = m.Name
}
t.Fatalf("order = %v, want %v", got, want)
}
}
}
// The active profile heads the chain no matter how low its own priority is —
// that's the documented precedence, and the UI must not imply otherwise.
func TestActiveProfileHeadsStatusList(t *testing.T) {
in := []LLMPoolMemberStatus{
{ProfileID: "1", Name: "loud", Priority: 999},
{ProfileID: "2", Name: "active", Priority: -5, Active: true},
}
sortPoolStatus(in)
if in[0].Name != "active" {
t.Fatalf("head = %q, want the active profile regardless of priority", in[0].Name)
}
}
File diff suppressed because it is too large. Load diff
@@ -0,0 +1,223 @@
"use client";
import * as React from "react";
import { ArrowDownIcon, ArrowUpIcon, XIcon } from "lucide-react";
import { Badge } from "@/components/ui/badge";
import { Button } from "@/components/ui/button";
import {
Combobox,
ComboboxChip,
ComboboxChips,
ComboboxChipsInput,
ComboboxContent,
ComboboxEmpty,
ComboboxItem,
ComboboxList,
ComboboxValue,
} from "@/components/ui/combobox";
import { Select, SelectContent, SelectGroup, SelectItem, SelectTrigger, SelectValue } from "@/components/ui/select";
import type { LLMProfile } from "@/lib/types";
interface TaskLLMProfileChainProps {
profiles: LLMProfile[];
value: string[];
onValueChange: (value: string[]) => void;
activeProfileId?: string;
onActiveProfileChange?: (value: string) => void;
disabled?: boolean;
inputId?: string;
portalContainer?: React.RefObject<HTMLElement | null>;
}
function ProfileRoleBadge({ index, currentIndex }: { index: number; currentIndex: number }) {
if (index === currentIndex) return <Badge variant="default">当前</Badge>;
if (index < currentIndex) {
return (
<Badge variant="outline" title="当前游标之前的配置不会被自动故障转移选中">
已跳过
</Badge>
);
}
return <Badge variant="secondary">备用 {index - currentIndex}</Badge>;
}
export function TaskLLMProfileChain({
profiles,
value,
onValueChange,
activeProfileId,
onActiveProfileChange,
disabled = false,
inputId,
portalContainer,
}: TaskLLMProfileChainProps) {
const profilesByID = React.useMemo(() => new Map(profiles.map((profile) => [profile.id, profile])), [profiles]);
const itemIDs = React.useMemo(() => profiles.map((profile) => profile.id), [profiles]);
const profilesUnavailable = profiles.length === 0;
const currentProfileID = activeProfileId && value.includes(activeProfileId) ? activeProfileId : value[0];
const currentIndex = value.indexOf(currentProfileID);
// api_key_hint is display-only and is empty for valid keys shorter than four
// characters. The task API remains authoritative for profile validation.
const profileLabel = React.useCallback(
(id: string) => {
const profile = profilesByID.get(id);
return profile ? `${profile.name} ${profile.model}` : `配置 #${id}`;
},
[profilesByID],
);
const move = (index: number, offset: -1 | 1) => {
const target = index + offset;
if (target < 0 || target >= value.length) return;
const next = [...value];
[next[index], next[target]] = [next[target], next[index]];
onValueChange(next);
};
const remove = (id: string) => {
const next = value.filter((profileID) => profileID !== id);
onValueChange(next);
if (activeProfileId === id) onActiveProfileChange?.(next[0] ?? "");
};
const handleValueChange = (next: string[]) => {
onValueChange(next);
if (activeProfileId && !next.includes(activeProfileId)) {
onActiveProfileChange?.(next[0] ?? "");
}
};
return (
<div className="flex flex-col gap-3">
<Combobox
items={itemIDs}
itemToStringValue={profileLabel}
multiple
value={value}
onValueChange={handleValueChange}
disabled={disabled}
>
<ComboboxChips>
<ComboboxValue>
{value.map((id) => (
<ComboboxChip key={id}>{profilesByID.get(id)?.name ?? `#${id}`}</ComboboxChip>
))}
</ComboboxValue>
<ComboboxChipsInput
id={inputId}
placeholder={profiles.length > 0 ? "搜索并添加 LLM 配置" : "暂无可用 LLM 配置"}
disabled={disabled ? true : profilesUnavailable}
/>
</ComboboxChips>
<ComboboxContent portalContainer={portalContainer}>
<ComboboxEmpty>没有匹配的 LLM 配置</ComboboxEmpty>
<ComboboxList>
{(id) => {
const profile = profilesByID.get(id);
return (
<ComboboxItem key={id} value={id} disabled={!profile}>
<div className="flex min-w-0 flex-1 flex-col">
<span className="flex min-w-0 items-center gap-2">
<span className="truncate">
{profile?.name ?? `配置 #${id}`}
{profile?.is_default ? "(激活)" : ""}
</span>
</span>
{profile && <span className="truncate text-muted-foreground text-xs">{profile.model}</span>}
</div>
</ComboboxItem>
);
}}
</ComboboxList>
</ComboboxContent>
</Combobox>
{value.length === 0 ? (
<p className="text-muted-foreground text-xs">未指定配置时,任务跟随 Agent 或全局激活配置。</p>
) : (
<div className="flex flex-col divide-y rounded-lg border">
{value.map((id, index) => {
const profile = profilesByID.get(id);
return (
<div
key={id}
className="grid min-w-0 grid-cols-[auto_minmax(0,1fr)_auto] items-center gap-2 px-2.5 py-2 sm:flex"
>
<span className="w-5 shrink-0 text-center text-muted-foreground text-xs tabular-nums">{index + 1}</span>
<div className="min-w-0 sm:flex-1">
<p className="truncate font-medium text-sm">{profile?.name ?? `配置 #${id}`}</p>
<div className="flex min-w-0 items-center gap-2">
<p className="truncate text-muted-foreground text-xs">{profile?.model ?? "配置已不可用"}</p>
</div>
</div>
<ProfileRoleBadge index={index} currentIndex={currentIndex} />
<div className="col-span-2 col-start-2 flex shrink-0 items-center gap-1 justify-self-end sm:col-auto sm:justify-self-auto">
<Button
type="button"
variant="ghost"
size="icon-sm"
aria-label="上移配置"
onClick={() => move(index, -1)}
disabled={disabled ? true : index === 0}
>
<ArrowUpIcon data-icon="inline-start" />
</Button>
<Button
type="button"
variant="ghost"
size="icon-sm"
aria-label="下移配置"
onClick={() => move(index, 1)}
disabled={disabled ? true : index === value.length - 1}
>
<ArrowDownIcon data-icon="inline-start" />
</Button>
<Button
type="button"
variant="ghost"
size="icon-sm"
aria-label="移除配置"
onClick={() => remove(id)}
disabled={disabled}
>
<XIcon data-icon="inline-start" />
</Button>
</div>
</div>
);
})}
</div>
)}
{onActiveProfileChange && value.length > 0 && (
<div className="flex flex-col items-stretch gap-2 sm:flex-row sm:items-center sm:justify-between sm:gap-3">
<span className="font-medium text-sm">当前配置</span>
<Select
value={activeProfileId && value.includes(activeProfileId) ? activeProfileId : value[0]}
onValueChange={onActiveProfileChange}
disabled={disabled}
>
<SelectTrigger size="sm" className="w-full sm:min-w-48 sm:max-w-64">
<SelectValue />
</SelectTrigger>
<SelectContent>
<SelectGroup>
{value.map((id) => {
const profile = profilesByID.get(id);
return (
<SelectItem key={id} value={id} disabled={!profile}>
{profile?.name ?? `配置 #${id}`}
</SelectItem>
);
})}
</SelectGroup>
</SelectContent>
</Select>
</div>
)}
</div>
);
}