feat(llm): add profile failover and pool health
This commit is contained in:
pphua committed
1 parent
d4b21f0bdb
commit
18926c129f
15 files changed
+2507
-761
No files matched your search
+27
-22
@@ -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) {
|
||||
|
||||
@@ -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
@@ -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
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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
|
||||
}
|
||||
@@ -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")
|
||||
}
|
||||
}
|
||||
+136
-39
@@ -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))
|
||||
|
||||
@@ -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)
|
||||
}
|
||||
}
|
||||
@@ -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)
|
||||
}
|
||||
}
|
||||
@@ -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
|
||||
}
|
||||
@@ -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>
|
||||
);
|
||||
}
|
||||
Reference in new issue
Block a user