BCAI/easyai-ai-gateway
3
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

feat: add priority demotion controls

Browse Source
This commit is contained in:
wangbo
2026-05-12 18:43:20 +08:00
parent 3917b84b5d
commit 98abd247d6
22 changed files with 917 additions and 85 deletions
Download Patch File Download Diff File Expand all files Collapse all files
apps/api/internal/store/candidates_test.go
+11
View File
@@ -59,3 +59,14 @@ func TestRuntimeCandidateSortingAvoidsFullCandidatesButKeepsFallback(t *testing.
t.Fatalf("expected full high-priority candidate to remain as avoided fallback, got %+v", candidates)
}
}
func TestDefaultRunnerPriorityDemotePolicyUsesAutoMode(t *testing.T) {
policy := defaultRunnerPriorityDemotePolicy()
if _, ok := policy["demoteStep"]; ok {
t.Fatal("priority demotion should be automatic and must not expose a demoteStep policy")
}
if policy["enabled"] != true {
t.Fatalf("expected default priority demotion to stay enabled, got %+v", policy["enabled"])
}
}
apps/api/internal/store/postgres.go
+22 -4
View File
@@ -3,6 +3,7 @@ package store
import (
"context"
"crypto/rand"
"database/sql"
"encoding/base64"
"encoding/json"
"errors"
@@ -67,6 +68,8 @@ type Platform struct {
AuthType string `json:"authType"`
Status string `json:"status"`
Priority int `json:"priority"`
DynamicPriority *int `json:"dynamicPriority,omitempty"`
EffectivePriority int `json:"effectivePriority"`
DefaultPricingMode string `json:"defaultPricingMode"`
DefaultDiscountFactor float64 `json:"defaultDiscountFactor"`
PricingRuleSetID string `json:"pricingRuleSetId,omitempty"`
@@ -508,13 +511,14 @@ type TaskParamPreprocessingLog struct {
func (s *Store) ListPlatforms(ctx context.Context) ([]Platform, error) {
rows, err := s.pool.Query(ctx, `
SELECT id::text, provider, platform_key, name, COALESCE(internal_name, ''), COALESCE(base_url, ''), auth_type, status, priority,
SELECT id::text, provider, platform_key, name, COALESCE(internal_name, ''), COALESCE(base_url, ''), auth_type, status,
priority, dynamic_priority, COALESCE(dynamic_priority, priority),
default_pricing_mode, default_discount_factor::float8, COALESCE(pricing_rule_set_id::text, ''),
config, credentials, retry_policy, rate_limit_policy,
COALESCE(to_char(cooldown_until AT TIME ZONE 'UTC', 'YYYY-MM-DD"T"HH24:MI:SS.MS"Z"'), ''),
created_at, updated_at
FROM integration_platforms
ORDER BY priority ASC, created_at DESC`)
ORDER BY COALESCE(dynamic_priority, priority) ASC, priority ASC, created_at DESC`)
if err != nil {
return nil, err
}
@@ -527,6 +531,7 @@ ORDER BY priority ASC, created_at DESC`)
var credentialsBytes []byte
var retryPolicyBytes []byte
var rateLimitPolicyBytes []byte
var dynamicPriority sql.NullInt64
if err := rows.Scan(
&platform.ID,
&platform.Provider,
@@ -537,6 +542,8 @@ ORDER BY priority ASC, created_at DESC`)
&platform.AuthType,
&platform.Status,
&platform.Priority,
&dynamicPriority,
&platform.EffectivePriority,
&platform.DefaultPricingMode,
&platform.DefaultDiscountFactor,
&platform.PricingRuleSetID,
@@ -550,6 +557,7 @@ ORDER BY priority ASC, created_at DESC`)
); err != nil {
return nil, err
}
platform.DynamicPriority = intPointerFromNull(dynamicPriority)
platform.Config = decodeObject(configBytes)
platform.CredentialsPreview = maskCredentialsPreview(credentialsBytes)
platform.RetryPolicy = decodeObject(retryPolicyBytes)
@@ -578,6 +586,7 @@ func (s *Store) CreatePlatform(ctx context.Context, input CreatePlatformInput) (
var credentialsResultBytes []byte
var retryPolicyBytes []byte
var rateLimitPolicyBytes []byte
var dynamicPriority sql.NullInt64
err := s.pool.QueryRow(ctx, `
INSERT INTO integration_platforms (
provider, platform_key, name, internal_name, base_url, auth_type, credentials, config,
@@ -588,7 +597,8 @@ VALUES (
$1, COALESCE(NULLIF($2, ''), 'platform_' || replace(gen_random_uuid()::text, '-', '')), $3, NULLIF($4, ''), $5, $6, $7, $8,
$9, $10, NULLIF($11, '')::uuid, $12, $13, $14
)
RETURNING id::text, provider, platform_key, name, COALESCE(internal_name, ''), COALESCE(base_url, ''), auth_type, status, priority,
RETURNING id::text, provider, platform_key, name, COALESCE(internal_name, ''), COALESCE(base_url, ''), auth_type, status,
priority, dynamic_priority, COALESCE(dynamic_priority, priority),
default_pricing_mode, default_discount_factor::float8, COALESCE(pricing_rule_set_id::text, ''),
config, credentials, retry_policy, rate_limit_policy,
COALESCE(to_char(cooldown_until AT TIME ZONE 'UTC', 'YYYY-MM-DD"T"HH24:MI:SS.MS"Z"'), ''),
@@ -606,6 +616,8 @@ RETURNING id::text, provider, platform_key, name, COALESCE(internal_name, ''), C
&platform.AuthType,
&platform.Status,
&platform.Priority,
&dynamicPriority,
&platform.EffectivePriority,
&platform.DefaultPricingMode,
&platform.DefaultDiscountFactor,
&platform.PricingRuleSetID,
@@ -620,6 +632,7 @@ RETURNING id::text, provider, platform_key, name, COALESCE(internal_name, ''), C
if err != nil {
return Platform{}, err
}
platform.DynamicPriority = intPointerFromNull(dynamicPriority)
platform.Config = decodeObject(configBytes)
platform.CredentialsPreview = maskCredentialsPreview(credentialsResultBytes)
platform.RetryPolicy = decodeObject(retryPolicyBytes)
@@ -650,6 +663,7 @@ func (s *Store) UpdatePlatform(ctx context.Context, id string, input CreatePlatf
var credentialsResultBytes []byte
var retryPolicyBytes []byte
var rateLimitPolicyBytes []byte
var dynamicPriority sql.NullInt64
err := s.pool.QueryRow(ctx, `
UPDATE integration_platforms
SET provider = $2,
@@ -672,7 +686,8 @@ SET provider = $2,
rate_limit_policy = $15,
updated_at = now()
WHERE id = $1::uuid
RETURNING id::text, provider, platform_key, name, COALESCE(internal_name, ''), COALESCE(base_url, ''), auth_type, status, priority,
RETURNING id::text, provider, platform_key, name, COALESCE(internal_name, ''), COALESCE(base_url, ''), auth_type, status,
priority, dynamic_priority, COALESCE(dynamic_priority, priority),
default_pricing_mode, default_discount_factor::float8, COALESCE(pricing_rule_set_id::text, ''),
config, credentials, retry_policy, rate_limit_policy,
COALESCE(to_char(cooldown_until AT TIME ZONE 'UTC', 'YYYY-MM-DD"T"HH24:MI:SS.MS"Z"'), ''),
@@ -702,6 +717,8 @@ RETURNING id::text, provider, platform_key, name, COALESCE(internal_name, ''), C
&platform.AuthType,
&platform.Status,
&platform.Priority,
&dynamicPriority,
&platform.EffectivePriority,
&platform.DefaultPricingMode,
&platform.DefaultDiscountFactor,
&platform.PricingRuleSetID,
@@ -716,6 +733,7 @@ RETURNING id::text, provider, platform_key, name, COALESCE(internal_name, ''), C
if err != nil {
return Platform{}, err
}
platform.DynamicPriority = intPointerFromNull(dynamicPriority)
platform.Config = decodeObject(configBytes)
platform.CredentialsPreview = maskCredentialsPreview(credentialsResultBytes)
platform.RetryPolicy = decodeObject(retryPolicyBytes)
apps/api/internal/store/rate_limit_status.go
+159 -20
View File
@@ -2,8 +2,11 @@ package store
import (
"context"
"database/sql"
"sort"
"strconv"
"strings"
"time"
)
type RateLimitMetricStatus struct {
@@ -17,30 +20,53 @@ type RateLimitMetricStatus struct {
}
type ModelRateLimitStatus struct {
PlatformModelID string `json:"platformModelId"`
PlatformID string `json:"platformId"`
PlatformName string `json:"platformName"`
Provider string `json:"provider"`
ModelName string `json:"modelName"`
ProviderModelName string `json:"providerModelName,omitempty"`
ModelAlias string `json:"modelAlias,omitempty"`
DisplayName string `json:"displayName"`
ModelType []string `json:"modelType"`
Enabled bool `json:"enabled"`
RateLimitPolicy map[string]any `json:"rateLimitPolicy,omitempty"`
PlatformCooldownUntil string `json:"platformCooldownUntil,omitempty"`
ModelCooldownUntil string `json:"modelCooldownUntil,omitempty"`
Concurrent RateLimitMetricStatus `json:"concurrent"`
QueuedTasks float64 `json:"queuedTasks"`
RPM RateLimitMetricStatus `json:"rpm"`
TPM RateLimitMetricStatus `json:"tpm"`
LoadRatio float64 `json:"loadRatio"`
PlatformModelID string `json:"platformModelId"`
PlatformID string `json:"platformId"`
PlatformName string `json:"platformName"`
Provider string `json:"provider"`
PlatformPriority int `json:"platformPriority"`
PlatformDynamicPriority *int `json:"platformDynamicPriority,omitempty"`
PlatformEffectivePriority int `json:"platformEffectivePriority"`
ModelName string `json:"modelName"`
ProviderModelName string `json:"providerModelName,omitempty"`
ModelAlias string `json:"modelAlias,omitempty"`
DisplayName string `json:"displayName"`
ModelType []string `json:"modelType"`
Enabled bool `json:"enabled"`
RateLimitPolicy map[string]any `json:"rateLimitPolicy,omitempty"`
PlatformCooldownUntil string `json:"platformCooldownUntil,omitempty"`
ModelCooldownUntil string `json:"modelCooldownUntil,omitempty"`
Concurrent RateLimitMetricStatus `json:"concurrent"`
QueuedTasks float64 `json:"queuedTasks"`
RPM RateLimitMetricStatus `json:"rpm"`
TPM RateLimitMetricStatus `json:"tpm"`
LoadRatio float64 `json:"loadRatio"`
RecentPriorityDemotions []PriorityDemotionRecord `json:"recentPriorityDemotions,omitempty"`
}
type PriorityDemotionRecord struct {
ID string `json:"id"`
TaskID string `json:"taskId"`
PlatformID string `json:"platformId"`
PlatformModelID string `json:"platformModelId,omitempty"`
Reason string `json:"reason,omitempty"`
ErrorCode string `json:"errorCode,omitempty"`
ErrorMessage string `json:"errorMessage,omitempty"`
Category string `json:"category,omitempty"`
StatusCode int `json:"statusCode,omitempty"`
PolicySource string `json:"policySource,omitempty"`
Policy string `json:"policy,omitempty"`
PolicyRule string `json:"policyRule,omitempty"`
MatchedValue string `json:"matchedValue,omitempty"`
DynamicPriority int `json:"dynamicPriority,omitempty"`
CreatedAt time.Time `json:"createdAt"`
}
func (s *Store) ListModelRateLimitStatuses(ctx context.Context) ([]ModelRateLimitStatus, error) {
rows, err := s.pool.Query(ctx, `
SELECT m.id::text, m.platform_id::text, p.name, p.provider,
m.model_name, COALESCE(NULLIF(m.provider_model_name, ''), m.model_name), COALESCE(m.model_alias, ''),
SELECT m.id::text, m.platform_id::text, p.name, p.provider,
p.priority, p.dynamic_priority, COALESCE(p.dynamic_priority, p.priority),
m.model_name, COALESCE(NULLIF(m.provider_model_name, ''), m.model_name), COALESCE(m.model_alias, ''),
m.model_type, m.display_name, m.enabled,
p.rate_limit_policy, COALESCE(rp.rate_limit_policy, '{}'::jsonb), COALESCE(NULLIF(m.runtime_policy_override, '{}'::jsonb), b.runtime_policy_override, '{}'::jsonb), m.rate_limit_policy,
COALESCE(to_char(p.cooldown_until AT TIME ZONE 'UTC', 'YYYY-MM-DD"T"HH24:MI:SS.MS"Z"'), ''),
@@ -119,6 +145,7 @@ ORDER BY p.priority ASC, m.model_name ASC`)
var runtimePolicyBytes []byte
var runtimeOverrideBytes []byte
var modelPolicyBytes []byte
var platformDynamicPriority sql.NullInt64
var platformCooldownUntil string
var modelCooldownUntil string
var concurrentCurrent float64
@@ -134,6 +161,9 @@ ORDER BY p.priority ASC, m.model_name ASC`)
&item.PlatformID,
&item.PlatformName,
&item.Provider,
&item.PlatformPriority,
&platformDynamicPriority,
&item.PlatformEffectivePriority,
&item.ModelName,
&item.ProviderModelName,
&item.ModelAlias,
@@ -157,6 +187,7 @@ ORDER BY p.priority ASC, m.model_name ASC`)
); err != nil {
return nil, err
}
item.PlatformDynamicPriority = intPointerFromNull(platformDynamicPriority)
item.ModelType = decodeStringArray(modelTypeBytes)
policy := effectiveModelRateLimitPolicy(
decodeObject(platformPolicyBytes),
@@ -177,6 +208,13 @@ ORDER BY p.priority ASC, m.model_name ASC`)
if err := rows.Err(); err != nil {
return nil, err
}
demotions, err := s.listRecentPriorityDemotionsByPlatform(ctx, items, 10)
if err != nil {
return nil, err
}
for index := range items {
items[index].RecentPriorityDemotions = demotions[items[index].PlatformID]
}
sort.SliceStable(items, func(i, j int) bool {
if items[i].LoadRatio == items[j].LoadRatio {
return strings.ToLower(items[i].DisplayName) < strings.ToLower(items[j].DisplayName)
@@ -186,6 +224,91 @@ ORDER BY p.priority ASC, m.model_name ASC`)
return items, nil
}
func (s *Store) listRecentPriorityDemotionsByPlatform(ctx context.Context, statuses []ModelRateLimitStatus, limit int) (map[string][]PriorityDemotionRecord, error) {
out := map[string][]PriorityDemotionRecord{}
if limit <= 0 || len(statuses) == 0 {
return out, nil
}
seen := map[string]bool{}
platformIDs := make([]string, 0, len(statuses))
for _, status := range statuses {
platformID := strings.TrimSpace(status.PlatformID)
if platformID == "" || seen[platformID] {
continue
}
seen[platformID] = true
platformIDs = append(platformIDs, platformID)
}
if len(platformIDs) == 0 {
return out, nil
}
rows, err := s.pool.Query(ctx, `
SELECT id::text, task_id::text, COALESCE(message, ''), payload, created_at
FROM (
SELECT e.*,
row_number() OVER (
PARTITION BY e.payload->>'platformId'
ORDER BY e.created_at DESC, e.seq DESC
) AS demotion_rank
FROM gateway_task_events e
WHERE e.event_type = 'task.policy.priority_demoted'
AND e.payload->>'platformId' = ANY($1::text[])
) ranked
WHERE demotion_rank <= $2
ORDER BY payload->>'platformId' ASC, created_at DESC`, platformIDs, limit)
if err != nil {
return nil, err
}
defer rows.Close()
for rows.Next() {
var id string
var taskID string
var message string
var payloadBytes []byte
var createdAt time.Time
if err := rows.Scan(&id, &taskID, &message, &payloadBytes, &createdAt); err != nil {
return nil, err
}
record := priorityDemotionRecordFromEventPayload(id, taskID, message, decodeObject(payloadBytes), createdAt)
if record.PlatformID == "" {
continue
}
out[record.PlatformID] = append(out[record.PlatformID], record)
}
return out, rows.Err()
}
func priorityDemotionRecordFromEventPayload(id string, taskID string, message string, payload map[string]any, createdAt time.Time) PriorityDemotionRecord {
errorMessage := stringValue(payload["errorMessage"])
if errorMessage == "" {
errorMessage = stringValue(payload["message"])
}
if errorMessage == "" {
errorMessage = strings.TrimSpace(message)
}
errorCode := stringValue(payload["errorCode"])
if errorCode == "" {
errorCode = stringValue(payload["code"])
}
return PriorityDemotionRecord{
ID: id,
TaskID: taskID,
PlatformID: stringValue(payload["platformId"]),
PlatformModelID: stringValue(payload["platformModelId"]),
Reason: stringValue(payload["reason"]),
ErrorCode: errorCode,
ErrorMessage: errorMessage,
Category: stringValue(payload["category"]),
StatusCode: intValue(payload["statusCode"]),
PolicySource: stringValue(payload["policySource"]),
Policy: stringValue(payload["policy"]),
PolicyRule: stringValue(payload["policyRule"]),
MatchedValue: stringValue(payload["matchedValue"]),
DynamicPriority: intValue(payload["dynamicPriority"]),
CreatedAt: createdAt,
}
}
func effectiveModelRateLimitPolicy(platformPolicy map[string]any, runtimePolicy map[string]any, runtimeOverride map[string]any, modelPolicy map[string]any) map[string]any {
policy := platformPolicy
if hasRateLimitRules(runtimePolicy) {
@@ -279,3 +402,19 @@ func floatValue(value any) float64 {
return 0
}
}
func intValue(value any) int {
switch typed := value.(type) {
case int:
return typed
case int64:
return int(typed)
case float64:
return int(typed)
case string:
parsed, _ := strconv.Atoi(strings.TrimSpace(typed))
return parsed
default:
return 0
}
}
apps/api/internal/store/rate_limit_status_test.go
+32 -1
View File
@@ -1,6 +1,9 @@
package store
import "testing"
import (
"testing"
"time"
)
func TestEffectiveModelRateLimitPolicyTreatsModelRulesAsAuthoritative(t *testing.T) {
policy := effectiveModelRateLimitPolicy(
@@ -30,3 +33,31 @@ func TestEffectiveModelRateLimitPolicyTreatsModelRulesAsAuthoritative(t *testing
t.Fatalf("expected missing model tpm limit to mean unlimited, got %v", got)
}
}
func TestPriorityDemotionRecordFromEventPayloadKeepsReason(t *testing.T) {
createdAt := time.Date(2026, 5, 12, 9, 30, 0, 0, time.UTC)
record := priorityDemotionRecordFromEventPayload("event-1", "task-1", "fallback message", map[string]any{
"platformId": "platform-1",
"platformModelId": "platform-model-1",
"reason": "priority_demote_policy",
"errorCode": "rate_limit",
"errorMessage": "upstream 429 rate limit",
"category": "rate_limit",
"statusCode": float64(429),
"policySource": "gateway_runner_policies.priority_demote_policy",
"policy": "priorityDemotePolicy",
"policyRule": "categories",
"matchedValue": "rate_limit",
"dynamicPriority": float64(1511),
}, createdAt)
if record.Reason != "priority_demote_policy" || record.ErrorMessage != "upstream 429 rate limit" {
t.Fatalf("expected demotion reason and error message to survive, got %+v", record)
}
if record.StatusCode != 429 || record.DynamicPriority != 1511 {
t.Fatalf("expected numeric demotion metadata, got %+v", record)
}
if !record.CreatedAt.Equal(createdAt) {
t.Fatalf("expected createdAt %s, got %s", createdAt, record.CreatedAt)
}
}
apps/api/internal/store/runner_policies.go
-1
View File
@@ -134,7 +134,6 @@ func defaultRunnerPolicy() RunnerPolicy {
func defaultRunnerPriorityDemotePolicy() map[string]any {
return map[string]any{
"enabled": true,
"demoteStep": 100,
"categories": []any{"network", "timeout", "stream_error", "rate_limit", "provider_5xx", "provider_overloaded"},
"codes": []any{"network", "timeout", "stream_read_error", "rate_limit", "server_error", "overloaded"},
"statusCodes": []any{408, 429, 500, 502, 503, 504},
apps/api/internal/store/runtime_policies.go
+65 -17
View File
@@ -2,8 +2,10 @@ package store
import (
"context"
"database/sql"
"encoding/json"
"strings"
"time"
"github.com/jackc/pgx/v5"
)
@@ -28,6 +30,14 @@ type runtimePolicyScanner interface {
Scan(dest ...any) error
}
type PlatformDynamicPriorityState struct {
PlatformID string `json:"platformId"`
Priority int `json:"priority"`
DynamicPriority *int `json:"dynamicPriority,omitempty"`
EffectivePriority int `json:"effectivePriority"`
UpdatedAt time.Time `json:"updatedAt"`
}
func (s *Store) ListRuntimePolicySets(ctx context.Context) ([]RuntimePolicySet, error) {
rows, err := s.pool.Query(ctx, `SELECT `+runtimePolicyColumns+` FROM model_runtime_policy_sets ORDER BY policy_key ASC`)
if err != nil {
@@ -149,26 +159,64 @@ WHERE id = $1::uuid`, platformModelID, cooldownSeconds)
return err
}
func (s *Store) DemoteCandidatePlatformPriority(ctx context.Context, platformID string, demoteStep int) error {
func (s *Store) DemoteCandidatePlatformPriority(ctx context.Context, platformID string) (int, error) {
if strings.TrimSpace(platformID) == "" {
return 0, nil
}
var dynamicPriority int
err := s.pool.QueryRow(ctx, `
UPDATE integration_platforms target
SET dynamic_priority = COALESCE((
SELECT MAX(COALESCE(peer.dynamic_priority, peer.priority))
FROM integration_platforms peer
WHERE peer.deleted_at IS NULL
), target.priority) + 1,
updated_at = now()
WHERE target.id = $1::uuid
RETURNING dynamic_priority`, platformID).Scan(&dynamicPriority)
return dynamicPriority, err
}
func (s *Store) UpdatePlatformDynamicPriority(ctx context.Context, platformID string, dynamicPriority *int) (PlatformDynamicPriorityState, error) {
if strings.TrimSpace(platformID) == "" {
return PlatformDynamicPriorityState{}, pgx.ErrNoRows
}
value := 0
reset := dynamicPriority == nil
if dynamicPriority != nil {
value = *dynamicPriority
}
return scanPlatformDynamicPriorityState(s.pool.QueryRow(ctx, `
UPDATE integration_platforms
SET dynamic_priority = CASE WHEN $2::boolean THEN priority ELSE $3::int END,
updated_at = now()
WHERE id = $1::uuid
AND deleted_at IS NULL
RETURNING id::text, priority, dynamic_priority, COALESCE(dynamic_priority, priority), updated_at`, platformID, reset, value))
}
func scanPlatformDynamicPriorityState(scanner runtimePolicyScanner) (PlatformDynamicPriorityState, error) {
var item PlatformDynamicPriorityState
var dynamicPriority sql.NullInt64
if err := scanner.Scan(
&item.PlatformID,
&item.Priority,
&dynamicPriority,
&item.EffectivePriority,
&item.UpdatedAt,
); err != nil {
return PlatformDynamicPriorityState{}, err
}
item.DynamicPriority = intPointerFromNull(dynamicPriority)
return item, nil
}
func intPointerFromNull(value sql.NullInt64) *int {
if !value.Valid {
return nil
}
if demoteStep <= 0 {
demoteStep = 100
}
_, err := s.pool.Exec(ctx, `
UPDATE integration_platforms target
SET dynamic_priority = GREATEST(
COALESCE(target.dynamic_priority, target.priority),
COALESCE((
SELECT MAX(COALESCE(peer.dynamic_priority, peer.priority))
FROM integration_platforms peer
WHERE peer.deleted_at IS NULL
), target.priority) + $2::int
),
updated_at = now()
WHERE target.id = $1::uuid`, platformID, demoteStep)
return err
converted := int(value.Int64)
return &converted
}
func scanRuntimePolicySet(scanner runtimePolicyScanner) (RuntimePolicySet, error) {