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