Sub2API v1.0 - AI API 网关（二开初始版本，基于上游 Wei-Shaw/sub2api）

backend/internal/service/grok_quota_fetcher.go

@@ -0,0 +1,348 @@
+package service
+
+import (
+	"encoding/json"
+	"fmt"
+	"net/http"
+	"strings"
+	"time"
+
+	"github.com/Wei-Shaw/sub2api/internal/pkg/xai"
+)
+
+const grokQuotaSnapshotExtraKey = "grok_usage_snapshot"
+
+type GrokQuotaFetcher struct{}
+
+func NewGrokQuotaFetcher() *GrokQuotaFetcher {
+	return &GrokQuotaFetcher{}
+}
+
+func (f *GrokQuotaFetcher) BuildUsageInfo(account *Account) *UsageInfo {
+	now := time.Now()
+	usage := &UsageInfo{
+		Source:             "passive",
+		UpdatedAt:          &now,
+		GrokFreeTokenLimit: xai.GrokFreeRolling24hTokenLimit,
+	}
+	if account == nil {
+		usage.ErrorCode = "quota_unknown"
+		usage.Error = "Grok quota is unknown until billing is probed or an upstream response includes xAI rate-limit headers"
+		return usage
+	}
+
+	billing, _ := grokBillingSnapshotFromExtra(account.Extra)
+	snapshot, err := grokQuotaSnapshotFromExtra(account.Extra)
+	activeProbeClearsForbidden := newerSuccessfulGrokActiveProbeClearsBillingForbidden(billing, snapshot)
+	if billing != nil {
+		usage.GrokBilling = billing
+		applyGrokBillingProgressWindows(usage, billing, now)
+		if billing.Plan != "" {
+			usage.SubscriptionTier = billing.Plan
+			usage.SubscriptionTierRaw = billing.Plan
+		}
+		if parsedAt, parseErr := time.Parse(time.RFC3339, billing.UpdatedAt); parseErr == nil {
+			usage.UpdatedAt = &parsedAt
+		}
+		if billing.FetchedAt != "" {
+			usage.GrokLastQuotaProbeAt = billing.FetchedAt
+		}
+		usage.GrokQuotaSnapshotState = "billing_observed"
+		usage.GrokLastStatusCode = billing.StatusCode
+		switch billing.StatusCode {
+		case 401:
+			usage.NeedsReauth = true
+			usage.ErrorCode = "unauthenticated"
+		case 403:
+			usage.IsForbidden = true
+			usage.ForbiddenType = "forbidden"
+			usage.ErrorCode = "forbidden"
+		case 429:
+			usage.ErrorCode = "rate_limited"
+		}
+		// Official weekly/monthly progress clears the "unknown until headers" state.
+		if usage.ErrorCode == "quota_unknown" && (usage.SevenDay != nil || usage.ThirtyDay != nil) {
+			usage.ErrorCode = ""
+			if strings.Contains(strings.ToLower(usage.Error), "unknown until") ||
+				strings.Contains(strings.ToLower(usage.Error), "no xai quota headers") {
+				usage.Error = ""
+			}
+		}
+	}
+
+	if err != nil || snapshot == nil {
+		applyGrokCredentialUsageFallback(usage, account, billing, nil)
+		if billing == nil {
+			usage.ErrorCode = "quota_unknown"
+			usage.Error = "Grok quota is unknown until billing is probed or an upstream response includes xAI rate-limit headers"
+		}
+		return usage
+	}
+
+	if parsedAt, parseErr := time.Parse(time.RFC3339, snapshot.UpdatedAt); parseErr == nil {
+		if billing == nil || usage.UpdatedAt == nil || parsedAt.After(*usage.UpdatedAt) {
+			usage.UpdatedAt = &parsedAt
+		}
+	}
+	usage.GrokRequestQuota = snapshot.Requests
+	usage.GrokTokenQuota = snapshot.Tokens
+	usage.GrokRetryAfterSeconds = snapshot.RetryAfterSeconds
+	if usage.SubscriptionTier == "" {
+		usage.SubscriptionTier = snapshot.SubscriptionTier
+		usage.SubscriptionTierRaw = snapshot.SubscriptionTier
+	}
+	if usage.GrokEntitlementStatus == "" {
+		usage.GrokEntitlementStatus = snapshot.EntitlementStatus
+	}
+	if usage.GrokLastQuotaProbeAt == "" {
+		usage.GrokLastQuotaProbeAt = snapshot.LastProbeAt
+	}
+	usage.GrokLastHeadersSeenAt = snapshot.LastHeadersSeenAt
+	if activeProbeClearsForbidden {
+		usage.IsForbidden = false
+		usage.ForbiddenType = ""
+		usage.ErrorCode = ""
+		usage.GrokLastQuotaProbeAt = snapshot.LastProbeAt
+		usage.GrokLastStatusCode = snapshot.StatusCode
+	} else if snapshot.StatusCode >= http.StatusBadRequest || usage.GrokLastStatusCode == 0 {
+		usage.GrokLastStatusCode = snapshot.StatusCode
+	}
+	if snapshot.HasObservedHeaders() {
+		if usage.GrokQuotaSnapshotState == "" {
+			usage.GrokQuotaSnapshotState = "observed"
+		}
+	} else if billing == nil {
+		usage.GrokQuotaSnapshotState = "no_headers"
+		usage.ErrorCode = "quota_unknown"
+		usage.Error = "No xAI quota headers observed on the latest Grok probe"
+	}
+
+	if usage.ErrorCode == "" {
+		switch snapshot.StatusCode {
+		case 401:
+			usage.NeedsReauth = true
+			usage.ErrorCode = "unauthenticated"
+		case 403:
+			usage.IsForbidden = true
+			usage.ForbiddenType = "forbidden"
+			usage.ErrorCode = "forbidden"
+			if usage.GrokEntitlementStatus == "" {
+				usage.GrokEntitlementStatus = "forbidden"
+			}
+		case 429:
+			usage.ErrorCode = "rate_limited"
+		}
+	}
+	if accountGrokNeedsReauth(account) {
+		usage.NeedsReauth = true
+		if usage.ErrorCode == "" {
+			usage.ErrorCode = "spending_limit"
+		}
+	}
+	applyGrokCredentialUsageFallback(usage, account, billing, snapshot)
+	if activeProbeClearsForbidden && strings.TrimSpace(snapshot.EntitlementStatus) == "" &&
+		strings.EqualFold(strings.TrimSpace(usage.GrokEntitlementStatus), "forbidden") {
+		usage.GrokEntitlementStatus = ""
+	}
+	return usage
+}
+
+func newerSuccessfulGrokActiveProbeClearsBillingForbidden(billing *xai.BillingSummary, snapshot *xai.QuotaSnapshot) bool {
+	if billing == nil || billing.StatusCode != http.StatusForbidden || snapshot == nil ||
+		snapshot.StatusCode != http.StatusOK || strings.TrimSpace(snapshot.ObservationSource) != "active_probe" {
+		return false
+	}
+
+	billingAt, billingOK := firstGrokObservationTime(billing.UpdatedAt, billing.FetchedAt)
+	probeAt, probeOK := firstGrokObservationTime(snapshot.LastProbeAt, snapshot.UpdatedAt)
+	// Both snapshots use second precision, so a billing request followed by the
+	// active probe in the same refresh can legitimately have equal timestamps.
+	return billingOK && probeOK && !probeAt.Before(billingAt)
+}
+
+func firstGrokObservationTime(values ...string) (time.Time, bool) {
+	for _, value := range values {
+		parsedAt, err := time.Parse(time.RFC3339, strings.TrimSpace(value))
+		if err == nil {
+			return parsedAt, true
+		}
+	}
+	return time.Time{}, false
+}
+
+func applyGrokCredentialUsageFallback(usage *UsageInfo, account *Account, billing *xai.BillingSummary, snapshot *xai.QuotaSnapshot) {
+	if usage == nil || account == nil {
+		return
+	}
+	if usage.GrokEntitlementStatus == "" {
+		usage.GrokEntitlementStatus = strings.TrimSpace(account.GetCredential("entitlement_status"))
+	}
+	applyGrokResolvedSubscriptionTier(usage, account, billing, snapshot)
+}
+
+func applyGrokResolvedSubscriptionTier(usage *UsageInfo, account *Account, billing *xai.BillingSummary, snapshot *xai.QuotaSnapshot) {
+	if usage == nil || account == nil {
+		return
+	}
+	if jwtTier := xai.SubscriptionTierFromJWT(account.GetCredential("access_token")); jwtTier != "" {
+		usage.SubscriptionTier = jwtTier
+		usage.SubscriptionTierRaw = jwtTier
+		return
+	}
+	signal := strings.TrimSpace(account.GetCredential("subscription_tier"))
+	if signal == "" && snapshot != nil {
+		signal = strings.TrimSpace(snapshot.SubscriptionTier)
+	}
+	if signal == "" && billing != nil {
+		signal = strings.TrimSpace(billing.Plan)
+	}
+	var limit *float64
+	if billing != nil {
+		limit = billing.MonthlyLimitCents
+	}
+	if plan := xai.CanonicalGrokPlan(limit, signal, snapshot); plan != "" {
+		usage.SubscriptionTier = plan
+		if usage.SubscriptionTierRaw == "" {
+			usage.SubscriptionTierRaw = firstNonEmpty(signal, plan)
+		}
+		return
+	}
+	if usage.SubscriptionTier == "" && signal != "" {
+		usage.SubscriptionTier = signal
+		usage.SubscriptionTierRaw = signal
+	}
+}
+
+func grokBillingSnapshotFromExtra(extra map[string]any) (*xai.BillingSummary, error) {
+	if extra == nil {
+		return nil, nil
+	}
+	raw, ok := extra[grokBillingExtraKey]
+	if !ok || raw == nil {
+		return nil, nil
+	}
+	switch snapshot := raw.(type) {
+	case *xai.BillingSummary:
+		return snapshot, nil
+	case xai.BillingSummary:
+		return &snapshot, nil
+	case map[string]any:
+		data, err := json.Marshal(snapshot)
+		if err != nil {
+			return nil, err
+		}
+		var out xai.BillingSummary
+		if err := json.Unmarshal(data, &out); err != nil {
+			return nil, err
+		}
+		return &out, nil
+	default:
+		data, err := json.Marshal(raw)
+		if err != nil {
+			return nil, fmt.Errorf("marshal grok billing snapshot: %w", err)
+		}
+		var out xai.BillingSummary
+		if err := json.Unmarshal(data, &out); err != nil {
+			return nil, err
+		}
+		return &out, nil
+	}
+}
+
+func stampGrokQuotaSnapshotForPlan(account *Account, snapshot *xai.QuotaSnapshot, model string) {
+	if snapshot == nil {
+		return
+	}
+	if strings.TrimSpace(snapshot.Model) == "" {
+		model = strings.TrimSpace(model)
+		if model != "" {
+			snapshot.Model = xai.ResolveGrokTextResponsesModelID(model)
+		}
+	}
+	var prev *xai.QuotaSnapshot
+	if account != nil {
+		prev, _ = grokQuotaSnapshotFromExtra(account.Extra)
+	}
+	snapshot.ApplyGrok45ResponsesPlanSignal(prev)
+}
+
+func grokQuotaSnapshotFromExtra(extra map[string]any) (*xai.QuotaSnapshot, error) {
+	if extra == nil {
+		return nil, nil
+	}
+	raw, ok := extra[grokQuotaSnapshotExtraKey]
+	if !ok || raw == nil {
+		return nil, nil
+	}
+	switch snapshot := raw.(type) {
+	case *xai.QuotaSnapshot:
+		return snapshot, nil
+	case xai.QuotaSnapshot:
+		return &snapshot, nil
+	case map[string]any:
+		data, err := json.Marshal(snapshot)
+		if err != nil {
+			return nil, err
+		}
+		var out xai.QuotaSnapshot
+		if err := json.Unmarshal(data, &out); err != nil {
+			return nil, err
+		}
+		return &out, nil
+	default:
+		data, err := json.Marshal(raw)
+		if err != nil {
+			return nil, fmt.Errorf("marshal grok quota snapshot: %w", err)
+		}
+		var out xai.QuotaSnapshot
+		if err := json.Unmarshal(data, &out); err != nil {
+			return nil, err
+		}
+		return &out, nil
+	}
+}
+
+// applyGrokBillingProgressWindows fills official weekly (seven_day) and monthly
+// (thirty_day) UsageProgress from a billing probe summary.
+func applyGrokBillingProgressWindows(usage *UsageInfo, billing *xai.BillingSummary, now time.Time) {
+	if usage == nil || billing == nil {
+		return
+	}
+	if billing.UsagePercent != nil {
+		seven := &UsageProgress{Utilization: *billing.UsagePercent}
+		if end, err := parseTime(strings.TrimSpace(billing.PeriodEnd)); err == nil {
+			seven.ResetsAt = &end
+			if sec := int(end.Sub(now).Seconds()); sec > 0 {
+				seven.RemainingSeconds = sec
+			}
+		}
+		if usage.SevenDay != nil {
+			seven.WindowStats = usage.SevenDay.WindowStats
+		}
+		usage.SevenDay = seven
+	}
+	var monthlyUtil *float64
+	if billing.UsedPercent != nil {
+		monthlyUtil = billing.UsedPercent
+	} else if billing.MonthlyLimitCents != nil && *billing.MonthlyLimitCents > 0 && billing.UsedCents != nil {
+		v := (*billing.UsedCents / *billing.MonthlyLimitCents) * 100
+		monthlyUtil = &v
+	}
+	if monthlyUtil != nil {
+		thirty := &UsageProgress{Utilization: *monthlyUtil}
+		endRaw := strings.TrimSpace(billing.BillingPeriodEnd)
+		if endRaw == "" && billing.PeriodType == "monthly" {
+			endRaw = strings.TrimSpace(billing.PeriodEnd)
+		}
+		if end, err := parseTime(endRaw); err == nil {
+			thirty.ResetsAt = &end
+			if sec := int(end.Sub(now).Seconds()); sec > 0 {
+				thirty.RemainingSeconds = sec
+			}
+		}
+		if usage.ThirtyDay != nil {
+			thirty.WindowStats = usage.ThirtyDay.WindowStats
+		}
+		usage.ThirtyDay = thirty
+	}
+}