grok_quota_service.go 21 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574
  1. package service
  2. import (
  3. "bytes"
  4. "context"
  5. "encoding/json"
  6. "io"
  7. "log/slog"
  8. "net/http"
  9. "strconv"
  10. "strings"
  11. "sync"
  12. "time"
  13. "github.com/Wei-Shaw/sub2api/internal/config"
  14. infraerrors "github.com/Wei-Shaw/sub2api/internal/pkg/errors"
  15. "github.com/Wei-Shaw/sub2api/internal/pkg/xai"
  16. "golang.org/x/sync/singleflight"
  17. )
  18. const (
  19. grokQuotaUpstreamTimeout = 20 * time.Second
  20. grokQuotaProbeInput = "hi"
  21. grokQuotaDefaultModel = grokDefaultResponsesModel
  22. grokBillingExtraKey = "grok_billing_snapshot"
  23. grokBillingMaxAttempts = 2
  24. grokBillingRetryDelay = 100 * time.Millisecond
  25. )
  26. type GrokQuotaProbeResult struct {
  27. Source string `json:"source"`
  28. Model string `json:"model,omitempty"`
  29. Billing *xai.BillingSummary `json:"billing,omitempty"`
  30. Snapshot *xai.QuotaSnapshot `json:"snapshot,omitempty"`
  31. LocalUsage24h *WindowStats `json:"local_usage_24h,omitempty"`
  32. LocalUsage7d *WindowStats `json:"local_usage_7d,omitempty"`
  33. LocalUsageMonthly *WindowStats `json:"local_usage_monthly,omitempty"`
  34. StatusCode int `json:"status_code,omitempty"`
  35. HeadersObserved bool `json:"headers_observed"`
  36. ResetSupported bool `json:"reset_supported"`
  37. FetchedAt int64 `json:"fetched_at"`
  38. Persisted bool `json:"persisted"`
  39. ProbeError string `json:"probe_error,omitempty"`
  40. }
  41. type GrokQuotaResetResult struct {
  42. Supported bool `json:"supported"`
  43. Code string `json:"code"`
  44. Message string `json:"message"`
  45. }
  46. type GrokQuotaService struct {
  47. accountRepo AccountRepository
  48. proxyRepo ProxyRepository
  49. tokenProvider *GrokTokenProvider
  50. httpUpstream HTTPUpstream
  51. usageLogRepo UsageLogRepository
  52. settingService *SettingService
  53. cfg *config.Config
  54. probeFlight singleflight.Group
  55. }
  56. func NewGrokQuotaService(
  57. accountRepo AccountRepository,
  58. proxyRepo ProxyRepository,
  59. tokenProvider *GrokTokenProvider,
  60. httpUpstream HTTPUpstream,
  61. cfg *config.Config,
  62. usageLogRepos ...UsageLogRepository,
  63. ) *GrokQuotaService {
  64. var usageLogRepo UsageLogRepository
  65. if len(usageLogRepos) > 0 {
  66. usageLogRepo = usageLogRepos[0]
  67. }
  68. return &GrokQuotaService{
  69. accountRepo: accountRepo,
  70. proxyRepo: proxyRepo,
  71. tokenProvider: tokenProvider,
  72. httpUpstream: httpUpstream,
  73. usageLogRepo: usageLogRepo,
  74. cfg: cfg,
  75. }
  76. }
  77. func (s *GrokQuotaService) SetSettingService(settingService *SettingService) {
  78. if s != nil {
  79. s.settingService = settingService
  80. }
  81. }
  82. // QueryQuota combines xAI billing data with an active quota-header probe for
  83. // Free accounts, whose billing response does not include usage_percent.
  84. func (s *GrokQuotaService) QueryQuota(ctx context.Context, accountID int64) (*GrokQuotaProbeResult, error) {
  85. billingResult, billingErr := s.ProbeBilling(ctx, accountID)
  86. if billingErr == nil && billingResult != nil && grokBillingHasAuthoritativeQuota(billingResult.Billing) {
  87. if acc, err := s.accountRepo.GetByID(ctx, accountID); err == nil {
  88. s.scheduleGrokObservedModelsSync(acc)
  89. }
  90. return billingResult, nil
  91. }
  92. probeResult, probeErr := s.ProbeUsage(ctx, accountID)
  93. if probeErr != nil {
  94. if billingResult != nil && billingResult.Billing != nil {
  95. billingResult.ProbeError = probeErr.Error()
  96. return billingResult, nil
  97. }
  98. return nil, probeErr
  99. }
  100. if probeResult == nil {
  101. if billingErr != nil {
  102. return nil, billingErr
  103. }
  104. return nil, infraerrors.New(http.StatusBadGateway, "GROK_QUOTA_PROBE_EMPTY", "Grok quota probe returned no result")
  105. }
  106. if billingResult != nil {
  107. probeResult.Source = "hybrid_probe"
  108. probeResult.Billing = billingResult.Billing
  109. probeResult.LocalUsage24h = billingResult.LocalUsage24h
  110. probeResult.LocalUsage7d = billingResult.LocalUsage7d
  111. probeResult.LocalUsageMonthly = billingResult.LocalUsageMonthly
  112. probeResult.Persisted = probeResult.Persisted || billingResult.Persisted
  113. }
  114. if acc, err := s.accountRepo.GetByID(ctx, accountID); err == nil {
  115. s.scheduleGrokObservedModelsSync(acc)
  116. }
  117. return probeResult, nil
  118. }
  119. func grokBillingHasAuthoritativeQuota(billing *xai.BillingSummary) bool {
  120. if billing == nil {
  121. return false
  122. }
  123. return billing.UsagePercent != nil ||
  124. billing.UsedPercent != nil ||
  125. (billing.MonthlyLimitCents != nil && *billing.MonthlyLimitCents > 0) ||
  126. strings.TrimSpace(billing.Plan) != ""
  127. }
  128. func (s *GrokQuotaService) ProbeUsage(ctx context.Context, accountID int64) (*GrokQuotaProbeResult, error) {
  129. return s.runProbeFlight(ctx, "active:"+strconv.FormatInt(accountID, 10), func(sharedCtx context.Context) (*GrokQuotaProbeResult, error) {
  130. return s.probeUsage(sharedCtx, accountID)
  131. })
  132. }
  133. func (s *GrokQuotaService) probeUsage(ctx context.Context, accountID int64) (*GrokQuotaProbeResult, error) {
  134. account, token, proxyURL, err := s.prepareProbe(ctx, accountID)
  135. if err != nil {
  136. return nil, err
  137. }
  138. probeModel := grokQuotaProbeModel()
  139. body, err := buildGrokQuotaProbeBody(probeModel)
  140. if err != nil {
  141. return nil, infraerrors.Newf(http.StatusBadRequest, "GROK_QUOTA_PROBE_BODY_ERROR", "failed to build probe body: %v", err)
  142. }
  143. targetURL, err := buildGrokResponsesURL(account, s.cfg, s.settingService)
  144. if err != nil {
  145. return nil, infraerrors.Newf(http.StatusBadRequest, "GROK_QUOTA_BASE_URL_INVALID", "invalid Grok base_url: %v", err)
  146. }
  147. callCtx, cancel := context.WithTimeout(ctx, grokQuotaUpstreamTimeout)
  148. defer cancel()
  149. req, err := http.NewRequestWithContext(callCtx, http.MethodPost, targetURL, bytes.NewReader(body))
  150. if err != nil {
  151. return nil, infraerrors.Newf(http.StatusInternalServerError, "GROK_QUOTA_PROBE_REQUEST_BUILD_FAILED", "failed to build upstream request: %v", err)
  152. }
  153. req.Header.Set("Authorization", "Bearer "+token)
  154. req.Header.Set("Content-Type", "application/json")
  155. req.Header.Set("Accept", "application/json, text/event-stream")
  156. if account.IsGrokOAuth() {
  157. applyGrokCLIHeaders(req.Header)
  158. }
  159. // 探测请求与真实转发保持同一套账号级请求头覆写,避免探测通过但转发失败。
  160. account.ApplyHeaderOverrides(req.Header)
  161. resp, err := s.httpUpstream.Do(req, proxyURL, account.ID, maxInt(account.Concurrency, 1))
  162. if err != nil {
  163. return nil, infraerrors.Newf(http.StatusBadGateway, "GROK_QUOTA_PROBE_REQUEST_FAILED", "upstream probe failed: %v", err)
  164. }
  165. defer func() { _ = resp.Body.Close() }()
  166. snapshot := xai.ObserveQuotaHeaders(resp.Header, resp.StatusCode, "active_probe")
  167. stampGrokQuotaSnapshotForPlan(account, snapshot, probeModel)
  168. resetAt, limited := grokRateLimitResetAtForAccount(account, snapshot, time.Now())
  169. if limited {
  170. normalizeGrokExhaustedWindowResets(snapshot, resetAt, time.Now())
  171. }
  172. // A failed probe must not erase a previously observed snapshot. 401/403 and
  173. // transport/server errors commonly carry no quota headers; only successful
  174. // responses, or 429 responses with useful rate-limit headers, are safe to
  175. // persist. A successful 200 with no headers is still persisted as an
  176. // explicit "no headers" observation so the UI can distinguish it from never
  177. // probed.
  178. persistErr := error(nil)
  179. persisted := false
  180. shouldPersist := resp.StatusCode < 400 || resp.StatusCode == http.StatusTooManyRequests
  181. if shouldPersist && (snapshot.HeadersObserved || resp.StatusCode == http.StatusOK) {
  182. persistErr = s.accountRepo.UpdateExtra(ctx, account.ID, map[string]any{
  183. grokQuotaSnapshotExtraKey: snapshot,
  184. })
  185. persisted = persistErr == nil
  186. }
  187. if limited {
  188. persistGrokRateLimit(ctx, s.accountRepo, account, resetAt)
  189. } else if isSuccessfulGrokRateLimitRecovery(account, snapshot) {
  190. clearGrokRateLimitAfterRecovery(ctx, s.accountRepo, account)
  191. }
  192. result := &GrokQuotaProbeResult{
  193. Source: "active_probe",
  194. Model: probeModel,
  195. Snapshot: snapshot,
  196. StatusCode: resp.StatusCode,
  197. HeadersObserved: snapshot.HeadersObserved,
  198. ResetSupported: false,
  199. FetchedAt: time.Now().Unix(),
  200. Persisted: persisted,
  201. }
  202. if resp.StatusCode == http.StatusTooManyRequests {
  203. return result, nil
  204. }
  205. if resp.StatusCode >= 400 {
  206. const reason = "GROK_QUOTA_PROBE_UPSTREAM_ERROR"
  207. _, _ = io.Copy(io.Discard, io.LimitReader(resp.Body, 4<<10))
  208. slog.Warn(
  209. "grok_quota_probe_failed",
  210. "account_id", account.ID,
  211. "model", probeModel,
  212. "status", resp.StatusCode,
  213. "reason", reason,
  214. )
  215. return nil, infraerrors.Newf(
  216. mapUpstreamStatus(resp.StatusCode),
  217. reason,
  218. "upstream returned %d for probe model %q",
  219. resp.StatusCode,
  220. probeModel,
  221. )
  222. }
  223. return result, nil
  224. }
  225. // ProbeBilling only calls the xAI billing endpoints. Account usage refreshes
  226. // use this method so opening the account list never consumes model quota.
  227. func (s *GrokQuotaService) ProbeBilling(ctx context.Context, accountID int64) (*GrokQuotaProbeResult, error) {
  228. return s.runProbeFlight(ctx, "billing:"+strconv.FormatInt(accountID, 10), func(sharedCtx context.Context) (*GrokQuotaProbeResult, error) {
  229. return s.probeBilling(sharedCtx, accountID)
  230. })
  231. }
  232. // ProbeMediaEligibility refreshes billing state and evaluates the persisted
  233. // account snapshot used by media scheduling. Probe failures remain fail-closed;
  234. // deterministic persisted states such as forbidden or Free are returned as
  235. // normal ineligibility decisions rather than transport errors.
  236. func (s *GrokQuotaService) ProbeMediaEligibility(ctx context.Context, accountID int64) (bool, string, error) {
  237. _, probeErr := s.ProbeBilling(ctx, accountID)
  238. account, err := s.loadGrokOAuthAccount(ctx, accountID)
  239. if err != nil {
  240. return false, "billing_probe_failed", err
  241. }
  242. eligible, reason := account.GrokMediaGenerationEligibility()
  243. if reason == "billing_unobserved" && probeErr != nil {
  244. return false, reason, probeErr
  245. }
  246. return eligible, reason, nil
  247. }
  248. func (s *GrokQuotaService) probeBilling(ctx context.Context, accountID int64) (*GrokQuotaProbeResult, error) {
  249. account, token, proxyURL, err := s.prepareProbe(ctx, accountID)
  250. if err != nil {
  251. return nil, err
  252. }
  253. probeCtx, cancel := context.WithTimeout(ctx, grokQuotaUpstreamTimeout)
  254. defer cancel()
  255. type billingResult struct {
  256. summary *xai.BillingSummary
  257. status int
  258. err error
  259. }
  260. var weekly, monthly billingResult
  261. var wg sync.WaitGroup
  262. wg.Add(2)
  263. go func() {
  264. defer wg.Done()
  265. weekly.summary, weekly.status, weekly.err = s.fetchBilling(probeCtx, account, token, proxyURL, true)
  266. }()
  267. go func() {
  268. defer wg.Done()
  269. monthly.summary, monthly.status, monthly.err = s.fetchBilling(probeCtx, account, token, proxyURL, false)
  270. }()
  271. wg.Wait()
  272. weeklyOK := weekly.summary != nil
  273. monthlyOK := monthly.summary != nil
  274. previous, _ := grokBillingSnapshotFromExtra(account.Extra)
  275. if !weeklyOK && !monthlyOK {
  276. probeErr := mergeGrokBillingProbeErrors(weekly.status, monthly.status, weekly.err, monthly.err)
  277. billing := xai.MergeBillingProbeResult(previous, nil, nil, false, false)
  278. if billing == nil {
  279. billing = &xai.BillingSummary{Partial: true, FailedWindows: []string{"weekly", "monthly"}}
  280. }
  281. billing.WeeklyStatusCode = weekly.status
  282. billing.MonthlyStatusCode = monthly.status
  283. billing = xai.StampBillingSummary(billing, preferBillingObservationStatus(weekly.status, monthly.status), "billing_probe")
  284. if persistErr := s.accountRepo.UpdateExtra(ctx, account.ID, map[string]any{grokBillingExtraKey: billing}); persistErr != nil {
  285. slog.Warn("grok_billing_failure_persist_failed", "account_id", account.ID, "error", persistErr)
  286. }
  287. return nil, probeErr
  288. }
  289. statusCode := preferSuccessfulBillingStatus(weekly.status, monthly.status, weeklyOK, monthlyOK)
  290. billing := xai.MergeBillingProbeResult(previous, weekly.summary, monthly.summary, weeklyOK, monthlyOK)
  291. billing.WeeklyStatusCode = weekly.status
  292. billing.MonthlyStatusCode = monthly.status
  293. billing = xai.StampBillingSummary(billing, statusCode, "billing_probe")
  294. persistErr := s.accountRepo.UpdateExtra(ctx, account.ID, map[string]any{
  295. grokBillingExtraKey: billing,
  296. })
  297. if persistErr != nil {
  298. slog.Warn("grok_billing_persist_failed", "account_id", account.ID, "error", persistErr)
  299. }
  300. now := time.Now().UTC()
  301. localUsage24h, localUsage7d, localUsageMonthly := grokLocalUsageForQuota(ctx, s.usageLogRepo, account.ID, billing, now)
  302. return &GrokQuotaProbeResult{
  303. Source: "billing_probe",
  304. Billing: billing,
  305. LocalUsage24h: localUsage24h,
  306. LocalUsage7d: localUsage7d,
  307. LocalUsageMonthly: localUsageMonthly,
  308. StatusCode: statusCode,
  309. FetchedAt: now.Unix(),
  310. Persisted: persistErr == nil,
  311. }, nil
  312. }
  313. func preferBillingObservationStatus(weeklyStatus, monthlyStatus int) int {
  314. if weeklyStatus == http.StatusForbidden || monthlyStatus == http.StatusForbidden {
  315. return http.StatusForbidden
  316. }
  317. if weeklyStatus != 0 {
  318. return weeklyStatus
  319. }
  320. return monthlyStatus
  321. }
  322. func (s *GrokQuotaService) runProbeFlight(
  323. ctx context.Context,
  324. key string,
  325. probe func(context.Context) (*GrokQuotaProbeResult, error),
  326. ) (*GrokQuotaProbeResult, error) {
  327. if s == nil {
  328. return nil, infraerrors.New(http.StatusInternalServerError, "GROK_QUOTA_NOT_CONFIGURED", "grok quota service is not configured")
  329. }
  330. resultCh := s.probeFlight.DoChan(key, func() (any, error) {
  331. sharedCtx, cancel := context.WithTimeout(context.Background(), grokQuotaUpstreamTimeout+5*time.Second)
  332. defer cancel()
  333. return probe(sharedCtx)
  334. })
  335. select {
  336. case <-ctx.Done():
  337. return nil, ctx.Err()
  338. case flightResult := <-resultCh:
  339. if flightResult.Err != nil {
  340. return nil, flightResult.Err
  341. }
  342. result, ok := flightResult.Val.(*GrokQuotaProbeResult)
  343. if !ok || result == nil {
  344. return nil, infraerrors.New(http.StatusInternalServerError, "GROK_QUOTA_PROBE_RESULT_INVALID", "invalid Grok quota probe result")
  345. }
  346. cloned := *result
  347. return &cloned, nil
  348. }
  349. }
  350. func (s *GrokQuotaService) fetchBilling(
  351. ctx context.Context,
  352. account *Account,
  353. token string,
  354. proxyURL string,
  355. weekly bool,
  356. ) (*xai.BillingSummary, int, error) {
  357. billingURL, err := buildGrokBillingURL(account, s.cfg, weekly)
  358. if err != nil {
  359. return nil, 0, infraerrors.Newf(http.StatusBadRequest, "GROK_QUOTA_BASE_URL_INVALID", "invalid Grok base_url: %v", err)
  360. }
  361. for attempt := 0; attempt < grokBillingMaxAttempts; attempt++ {
  362. req, err := http.NewRequestWithContext(ctx, http.MethodGet, billingURL, nil)
  363. if err != nil {
  364. return nil, 0, infraerrors.Newf(http.StatusInternalServerError, "GROK_QUOTA_PROBE_REQUEST_BUILD_FAILED", "failed to build billing request: %v", err)
  365. }
  366. xai.ApplyCLIBillingHeaders(req, token)
  367. // billing 探测与真实转发保持同一套账号级请求头覆写。
  368. account.ApplyHeaderOverrides(req.Header)
  369. resp, requestErr := s.httpUpstream.Do(req, proxyURL, account.ID, maxInt(account.Concurrency, 2))
  370. statusCode := 0
  371. var bodyBytes []byte
  372. if requestErr == nil {
  373. statusCode = resp.StatusCode
  374. bodyBytes, _ = io.ReadAll(io.LimitReader(resp.Body, 1<<20))
  375. _ = resp.Body.Close()
  376. }
  377. shouldRetry := requestErr != nil || isRetryableGrokBillingStatus(statusCode)
  378. if shouldRetry && attempt+1 < grokBillingMaxAttempts {
  379. timer := time.NewTimer(grokBillingRetryDelay)
  380. select {
  381. case <-timer.C:
  382. case <-ctx.Done():
  383. timer.Stop()
  384. return nil, statusCode, infraerrors.Newf(http.StatusBadGateway, "GROK_QUOTA_PROBE_REQUEST_FAILED", "billing request failed: %v", ctx.Err())
  385. }
  386. continue
  387. }
  388. if requestErr != nil {
  389. return nil, 0, infraerrors.Newf(http.StatusBadGateway, "GROK_QUOTA_PROBE_REQUEST_FAILED", "billing request failed: %v", requestErr)
  390. }
  391. if statusCode == http.StatusTooManyRequests {
  392. return nil, statusCode, nil
  393. }
  394. if statusCode >= 400 {
  395. bodyText := truncate(strings.TrimSpace(string(bodyBytes)), 240)
  396. slog.Warn("grok_quota_billing_failed", "account_id", account.ID, "weekly", weekly, "status", statusCode, "body", bodyText)
  397. return nil, statusCode, infraerrors.Newf(mapUpstreamStatus(statusCode), "GROK_QUOTA_PROBE_UPSTREAM_ERROR", "billing returned %d: %s", statusCode, bodyText)
  398. }
  399. payload, err := xai.ParseBillingPayload(bodyBytes)
  400. if err != nil {
  401. return nil, statusCode, infraerrors.Newf(http.StatusBadGateway, "GROK_QUOTA_BILLING_PARSE_ERROR", "failed to parse billing body: %v", err)
  402. }
  403. return xai.BuildBillingSummary(payload.Config), statusCode, nil
  404. }
  405. return nil, 0, infraerrors.New(http.StatusBadGateway, "GROK_QUOTA_PROBE_REQUEST_FAILED", "billing request failed")
  406. }
  407. func isRetryableGrokBillingStatus(statusCode int) bool {
  408. switch statusCode {
  409. case http.StatusBadGateway, http.StatusServiceUnavailable, http.StatusGatewayTimeout:
  410. return true
  411. default:
  412. return false
  413. }
  414. }
  415. func mergeGrokBillingProbeErrors(weeklyStatus, monthlyStatus int, weeklyErr, monthlyErr error) error {
  416. weeklyKey := grokBillingProbeErrorKey(weeklyStatus, weeklyErr)
  417. monthlyKey := grokBillingProbeErrorKey(monthlyStatus, monthlyErr)
  418. if weeklyKey == monthlyKey {
  419. switch {
  420. case weeklyErr != nil:
  421. return weeklyErr
  422. case monthlyErr != nil:
  423. return monthlyErr
  424. case weeklyStatus == http.StatusTooManyRequests:
  425. return infraerrors.New(http.StatusTooManyRequests, "GROK_QUOTA_PROBE_UPSTREAM_ERROR", "billing rate limited")
  426. case weeklyStatus != 0 && weeklyStatus != http.StatusOK:
  427. return infraerrors.New(mapUpstreamStatus(weeklyStatus), "GROK_QUOTA_PROBE_UPSTREAM_ERROR", "xAI billing endpoints returned the same upstream error")
  428. default:
  429. return infraerrors.New(http.StatusBadGateway, "GROK_QUOTA_BILLING_EMPTY", "xAI billing endpoints returned no quota data")
  430. }
  431. }
  432. slog.Warn("grok_quota_probe_parts_failed", "weekly_status", weeklyStatus, "weekly_error", weeklyErr, "monthly_status", monthlyStatus, "monthly_error", monthlyErr)
  433. return infraerrors.New(http.StatusBadGateway, "GROK_QUOTA_PROBE_PARTS_FAILED", "weekly and monthly billing probes failed differently").WithMetadata(map[string]string{
  434. "weekly_status": strconv.Itoa(weeklyStatus), "monthly_status": strconv.Itoa(monthlyStatus),
  435. })
  436. }
  437. func grokBillingProbeErrorKey(status int, err error) string {
  438. if err != nil {
  439. return strconv.Itoa(status) + ":" + strconv.Itoa(infraerrors.Code(err)) + ":" + infraerrors.Reason(err)
  440. }
  441. return strconv.Itoa(status) + ":empty"
  442. }
  443. func preferSuccessfulBillingStatus(weeklyStatus, monthlyStatus int, weeklyOK, monthlyOK bool) int {
  444. if weeklyOK && weeklyStatus >= 200 && weeklyStatus < 300 {
  445. return weeklyStatus
  446. }
  447. if monthlyOK && monthlyStatus >= 200 && monthlyStatus < 300 {
  448. return monthlyStatus
  449. }
  450. if weeklyStatus != 0 {
  451. return weeklyStatus
  452. }
  453. return monthlyStatus
  454. }
  455. func (s *GrokQuotaService) ResetQuota(ctx context.Context, accountID int64) (*GrokQuotaResetResult, error) {
  456. if _, err := s.loadGrokOAuthAccount(ctx, accountID); err != nil {
  457. return nil, err
  458. }
  459. return nil, infraerrors.New(http.StatusNotImplemented, "GROK_QUOTA_RESET_UNSUPPORTED", "xAI does not expose a Grok subscription quota reset endpoint for OAuth accounts")
  460. }
  461. func (s *GrokQuotaService) prepareProbe(ctx context.Context, accountID int64) (*Account, string, string, error) {
  462. if s == nil || s.tokenProvider == nil || s.httpUpstream == nil {
  463. return nil, "", "", infraerrors.New(http.StatusInternalServerError, "GROK_QUOTA_NOT_CONFIGURED", "grok quota service is not configured")
  464. }
  465. account, err := s.loadGrokOAuthAccount(ctx, accountID)
  466. if err != nil {
  467. return nil, "", "", err
  468. }
  469. proxyURL := s.resolveProxyURL(ctx, account)
  470. token, err := s.tokenProvider.GetAccessToken(ctx, account)
  471. if err != nil {
  472. return nil, "", "", infraerrors.Newf(http.StatusBadGateway, "GROK_QUOTA_TOKEN_UNAVAILABLE", "failed to acquire access token: %v", err)
  473. }
  474. if strings.TrimSpace(token) == "" {
  475. return nil, "", "", infraerrors.New(http.StatusBadGateway, "GROK_QUOTA_TOKEN_UNAVAILABLE", "access token is empty")
  476. }
  477. return account, token, proxyURL, nil
  478. }
  479. func (s *GrokQuotaService) resolveProxyURL(ctx context.Context, account *Account) string {
  480. if account == nil || account.ProxyID == nil {
  481. return ""
  482. }
  483. switch {
  484. case account.Proxy != nil:
  485. return account.Proxy.URL()
  486. case s != nil && s.proxyRepo != nil:
  487. if proxy, err := s.proxyRepo.GetByID(ctx, *account.ProxyID); err == nil && proxy != nil {
  488. account.Proxy = proxy
  489. return proxy.URL()
  490. }
  491. }
  492. return ""
  493. }
  494. func (s *GrokQuotaService) loadGrokOAuthAccount(ctx context.Context, accountID int64) (*Account, error) {
  495. if s == nil || s.accountRepo == nil {
  496. return nil, infraerrors.New(http.StatusInternalServerError, "GROK_QUOTA_NOT_CONFIGURED", "grok quota service is not configured")
  497. }
  498. account, err := s.accountRepo.GetByID(ctx, accountID)
  499. if err != nil {
  500. return nil, infraerrors.Newf(http.StatusNotFound, "GROK_QUOTA_ACCOUNT_NOT_FOUND", "account not found: %v", err)
  501. }
  502. if account == nil {
  503. return nil, infraerrors.New(http.StatusNotFound, "GROK_QUOTA_ACCOUNT_NOT_FOUND", "account not found")
  504. }
  505. if account.Platform != PlatformGrok {
  506. return nil, infraerrors.New(http.StatusBadRequest, "GROK_QUOTA_INVALID_PLATFORM", "account is not a Grok account")
  507. }
  508. if account.Type != AccountTypeOAuth {
  509. return nil, infraerrors.New(http.StatusBadRequest, "GROK_QUOTA_INVALID_TYPE", "account is not an OAuth account")
  510. }
  511. return account, nil
  512. }
  513. func grokQuotaProbeModel() string {
  514. return grokQuotaDefaultModel
  515. }
  516. func buildGrokQuotaProbeBody(model string) ([]byte, error) {
  517. model = strings.TrimSpace(model)
  518. if model == "" {
  519. model = grokQuotaDefaultModel
  520. }
  521. return json.Marshal(map[string]any{
  522. "model": model,
  523. "input": grokQuotaProbeInput,
  524. "stream": true,
  525. })
  526. }
  527. func maxInt(a, b int) int {
  528. if a > b {
  529. return a
  530. }
  531. return b
  532. }