easyai-ai-gateway/apps/api/internal/store/runner_policies.go

package store

import (
	"context"
	"encoding/json"
	"strings"
	"time"

	"github.com/jackc/pgx/v5"
)

const runnerPolicyColumns = `
id::text, policy_key, name, COALESCE(description, ''), failover_policy,
hard_stop_policy, priority_demote_policy, metadata, status, created_at, updated_at`

type RunnerPolicyInput struct {
	PolicyKey            string         `json:"policyKey"`
	Name                 string         `json:"name"`
	Description          string         `json:"description"`
	FailoverPolicy       map[string]any `json:"failoverPolicy"`
	HardStopPolicy       map[string]any `json:"hardStopPolicy"`
	PriorityDemotePolicy map[string]any `json:"priorityDemotePolicy"`
	Metadata             map[string]any `json:"metadata"`
	Status               string         `json:"status"`
}

type runnerPolicyScanner interface {
	Scan(dest ...any) error
}

func (s *Store) GetActiveRunnerPolicy(ctx context.Context) (RunnerPolicy, error) {
	item, err := scanRunnerPolicy(s.pool.QueryRow(ctx, `
SELECT `+runnerPolicyColumns+`
FROM gateway_runner_policies
ORDER BY CASE WHEN policy_key = 'default-runner-v1' THEN 0 ELSE 1 END,
         CASE WHEN status = 'active' THEN 0 ELSE 1 END,
         updated_at DESC
LIMIT 1`))
	if err != nil {
		if err == pgx.ErrNoRows || IsUndefinedDatabaseObject(err) {
			return defaultRunnerPolicy(), nil
		}
		return RunnerPolicy{}, err
	}
	return item, nil
}

func (s *Store) UpsertDefaultRunnerPolicy(ctx context.Context, input RunnerPolicyInput) (RunnerPolicy, error) {
	input = normalizeRunnerPolicyInput(input)
	failoverPolicy, _ := json.Marshal(emptyObjectIfNil(input.FailoverPolicy))
	hardStopPolicy, _ := json.Marshal(emptyObjectIfNil(input.HardStopPolicy))
	priorityDemotePolicy, _ := json.Marshal(emptyObjectIfNil(input.PriorityDemotePolicy))
	metadata, _ := json.Marshal(emptyObjectIfNil(input.Metadata))
	return scanRunnerPolicy(s.pool.QueryRow(ctx, `
INSERT INTO gateway_runner_policies (
  policy_key, name, description, failover_policy, hard_stop_policy, priority_demote_policy, metadata, status
)
VALUES ($1, $2, NULLIF($3, ''), $4, $5, $6, $7, $8)
ON CONFLICT (policy_key) DO UPDATE
SET name = EXCLUDED.name,
    description = EXCLUDED.description,
    failover_policy = EXCLUDED.failover_policy,
    hard_stop_policy = EXCLUDED.hard_stop_policy,
    priority_demote_policy = EXCLUDED.priority_demote_policy,
    metadata = EXCLUDED.metadata,
    status = EXCLUDED.status,
    updated_at = now()
RETURNING `+runnerPolicyColumns,
		input.PolicyKey, input.Name, input.Description, failoverPolicy, hardStopPolicy, priorityDemotePolicy, metadata, input.Status,
	))
}

func scanRunnerPolicy(scanner runnerPolicyScanner) (RunnerPolicy, error) {
	var item RunnerPolicy
	var failoverPolicy []byte
	var hardStopPolicy []byte
	var priorityDemotePolicy []byte
	var metadata []byte
	if err := scanner.Scan(
		&item.ID,
		&item.PolicyKey,
		&item.Name,
		&item.Description,
		&failoverPolicy,
		&hardStopPolicy,
		&priorityDemotePolicy,
		&metadata,
		&item.Status,
		&item.CreatedAt,
		&item.UpdatedAt,
	); err != nil {
		return RunnerPolicy{}, err
	}
	item.FailoverPolicy = decodeObject(failoverPolicy)
	item.HardStopPolicy = decodeObject(hardStopPolicy)
	item.PriorityDemotePolicy = decodeObject(priorityDemotePolicy)
	item.Metadata = decodeObject(metadata)
	return item, nil
}

func normalizeRunnerPolicyInput(input RunnerPolicyInput) RunnerPolicyInput {
	input.PolicyKey = strings.TrimSpace(input.PolicyKey)
	if input.PolicyKey == "" {
		input.PolicyKey = "default-runner-v1"
	}
	input.Name = strings.TrimSpace(input.Name)
	if input.Name == "" {
		input.Name = "默认全局调度策略"
	}
	input.Description = strings.TrimSpace(input.Description)
	input.Status = strings.TrimSpace(input.Status)
	if input.Status == "" {
		input.Status = "active"
	}
	return input
}

func defaultRunnerPolicy() RunnerPolicy {
	now := time.Now()
	return RunnerPolicy{
		PolicyKey:            "default-runner-v1",
		Name:                 "默认全局调度策略",
		Description:          "控制多个候选平台之间的故障切换；模型运行策略只可覆盖 failoverPolicy，不能覆盖 hardStopPolicy。",
		FailoverPolicy:       defaultRunnerFailoverPolicy(),
		HardStopPolicy:       defaultRunnerHardStopPolicy(),
		PriorityDemotePolicy: defaultRunnerPriorityDemotePolicy(),
		Metadata:             map[string]any{"source": "code-default"},
		Status:               "active",
		CreatedAt:            now,
		UpdatedAt:            now,
	}
}

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},
		"keywords":    []any{"timeout", "network", "rate_limit", "overloaded", "temporarily_unavailable", "server_error", "429", "5xx"},
	}
}

func defaultRunnerFailoverPolicy() map[string]any {
	return map[string]any{
		"enabled":            true,
		"maxPlatforms":       99,
		"maxDurationSeconds": 600,
		"allowCategories":    []any{"network", "timeout", "stream_error", "rate_limit", "provider_5xx", "provider_overloaded", "auth_error"},
		"denyCategories":     []any{"request_error", "unsupported_model", "user_permission", "insufficient_balance"},
		"allowCodes":         []any{"auth_failed", "invalid_api_key", "missing_credentials"},
		"allowKeywords":      []any{"timeout", "network", "rate_limit", "overloaded", "temporarily_unavailable", "server_error", "auth_failed", "invalid_api_key", "missing_credentials", "unauthorized", "forbidden", "429", "5xx"},
		"denyKeywords":       []any{"invalid_parameter", "missing required", "bad request"},
		"allowStatusCodes":   []any{401, 403, 408, 429, 500, 502, 503, 504},
		"denyStatusCodes":    []any{},
	}
}

func defaultRunnerHardStopPolicy() map[string]any {
	return map[string]any{
		"enabled":     true,
		"categories":  []any{"request_error", "unsupported_model", "user_permission", "insufficient_balance"},
		"codes":       []any{"bad_request", "invalid_request", "invalid_parameter", "missing_required", "unsupported_kind", "unsupported_model", "insufficient_balance", "permission_denied"},
		"statusCodes": []any{},
		"keywords":    []any{"invalid_parameter", "missing required", "bad request", "insufficient balance"},
	}
}