upstream_billing_probe.go 44 KB

12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758596061626364656667686970717273747576777879808182838485868788899091929394959697989910010110210310410510610710810911011111211311411511611711811912012112212312412512612712812913013113213313413513613713813914014114214314414514614714814915015115215315415515615715815916016116216316416516616716816917017117217317417517617717817918018118218318418518618718818919019119219319419519619719819920020120220320420520620720820921021121221321421521621721821922022122222322422522622722822923023123223323423523623723823924024124224324424524624724824925025125225325425525625725825926026126226326426526626726826927027127227327427527627727827928028128228328428528628728828929029129229329429529629729829930030130230330430530630730830931031131231331431531631731831932032132232332432532632732832933033133233333433533633733833934034134234334434534634734834935035135235335435535635735835936036136236336436536636736836937037137237337437537637737837938038138238338438538638738838939039139239339439539639739839940040140240340440540640740840941041141241341441541641741841942042142242342442542642742842943043143243343443543643743843944044144244344444544644744844945045145245345445545645745845946046146246346446546646746846947047147247347447547647747847948048148248348448548648748848949049149249349449549649749849950050150250350450550650750850951051151251351451551651751851952052152252352452552652752852953053153253353453553653753853954054154254354454554654754854955055155255355455555655755855956056156256356456556656756856957057157257357457557657757857958058158258358458558658758858959059159259359459559659759859960060160260360460560660760860961061161261361461561661761861962062162262362462562662762862963063163263363463563663763863964064164264364464564664764864965065165265365465565665765865966066166266366466566666766866967067167267367467567667767867968068168268368468568668768868969069169269369469569669769869970070170270370470570670770870971071171271371471571671771871972072172272372472572672772872973073173273373473573673773873974074174274374474574674774874975075175275375475575675775875976076176276376476576676776876977077177277377477577677777877978078178278378478578678778878979079179279379479579679779879980080180280380480580680780880981081181281381481581681781881982082182282382482582682782882983083183283383483583683783883984084184284384484584684784884985085185285385485585685785885986086186286386486586686786886987087187287387487587687787887988088188288388488588688788888989089189289389489589689789889990090190290390490590690790890991091191291391491591691791891992092192292392492592692792892993093193293393493593693793893994094194294394494594694794894995095195295395495595695795895996096196296396496596696796896997097197297397497597697797897998098198298398498598698798898999099199299399499599699799899910001001100210031004100510061007100810091010101110121013101410151016101710181019102010211022102310241025102610271028102910301031103210331034103510361037103810391040104110421043104410451046104710481049105010511052105310541055105610571058105910601061106210631064106510661067106810691070107110721073107410751076107710781079108010811082108310841085108610871088108910901091109210931094109510961097109810991100110111021103110411051106110711081109111011111112111311141115111611171118111911201121112211231124112511261127112811291130113111321133113411351136113711381139114011411142114311441145114611471148114911501151115211531154115511561157115811591160116111621163116411651166116711681169117011711172117311741175117611771178
  1. package service
  2. import (
  3. "bytes"
  4. "context"
  5. "database/sql"
  6. "encoding/json"
  7. "errors"
  8. "fmt"
  9. "io"
  10. "log/slog"
  11. "math"
  12. "math/rand/v2"
  13. "net/http"
  14. "net/url"
  15. "sort"
  16. "strconv"
  17. "strings"
  18. "sync"
  19. "time"
  20. infraerrors "github.com/Wei-Shaw/sub2api/internal/pkg/errors"
  21. "github.com/Wei-Shaw/sub2api/internal/pkg/logger"
  22. "github.com/Wei-Shaw/sub2api/internal/pkg/tlsfingerprint"
  23. "github.com/google/uuid"
  24. "golang.org/x/sync/errgroup"
  25. "golang.org/x/sync/singleflight"
  26. )
  27. const (
  28. // These values live in accounts.extra so PR2 does not require a schema migration.
  29. UpstreamBillingProbeExtraKey = "upstream_billing_probe"
  30. UpstreamBillingProbeEnabledExtraKey = "upstream_billing_probe_enabled"
  31. UpstreamBillingRateSyncEnabledExtraKey = "upstream_billing_rate_sync_enabled"
  32. upstreamBillingProbeDefaultIntervalMinutes = 30
  33. upstreamBillingProbeMinIntervalMinutes = 5
  34. upstreamBillingProbeMaxIntervalMinutes = 24 * 60
  35. upstreamBillingProbeCycleInterval = time.Minute
  36. upstreamBillingProbeRequestTimeout = 10 * time.Second
  37. upstreamBillingProbeMaxBodyBytes = 64 * 1024
  38. upstreamBillingProbeMaxPerCycle = 20
  39. upstreamBillingProbeConcurrency = 4
  40. upstreamBillingProbeMaxDelay = 24 * time.Hour
  41. // unsupported 账号的重探间隔倍数:上游不是 sub2api 中转就不会突然长出
  42. // /v1/sub2api/billing,按常规 interval 重排只会持续占满每周期
  43. // upstreamBillingProbeMaxPerCycle 个名额。
  44. upstreamBillingProbeUnsupportedDelayFactor = 8
  45. upstreamBillingProbeAccountRateScale = 10000.0
  46. upstreamBillingProbeLeaderLockKey = "upstream:billing:probe:leader"
  47. upstreamBillingProbeLeaderLockTTL = 2 * time.Minute
  48. )
  49. // UpstreamBillingProbeMaxBatchSize limits one manual batch and one runner cycle.
  50. const UpstreamBillingProbeMaxBatchSize = upstreamBillingProbeMaxPerCycle
  51. // upstreamBillingRateSyncMaxMultiplier bounds the value the automatic
  52. // write-back may push into accounts.rate_multiplier.
  53. //
  54. // No other code path bounds that column from above — admins may type any
  55. // non-negative number and the only ceiling is the DECIMAL(10,4) column itself
  56. // (999999.9999). That ceiling is meaningless as a guard: rate_multiplier
  57. // scales the per-request account cost that feeds quota_used, so a single
  58. // declared 999999 would exhaust any account quota on the first request and
  59. // poison cost reporting. 100 is picked as a deliberately generous bound: it is
  60. // two orders of magnitude above the 1.0 default and far above any plausible
  61. // upstream resale markup, so no legitimate declaration is rejected while an
  62. // absurd or hostile one cannot reach the quota control plane unattended.
  63. // It only constrains the automatic path; manual edits keep their old range.
  64. const upstreamBillingRateSyncMaxMultiplier = 100.0
  65. var (
  66. ErrUpstreamBillingProbeUnavailable = infraerrors.ServiceUnavailable(
  67. "UPSTREAM_BILLING_PROBE_UNAVAILABLE", "upstream billing probe is unavailable",
  68. )
  69. ErrUpstreamBillingProbeAccountInvalid = infraerrors.BadRequest(
  70. "UPSTREAM_BILLING_PROBE_ACCOUNT_INVALID", "account is not an API key account",
  71. )
  72. ErrUpstreamBillingProbeIdentityChanged = infraerrors.Conflict(
  73. "UPSTREAM_BILLING_PROBE_IDENTITY_CHANGED", "account identity changed during upstream billing probe; retry the probe",
  74. )
  75. ErrUpstreamBillingRateSyncBulkConflict = infraerrors.Conflict(
  76. "UPSTREAM_BILLING_RATE_SYNC_BULK_CONFLICT",
  77. "account rate multiplier cannot be changed in bulk while upstream billing rate sync is enabled",
  78. )
  79. ErrUpstreamBillingRateSyncConflict = infraerrors.Conflict(
  80. "UPSTREAM_BILLING_RATE_SYNC_CONFLICT",
  81. "account rate multiplier cannot be changed while upstream billing rate sync is enabled",
  82. )
  83. )
  84. const (
  85. UpstreamBillingProbeStatusOK = "ok"
  86. UpstreamBillingProbeStatusUnsupported = "unsupported"
  87. UpstreamBillingProbeStatusFailed = "failed"
  88. )
  89. // UpstreamBillingProbeSettings controls the periodic probe runner.
  90. type UpstreamBillingProbeSettings struct {
  91. Enabled bool `json:"enabled"`
  92. IntervalMinutes int `json:"interval_minutes"`
  93. }
  94. // UpstreamBillingProbeSnapshot is persisted in accounts.extra. Data is kept as
  95. // a sanitized map so future response fields do not require a database change.
  96. type UpstreamBillingProbeSnapshot struct {
  97. Status string `json:"status"`
  98. Data map[string]any `json:"data,omitempty"`
  99. ReceivedAt *time.Time `json:"received_at,omitempty"`
  100. FreshUntil *time.Time `json:"fresh_until,omitempty"`
  101. LastAttemptAt time.Time `json:"last_attempt_at"`
  102. NextProbeAt time.Time `json:"next_probe_at"`
  103. FailureCount int `json:"failure_count,omitempty"`
  104. HTTPStatus int `json:"http_status,omitempty"`
  105. LastError string `json:"last_error,omitempty"`
  106. // SyncedRateMultiplier records the value this probe wrote into
  107. // accounts.rate_multiplier. It is only set when the account opted into rate
  108. // sync and the declared value passed the write-back range check, so the
  109. // stored snapshot always answers "did this probe move the account rate, and
  110. // to what" without a separate history table.
  111. SyncedRateMultiplier *float64 `json:"synced_rate_multiplier,omitempty"`
  112. }
  113. // UpstreamBillingProbeResult is returned by manual probe endpoints.
  114. type UpstreamBillingProbeResult struct {
  115. AccountID int64 `json:"account_id"`
  116. Snapshot *UpstreamBillingProbeSnapshot `json:"snapshot,omitempty"`
  117. Error string `json:"error,omitempty"`
  118. }
  119. // UpstreamBillingRateSnapshotItem is the compact representation used by the
  120. // account table's background refresh. It intentionally excludes credentials,
  121. // runtime counters, and usage data from the response.
  122. type UpstreamBillingRateSnapshotItem struct {
  123. AccountID int64 `json:"account_id"`
  124. Snapshot *UpstreamBillingProbeSnapshot `json:"snapshot"`
  125. }
  126. // BuildUpstreamBillingRateSnapshotItems projects account rows into the
  127. // read-only payload used by the rate refresh endpoint. Decode snapshots here
  128. // so malformed or legacy extra data is handled consistently with probe logic.
  129. func BuildUpstreamBillingRateSnapshotItems(accounts []Account) []UpstreamBillingRateSnapshotItem {
  130. items := make([]UpstreamBillingRateSnapshotItem, 0, len(accounts))
  131. for _, account := range accounts {
  132. var snapshot *UpstreamBillingProbeSnapshot
  133. // The billing endpoint is supported by every API-key platform; limiting
  134. // this projection to OpenAI would make the background refresh erase the
  135. // other platforms' persisted snapshots from the table.
  136. if account.Type == AccountTypeAPIKey {
  137. snapshot = decodeUpstreamBillingProbeSnapshot(account.Extra)
  138. }
  139. items = append(items, UpstreamBillingRateSnapshotItem{
  140. AccountID: account.ID,
  141. Snapshot: snapshot,
  142. })
  143. }
  144. return items
  145. }
  146. type upstreamBillingProbeResponse struct {
  147. Object string `json:"object"`
  148. SchemaVersion int `json:"schema_version"`
  149. BillingScope string `json:"billing_scope"`
  150. GroupRateMultiplier *float64 `json:"group_rate_multiplier"`
  151. UserRateMultiplier *float64 `json:"user_rate_multiplier"`
  152. ResolvedRateMultiplier *float64 `json:"resolved_rate_multiplier"`
  153. PeakRateEnabled *bool `json:"peak_rate_enabled"`
  154. PeakStart *string `json:"peak_start"`
  155. PeakEnd *string `json:"peak_end"`
  156. PeakRateMultiplier *float64 `json:"peak_rate_multiplier"`
  157. AppliedPeakMultiplier *float64 `json:"applied_peak_multiplier"`
  158. EffectiveRateMultiplier *float64 `json:"effective_rate_multiplier"`
  159. Timezone *string `json:"timezone"`
  160. ObservedAt string `json:"observed_at"`
  161. }
  162. // GetUpstreamBillingProbeSettings returns defaults when the setting is absent.
  163. func (s *SettingService) GetUpstreamBillingProbeSettings(ctx context.Context) (*UpstreamBillingProbeSettings, error) {
  164. defaults := defaultUpstreamBillingProbeSettings()
  165. if s == nil || s.settingRepo == nil {
  166. return defaults, nil
  167. }
  168. value, err := s.settingRepo.GetValue(ctx, SettingKeyUpstreamBillingProbeSettings)
  169. if err != nil {
  170. if errors.Is(err, ErrSettingNotFound) {
  171. return defaults, nil
  172. }
  173. return nil, fmt.Errorf("get upstream billing probe settings: %w", err)
  174. }
  175. if strings.TrimSpace(value) == "" {
  176. return defaults, nil
  177. }
  178. settings := *defaults
  179. if err := json.Unmarshal([]byte(value), &settings); err != nil {
  180. return nil, fmt.Errorf("parse upstream billing probe settings: %w", err)
  181. }
  182. if settings.IntervalMinutes == 0 {
  183. settings.IntervalMinutes = defaults.IntervalMinutes
  184. }
  185. normalizeUpstreamBillingProbeSettings(&settings)
  186. return &settings, nil
  187. }
  188. // SetUpstreamBillingProbeSettings validates and persists the runner settings.
  189. func (s *SettingService) SetUpstreamBillingProbeSettings(ctx context.Context, settings *UpstreamBillingProbeSettings) error {
  190. if s == nil || s.settingRepo == nil {
  191. return fmt.Errorf("setting repository is unavailable")
  192. }
  193. if settings == nil {
  194. return infraerrors.BadRequest("INVALID_UPSTREAM_BILLING_PROBE_SETTINGS", "settings cannot be nil")
  195. }
  196. if settings.IntervalMinutes < upstreamBillingProbeMinIntervalMinutes || settings.IntervalMinutes > upstreamBillingProbeMaxIntervalMinutes {
  197. return infraerrors.BadRequest(
  198. "INVALID_UPSTREAM_BILLING_PROBE_INTERVAL",
  199. fmt.Sprintf("interval_minutes must be between %d and %d", upstreamBillingProbeMinIntervalMinutes, upstreamBillingProbeMaxIntervalMinutes),
  200. )
  201. }
  202. normalizeUpstreamBillingProbeSettings(settings)
  203. data, err := json.Marshal(settings)
  204. if err != nil {
  205. return fmt.Errorf("marshal upstream billing probe settings: %w", err)
  206. }
  207. return s.settingRepo.Set(ctx, SettingKeyUpstreamBillingProbeSettings, string(data))
  208. }
  209. func defaultUpstreamBillingProbeSettings() *UpstreamBillingProbeSettings {
  210. return &UpstreamBillingProbeSettings{Enabled: true, IntervalMinutes: upstreamBillingProbeDefaultIntervalMinutes}
  211. }
  212. func normalizeUpstreamBillingProbeSettings(settings *UpstreamBillingProbeSettings) {
  213. if settings.IntervalMinutes < upstreamBillingProbeMinIntervalMinutes {
  214. settings.IntervalMinutes = upstreamBillingProbeMinIntervalMinutes
  215. }
  216. if settings.IntervalMinutes > upstreamBillingProbeMaxIntervalMinutes {
  217. settings.IntervalMinutes = upstreamBillingProbeMaxIntervalMinutes
  218. }
  219. }
  220. // UpstreamBillingProbeService discovers a remote Sub2API billing snapshot.
  221. type UpstreamBillingProbeService struct {
  222. accountRepo AccountRepository
  223. accountTestService *AccountTestService
  224. settingService *SettingService
  225. parentCtx context.Context
  226. parentCancel context.CancelFunc
  227. wg sync.WaitGroup
  228. mu sync.Mutex
  229. started bool
  230. stopped bool
  231. cycleMu sync.Mutex
  232. probeGroup singleflight.Group
  233. probeSlots chan struct{}
  234. now func() time.Time
  235. lockCache LeaderLockCache
  236. db *sql.DB
  237. instanceID string
  238. }
  239. type upstreamBillingProbeSnapshotWriter interface {
  240. UpdateUpstreamBillingProbeSnapshot(context.Context, *Account, *UpstreamBillingProbeSnapshot, *float64) error
  241. }
  242. type upstreamBillingProbeDueAccountLister interface {
  243. ListDueUpstreamBillingProbeAccounts(context.Context, time.Time, int) ([]Account, error)
  244. }
  245. func NewUpstreamBillingProbeService(
  246. accountRepo AccountRepository,
  247. accountTestService *AccountTestService,
  248. settingService *SettingService,
  249. ) *UpstreamBillingProbeService {
  250. ctx, cancel := context.WithCancel(context.Background())
  251. return &UpstreamBillingProbeService{
  252. accountRepo: accountRepo,
  253. accountTestService: accountTestService,
  254. settingService: settingService,
  255. parentCtx: ctx,
  256. parentCancel: cancel,
  257. probeSlots: make(chan struct{}, upstreamBillingProbeConcurrency),
  258. now: time.Now,
  259. instanceID: uuid.NewString(),
  260. }
  261. }
  262. func (s *UpstreamBillingProbeService) SetLeaderLock(lockCache LeaderLockCache, db *sql.DB) {
  263. if s == nil {
  264. return
  265. }
  266. s.lockCache = lockCache
  267. s.db = db
  268. }
  269. // ProvideUpstreamBillingProbeService starts the process-wide periodic runner.
  270. func ProvideUpstreamBillingProbeService(
  271. accountRepo AccountRepository,
  272. accountTestService *AccountTestService,
  273. settingService *SettingService,
  274. lockCache LeaderLockCache,
  275. db *sql.DB,
  276. ) *UpstreamBillingProbeService {
  277. svc := NewUpstreamBillingProbeService(accountRepo, accountTestService, settingService)
  278. svc.SetLeaderLock(lockCache, db)
  279. svc.Start()
  280. return svc
  281. }
  282. func (s *UpstreamBillingProbeService) Start() {
  283. if s == nil {
  284. return
  285. }
  286. s.mu.Lock()
  287. if s.started || s.stopped {
  288. s.mu.Unlock()
  289. return
  290. }
  291. s.started = true
  292. s.wg.Add(1)
  293. s.mu.Unlock()
  294. go s.runLoop()
  295. }
  296. func (s *UpstreamBillingProbeService) Stop() {
  297. if s == nil {
  298. return
  299. }
  300. s.mu.Lock()
  301. if s.stopped {
  302. s.mu.Unlock()
  303. return
  304. }
  305. s.stopped = true
  306. s.parentCancel()
  307. s.mu.Unlock()
  308. s.wg.Wait()
  309. }
  310. func (s *UpstreamBillingProbeService) runLoop() {
  311. defer s.wg.Done()
  312. _ = s.RunDue(s.parentCtx)
  313. ticker := time.NewTicker(upstreamBillingProbeCycleInterval)
  314. defer ticker.Stop()
  315. for {
  316. select {
  317. case <-s.parentCtx.Done():
  318. return
  319. case <-ticker.C:
  320. if err := s.RunDue(s.parentCtx); err != nil {
  321. logger.LegacyPrintf("service.upstream_billing_probe", "run_due_failed: err=%v", err)
  322. }
  323. }
  324. }
  325. }
  326. // RunDue executes at most one bounded batch of due accounts.
  327. func (s *UpstreamBillingProbeService) RunDue(ctx context.Context) error {
  328. if s == nil || s.accountRepo == nil {
  329. return nil
  330. }
  331. s.cycleMu.Lock()
  332. defer s.cycleMu.Unlock()
  333. settings, err := s.getSettings(ctx)
  334. if err != nil {
  335. return err
  336. }
  337. if !settings.Enabled {
  338. return nil
  339. }
  340. runRelease, acquired, lockErr := s.tryAcquireLeaderLock(ctx, upstreamBillingProbeLeaderLockKey)
  341. if lockErr != nil {
  342. return fmt.Errorf("acquire upstream billing probe leader lock: %w", lockErr)
  343. }
  344. if !acquired {
  345. return nil
  346. }
  347. defer runRelease()
  348. lockNow := time.Now()
  349. cadenceRelease, acquired, lockErr := s.tryAcquireLeaderLock(ctx, upstreamBillingProbeLeaderLockKeyAt(lockNow))
  350. if lockErr != nil {
  351. return fmt.Errorf("acquire upstream billing probe cadence lock: %w", lockErr)
  352. }
  353. if !acquired {
  354. return nil
  355. }
  356. defer releaseUpstreamBillingProbeLeaderLock(cadenceRelease, lockNow.Truncate(upstreamBillingProbeCycleInterval).Add(upstreamBillingProbeCycleInterval))
  357. now := s.currentTime()
  358. accounts, err := s.listDueAccounts(ctx, now)
  359. if err != nil {
  360. return fmt.Errorf("list enabled upstream billing probes: %w", err)
  361. }
  362. due := make([]Account, 0, len(accounts))
  363. for i := range accounts {
  364. account := accounts[i]
  365. if !isUpstreamBillingProbeAccount(&account) || !account.IsActive() || !upstreamBillingProbeEnabled(&account) {
  366. continue
  367. }
  368. snapshot := decodeUpstreamBillingProbeSnapshot(account.Extra)
  369. if snapshot != nil && !snapshot.NextProbeAt.IsZero() && now.Before(snapshot.NextProbeAt) {
  370. continue
  371. }
  372. due = append(due, account)
  373. }
  374. sort.SliceStable(due, func(i, j int) bool {
  375. left := decodeUpstreamBillingProbeSnapshot(due[i].Extra)
  376. right := decodeUpstreamBillingProbeSnapshot(due[j].Extra)
  377. leftUnset := left == nil || left.NextProbeAt.IsZero()
  378. rightUnset := right == nil || right.NextProbeAt.IsZero()
  379. if leftUnset && rightUnset {
  380. return due[i].ID < due[j].ID
  381. }
  382. if leftUnset {
  383. return true
  384. }
  385. if rightUnset {
  386. return false
  387. }
  388. return left.NextProbeAt.Before(right.NextProbeAt)
  389. })
  390. if len(due) > upstreamBillingProbeMaxPerCycle {
  391. due = due[:upstreamBillingProbeMaxPerCycle]
  392. }
  393. var group errgroup.Group
  394. for i := range due {
  395. accountID := due[i].ID
  396. group.Go(func() error {
  397. if _, probeErr := s.probeScheduledAccount(ctx, accountID, settings.IntervalMinutes); probeErr != nil {
  398. logger.LegacyPrintf("service.upstream_billing_probe", "probe_due_failed: account_id=%d err=%v", accountID, probeErr)
  399. }
  400. return nil
  401. })
  402. }
  403. return group.Wait()
  404. }
  405. func (s *UpstreamBillingProbeService) listDueAccounts(ctx context.Context, now time.Time) ([]Account, error) {
  406. if lister, ok := s.accountRepo.(upstreamBillingProbeDueAccountLister); ok {
  407. return lister.ListDueUpstreamBillingProbeAccounts(ctx, now, upstreamBillingProbeMaxPerCycle)
  408. }
  409. // Non-production repositories and older adapters keep the generic path. The
  410. // runner still truncates before issuing network requests.
  411. return s.accountRepo.FindByExtraField(ctx, UpstreamBillingProbeEnabledExtraKey, true)
  412. }
  413. func (s *UpstreamBillingProbeService) getSettings(ctx context.Context) (*UpstreamBillingProbeSettings, error) {
  414. if s.settingService == nil {
  415. return defaultUpstreamBillingProbeSettings(), nil
  416. }
  417. return s.settingService.GetUpstreamBillingProbeSettings(ctx)
  418. }
  419. func (s *UpstreamBillingProbeService) GetSettings(ctx context.Context) (*UpstreamBillingProbeSettings, error) {
  420. return s.getSettings(ctx)
  421. }
  422. func (s *UpstreamBillingProbeService) UpdateSettings(ctx context.Context, settings *UpstreamBillingProbeSettings) error {
  423. if s == nil || s.settingService == nil {
  424. return ErrUpstreamBillingProbeUnavailable
  425. }
  426. return s.settingService.SetUpstreamBillingProbeSettings(ctx, settings)
  427. }
  428. // ProbeAccount performs one manual or scheduled probe. Manual calls ignore both switches.
  429. func (s *UpstreamBillingProbeService) ProbeAccount(ctx context.Context, accountID int64) (*UpstreamBillingProbeSnapshot, error) {
  430. if s == nil || s.accountRepo == nil {
  431. return nil, ErrUpstreamBillingProbeUnavailable
  432. }
  433. settings, err := s.getSettings(ctx)
  434. if err != nil {
  435. return nil, err
  436. }
  437. return s.probeAccount(ctx, accountID, settings.IntervalMinutes)
  438. }
  439. func (s *UpstreamBillingProbeService) probeAccount(ctx context.Context, accountID int64, intervalMinutes int) (*UpstreamBillingProbeSnapshot, error) {
  440. return s.probeAccountWithMode(ctx, accountID, intervalMinutes, false)
  441. }
  442. func (s *UpstreamBillingProbeService) probeScheduledAccount(ctx context.Context, accountID int64, intervalMinutes int) (*UpstreamBillingProbeSnapshot, error) {
  443. return s.probeAccountWithMode(ctx, accountID, intervalMinutes, true)
  444. }
  445. func (s *UpstreamBillingProbeService) probeAccountWithMode(ctx context.Context, accountID int64, intervalMinutes int, requireEnabled bool) (*UpstreamBillingProbeSnapshot, error) {
  446. key := strconv.FormatInt(accountID, 10)
  447. value, err, _ := s.probeGroup.Do(key, func() (any, error) {
  448. select {
  449. case s.probeSlots <- struct{}{}:
  450. defer func() { <-s.probeSlots }()
  451. case <-ctx.Done():
  452. return nil, ctx.Err()
  453. }
  454. account, loadErr := s.accountRepo.GetByID(ctx, accountID)
  455. if loadErr != nil {
  456. return nil, loadErr
  457. }
  458. if !isUpstreamBillingProbeAccount(account) {
  459. return nil, ErrUpstreamBillingProbeAccountInvalid
  460. }
  461. if requireEnabled {
  462. if !account.IsActive() || !upstreamBillingProbeEnabled(account) {
  463. return nil, nil
  464. }
  465. if snapshot := decodeUpstreamBillingProbeSnapshot(account.Extra); snapshot != nil &&
  466. !snapshot.NextProbeAt.IsZero() && s.currentTime().Before(snapshot.NextProbeAt) {
  467. return nil, nil
  468. }
  469. }
  470. return s.probeLoadedAccount(ctx, account, intervalMinutes)
  471. })
  472. if err != nil {
  473. return nil, err
  474. }
  475. if value == nil {
  476. return nil, nil
  477. }
  478. snapshot, ok := value.(*UpstreamBillingProbeSnapshot)
  479. if !ok {
  480. return nil, fmt.Errorf("invalid upstream billing probe result")
  481. }
  482. return snapshot, nil
  483. }
  484. // ProbeAccounts performs a bounded manual batch with the same concurrency limit as the runner.
  485. func (s *UpstreamBillingProbeService) ProbeAccounts(ctx context.Context, accountIDs []int64) []UpstreamBillingProbeResult {
  486. if len(accountIDs) > upstreamBillingProbeMaxPerCycle {
  487. accountIDs = accountIDs[:upstreamBillingProbeMaxPerCycle]
  488. }
  489. results := make([]UpstreamBillingProbeResult, len(accountIDs))
  490. if s == nil || s.accountRepo == nil {
  491. for i, accountID := range accountIDs {
  492. results[i] = UpstreamBillingProbeResult{AccountID: accountID, Error: ErrUpstreamBillingProbeUnavailable.Error()}
  493. }
  494. return results
  495. }
  496. settings, settingsErr := s.getSettings(ctx)
  497. if settingsErr != nil {
  498. for i, accountID := range accountIDs {
  499. results[i] = UpstreamBillingProbeResult{AccountID: accountID, Error: safeProbeError(settingsErr)}
  500. }
  501. return results
  502. }
  503. var group errgroup.Group
  504. for i, accountID := range accountIDs {
  505. i, accountID := i, accountID
  506. results[i].AccountID = accountID
  507. group.Go(func() error {
  508. snapshot, err := s.probeAccount(ctx, accountID, settings.IntervalMinutes)
  509. if err != nil {
  510. results[i].Error = safeProbeError(err)
  511. return nil
  512. }
  513. results[i].Snapshot = snapshot
  514. return nil
  515. })
  516. }
  517. _ = group.Wait()
  518. return results
  519. }
  520. func upstreamBillingProbeLeaderLockKeyAt(now time.Time) string {
  521. return fmt.Sprintf("%s:%d", upstreamBillingProbeLeaderLockKey, now.Unix()/int64(upstreamBillingProbeCycleInterval/time.Second))
  522. }
  523. func (s *UpstreamBillingProbeService) tryAcquireLeaderLock(ctx context.Context, key string) (func(), bool, error) {
  524. lockCtx, cancel := context.WithTimeout(ctx, 2*time.Second)
  525. defer cancel()
  526. if s.lockCache != nil {
  527. acquired, err := s.lockCache.TryAcquireLeaderLock(lockCtx, key, s.instanceID, upstreamBillingProbeLeaderLockTTL)
  528. if err != nil {
  529. return nil, false, err
  530. }
  531. if !acquired {
  532. return nil, false, nil
  533. }
  534. return func() {
  535. releaseCtx, releaseCancel := context.WithTimeout(context.Background(), 2*time.Second)
  536. defer releaseCancel()
  537. _ = s.lockCache.ReleaseLeaderLock(releaseCtx, key, s.instanceID)
  538. }, true, nil
  539. }
  540. if s.db != nil {
  541. return tryAcquireDBAdvisoryLockWithError(lockCtx, s.db, hashAdvisoryLockID(key))
  542. }
  543. return func() {}, true, nil
  544. }
  545. func releaseUpstreamBillingProbeLeaderLock(release func(), releaseAt time.Time) {
  546. delay := time.Until(releaseAt)
  547. if delay <= 0 {
  548. release()
  549. return
  550. }
  551. time.AfterFunc(delay, release)
  552. }
  553. func (s *UpstreamBillingProbeService) SetAccountEnabled(ctx context.Context, accountID int64, enabled bool) error {
  554. if s == nil || s.accountRepo == nil {
  555. return ErrUpstreamBillingProbeUnavailable
  556. }
  557. account, err := s.accountRepo.GetByID(ctx, accountID)
  558. if err != nil {
  559. return err
  560. }
  561. if !isUpstreamBillingProbeAccount(account) {
  562. return ErrUpstreamBillingProbeAccountInvalid
  563. }
  564. updates := map[string]any{UpstreamBillingProbeEnabledExtraKey: enabled}
  565. if !enabled {
  566. updates[UpstreamBillingRateSyncEnabledExtraKey] = false
  567. }
  568. return s.accountRepo.UpdateExtra(ctx, accountID, updates)
  569. }
  570. func (s *UpstreamBillingProbeService) probeLoadedAccount(ctx context.Context, account *Account, intervalMinutes int) (*UpstreamBillingProbeSnapshot, error) {
  571. now := s.currentTime().UTC()
  572. if s.accountTestService == nil || s.accountTestService.httpUpstream == nil {
  573. return s.persistProbeFailure(ctx, account, intervalMinutes, now, 0, "transport_unavailable", 0)
  574. }
  575. // 平台放宽后取数直读 credentials:所有 API-key 平台的密钥与自定义上游
  576. // 统一存放在 credentials.api_key / credentials.base_url。
  577. apiKey := account.GetCredential("api_key")
  578. if apiKey == "" {
  579. return s.persistProbeFailure(ctx, account, intervalMinutes, now, 0, "missing_api_key", 0)
  580. }
  581. baseURL := account.GetCredential("base_url")
  582. if account.IsCNProvider() && account.IsAdaptiveAPIProtocol() {
  583. baseURL = account.GetCNProtocolBaseURL(APIProtocolChatCompletions)
  584. }
  585. if account.Platform == PlatformOpenAI {
  586. if baseURL == "" {
  587. // 保持官方语义:OpenAI 账号无自定义 base 时探官方域(404 → unsupported)。
  588. baseURL = "https://api.openai.com"
  589. }
  590. } else if upstreamBillingProbeTargetIsOfficialAPI(baseURL) {
  591. // 其他平台 base_url 为空或指向官方 API 根域(前端创建时会把空值
  592. // 填成官方默认域,且提供 us-east-1.api.x.ai 等官方区域预设)⇒
  593. // 必无 /v1/sub2api/billing;不发请求,直接记 unsupported,避免
  594. // 拿账号 Key 周期性请求官方域的不存在路径。
  595. return s.persistProbeFailure(ctx, account, intervalMinutes, now, 0, "unsupported", 0)
  596. }
  597. normalizedBaseURL, err := s.accountTestService.validateUpstreamBaseURL(baseURL)
  598. if err != nil {
  599. return s.persistProbeFailure(ctx, account, intervalMinutes, now, 0, "invalid_base_url", 0)
  600. }
  601. proxyURL := ""
  602. if account.ProxyID != nil {
  603. if account.Proxy == nil {
  604. return s.persistProbeFailure(ctx, account, intervalMinutes, now, 0, "proxy_unavailable", 0)
  605. }
  606. if account.Proxy.ID != *account.ProxyID {
  607. return nil, ErrUpstreamBillingProbeIdentityChanged
  608. }
  609. proxyURL = account.Proxy.URL()
  610. }
  611. probeURL := buildOpenAIEndpointURL(normalizedBaseURL, "/v1/sub2api/billing")
  612. probeCtx, cancel := context.WithTimeout(ctx, upstreamBillingProbeRequestTimeout)
  613. defer cancel()
  614. req, err := http.NewRequestWithContext(probeCtx, http.MethodGet, probeURL, bytes.NewReader(nil))
  615. if err != nil {
  616. return s.persistProbeFailure(ctx, account, intervalMinutes, now, 0, "request_build_failed", 0)
  617. }
  618. // OpenAI 账号保持官方 openai 传输画像;其他平台探测走默认画像。
  619. profile := HTTPUpstreamProfileDefault
  620. if account.Platform == PlatformOpenAI {
  621. profile = HTTPUpstreamProfileOpenAI
  622. }
  623. reqCtx := WithHTTPUpstreamProfile(req.Context(), profile)
  624. req = req.WithContext(WithHTTPUpstreamRedirectsDisabled(reqCtx))
  625. req.Header.Set("Accept", "application/json")
  626. req.Header.Set("Authorization", "Bearer "+apiKey)
  627. account.ApplyHeaderOverrides(req.Header)
  628. var tlsProfile *tlsfingerprint.Profile
  629. if s.accountTestService.tlsFPProfileService != nil {
  630. tlsProfile = s.accountTestService.tlsFPProfileService.ResolveTLSProfile(account)
  631. }
  632. resp, err := s.accountTestService.httpUpstream.DoWithTLS(req, proxyURL, account.ID, account.Concurrency, tlsProfile)
  633. if err != nil {
  634. return s.persistProbeFailure(ctx, account, intervalMinutes, now, 0, "request_failed", 0)
  635. }
  636. if resp == nil || resp.Body == nil {
  637. return s.persistProbeFailure(ctx, account, intervalMinutes, now, 0, "empty_response", 0)
  638. }
  639. defer func() { _ = resp.Body.Close() }()
  640. body, readErr := io.ReadAll(io.LimitReader(resp.Body, upstreamBillingProbeMaxBodyBytes+1))
  641. if readErr != nil {
  642. return s.persistProbeFailure(ctx, account, intervalMinutes, now, resp.StatusCode, "response_read_failed", retryAfter(resp.Header, now))
  643. }
  644. if len(body) > upstreamBillingProbeMaxBodyBytes {
  645. return s.persistProbeFailure(ctx, account, intervalMinutes, now, resp.StatusCode, "response_too_large", retryAfter(resp.Header, now))
  646. }
  647. if resp.StatusCode == http.StatusNotFound || resp.StatusCode == http.StatusMethodNotAllowed {
  648. return s.persistProbeFailure(ctx, account, intervalMinutes, now, resp.StatusCode, "unsupported", retryAfter(resp.Header, now))
  649. }
  650. if resp.StatusCode < 200 || resp.StatusCode >= 300 {
  651. return s.persistProbeFailure(ctx, account, intervalMinutes, now, resp.StatusCode, "http_error", retryAfter(resp.Header, now))
  652. }
  653. data, err := parseUpstreamBillingProbeResponse(body)
  654. if err != nil {
  655. return s.persistProbeFailure(ctx, account, intervalMinutes, now, resp.StatusCode, "invalid_response", retryAfter(resp.Header, now))
  656. }
  657. snapshot := &UpstreamBillingProbeSnapshot{
  658. Status: UpstreamBillingProbeStatusOK,
  659. Data: data,
  660. ReceivedAt: probeTimePtr(now),
  661. FreshUntil: probeTimePtr(now.Add(2 * time.Duration(intervalMinutes) * time.Minute)),
  662. LastAttemptAt: now,
  663. NextProbeAt: now.Add(nextProbeDelay(intervalMinutes, 0)),
  664. HTTPStatus: resp.StatusCode,
  665. }
  666. // 账号级值域与精度只在真要写回时才有影响:只观察上游声明、未开启同步的
  667. // 账号不因声明值不适配 accounts.rate_multiplier 而被记成探测失败并进入
  668. // 指数退避——探测本身成功了,原始声明照常存进快照供展示。
  669. var syncRate *float64
  670. previousRate := account.BillingRateMultiplier()
  671. if upstreamBillingRateSyncEnabled(account) {
  672. if value, valid := upstreamBillingProbeSyncRate(data); valid {
  673. syncRate = &value
  674. snapshot.SyncedRateMultiplier = &value
  675. } else {
  676. declared, _ := resolveAccountExtraNumber(data, "resolved_rate_multiplier")
  677. slog.Warn("upstream_billing_rate_sync_rejected",
  678. "source", "upstream_billing_probe",
  679. "account_id", account.ID,
  680. "declared_resolved_rate_multiplier", declared,
  681. "max_rate_multiplier", upstreamBillingRateSyncMaxMultiplier,
  682. "current_rate_multiplier", previousRate,
  683. )
  684. }
  685. }
  686. if err := s.updateSnapshot(ctx, account, snapshot, syncRate); err != nil {
  687. return nil, err
  688. }
  689. if syncRate != nil {
  690. // 写回是后台任务的裸 SQL,不经过管理端路由,因此不会产生 audit_logs 行。
  691. // old_rate_multiplier 是本次探测开始时读到的值(写回的 CAS 不比对该列)。
  692. slog.Info("upstream_billing_rate_sync_applied",
  693. "source", "upstream_billing_probe",
  694. "account_id", account.ID,
  695. "old_rate_multiplier", previousRate,
  696. "new_rate_multiplier", *syncRate,
  697. )
  698. }
  699. return snapshot, nil
  700. }
  701. func (s *UpstreamBillingProbeService) persistProbeFailure(
  702. ctx context.Context,
  703. account *Account,
  704. intervalMinutes int,
  705. now time.Time,
  706. statusCode int,
  707. reason string,
  708. retryAfterDuration time.Duration,
  709. ) (*UpstreamBillingProbeSnapshot, error) {
  710. previous := decodeUpstreamBillingProbeSnapshot(account.Extra)
  711. failureCount := 1
  712. if previous != nil {
  713. failureCount = previous.FailureCount + 1
  714. }
  715. status := UpstreamBillingProbeStatusFailed
  716. delay := nextProbeDelay(intervalMinutes, retryAfterDuration)
  717. if reason == "unsupported" {
  718. status = UpstreamBillingProbeStatusUnsupported
  719. delay = unsupportedProbeDelay(intervalMinutes, retryAfterDuration)
  720. }
  721. snapshot := &UpstreamBillingProbeSnapshot{
  722. Status: status,
  723. LastAttemptAt: now,
  724. NextProbeAt: now.Add(delay),
  725. FailureCount: failureCount,
  726. HTTPStatus: statusCode,
  727. LastError: reason,
  728. }
  729. if previous != nil {
  730. snapshot.Data = previous.Data
  731. snapshot.ReceivedAt = previous.ReceivedAt
  732. snapshot.FreshUntil = previous.FreshUntil
  733. if snapshot.FreshUntil == nil && previous.Status == UpstreamBillingProbeStatusOK && previous.ReceivedAt != nil {
  734. snapshot.FreshUntil = probeTimePtr(previous.ReceivedAt.Add(2 * time.Duration(intervalMinutes) * time.Minute))
  735. }
  736. }
  737. if err := s.updateSnapshot(ctx, account, snapshot, nil); err != nil {
  738. return nil, err
  739. }
  740. return snapshot, nil
  741. }
  742. func (s *UpstreamBillingProbeService) updateSnapshot(
  743. ctx context.Context,
  744. account *Account,
  745. snapshot *UpstreamBillingProbeSnapshot,
  746. rateMultiplier *float64,
  747. ) error {
  748. writer, ok := s.accountRepo.(upstreamBillingProbeSnapshotWriter)
  749. if !ok {
  750. return ErrUpstreamBillingProbeUnavailable
  751. }
  752. return writer.UpdateUpstreamBillingProbeSnapshot(ctx, account, snapshot, rateMultiplier)
  753. }
  754. func parseUpstreamBillingProbeResponse(body []byte) (map[string]any, error) {
  755. var response upstreamBillingProbeResponse
  756. if err := json.Unmarshal(body, &response); err != nil {
  757. return nil, err
  758. }
  759. if response.Object != "sub2api.key_billing" || response.SchemaVersion != 1 || response.BillingScope != "token" {
  760. return nil, fmt.Errorf("unexpected billing response schema")
  761. }
  762. if response.GroupRateMultiplier == nil || response.ResolvedRateMultiplier == nil ||
  763. response.PeakRateEnabled == nil || response.EffectiveRateMultiplier == nil {
  764. return nil, fmt.Errorf("incomplete billing response")
  765. }
  766. for _, value := range []float64{
  767. *response.GroupRateMultiplier,
  768. *response.ResolvedRateMultiplier,
  769. *response.EffectiveRateMultiplier,
  770. } {
  771. if value < 0 || math.IsNaN(value) || math.IsInf(value, 0) {
  772. return nil, fmt.Errorf("invalid billing multiplier")
  773. }
  774. }
  775. if response.UserRateMultiplier != nil && (*response.UserRateMultiplier < 0 || math.IsNaN(*response.UserRateMultiplier) || math.IsInf(*response.UserRateMultiplier, 0)) {
  776. return nil, fmt.Errorf("invalid user billing multiplier")
  777. }
  778. expectedResolved := *response.GroupRateMultiplier
  779. if response.UserRateMultiplier != nil {
  780. expectedResolved = *response.UserRateMultiplier
  781. }
  782. if !equalBillingMultiplier(*response.ResolvedRateMultiplier, expectedResolved) {
  783. return nil, fmt.Errorf("inconsistent resolved billing multiplier")
  784. }
  785. observedAt, err := time.Parse(time.RFC3339Nano, response.ObservedAt)
  786. if err != nil || observedAt.IsZero() {
  787. return nil, fmt.Errorf("invalid observed_at")
  788. }
  789. data := map[string]any{
  790. "object": response.Object,
  791. "schema_version": response.SchemaVersion,
  792. "billing_scope": response.BillingScope,
  793. "group_rate_multiplier": *response.GroupRateMultiplier,
  794. "resolved_rate_multiplier": *response.ResolvedRateMultiplier,
  795. "peak_rate_enabled": *response.PeakRateEnabled,
  796. "effective_rate_multiplier": *response.EffectiveRateMultiplier,
  797. "observed_at": observedAt.UTC().Format(time.RFC3339Nano),
  798. }
  799. if response.UserRateMultiplier != nil {
  800. data["user_rate_multiplier"] = *response.UserRateMultiplier
  801. }
  802. if *response.PeakRateEnabled {
  803. if response.PeakStart == nil || response.PeakEnd == nil || response.Timezone == nil ||
  804. response.PeakRateMultiplier == nil || response.AppliedPeakMultiplier == nil ||
  805. *response.PeakStart == "" || *response.PeakEnd == "" || *response.Timezone == "" ||
  806. *response.PeakRateMultiplier < 0 || *response.AppliedPeakMultiplier < 0 ||
  807. math.IsNaN(*response.PeakRateMultiplier) || math.IsInf(*response.PeakRateMultiplier, 0) ||
  808. math.IsNaN(*response.AppliedPeakMultiplier) || math.IsInf(*response.AppliedPeakMultiplier, 0) {
  809. return nil, fmt.Errorf("incomplete peak billing response")
  810. }
  811. data["peak_start"] = *response.PeakStart
  812. data["peak_end"] = *response.PeakEnd
  813. data["peak_rate_multiplier"] = *response.PeakRateMultiplier
  814. data["applied_peak_multiplier"] = *response.AppliedPeakMultiplier
  815. data["timezone"] = *response.Timezone
  816. }
  817. appliedPeak, ok := upstreamBillingPeakMultiplierAt(data, observedAt)
  818. if !ok {
  819. return nil, fmt.Errorf("invalid peak billing response")
  820. }
  821. if response.PeakRateEnabled != nil && *response.PeakRateEnabled {
  822. if !equalBillingMultiplier(*response.AppliedPeakMultiplier, appliedPeak) {
  823. return nil, fmt.Errorf("inconsistent applied peak multiplier")
  824. }
  825. } else if response.AppliedPeakMultiplier != nil && !equalBillingMultiplier(*response.AppliedPeakMultiplier, 1) {
  826. return nil, fmt.Errorf("inconsistent applied peak multiplier")
  827. }
  828. if !equalBillingMultiplier(*response.EffectiveRateMultiplier, *response.ResolvedRateMultiplier*appliedPeak) {
  829. return nil, fmt.Errorf("inconsistent effective billing multiplier")
  830. }
  831. return data, nil
  832. }
  833. func upstreamBillingRateAt(data map[string]any, now time.Time) (float64, bool) {
  834. if scope, _ := data["billing_scope"].(string); scope != "token" {
  835. return 0, false
  836. }
  837. base, ok := resolveAccountExtraNumber(data, "resolved_rate_multiplier")
  838. if !ok || base < 0 || math.IsNaN(base) || math.IsInf(base, 0) {
  839. return 0, false
  840. }
  841. appliedPeak, ok := upstreamBillingPeakMultiplierAt(data, now)
  842. if !ok {
  843. return 0, false
  844. }
  845. base *= appliedPeak
  846. if math.IsNaN(base) || math.IsInf(base, 0) {
  847. return 0, false
  848. }
  849. return base, true
  850. }
  851. // upstreamBillingProbeSyncRate converts the declared multiplier into the value
  852. // the automatic write-back may store in accounts.rate_multiplier, at the
  853. // precision that column supports (DECIMAL(10,4)).
  854. //
  855. // It reads resolved_rate_multiplier, not effective_rate_multiplier: the
  856. // effective value folds in the peak coefficient that happened to apply at the
  857. // instant of the probe, so writing it would freeze one probe cycle's peak (or
  858. // off-peak) factor into a static column, while display and scheduling
  859. // recompute the peak factor for the current time through upstreamBillingRateAt.
  860. //
  861. // The accepted range is deliberately narrower than the column:
  862. // - 0 is rejected. accountCost multiplies the request cost by this value, so
  863. // an upstream-declared 0 would stop quota_used from ever growing and every
  864. // admin-configured account quota and cost alert would silently stop
  865. // working. Admins may still set 0 by hand; only the automatic path refuses.
  866. // - anything above upstreamBillingRateSyncMaxMultiplier is rejected.
  867. //
  868. // A rejected declaration leaves the current multiplier untouched; the probe
  869. // still records an OK snapshot carrying the raw declaration for display.
  870. func upstreamBillingProbeSyncRate(data map[string]any) (float64, bool) {
  871. value, ok := resolveAccountExtraNumber(data, "resolved_rate_multiplier")
  872. if !ok || math.IsNaN(value) || math.IsInf(value, 0) {
  873. return 0, false
  874. }
  875. rounded := math.Round(value*upstreamBillingProbeAccountRateScale) / upstreamBillingProbeAccountRateScale
  876. if rounded <= 0 || rounded > upstreamBillingRateSyncMaxMultiplier {
  877. return 0, false
  878. }
  879. return rounded, true
  880. }
  881. func upstreamBillingPeakMultiplierAt(data map[string]any, now time.Time) (float64, bool) {
  882. peakEnabled, ok := data["peak_rate_enabled"].(bool)
  883. if !ok {
  884. return 0, false
  885. }
  886. if !peakEnabled {
  887. return 1, true
  888. }
  889. start, startOK := data["peak_start"].(string)
  890. end, endOK := data["peak_end"].(string)
  891. timezoneName, timezoneOK := data["timezone"].(string)
  892. peakMultiplier, multiplierOK := resolveAccountExtraNumber(data, "peak_rate_multiplier")
  893. startMinute, validStart := parseMinutes(start)
  894. endMinute, validEnd := parseMinutes(end)
  895. if !startOK || !endOK || !timezoneOK || !multiplierOK || !validStart || !validEnd ||
  896. startMinute >= endMinute || peakMultiplier < 0 || math.IsNaN(peakMultiplier) || math.IsInf(peakMultiplier, 0) {
  897. return 0, false
  898. }
  899. location, err := time.LoadLocation(timezoneName)
  900. if err != nil {
  901. return 0, false
  902. }
  903. local := now.In(location)
  904. minute := local.Hour()*60 + local.Minute()
  905. if minute >= startMinute && minute < endMinute {
  906. return peakMultiplier, true
  907. }
  908. return 1, true
  909. }
  910. func equalBillingMultiplier(left, right float64) bool {
  911. if math.IsNaN(left) || math.IsNaN(right) || math.IsInf(left, 0) || math.IsInf(right, 0) {
  912. return false
  913. }
  914. scale := math.Max(1, math.Max(math.Abs(left), math.Abs(right)))
  915. return math.Abs(left-right) <= 1e-9*scale
  916. }
  917. func decodeUpstreamBillingProbeSnapshot(extra map[string]any) *UpstreamBillingProbeSnapshot {
  918. if extra == nil {
  919. return nil
  920. }
  921. value, ok := extra[UpstreamBillingProbeExtraKey]
  922. if !ok {
  923. return nil
  924. }
  925. raw, err := json.Marshal(value)
  926. if err != nil {
  927. return nil
  928. }
  929. var snapshot UpstreamBillingProbeSnapshot
  930. if err := json.Unmarshal(raw, &snapshot); err != nil || snapshot.Status == "" {
  931. return nil
  932. }
  933. if snapshot.Status != UpstreamBillingProbeStatusOK &&
  934. snapshot.Status != UpstreamBillingProbeStatusUnsupported &&
  935. snapshot.Status != UpstreamBillingProbeStatusFailed {
  936. return nil
  937. }
  938. return &snapshot
  939. }
  940. // IsUpstreamBillingProbeIdentity reports whether an account identity may opt
  941. // in to the upstream billing probe. `/v1/sub2api/billing` is a key-scoped
  942. // sub2api convention shared by the supported API-key platforms (including the
  943. // CN providers, whose official-domain accounts are short-circuited to
  944. // "unsupported" by upstreamBillingProbeTargetIsOfficialAPI).
  945. // Non-sub2api upstreams return 404 and the snapshot records "unsupported".
  946. // Only AccountTypeAPIKey is in scope. OAuth/Bedrock hold no static API key to
  947. // present at all; AccountTypeUpstream (antigravity relay accounts) does carry
  948. // a base_url plus a static api_key, but it is deliberately left out of the
  949. // current supported set. New antigravity relay accounts are created with
  950. // type=apikey by the admin form, so only pre-existing type=upstream rows
  951. // cannot turn the probe on.
  952. func IsUpstreamBillingProbeIdentity(platform, accountType string) bool {
  953. if accountType != AccountTypeAPIKey {
  954. return false
  955. }
  956. switch platform {
  957. case PlatformOpenAI, PlatformAnthropic, PlatformGemini, PlatformAntigravity, PlatformGrok,
  958. PlatformKimi, PlatformZhipu, PlatformDeepseek, PlatformMiniMax:
  959. return true
  960. default:
  961. return false
  962. }
  963. }
  964. func isUpstreamBillingProbeAccount(account *Account) bool {
  965. return account != nil && IsUpstreamBillingProbeIdentity(account.Platform, account.Type)
  966. }
  967. // upstreamBillingProbeOfficialAPIDomains lists the root domains of official
  968. // provider APIs. The create form fills empty base_url values with official
  969. // defaults (and offers official regional presets like us-east-1.api.x.ai),
  970. // so probing them would send the account key to an official API path that
  971. // cannot exist. Matching is by registrable root domain — exact host or any
  972. // subdomain, after stripping the port and a trailing DNS dot — because no
  973. // third-party sub2api relay can live under these domains, while custom
  974. // relays (the only targets that can answer /v1/sub2api/billing) always do
  975. // probe. OpenAI-platform accounts never reach this check: they keep the
  976. // upstream-official behavior of probing api.openai.com.
  977. // ollama.com is a first-class configuration here (Ollama Cloud accounts are
  978. // platform openai/anthropic with base_url https://ollama.com/v1), and it is
  979. // an official provider API just like the rest, so it belongs on this list.
  980. // CN provider domains (moonshot.cn / kimi.com / bigmodel.cn / deepseek.com)
  981. // serve the same role: official APIs that can never host /v1/sub2api/billing,
  982. // so their accounts short-circuit to "unsupported" without a request.
  983. var upstreamBillingProbeOfficialAPIDomains = []string{
  984. "anthropic.com",
  985. "googleapis.com",
  986. "x.ai",
  987. "grok.com",
  988. "openai.com",
  989. "ollama.com",
  990. "moonshot.cn",
  991. "kimi.com",
  992. "bigmodel.cn",
  993. "deepseek.com",
  994. }
  995. func upstreamBillingProbeTargetIsOfficialAPI(baseURL string) bool {
  996. baseURL = strings.TrimSpace(baseURL)
  997. if baseURL == "" {
  998. return true
  999. }
  1000. parsed, err := url.Parse(baseURL)
  1001. if err != nil {
  1002. return false
  1003. }
  1004. host := strings.TrimSuffix(strings.ToLower(parsed.Hostname()), ".")
  1005. if host == "" {
  1006. return true
  1007. }
  1008. for _, domain := range upstreamBillingProbeOfficialAPIDomains {
  1009. if host == domain || strings.HasSuffix(host, "."+domain) {
  1010. return true
  1011. }
  1012. }
  1013. return false
  1014. }
  1015. func upstreamBillingProbeEnabled(account *Account) bool {
  1016. if account == nil || account.Extra == nil {
  1017. return false
  1018. }
  1019. enabled, ok := account.Extra[UpstreamBillingProbeEnabledExtraKey].(bool)
  1020. return ok && enabled
  1021. }
  1022. // upstreamBillingRateSyncEnabled is the probe-side pre-filter deciding whether
  1023. // a rate is even proposed for write-back. It is a necessary condition, not the
  1024. // authority: the repository CAS re-checks both switches against the row it
  1025. // updates, so a switch flipped between load and write can never sneak a rate in.
  1026. func upstreamBillingRateSyncEnabled(account *Account) bool {
  1027. if account == nil || account.Extra == nil {
  1028. return false
  1029. }
  1030. enabled, ok := account.Extra[UpstreamBillingRateSyncEnabledExtraKey].(bool)
  1031. return ok && enabled && upstreamBillingProbeEnabled(account)
  1032. }
  1033. func (s *UpstreamBillingProbeService) currentTime() time.Time {
  1034. if s != nil && s.now != nil {
  1035. return s.now()
  1036. }
  1037. return time.Now()
  1038. }
  1039. func nextProbeDelay(intervalMinutes int, retryAfterDuration time.Duration) time.Duration {
  1040. interval := time.Duration(intervalMinutes) * time.Minute
  1041. if interval < upstreamBillingProbeMinIntervalMinutes*time.Minute {
  1042. interval = upstreamBillingProbeMinIntervalMinutes * time.Minute
  1043. }
  1044. if interval > upstreamBillingProbeMaxDelay {
  1045. interval = upstreamBillingProbeMaxDelay
  1046. }
  1047. jitterRange := interval / 5
  1048. if jitterRange > 5*time.Minute {
  1049. jitterRange = 5 * time.Minute
  1050. }
  1051. if jitterRange > 0 {
  1052. interval += time.Duration(rand.Int64N(int64(jitterRange)*2+1)) - jitterRange
  1053. }
  1054. if retryAfterDuration > interval {
  1055. // Retry-After is an explicit upstream instruction; do not shorten it
  1056. // with the local maximum delay.
  1057. return retryAfterDuration
  1058. }
  1059. if interval > upstreamBillingProbeMaxDelay {
  1060. return upstreamBillingProbeMaxDelay
  1061. }
  1062. return interval
  1063. }
  1064. // unsupportedProbeDelay 拉长 unsupported 账号的重探间隔,让无效候选自然退出
  1065. // 热队列,不再和真正接入 sub2api 的中转账号抢每周期的探测名额。
  1066. // 仍按 upstreamBillingProbeMaxDelay 封顶,保证上游后来接入 sub2api 时最迟一天
  1067. // 内会被重新发现;base 本身已达上限(例如 Retry-After 明确要求更久)时原样返回,
  1068. // 不缩短上游指令。
  1069. func unsupportedProbeDelay(intervalMinutes int, retryAfterDuration time.Duration) time.Duration {
  1070. base := nextProbeDelay(intervalMinutes, retryAfterDuration)
  1071. if base >= upstreamBillingProbeMaxDelay {
  1072. return base
  1073. }
  1074. stretched := base * upstreamBillingProbeUnsupportedDelayFactor
  1075. if stretched > upstreamBillingProbeMaxDelay {
  1076. return upstreamBillingProbeMaxDelay
  1077. }
  1078. return stretched
  1079. }
  1080. func retryAfter(header http.Header, now time.Time) time.Duration {
  1081. value := strings.TrimSpace(header.Get("Retry-After"))
  1082. if value == "" {
  1083. return 0
  1084. }
  1085. if seconds, err := strconv.Atoi(value); err == nil && seconds > 0 {
  1086. return time.Duration(seconds) * time.Second
  1087. }
  1088. if at, err := http.ParseTime(value); err == nil {
  1089. if delay := at.Sub(now); delay > 0 {
  1090. return delay
  1091. }
  1092. }
  1093. return 0
  1094. }
  1095. func probeTimePtr(value time.Time) *time.Time {
  1096. return &value
  1097. }
  1098. func safeProbeError(err error) string {
  1099. if err == nil {
  1100. return ""
  1101. }
  1102. if errors.Is(err, ErrUpstreamBillingProbeAccountInvalid) {
  1103. return ErrUpstreamBillingProbeAccountInvalid.Error()
  1104. }
  1105. if errors.Is(err, ErrUpstreamBillingProbeUnavailable) {
  1106. return ErrUpstreamBillingProbeUnavailable.Error()
  1107. }
  1108. return "probe_failed"
  1109. }