Files
sub2api/backend/internal/service/upstream_billing_probe.go
T
李建琦 6d655c9903
Release / update-version (push) Has been cancelled
Release / build-frontend (push) Has been cancelled
Release / release (push) Has been cancelled
Release / sync-version-file (push) Has been cancelled
CI / shell (push) Canceled after 0s
CI / test (push) Canceled after 0s
CI / frontend (push) Canceled after 0s
CI / golangci-lint (push) Canceled after 0s
Security Scan / backend-security (push) Canceled after 0s
Security Scan / frontend-security (push) Canceled after 0s
Sub2API v1.0 - AI API 网关(二开初始版本,基于上游 Wei-Shaw/sub2api)
2026-08-21 18:30:13 +08:00

1150 lines
43 KiB
Go

package service
import (
"bytes"
"context"
"database/sql"
"encoding/json"
"errors"
"fmt"
"io"
"log/slog"
"math"
"math/rand/v2"
"net/http"
"net/url"
"sort"
"strconv"
"strings"
"sync"
"time"
infraerrors "github.com/Wei-Shaw/sub2api/internal/pkg/errors"
"github.com/Wei-Shaw/sub2api/internal/pkg/logger"
"github.com/Wei-Shaw/sub2api/internal/pkg/tlsfingerprint"
"github.com/google/uuid"
"golang.org/x/sync/errgroup"
"golang.org/x/sync/singleflight"
)
const (
// These values live in accounts.extra so PR2 does not require a schema migration.
UpstreamBillingProbeExtraKey = "upstream_billing_probe"
UpstreamBillingProbeEnabledExtraKey = "upstream_billing_probe_enabled"
UpstreamBillingRateSyncEnabledExtraKey = "upstream_billing_rate_sync_enabled"
upstreamBillingProbeDefaultIntervalMinutes = 30
upstreamBillingProbeMinIntervalMinutes = 5
upstreamBillingProbeMaxIntervalMinutes = 24 * 60
upstreamBillingProbeCycleInterval = time.Minute
upstreamBillingProbeRequestTimeout = 10 * time.Second
upstreamBillingProbeMaxBodyBytes = 64 * 1024
upstreamBillingProbeMaxPerCycle = 20
upstreamBillingProbeConcurrency = 4
upstreamBillingProbeMaxDelay = 24 * time.Hour
// unsupported 账号的重探间隔倍数:上游不是 sub2api 中转就不会突然长出
// /v1/sub2api/billing,按常规 interval 重排只会持续占满每周期
// upstreamBillingProbeMaxPerCycle 个名额。
upstreamBillingProbeUnsupportedDelayFactor = 8
upstreamBillingProbeAccountRateScale = 10000.0
upstreamBillingProbeLeaderLockKey = "upstream:billing:probe:leader"
upstreamBillingProbeLeaderLockTTL = 2 * time.Minute
)
// UpstreamBillingProbeMaxBatchSize limits one manual batch and one runner cycle.
const UpstreamBillingProbeMaxBatchSize = upstreamBillingProbeMaxPerCycle
// upstreamBillingRateSyncMaxMultiplier bounds the value the automatic
// write-back may push into accounts.rate_multiplier.
//
// No other code path bounds that column from above — admins may type any
// non-negative number and the only ceiling is the DECIMAL(10,4) column itself
// (999999.9999). That ceiling is meaningless as a guard: rate_multiplier
// scales the per-request account cost that feeds quota_used, so a single
// declared 999999 would exhaust any account quota on the first request and
// poison cost reporting. 100 is picked as a deliberately generous bound: it is
// two orders of magnitude above the 1.0 default and far above any plausible
// upstream resale markup, so no legitimate declaration is rejected while an
// absurd or hostile one cannot reach the quota control plane unattended.
// It only constrains the automatic path; manual edits keep their old range.
const upstreamBillingRateSyncMaxMultiplier = 100.0
var (
ErrUpstreamBillingProbeUnavailable = infraerrors.ServiceUnavailable(
"UPSTREAM_BILLING_PROBE_UNAVAILABLE", "upstream billing probe is unavailable",
)
ErrUpstreamBillingProbeAccountInvalid = infraerrors.BadRequest(
"UPSTREAM_BILLING_PROBE_ACCOUNT_INVALID", "account is not an API key account",
)
ErrUpstreamBillingProbeIdentityChanged = infraerrors.Conflict(
"UPSTREAM_BILLING_PROBE_IDENTITY_CHANGED", "account identity changed during upstream billing probe; retry the probe",
)
ErrUpstreamBillingRateSyncBulkConflict = infraerrors.Conflict(
"UPSTREAM_BILLING_RATE_SYNC_BULK_CONFLICT",
"account rate multiplier cannot be changed in bulk while upstream billing rate sync is enabled",
)
ErrUpstreamBillingRateSyncConflict = infraerrors.Conflict(
"UPSTREAM_BILLING_RATE_SYNC_CONFLICT",
"account rate multiplier cannot be changed while upstream billing rate sync is enabled",
)
)
const (
UpstreamBillingProbeStatusOK = "ok"
UpstreamBillingProbeStatusUnsupported = "unsupported"
UpstreamBillingProbeStatusFailed = "failed"
)
// UpstreamBillingProbeSettings controls the periodic probe runner.
type UpstreamBillingProbeSettings struct {
Enabled bool `json:"enabled"`
IntervalMinutes int `json:"interval_minutes"`
}
// UpstreamBillingProbeSnapshot is persisted in accounts.extra. Data is kept as
// a sanitized map so future response fields do not require a database change.
type UpstreamBillingProbeSnapshot struct {
Status string `json:"status"`
Data map[string]any `json:"data,omitempty"`
ReceivedAt *time.Time `json:"received_at,omitempty"`
FreshUntil *time.Time `json:"fresh_until,omitempty"`
LastAttemptAt time.Time `json:"last_attempt_at"`
NextProbeAt time.Time `json:"next_probe_at"`
FailureCount int `json:"failure_count,omitempty"`
HTTPStatus int `json:"http_status,omitempty"`
LastError string `json:"last_error,omitempty"`
// SyncedRateMultiplier records the value this probe wrote into
// accounts.rate_multiplier. It is only set when the account opted into rate
// sync and the declared value passed the write-back range check, so the
// stored snapshot always answers "did this probe move the account rate, and
// to what" without a separate history table.
SyncedRateMultiplier *float64 `json:"synced_rate_multiplier,omitempty"`
}
// UpstreamBillingProbeResult is returned by manual probe endpoints.
type UpstreamBillingProbeResult struct {
AccountID int64 `json:"account_id"`
Snapshot *UpstreamBillingProbeSnapshot `json:"snapshot,omitempty"`
Error string `json:"error,omitempty"`
}
type upstreamBillingProbeResponse struct {
Object string `json:"object"`
SchemaVersion int `json:"schema_version"`
BillingScope string `json:"billing_scope"`
GroupRateMultiplier *float64 `json:"group_rate_multiplier"`
UserRateMultiplier *float64 `json:"user_rate_multiplier"`
ResolvedRateMultiplier *float64 `json:"resolved_rate_multiplier"`
PeakRateEnabled *bool `json:"peak_rate_enabled"`
PeakStart *string `json:"peak_start"`
PeakEnd *string `json:"peak_end"`
PeakRateMultiplier *float64 `json:"peak_rate_multiplier"`
AppliedPeakMultiplier *float64 `json:"applied_peak_multiplier"`
EffectiveRateMultiplier *float64 `json:"effective_rate_multiplier"`
Timezone *string `json:"timezone"`
ObservedAt string `json:"observed_at"`
}
// GetUpstreamBillingProbeSettings returns defaults when the setting is absent.
func (s *SettingService) GetUpstreamBillingProbeSettings(ctx context.Context) (*UpstreamBillingProbeSettings, error) {
defaults := defaultUpstreamBillingProbeSettings()
if s == nil || s.settingRepo == nil {
return defaults, nil
}
value, err := s.settingRepo.GetValue(ctx, SettingKeyUpstreamBillingProbeSettings)
if err != nil {
if errors.Is(err, ErrSettingNotFound) {
return defaults, nil
}
return nil, fmt.Errorf("get upstream billing probe settings: %w", err)
}
if strings.TrimSpace(value) == "" {
return defaults, nil
}
settings := *defaults
if err := json.Unmarshal([]byte(value), &settings); err != nil {
return nil, fmt.Errorf("parse upstream billing probe settings: %w", err)
}
if settings.IntervalMinutes == 0 {
settings.IntervalMinutes = defaults.IntervalMinutes
}
normalizeUpstreamBillingProbeSettings(&settings)
return &settings, nil
}
// SetUpstreamBillingProbeSettings validates and persists the runner settings.
func (s *SettingService) SetUpstreamBillingProbeSettings(ctx context.Context, settings *UpstreamBillingProbeSettings) error {
if s == nil || s.settingRepo == nil {
return fmt.Errorf("setting repository is unavailable")
}
if settings == nil {
return infraerrors.BadRequest("INVALID_UPSTREAM_BILLING_PROBE_SETTINGS", "settings cannot be nil")
}
if settings.IntervalMinutes < upstreamBillingProbeMinIntervalMinutes || settings.IntervalMinutes > upstreamBillingProbeMaxIntervalMinutes {
return infraerrors.BadRequest(
"INVALID_UPSTREAM_BILLING_PROBE_INTERVAL",
fmt.Sprintf("interval_minutes must be between %d and %d", upstreamBillingProbeMinIntervalMinutes, upstreamBillingProbeMaxIntervalMinutes),
)
}
normalizeUpstreamBillingProbeSettings(settings)
data, err := json.Marshal(settings)
if err != nil {
return fmt.Errorf("marshal upstream billing probe settings: %w", err)
}
return s.settingRepo.Set(ctx, SettingKeyUpstreamBillingProbeSettings, string(data))
}
func defaultUpstreamBillingProbeSettings() *UpstreamBillingProbeSettings {
return &UpstreamBillingProbeSettings{Enabled: true, IntervalMinutes: upstreamBillingProbeDefaultIntervalMinutes}
}
func normalizeUpstreamBillingProbeSettings(settings *UpstreamBillingProbeSettings) {
if settings.IntervalMinutes < upstreamBillingProbeMinIntervalMinutes {
settings.IntervalMinutes = upstreamBillingProbeMinIntervalMinutes
}
if settings.IntervalMinutes > upstreamBillingProbeMaxIntervalMinutes {
settings.IntervalMinutes = upstreamBillingProbeMaxIntervalMinutes
}
}
// UpstreamBillingProbeService discovers a remote Sub2API billing snapshot.
type UpstreamBillingProbeService struct {
accountRepo AccountRepository
accountTestService *AccountTestService
settingService *SettingService
parentCtx context.Context
parentCancel context.CancelFunc
wg sync.WaitGroup
mu sync.Mutex
started bool
stopped bool
cycleMu sync.Mutex
probeGroup singleflight.Group
probeSlots chan struct{}
now func() time.Time
lockCache LeaderLockCache
db *sql.DB
instanceID string
}
type upstreamBillingProbeSnapshotWriter interface {
UpdateUpstreamBillingProbeSnapshot(context.Context, *Account, *UpstreamBillingProbeSnapshot, *float64) error
}
type upstreamBillingProbeDueAccountLister interface {
ListDueUpstreamBillingProbeAccounts(context.Context, time.Time, int) ([]Account, error)
}
func NewUpstreamBillingProbeService(
accountRepo AccountRepository,
accountTestService *AccountTestService,
settingService *SettingService,
) *UpstreamBillingProbeService {
ctx, cancel := context.WithCancel(context.Background())
return &UpstreamBillingProbeService{
accountRepo: accountRepo,
accountTestService: accountTestService,
settingService: settingService,
parentCtx: ctx,
parentCancel: cancel,
probeSlots: make(chan struct{}, upstreamBillingProbeConcurrency),
now: time.Now,
instanceID: uuid.NewString(),
}
}
func (s *UpstreamBillingProbeService) SetLeaderLock(lockCache LeaderLockCache, db *sql.DB) {
if s == nil {
return
}
s.lockCache = lockCache
s.db = db
}
// ProvideUpstreamBillingProbeService starts the process-wide periodic runner.
func ProvideUpstreamBillingProbeService(
accountRepo AccountRepository,
accountTestService *AccountTestService,
settingService *SettingService,
lockCache LeaderLockCache,
db *sql.DB,
) *UpstreamBillingProbeService {
svc := NewUpstreamBillingProbeService(accountRepo, accountTestService, settingService)
svc.SetLeaderLock(lockCache, db)
svc.Start()
return svc
}
func (s *UpstreamBillingProbeService) Start() {
if s == nil {
return
}
s.mu.Lock()
if s.started || s.stopped {
s.mu.Unlock()
return
}
s.started = true
s.wg.Add(1)
s.mu.Unlock()
go s.runLoop()
}
func (s *UpstreamBillingProbeService) Stop() {
if s == nil {
return
}
s.mu.Lock()
if s.stopped {
s.mu.Unlock()
return
}
s.stopped = true
s.parentCancel()
s.mu.Unlock()
s.wg.Wait()
}
func (s *UpstreamBillingProbeService) runLoop() {
defer s.wg.Done()
_ = s.RunDue(s.parentCtx)
ticker := time.NewTicker(upstreamBillingProbeCycleInterval)
defer ticker.Stop()
for {
select {
case <-s.parentCtx.Done():
return
case <-ticker.C:
if err := s.RunDue(s.parentCtx); err != nil {
logger.LegacyPrintf("service.upstream_billing_probe", "run_due_failed: err=%v", err)
}
}
}
}
// RunDue executes at most one bounded batch of due accounts.
func (s *UpstreamBillingProbeService) RunDue(ctx context.Context) error {
if s == nil || s.accountRepo == nil {
return nil
}
s.cycleMu.Lock()
defer s.cycleMu.Unlock()
settings, err := s.getSettings(ctx)
if err != nil {
return err
}
if !settings.Enabled {
return nil
}
runRelease, acquired, lockErr := s.tryAcquireLeaderLock(ctx, upstreamBillingProbeLeaderLockKey)
if lockErr != nil {
return fmt.Errorf("acquire upstream billing probe leader lock: %w", lockErr)
}
if !acquired {
return nil
}
defer runRelease()
lockNow := time.Now()
cadenceRelease, acquired, lockErr := s.tryAcquireLeaderLock(ctx, upstreamBillingProbeLeaderLockKeyAt(lockNow))
if lockErr != nil {
return fmt.Errorf("acquire upstream billing probe cadence lock: %w", lockErr)
}
if !acquired {
return nil
}
defer releaseUpstreamBillingProbeLeaderLock(cadenceRelease, lockNow.Truncate(upstreamBillingProbeCycleInterval).Add(upstreamBillingProbeCycleInterval))
now := s.currentTime()
accounts, err := s.listDueAccounts(ctx, now)
if err != nil {
return fmt.Errorf("list enabled upstream billing probes: %w", err)
}
due := make([]Account, 0, len(accounts))
for i := range accounts {
account := accounts[i]
if !isUpstreamBillingProbeAccount(&account) || !account.IsActive() || !upstreamBillingProbeEnabled(&account) {
continue
}
snapshot := decodeUpstreamBillingProbeSnapshot(account.Extra)
if snapshot != nil && !snapshot.NextProbeAt.IsZero() && now.Before(snapshot.NextProbeAt) {
continue
}
due = append(due, account)
}
sort.SliceStable(due, func(i, j int) bool {
left := decodeUpstreamBillingProbeSnapshot(due[i].Extra)
right := decodeUpstreamBillingProbeSnapshot(due[j].Extra)
leftUnset := left == nil || left.NextProbeAt.IsZero()
rightUnset := right == nil || right.NextProbeAt.IsZero()
if leftUnset && rightUnset {
return due[i].ID < due[j].ID
}
if leftUnset {
return true
}
if rightUnset {
return false
}
return left.NextProbeAt.Before(right.NextProbeAt)
})
if len(due) > upstreamBillingProbeMaxPerCycle {
due = due[:upstreamBillingProbeMaxPerCycle]
}
var group errgroup.Group
for i := range due {
accountID := due[i].ID
group.Go(func() error {
if _, probeErr := s.probeScheduledAccount(ctx, accountID, settings.IntervalMinutes); probeErr != nil {
logger.LegacyPrintf("service.upstream_billing_probe", "probe_due_failed: account_id=%d err=%v", accountID, probeErr)
}
return nil
})
}
return group.Wait()
}
func (s *UpstreamBillingProbeService) listDueAccounts(ctx context.Context, now time.Time) ([]Account, error) {
if lister, ok := s.accountRepo.(upstreamBillingProbeDueAccountLister); ok {
return lister.ListDueUpstreamBillingProbeAccounts(ctx, now, upstreamBillingProbeMaxPerCycle)
}
// Non-production repositories and older adapters keep the generic path. The
// runner still truncates before issuing network requests.
return s.accountRepo.FindByExtraField(ctx, UpstreamBillingProbeEnabledExtraKey, true)
}
func (s *UpstreamBillingProbeService) getSettings(ctx context.Context) (*UpstreamBillingProbeSettings, error) {
if s.settingService == nil {
return defaultUpstreamBillingProbeSettings(), nil
}
return s.settingService.GetUpstreamBillingProbeSettings(ctx)
}
func (s *UpstreamBillingProbeService) GetSettings(ctx context.Context) (*UpstreamBillingProbeSettings, error) {
return s.getSettings(ctx)
}
func (s *UpstreamBillingProbeService) UpdateSettings(ctx context.Context, settings *UpstreamBillingProbeSettings) error {
if s == nil || s.settingService == nil {
return ErrUpstreamBillingProbeUnavailable
}
return s.settingService.SetUpstreamBillingProbeSettings(ctx, settings)
}
// ProbeAccount performs one manual or scheduled probe. Manual calls ignore both switches.
func (s *UpstreamBillingProbeService) ProbeAccount(ctx context.Context, accountID int64) (*UpstreamBillingProbeSnapshot, error) {
if s == nil || s.accountRepo == nil {
return nil, ErrUpstreamBillingProbeUnavailable
}
settings, err := s.getSettings(ctx)
if err != nil {
return nil, err
}
return s.probeAccount(ctx, accountID, settings.IntervalMinutes)
}
func (s *UpstreamBillingProbeService) probeAccount(ctx context.Context, accountID int64, intervalMinutes int) (*UpstreamBillingProbeSnapshot, error) {
return s.probeAccountWithMode(ctx, accountID, intervalMinutes, false)
}
func (s *UpstreamBillingProbeService) probeScheduledAccount(ctx context.Context, accountID int64, intervalMinutes int) (*UpstreamBillingProbeSnapshot, error) {
return s.probeAccountWithMode(ctx, accountID, intervalMinutes, true)
}
func (s *UpstreamBillingProbeService) probeAccountWithMode(ctx context.Context, accountID int64, intervalMinutes int, requireEnabled bool) (*UpstreamBillingProbeSnapshot, error) {
key := strconv.FormatInt(accountID, 10)
value, err, _ := s.probeGroup.Do(key, func() (any, error) {
select {
case s.probeSlots <- struct{}{}:
defer func() { <-s.probeSlots }()
case <-ctx.Done():
return nil, ctx.Err()
}
account, loadErr := s.accountRepo.GetByID(ctx, accountID)
if loadErr != nil {
return nil, loadErr
}
if !isUpstreamBillingProbeAccount(account) {
return nil, ErrUpstreamBillingProbeAccountInvalid
}
if requireEnabled {
if !account.IsActive() || !upstreamBillingProbeEnabled(account) {
return nil, nil
}
if snapshot := decodeUpstreamBillingProbeSnapshot(account.Extra); snapshot != nil &&
!snapshot.NextProbeAt.IsZero() && s.currentTime().Before(snapshot.NextProbeAt) {
return nil, nil
}
}
return s.probeLoadedAccount(ctx, account, intervalMinutes)
})
if err != nil {
return nil, err
}
if value == nil {
return nil, nil
}
snapshot, ok := value.(*UpstreamBillingProbeSnapshot)
if !ok {
return nil, fmt.Errorf("invalid upstream billing probe result")
}
return snapshot, nil
}
// ProbeAccounts performs a bounded manual batch with the same concurrency limit as the runner.
func (s *UpstreamBillingProbeService) ProbeAccounts(ctx context.Context, accountIDs []int64) []UpstreamBillingProbeResult {
if len(accountIDs) > upstreamBillingProbeMaxPerCycle {
accountIDs = accountIDs[:upstreamBillingProbeMaxPerCycle]
}
results := make([]UpstreamBillingProbeResult, len(accountIDs))
if s == nil || s.accountRepo == nil {
for i, accountID := range accountIDs {
results[i] = UpstreamBillingProbeResult{AccountID: accountID, Error: ErrUpstreamBillingProbeUnavailable.Error()}
}
return results
}
settings, settingsErr := s.getSettings(ctx)
if settingsErr != nil {
for i, accountID := range accountIDs {
results[i] = UpstreamBillingProbeResult{AccountID: accountID, Error: safeProbeError(settingsErr)}
}
return results
}
var group errgroup.Group
for i, accountID := range accountIDs {
i, accountID := i, accountID
results[i].AccountID = accountID
group.Go(func() error {
snapshot, err := s.probeAccount(ctx, accountID, settings.IntervalMinutes)
if err != nil {
results[i].Error = safeProbeError(err)
return nil
}
results[i].Snapshot = snapshot
return nil
})
}
_ = group.Wait()
return results
}
func upstreamBillingProbeLeaderLockKeyAt(now time.Time) string {
return fmt.Sprintf("%s:%d", upstreamBillingProbeLeaderLockKey, now.Unix()/int64(upstreamBillingProbeCycleInterval/time.Second))
}
func (s *UpstreamBillingProbeService) tryAcquireLeaderLock(ctx context.Context, key string) (func(), bool, error) {
lockCtx, cancel := context.WithTimeout(ctx, 2*time.Second)
defer cancel()
if s.lockCache != nil {
acquired, err := s.lockCache.TryAcquireLeaderLock(lockCtx, key, s.instanceID, upstreamBillingProbeLeaderLockTTL)
if err != nil {
return nil, false, err
}
if !acquired {
return nil, false, nil
}
return func() {
releaseCtx, releaseCancel := context.WithTimeout(context.Background(), 2*time.Second)
defer releaseCancel()
_ = s.lockCache.ReleaseLeaderLock(releaseCtx, key, s.instanceID)
}, true, nil
}
if s.db != nil {
return tryAcquireDBAdvisoryLockWithError(lockCtx, s.db, hashAdvisoryLockID(key))
}
return func() {}, true, nil
}
func releaseUpstreamBillingProbeLeaderLock(release func(), releaseAt time.Time) {
delay := time.Until(releaseAt)
if delay <= 0 {
release()
return
}
time.AfterFunc(delay, release)
}
func (s *UpstreamBillingProbeService) SetAccountEnabled(ctx context.Context, accountID int64, enabled bool) error {
if s == nil || s.accountRepo == nil {
return ErrUpstreamBillingProbeUnavailable
}
account, err := s.accountRepo.GetByID(ctx, accountID)
if err != nil {
return err
}
if !isUpstreamBillingProbeAccount(account) {
return ErrUpstreamBillingProbeAccountInvalid
}
updates := map[string]any{UpstreamBillingProbeEnabledExtraKey: enabled}
if !enabled {
updates[UpstreamBillingRateSyncEnabledExtraKey] = false
}
return s.accountRepo.UpdateExtra(ctx, accountID, updates)
}
func (s *UpstreamBillingProbeService) probeLoadedAccount(ctx context.Context, account *Account, intervalMinutes int) (*UpstreamBillingProbeSnapshot, error) {
now := s.currentTime().UTC()
if s.accountTestService == nil || s.accountTestService.httpUpstream == nil {
return s.persistProbeFailure(ctx, account, intervalMinutes, now, 0, "transport_unavailable", 0)
}
// 平台放宽后取数直读 credentials:所有 API-key 平台的密钥与自定义上游
// 统一存放在 credentials.api_key / credentials.base_url。
apiKey := account.GetCredential("api_key")
if apiKey == "" {
return s.persistProbeFailure(ctx, account, intervalMinutes, now, 0, "missing_api_key", 0)
}
baseURL := account.GetCredential("base_url")
if account.IsCNProvider() && account.IsAdaptiveAPIProtocol() {
baseURL = account.GetCNProtocolBaseURL(APIProtocolChatCompletions)
}
if account.Platform == PlatformOpenAI {
if baseURL == "" {
// 保持官方语义:OpenAI 账号无自定义 base 时探官方域(404 → unsupported)。
baseURL = "https://api.openai.com"
}
} else if upstreamBillingProbeTargetIsOfficialAPI(baseURL) {
// 其他平台 base_url 为空或指向官方 API 根域(前端创建时会把空值
// 填成官方默认域,且提供 us-east-1.api.x.ai 等官方区域预设)⇒
// 必无 /v1/sub2api/billing;不发请求,直接记 unsupported,避免
// 拿账号 Key 周期性请求官方域的不存在路径。
return s.persistProbeFailure(ctx, account, intervalMinutes, now, 0, "unsupported", 0)
}
normalizedBaseURL, err := s.accountTestService.validateUpstreamBaseURL(baseURL)
if err != nil {
return s.persistProbeFailure(ctx, account, intervalMinutes, now, 0, "invalid_base_url", 0)
}
proxyURL := ""
if account.ProxyID != nil {
if account.Proxy == nil {
return s.persistProbeFailure(ctx, account, intervalMinutes, now, 0, "proxy_unavailable", 0)
}
if account.Proxy.ID != *account.ProxyID {
return nil, ErrUpstreamBillingProbeIdentityChanged
}
proxyURL = account.Proxy.URL()
}
probeURL := buildOpenAIEndpointURL(normalizedBaseURL, "/v1/sub2api/billing")
probeCtx, cancel := context.WithTimeout(ctx, upstreamBillingProbeRequestTimeout)
defer cancel()
req, err := http.NewRequestWithContext(probeCtx, http.MethodGet, probeURL, bytes.NewReader(nil))
if err != nil {
return s.persistProbeFailure(ctx, account, intervalMinutes, now, 0, "request_build_failed", 0)
}
// OpenAI 账号保持官方 openai 传输画像;其他平台探测走默认画像。
profile := HTTPUpstreamProfileDefault
if account.Platform == PlatformOpenAI {
profile = HTTPUpstreamProfileOpenAI
}
reqCtx := WithHTTPUpstreamProfile(req.Context(), profile)
req = req.WithContext(WithHTTPUpstreamRedirectsDisabled(reqCtx))
req.Header.Set("Accept", "application/json")
req.Header.Set("Authorization", "Bearer "+apiKey)
account.ApplyHeaderOverrides(req.Header)
var tlsProfile *tlsfingerprint.Profile
if s.accountTestService.tlsFPProfileService != nil {
tlsProfile = s.accountTestService.tlsFPProfileService.ResolveTLSProfile(account)
}
resp, err := s.accountTestService.httpUpstream.DoWithTLS(req, proxyURL, account.ID, account.Concurrency, tlsProfile)
if err != nil {
return s.persistProbeFailure(ctx, account, intervalMinutes, now, 0, "request_failed", 0)
}
if resp == nil || resp.Body == nil {
return s.persistProbeFailure(ctx, account, intervalMinutes, now, 0, "empty_response", 0)
}
defer func() { _ = resp.Body.Close() }()
body, readErr := io.ReadAll(io.LimitReader(resp.Body, upstreamBillingProbeMaxBodyBytes+1))
if readErr != nil {
return s.persistProbeFailure(ctx, account, intervalMinutes, now, resp.StatusCode, "response_read_failed", retryAfter(resp.Header, now))
}
if len(body) > upstreamBillingProbeMaxBodyBytes {
return s.persistProbeFailure(ctx, account, intervalMinutes, now, resp.StatusCode, "response_too_large", retryAfter(resp.Header, now))
}
if resp.StatusCode == http.StatusNotFound || resp.StatusCode == http.StatusMethodNotAllowed {
return s.persistProbeFailure(ctx, account, intervalMinutes, now, resp.StatusCode, "unsupported", retryAfter(resp.Header, now))
}
if resp.StatusCode < 200 || resp.StatusCode >= 300 {
return s.persistProbeFailure(ctx, account, intervalMinutes, now, resp.StatusCode, "http_error", retryAfter(resp.Header, now))
}
data, err := parseUpstreamBillingProbeResponse(body)
if err != nil {
return s.persistProbeFailure(ctx, account, intervalMinutes, now, resp.StatusCode, "invalid_response", retryAfter(resp.Header, now))
}
snapshot := &UpstreamBillingProbeSnapshot{
Status: UpstreamBillingProbeStatusOK,
Data: data,
ReceivedAt: probeTimePtr(now),
FreshUntil: probeTimePtr(now.Add(2 * time.Duration(intervalMinutes) * time.Minute)),
LastAttemptAt: now,
NextProbeAt: now.Add(nextProbeDelay(intervalMinutes, 0)),
HTTPStatus: resp.StatusCode,
}
// 账号级值域与精度只在真要写回时才有影响:只观察上游声明、未开启同步的
// 账号不因声明值不适配 accounts.rate_multiplier 而被记成探测失败并进入
// 指数退避——探测本身成功了,原始声明照常存进快照供展示。
var syncRate *float64
previousRate := account.BillingRateMultiplier()
if upstreamBillingRateSyncEnabled(account) {
if value, valid := upstreamBillingProbeSyncRate(data); valid {
syncRate = &value
snapshot.SyncedRateMultiplier = &value
} else {
declared, _ := resolveAccountExtraNumber(data, "resolved_rate_multiplier")
slog.Warn("upstream_billing_rate_sync_rejected",
"source", "upstream_billing_probe",
"account_id", account.ID,
"declared_resolved_rate_multiplier", declared,
"max_rate_multiplier", upstreamBillingRateSyncMaxMultiplier,
"current_rate_multiplier", previousRate,
)
}
}
if err := s.updateSnapshot(ctx, account, snapshot, syncRate); err != nil {
return nil, err
}
if syncRate != nil {
// 写回是后台任务的裸 SQL,不经过管理端路由,因此不会产生 audit_logs 行。
// old_rate_multiplier 是本次探测开始时读到的值(写回的 CAS 不比对该列)。
slog.Info("upstream_billing_rate_sync_applied",
"source", "upstream_billing_probe",
"account_id", account.ID,
"old_rate_multiplier", previousRate,
"new_rate_multiplier", *syncRate,
)
}
return snapshot, nil
}
func (s *UpstreamBillingProbeService) persistProbeFailure(
ctx context.Context,
account *Account,
intervalMinutes int,
now time.Time,
statusCode int,
reason string,
retryAfterDuration time.Duration,
) (*UpstreamBillingProbeSnapshot, error) {
previous := decodeUpstreamBillingProbeSnapshot(account.Extra)
failureCount := 1
if previous != nil {
failureCount = previous.FailureCount + 1
}
status := UpstreamBillingProbeStatusFailed
delay := nextProbeDelay(intervalMinutes, retryAfterDuration)
if reason == "unsupported" {
status = UpstreamBillingProbeStatusUnsupported
delay = unsupportedProbeDelay(intervalMinutes, retryAfterDuration)
}
snapshot := &UpstreamBillingProbeSnapshot{
Status: status,
LastAttemptAt: now,
NextProbeAt: now.Add(delay),
FailureCount: failureCount,
HTTPStatus: statusCode,
LastError: reason,
}
if previous != nil {
snapshot.Data = previous.Data
snapshot.ReceivedAt = previous.ReceivedAt
snapshot.FreshUntil = previous.FreshUntil
if snapshot.FreshUntil == nil && previous.Status == UpstreamBillingProbeStatusOK && previous.ReceivedAt != nil {
snapshot.FreshUntil = probeTimePtr(previous.ReceivedAt.Add(2 * time.Duration(intervalMinutes) * time.Minute))
}
}
if err := s.updateSnapshot(ctx, account, snapshot, nil); err != nil {
return nil, err
}
return snapshot, nil
}
func (s *UpstreamBillingProbeService) updateSnapshot(
ctx context.Context,
account *Account,
snapshot *UpstreamBillingProbeSnapshot,
rateMultiplier *float64,
) error {
writer, ok := s.accountRepo.(upstreamBillingProbeSnapshotWriter)
if !ok {
return ErrUpstreamBillingProbeUnavailable
}
return writer.UpdateUpstreamBillingProbeSnapshot(ctx, account, snapshot, rateMultiplier)
}
func parseUpstreamBillingProbeResponse(body []byte) (map[string]any, error) {
var response upstreamBillingProbeResponse
if err := json.Unmarshal(body, &response); err != nil {
return nil, err
}
if response.Object != "sub2api.key_billing" || response.SchemaVersion != 1 || response.BillingScope != "token" {
return nil, fmt.Errorf("unexpected billing response schema")
}
if response.GroupRateMultiplier == nil || response.ResolvedRateMultiplier == nil ||
response.PeakRateEnabled == nil || response.EffectiveRateMultiplier == nil {
return nil, fmt.Errorf("incomplete billing response")
}
for _, value := range []float64{
*response.GroupRateMultiplier,
*response.ResolvedRateMultiplier,
*response.EffectiveRateMultiplier,
} {
if value < 0 || math.IsNaN(value) || math.IsInf(value, 0) {
return nil, fmt.Errorf("invalid billing multiplier")
}
}
if response.UserRateMultiplier != nil && (*response.UserRateMultiplier < 0 || math.IsNaN(*response.UserRateMultiplier) || math.IsInf(*response.UserRateMultiplier, 0)) {
return nil, fmt.Errorf("invalid user billing multiplier")
}
expectedResolved := *response.GroupRateMultiplier
if response.UserRateMultiplier != nil {
expectedResolved = *response.UserRateMultiplier
}
if !equalBillingMultiplier(*response.ResolvedRateMultiplier, expectedResolved) {
return nil, fmt.Errorf("inconsistent resolved billing multiplier")
}
observedAt, err := time.Parse(time.RFC3339Nano, response.ObservedAt)
if err != nil || observedAt.IsZero() {
return nil, fmt.Errorf("invalid observed_at")
}
data := map[string]any{
"object": response.Object,
"schema_version": response.SchemaVersion,
"billing_scope": response.BillingScope,
"group_rate_multiplier": *response.GroupRateMultiplier,
"resolved_rate_multiplier": *response.ResolvedRateMultiplier,
"peak_rate_enabled": *response.PeakRateEnabled,
"effective_rate_multiplier": *response.EffectiveRateMultiplier,
"observed_at": observedAt.UTC().Format(time.RFC3339Nano),
}
if response.UserRateMultiplier != nil {
data["user_rate_multiplier"] = *response.UserRateMultiplier
}
if *response.PeakRateEnabled {
if response.PeakStart == nil || response.PeakEnd == nil || response.Timezone == nil ||
response.PeakRateMultiplier == nil || response.AppliedPeakMultiplier == nil ||
*response.PeakStart == "" || *response.PeakEnd == "" || *response.Timezone == "" ||
*response.PeakRateMultiplier < 0 || *response.AppliedPeakMultiplier < 0 ||
math.IsNaN(*response.PeakRateMultiplier) || math.IsInf(*response.PeakRateMultiplier, 0) ||
math.IsNaN(*response.AppliedPeakMultiplier) || math.IsInf(*response.AppliedPeakMultiplier, 0) {
return nil, fmt.Errorf("incomplete peak billing response")
}
data["peak_start"] = *response.PeakStart
data["peak_end"] = *response.PeakEnd
data["peak_rate_multiplier"] = *response.PeakRateMultiplier
data["applied_peak_multiplier"] = *response.AppliedPeakMultiplier
data["timezone"] = *response.Timezone
}
appliedPeak, ok := upstreamBillingPeakMultiplierAt(data, observedAt)
if !ok {
return nil, fmt.Errorf("invalid peak billing response")
}
if response.PeakRateEnabled != nil && *response.PeakRateEnabled {
if !equalBillingMultiplier(*response.AppliedPeakMultiplier, appliedPeak) {
return nil, fmt.Errorf("inconsistent applied peak multiplier")
}
} else if response.AppliedPeakMultiplier != nil && !equalBillingMultiplier(*response.AppliedPeakMultiplier, 1) {
return nil, fmt.Errorf("inconsistent applied peak multiplier")
}
if !equalBillingMultiplier(*response.EffectiveRateMultiplier, *response.ResolvedRateMultiplier*appliedPeak) {
return nil, fmt.Errorf("inconsistent effective billing multiplier")
}
return data, nil
}
func upstreamBillingRateAt(data map[string]any, now time.Time) (float64, bool) {
if scope, _ := data["billing_scope"].(string); scope != "token" {
return 0, false
}
base, ok := resolveAccountExtraNumber(data, "resolved_rate_multiplier")
if !ok || base < 0 || math.IsNaN(base) || math.IsInf(base, 0) {
return 0, false
}
appliedPeak, ok := upstreamBillingPeakMultiplierAt(data, now)
if !ok {
return 0, false
}
base *= appliedPeak
if math.IsNaN(base) || math.IsInf(base, 0) {
return 0, false
}
return base, true
}
// upstreamBillingProbeSyncRate converts the declared multiplier into the value
// the automatic write-back may store in accounts.rate_multiplier, at the
// precision that column supports (DECIMAL(10,4)).
//
// It reads resolved_rate_multiplier, not effective_rate_multiplier: the
// effective value folds in the peak coefficient that happened to apply at the
// instant of the probe, so writing it would freeze one probe cycle's peak (or
// off-peak) factor into a static column, while display and scheduling
// recompute the peak factor for the current time through upstreamBillingRateAt.
//
// The accepted range is deliberately narrower than the column:
// - 0 is rejected. accountCost multiplies the request cost by this value, so
// an upstream-declared 0 would stop quota_used from ever growing and every
// admin-configured account quota and cost alert would silently stop
// working. Admins may still set 0 by hand; only the automatic path refuses.
// - anything above upstreamBillingRateSyncMaxMultiplier is rejected.
//
// A rejected declaration leaves the current multiplier untouched; the probe
// still records an OK snapshot carrying the raw declaration for display.
func upstreamBillingProbeSyncRate(data map[string]any) (float64, bool) {
value, ok := resolveAccountExtraNumber(data, "resolved_rate_multiplier")
if !ok || math.IsNaN(value) || math.IsInf(value, 0) {
return 0, false
}
rounded := math.Round(value*upstreamBillingProbeAccountRateScale) / upstreamBillingProbeAccountRateScale
if rounded <= 0 || rounded > upstreamBillingRateSyncMaxMultiplier {
return 0, false
}
return rounded, true
}
func upstreamBillingPeakMultiplierAt(data map[string]any, now time.Time) (float64, bool) {
peakEnabled, ok := data["peak_rate_enabled"].(bool)
if !ok {
return 0, false
}
if !peakEnabled {
return 1, true
}
start, startOK := data["peak_start"].(string)
end, endOK := data["peak_end"].(string)
timezoneName, timezoneOK := data["timezone"].(string)
peakMultiplier, multiplierOK := resolveAccountExtraNumber(data, "peak_rate_multiplier")
startMinute, validStart := parseMinutes(start)
endMinute, validEnd := parseMinutes(end)
if !startOK || !endOK || !timezoneOK || !multiplierOK || !validStart || !validEnd ||
startMinute >= endMinute || peakMultiplier < 0 || math.IsNaN(peakMultiplier) || math.IsInf(peakMultiplier, 0) {
return 0, false
}
location, err := time.LoadLocation(timezoneName)
if err != nil {
return 0, false
}
local := now.In(location)
minute := local.Hour()*60 + local.Minute()
if minute >= startMinute && minute < endMinute {
return peakMultiplier, true
}
return 1, true
}
func equalBillingMultiplier(left, right float64) bool {
if math.IsNaN(left) || math.IsNaN(right) || math.IsInf(left, 0) || math.IsInf(right, 0) {
return false
}
scale := math.Max(1, math.Max(math.Abs(left), math.Abs(right)))
return math.Abs(left-right) <= 1e-9*scale
}
func decodeUpstreamBillingProbeSnapshot(extra map[string]any) *UpstreamBillingProbeSnapshot {
if extra == nil {
return nil
}
value, ok := extra[UpstreamBillingProbeExtraKey]
if !ok {
return nil
}
raw, err := json.Marshal(value)
if err != nil {
return nil
}
var snapshot UpstreamBillingProbeSnapshot
if err := json.Unmarshal(raw, &snapshot); err != nil || snapshot.Status == "" {
return nil
}
if snapshot.Status != UpstreamBillingProbeStatusOK &&
snapshot.Status != UpstreamBillingProbeStatusUnsupported &&
snapshot.Status != UpstreamBillingProbeStatusFailed {
return nil
}
return &snapshot
}
// IsUpstreamBillingProbeIdentity reports whether an account identity may opt
// in to the upstream billing probe. `/v1/sub2api/billing` is a key-scoped
// sub2api convention shared by the supported API-key platforms (including the
// CN providers, whose official-domain accounts are short-circuited to
// "unsupported" by upstreamBillingProbeTargetIsOfficialAPI).
// Non-sub2api upstreams return 404 and the snapshot records "unsupported".
// Only AccountTypeAPIKey is in scope. OAuth/Bedrock hold no static API key to
// present at all; AccountTypeUpstream (antigravity relay accounts) does carry
// a base_url plus a static api_key, but it is deliberately left out of the
// current supported set. New antigravity relay accounts are created with
// type=apikey by the admin form, so only pre-existing type=upstream rows
// cannot turn the probe on.
func IsUpstreamBillingProbeIdentity(platform, accountType string) bool {
if accountType != AccountTypeAPIKey {
return false
}
switch platform {
case PlatformOpenAI, PlatformAnthropic, PlatformGemini, PlatformAntigravity, PlatformGrok,
PlatformKimi, PlatformZhipu, PlatformDeepseek:
return true
default:
return false
}
}
func isUpstreamBillingProbeAccount(account *Account) bool {
return account != nil && IsUpstreamBillingProbeIdentity(account.Platform, account.Type)
}
// upstreamBillingProbeOfficialAPIDomains lists the root domains of official
// provider APIs. The create form fills empty base_url values with official
// defaults (and offers official regional presets like us-east-1.api.x.ai),
// so probing them would send the account key to an official API path that
// cannot exist. Matching is by registrable root domain — exact host or any
// subdomain, after stripping the port and a trailing DNS dot — because no
// third-party sub2api relay can live under these domains, while custom
// relays (the only targets that can answer /v1/sub2api/billing) always do
// probe. OpenAI-platform accounts never reach this check: they keep the
// upstream-official behavior of probing api.openai.com.
// ollama.com is a first-class configuration here (Ollama Cloud accounts are
// platform openai/anthropic with base_url https://ollama.com/v1), and it is
// an official provider API just like the rest, so it belongs on this list.
// CN provider domains (moonshot.cn / kimi.com / bigmodel.cn / deepseek.com)
// serve the same role: official APIs that can never host /v1/sub2api/billing,
// so their accounts short-circuit to "unsupported" without a request.
var upstreamBillingProbeOfficialAPIDomains = []string{
"anthropic.com",
"googleapis.com",
"x.ai",
"grok.com",
"openai.com",
"ollama.com",
"moonshot.cn",
"kimi.com",
"bigmodel.cn",
"deepseek.com",
}
func upstreamBillingProbeTargetIsOfficialAPI(baseURL string) bool {
baseURL = strings.TrimSpace(baseURL)
if baseURL == "" {
return true
}
parsed, err := url.Parse(baseURL)
if err != nil {
return false
}
host := strings.TrimSuffix(strings.ToLower(parsed.Hostname()), ".")
if host == "" {
return true
}
for _, domain := range upstreamBillingProbeOfficialAPIDomains {
if host == domain || strings.HasSuffix(host, "."+domain) {
return true
}
}
return false
}
func upstreamBillingProbeEnabled(account *Account) bool {
if account == nil || account.Extra == nil {
return false
}
enabled, ok := account.Extra[UpstreamBillingProbeEnabledExtraKey].(bool)
return ok && enabled
}
// upstreamBillingRateSyncEnabled is the probe-side pre-filter deciding whether
// a rate is even proposed for write-back. It is a necessary condition, not the
// authority: the repository CAS re-checks both switches against the row it
// updates, so a switch flipped between load and write can never sneak a rate in.
func upstreamBillingRateSyncEnabled(account *Account) bool {
if account == nil || account.Extra == nil {
return false
}
enabled, ok := account.Extra[UpstreamBillingRateSyncEnabledExtraKey].(bool)
return ok && enabled && upstreamBillingProbeEnabled(account)
}
func (s *UpstreamBillingProbeService) currentTime() time.Time {
if s != nil && s.now != nil {
return s.now()
}
return time.Now()
}
func nextProbeDelay(intervalMinutes int, retryAfterDuration time.Duration) time.Duration {
interval := time.Duration(intervalMinutes) * time.Minute
if interval < upstreamBillingProbeMinIntervalMinutes*time.Minute {
interval = upstreamBillingProbeMinIntervalMinutes * time.Minute
}
if interval > upstreamBillingProbeMaxDelay {
interval = upstreamBillingProbeMaxDelay
}
jitterRange := interval / 5
if jitterRange > 5*time.Minute {
jitterRange = 5 * time.Minute
}
if jitterRange > 0 {
interval += time.Duration(rand.Int64N(int64(jitterRange)*2+1)) - jitterRange
}
if retryAfterDuration > interval {
// Retry-After is an explicit upstream instruction; do not shorten it
// with the local maximum delay.
return retryAfterDuration
}
if interval > upstreamBillingProbeMaxDelay {
return upstreamBillingProbeMaxDelay
}
return interval
}
// unsupportedProbeDelay 拉长 unsupported 账号的重探间隔,让无效候选自然退出
// 热队列,不再和真正接入 sub2api 的中转账号抢每周期的探测名额。
// 仍按 upstreamBillingProbeMaxDelay 封顶,保证上游后来接入 sub2api 时最迟一天
// 内会被重新发现;base 本身已达上限(例如 Retry-After 明确要求更久)时原样返回,
// 不缩短上游指令。
func unsupportedProbeDelay(intervalMinutes int, retryAfterDuration time.Duration) time.Duration {
base := nextProbeDelay(intervalMinutes, retryAfterDuration)
if base >= upstreamBillingProbeMaxDelay {
return base
}
stretched := base * upstreamBillingProbeUnsupportedDelayFactor
if stretched > upstreamBillingProbeMaxDelay {
return upstreamBillingProbeMaxDelay
}
return stretched
}
func retryAfter(header http.Header, now time.Time) time.Duration {
value := strings.TrimSpace(header.Get("Retry-After"))
if value == "" {
return 0
}
if seconds, err := strconv.Atoi(value); err == nil && seconds > 0 {
return time.Duration(seconds) * time.Second
}
if at, err := http.ParseTime(value); err == nil {
if delay := at.Sub(now); delay > 0 {
return delay
}
}
return 0
}
func probeTimePtr(value time.Time) *time.Time {
return &value
}
func safeProbeError(err error) string {
if err == nil {
return ""
}
if errors.Is(err, ErrUpstreamBillingProbeAccountInvalid) {
return ErrUpstreamBillingProbeAccountInvalid.Error()
}
if errors.Is(err, ErrUpstreamBillingProbeUnavailable) {
return ErrUpstreamBillingProbeUnavailable.Error()
}
return "probe_failed"
}