mirror of
https://github.com/Wei-Shaw/sub2api.git
synced 2026-09-24 16:05:44 +08:00
feat(grok): improve free quota probing and usage
This commit is contained in:
superman2003
@@ -0,0 +1,372 @@
|
||||
package xai
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"math"
|
||||
"net/http"
|
||||
"strconv"
|
||||
"strings"
|
||||
"time"
|
||||
)
|
||||
|
||||
const (
|
||||
// CLI client identity required by cli-chat-proxy billing endpoints.
|
||||
CLITokenAuthHeader = "x-xai-token-auth"
|
||||
CLITokenAuthValue = "xai-grok-cli"
|
||||
CLIClientVersionHeader = "x-grok-client-version"
|
||||
// Keep in sync with https://x.ai/cli/stable.
|
||||
CLIClientVersion = "0.2.93"
|
||||
CLIUserAgent = "grok-pager/" + CLIClientVersion + " grok-shell/" + CLIClientVersion + " (macos; aarch64)"
|
||||
|
||||
BillingWeeklyPath = "/billing?format=credits"
|
||||
BillingMonthlyPath = "/billing"
|
||||
|
||||
SuperGrokLimitCents = 15_000 // $150.00
|
||||
SuperGrokHeavyLimitCents = 150_000 // $1,500.00
|
||||
)
|
||||
|
||||
// BillingPeriod describes the current weekly/monthly window.
|
||||
type BillingPeriod struct {
|
||||
Type string `json:"type,omitempty"`
|
||||
Start string `json:"start,omitempty"`
|
||||
End string `json:"end,omitempty"`
|
||||
}
|
||||
|
||||
// BillingProductUsage is per-product usage inside the weekly credits window.
|
||||
type BillingProductUsage struct {
|
||||
Product string `json:"product,omitempty"`
|
||||
UsagePercent *float64 `json:"usagePercent,omitempty"`
|
||||
}
|
||||
|
||||
// BillingConfig is the nested config object from /v1/billing responses.
|
||||
type BillingConfig struct {
|
||||
CurrentPeriod *BillingPeriod `json:"currentPeriod,omitempty"`
|
||||
CreditUsagePercent *float64 `json:"creditUsagePercent,omitempty"`
|
||||
ProductUsage []BillingProductUsage `json:"productUsage,omitempty"`
|
||||
MonthlyLimit json.RawMessage `json:"monthlyLimit,omitempty"`
|
||||
Used json.RawMessage `json:"used,omitempty"`
|
||||
BillingPeriodStart string `json:"billingPeriodStart,omitempty"`
|
||||
BillingPeriodEnd string `json:"billingPeriodEnd,omitempty"`
|
||||
}
|
||||
|
||||
// BillingPayload is the top-level body from /v1/billing.
|
||||
type BillingPayload struct {
|
||||
Config *BillingConfig `json:"config,omitempty"`
|
||||
}
|
||||
|
||||
// BillingProductSummary is a normalized product usage row for UI.
|
||||
type BillingProductSummary struct {
|
||||
Product string `json:"product"`
|
||||
UsagePercent *float64 `json:"usage_percent,omitempty"`
|
||||
}
|
||||
|
||||
// BillingSummary is the merged weekly + monthly billing view.
|
||||
type BillingSummary struct {
|
||||
PeriodType string `json:"period_type,omitempty"` // weekly | monthly | unknown
|
||||
UsagePercent *float64 `json:"usage_percent,omitempty"`
|
||||
PeriodStart string `json:"period_start,omitempty"`
|
||||
PeriodEnd string `json:"period_end,omitempty"`
|
||||
ProductUsage []BillingProductSummary `json:"product_usage,omitempty"`
|
||||
MonthlyLimitCents *float64 `json:"monthly_limit_cents,omitempty"`
|
||||
UsedCents *float64 `json:"used_cents,omitempty"`
|
||||
IncludedUsedCents *float64 `json:"included_used_cents,omitempty"`
|
||||
BillingPeriodStart string `json:"billing_period_start,omitempty"`
|
||||
BillingPeriodEnd string `json:"billing_period_end,omitempty"`
|
||||
UsedPercent *float64 `json:"used_percent,omitempty"`
|
||||
Plan string `json:"plan,omitempty"` // SuperGrok | SuperGrok Heavy | ""
|
||||
StatusCode int `json:"status_code,omitempty"`
|
||||
Source string `json:"source,omitempty"`
|
||||
FetchedAt string `json:"fetched_at,omitempty"`
|
||||
UpdatedAt string `json:"updated_at,omitempty"`
|
||||
WeeklyUpdatedAt string `json:"weekly_updated_at,omitempty"`
|
||||
MonthlyUpdatedAt string `json:"monthly_updated_at,omitempty"`
|
||||
Partial bool `json:"partial,omitempty"`
|
||||
FailedWindows []string `json:"failed_windows,omitempty"`
|
||||
}
|
||||
|
||||
// BuildBillingURL builds weekly or monthly billing URL against the CLI chat proxy.
|
||||
func BuildBillingURL(formatCredits bool) string {
|
||||
base := strings.TrimRight(DefaultCLIBaseURL, "/")
|
||||
if formatCredits {
|
||||
return base + BillingWeeklyPath
|
||||
}
|
||||
return base + BillingMonthlyPath
|
||||
}
|
||||
|
||||
// ApplyCLIBillingHeaders sets Authorization + CLI identity headers for billing GETs.
|
||||
func ApplyCLIBillingHeaders(req *http.Request, accessToken string) {
|
||||
if req == nil {
|
||||
return
|
||||
}
|
||||
token := strings.TrimSpace(accessToken)
|
||||
if token != "" {
|
||||
req.Header.Set("Authorization", "Bearer "+token)
|
||||
}
|
||||
req.Header.Set("Accept", "application/json")
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
req.Header.Set(CLITokenAuthHeader, CLITokenAuthValue)
|
||||
req.Header.Set(CLIClientVersionHeader, CLIClientVersion)
|
||||
req.Header.Set("User-Agent", CLIUserAgent)
|
||||
}
|
||||
|
||||
// ParseBillingPayload unmarshals a billing API response body.
|
||||
func ParseBillingPayload(body []byte) (*BillingPayload, error) {
|
||||
if len(body) == 0 {
|
||||
return nil, fmt.Errorf("empty billing body")
|
||||
}
|
||||
var payload BillingPayload
|
||||
if err := json.Unmarshal(body, &payload); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return &payload, nil
|
||||
}
|
||||
|
||||
// BuildBillingSummary normalizes a billing config into a UI-friendly summary.
|
||||
func BuildBillingSummary(config *BillingConfig) *BillingSummary {
|
||||
if config == nil {
|
||||
return nil
|
||||
}
|
||||
summary := &BillingSummary{}
|
||||
period := config.CurrentPeriod
|
||||
periodType := resolvePeriodType(period)
|
||||
creditUsage := cloneFloat(config.CreditUsagePercent)
|
||||
|
||||
periodStart := ""
|
||||
periodEnd := ""
|
||||
if period != nil {
|
||||
periodStart = strings.TrimSpace(period.Start)
|
||||
periodEnd = strings.TrimSpace(period.End)
|
||||
}
|
||||
if periodStart == "" {
|
||||
periodStart = strings.TrimSpace(config.BillingPeriodStart)
|
||||
}
|
||||
if periodEnd == "" {
|
||||
periodEnd = strings.TrimSpace(config.BillingPeriodEnd)
|
||||
}
|
||||
|
||||
products := make([]BillingProductSummary, 0, len(config.ProductUsage))
|
||||
for _, item := range config.ProductUsage {
|
||||
product := strings.TrimSpace(item.Product)
|
||||
if product == "" {
|
||||
continue
|
||||
}
|
||||
products = append(products, BillingProductSummary{
|
||||
Product: product,
|
||||
UsagePercent: cloneFloat(item.UsagePercent),
|
||||
})
|
||||
}
|
||||
|
||||
monthlyLimit := parseCentValue(config.MonthlyLimit)
|
||||
used := parseCentValue(config.Used)
|
||||
billingStart := strings.TrimSpace(config.BillingPeriodStart)
|
||||
billingEnd := strings.TrimSpace(config.BillingPeriodEnd)
|
||||
|
||||
var includedUsed *float64
|
||||
if used != nil {
|
||||
if monthlyLimit != nil && *monthlyLimit > 0 {
|
||||
v := math.Min(*used, *monthlyLimit)
|
||||
includedUsed = &v
|
||||
} else {
|
||||
includedUsed = cloneFloat(used)
|
||||
}
|
||||
}
|
||||
|
||||
var usedPercent *float64
|
||||
if monthlyLimit != nil && *monthlyLimit > 0 && includedUsed != nil {
|
||||
v := (*includedUsed / *monthlyLimit) * 100
|
||||
usedPercent = &v
|
||||
}
|
||||
|
||||
hasWeekly := creditUsage != nil || periodType == "weekly" || len(products) > 0
|
||||
hasMonthly := monthlyLimit != nil || used != nil || (!hasWeekly && billingEnd != "")
|
||||
if !hasWeekly && !hasMonthly {
|
||||
return nil
|
||||
}
|
||||
|
||||
if hasWeekly {
|
||||
if periodType == "unknown" {
|
||||
periodType = "weekly"
|
||||
}
|
||||
summary.PeriodType = periodType
|
||||
summary.UsagePercent = creditUsage
|
||||
summary.PeriodStart = periodStart
|
||||
summary.PeriodEnd = periodEnd
|
||||
} else {
|
||||
// Monthly-only: do not put monthly % into UsagePercent (weekly bar field).
|
||||
// Frontend weekly bar only renders when PeriodType == weekly.
|
||||
summary.PeriodType = "monthly"
|
||||
summary.PeriodStart = billingStart
|
||||
summary.PeriodEnd = billingEnd
|
||||
}
|
||||
summary.ProductUsage = products
|
||||
summary.MonthlyLimitCents = monthlyLimit
|
||||
summary.UsedCents = used
|
||||
summary.IncludedUsedCents = includedUsed
|
||||
if hasMonthly {
|
||||
summary.BillingPeriodStart = billingStart
|
||||
summary.BillingPeriodEnd = billingEnd
|
||||
}
|
||||
summary.UsedPercent = usedPercent
|
||||
summary.Plan = resolvePlan(monthlyLimit)
|
||||
return summary
|
||||
}
|
||||
|
||||
// MergeBillingProbeResult updates successful billing domains while retaining
|
||||
// the previous value for any domain that could not be refreshed.
|
||||
func MergeBillingProbeResult(previous, weekly, monthly *BillingSummary, weeklyOK, monthlyOK bool) *BillingSummary {
|
||||
var out BillingSummary
|
||||
if previous != nil {
|
||||
out = *previous
|
||||
previousUpdatedAt := previous.UpdatedAt
|
||||
if previousUpdatedAt == "" {
|
||||
previousUpdatedAt = previous.FetchedAt
|
||||
}
|
||||
if out.WeeklyUpdatedAt == "" && (out.UsagePercent != nil || len(out.ProductUsage) > 0) {
|
||||
out.WeeklyUpdatedAt = previousUpdatedAt
|
||||
}
|
||||
if out.MonthlyUpdatedAt == "" && (out.MonthlyLimitCents != nil || out.UsedPercent != nil) {
|
||||
out.MonthlyUpdatedAt = previousUpdatedAt
|
||||
}
|
||||
}
|
||||
now := time.Now().UTC().Format(time.RFC3339)
|
||||
|
||||
if weeklyOK && weekly != nil {
|
||||
out.PeriodType = weekly.PeriodType
|
||||
out.UsagePercent = weekly.UsagePercent
|
||||
out.PeriodStart = weekly.PeriodStart
|
||||
out.PeriodEnd = weekly.PeriodEnd
|
||||
out.ProductUsage = weekly.ProductUsage
|
||||
out.WeeklyUpdatedAt = now
|
||||
}
|
||||
if monthlyOK && monthly != nil {
|
||||
if out.PeriodType == "" {
|
||||
out.PeriodType = "monthly"
|
||||
}
|
||||
out.MonthlyLimitCents = monthly.MonthlyLimitCents
|
||||
out.UsedCents = monthly.UsedCents
|
||||
out.IncludedUsedCents = monthly.IncludedUsedCents
|
||||
out.BillingPeriodStart = monthly.BillingPeriodStart
|
||||
out.BillingPeriodEnd = monthly.BillingPeriodEnd
|
||||
out.UsedPercent = monthly.UsedPercent
|
||||
out.Plan = monthly.Plan
|
||||
out.MonthlyUpdatedAt = now
|
||||
}
|
||||
|
||||
out.Partial = !weeklyOK || !monthlyOK
|
||||
out.FailedWindows = nil
|
||||
if !weeklyOK {
|
||||
out.FailedWindows = append(out.FailedWindows, "weekly")
|
||||
}
|
||||
if !monthlyOK {
|
||||
out.FailedWindows = append(out.FailedWindows, "monthly")
|
||||
}
|
||||
if !weeklyOK && !monthlyOK && previous == nil {
|
||||
return nil
|
||||
}
|
||||
return &out
|
||||
}
|
||||
|
||||
// StampBillingSummary sets fetch metadata.
|
||||
func StampBillingSummary(summary *BillingSummary, statusCode int, source string) *BillingSummary {
|
||||
if summary == nil {
|
||||
return nil
|
||||
}
|
||||
now := time.Now().UTC().Format(time.RFC3339)
|
||||
summary.StatusCode = statusCode
|
||||
summary.Source = source
|
||||
summary.FetchedAt = now
|
||||
summary.UpdatedAt = now
|
||||
return summary
|
||||
}
|
||||
|
||||
func resolvePeriodType(period *BillingPeriod) string {
|
||||
if period == nil {
|
||||
return "unknown"
|
||||
}
|
||||
raw := strings.ToLower(strings.TrimSpace(period.Type))
|
||||
if strings.Contains(raw, "weekly") {
|
||||
return "weekly"
|
||||
}
|
||||
if strings.Contains(raw, "monthly") {
|
||||
return "monthly"
|
||||
}
|
||||
return "unknown"
|
||||
}
|
||||
|
||||
func resolvePlan(monthlyLimitCents *float64) string {
|
||||
if monthlyLimitCents == nil {
|
||||
return ""
|
||||
}
|
||||
// Allow small float noise.
|
||||
limit := math.Round(*monthlyLimitCents)
|
||||
switch limit {
|
||||
case SuperGrokLimitCents:
|
||||
return "SuperGrok"
|
||||
case SuperGrokHeavyLimitCents:
|
||||
return "SuperGrok Heavy"
|
||||
default:
|
||||
return ""
|
||||
}
|
||||
}
|
||||
|
||||
func parseCentValue(raw json.RawMessage) *float64 {
|
||||
if len(raw) == 0 || string(raw) == "null" {
|
||||
return nil
|
||||
}
|
||||
// Object form: {"val": 123}
|
||||
var obj struct {
|
||||
Val any `json:"val"`
|
||||
}
|
||||
if err := json.Unmarshal(raw, &obj); err == nil && obj.Val != nil {
|
||||
return anyToFloat(obj.Val)
|
||||
}
|
||||
// Bare number / string
|
||||
var n any
|
||||
if err := json.Unmarshal(raw, &n); err != nil {
|
||||
return nil
|
||||
}
|
||||
return anyToFloat(n)
|
||||
}
|
||||
|
||||
func anyToFloat(v any) *float64 {
|
||||
switch n := v.(type) {
|
||||
case float64:
|
||||
return &n
|
||||
case float32:
|
||||
f := float64(n)
|
||||
return &f
|
||||
case int:
|
||||
f := float64(n)
|
||||
return &f
|
||||
case int64:
|
||||
f := float64(n)
|
||||
return &f
|
||||
case json.Number:
|
||||
f, err := n.Float64()
|
||||
if err != nil {
|
||||
return nil
|
||||
}
|
||||
return &f
|
||||
case string:
|
||||
s := strings.TrimSpace(n)
|
||||
if s == "" {
|
||||
return nil
|
||||
}
|
||||
f, err := strconv.ParseFloat(s, 64)
|
||||
if err != nil {
|
||||
return nil
|
||||
}
|
||||
return &f
|
||||
default:
|
||||
return nil
|
||||
}
|
||||
}
|
||||
|
||||
func cloneFloat(v *float64) *float64 {
|
||||
if v == nil {
|
||||
return nil
|
||||
}
|
||||
f := *v
|
||||
return &f
|
||||
}
|
||||
@@ -0,0 +1,127 @@
|
||||
package xai
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"net/http"
|
||||
"testing"
|
||||
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
func TestBuildBillingURL(t *testing.T) {
|
||||
t.Parallel()
|
||||
require.Equal(t, "https://cli-chat-proxy.grok.com/v1/billing?format=credits", BuildBillingURL(true))
|
||||
require.Equal(t, "https://cli-chat-proxy.grok.com/v1/billing", BuildBillingURL(false))
|
||||
}
|
||||
|
||||
func TestApplyCLIBillingHeaders(t *testing.T) {
|
||||
t.Parallel()
|
||||
req, err := http.NewRequest(http.MethodGet, BuildBillingURL(true), nil)
|
||||
require.NoError(t, err)
|
||||
|
||||
ApplyCLIBillingHeaders(req, " token ")
|
||||
|
||||
require.Equal(t, "Bearer token", req.Header.Get("Authorization"))
|
||||
require.Equal(t, CLITokenAuthValue, req.Header.Get(CLITokenAuthHeader))
|
||||
require.Equal(t, CLIClientVersion, req.Header.Get(CLIClientVersionHeader))
|
||||
require.Equal(t, "grok-pager/"+CLIClientVersion+" grok-shell/"+CLIClientVersion+" (macos; aarch64)", req.UserAgent())
|
||||
}
|
||||
|
||||
func TestBuildBillingSummaryWeeklyAndMonthly(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
weeklyBody := []byte(`{
|
||||
"config": {
|
||||
"currentPeriod": {"type":"WEEKLY","start":"2026-07-09T03:25:00Z","end":"2026-07-16T03:25:00Z"},
|
||||
"creditUsagePercent": 2.0,
|
||||
"productUsage": [{"product":"Api","usagePercent":2.0}]
|
||||
}
|
||||
}`)
|
||||
monthlyBody := []byte(`{
|
||||
"config": {
|
||||
"monthlyLimit": {"val": 15000},
|
||||
"used": {"val": 78},
|
||||
"billingPeriodStart": "2026-07-01T00:00:00Z",
|
||||
"billingPeriodEnd": "2026-08-01T00:00:00Z"
|
||||
}
|
||||
}`)
|
||||
|
||||
weeklyPayload, err := ParseBillingPayload(weeklyBody)
|
||||
require.NoError(t, err)
|
||||
monthlyPayload, err := ParseBillingPayload(monthlyBody)
|
||||
require.NoError(t, err)
|
||||
|
||||
weekly := BuildBillingSummary(weeklyPayload.Config)
|
||||
monthly := BuildBillingSummary(monthlyPayload.Config)
|
||||
require.NotNil(t, weekly)
|
||||
require.NotNil(t, monthly)
|
||||
require.Equal(t, "weekly", weekly.PeriodType)
|
||||
require.InDelta(t, 2.0, *weekly.UsagePercent, 1e-9)
|
||||
require.Equal(t, "Api", weekly.ProductUsage[0].Product)
|
||||
require.Equal(t, "SuperGrok", monthly.Plan)
|
||||
require.InDelta(t, 15000, *monthly.MonthlyLimitCents, 1e-9)
|
||||
require.InDelta(t, 78, *monthly.UsedCents, 1e-9)
|
||||
require.InDelta(t, 0.52, *monthly.UsedPercent, 1e-2)
|
||||
|
||||
merged := MergeBillingProbeResult(nil, weekly, monthly, true, true)
|
||||
require.Equal(t, "weekly", merged.PeriodType)
|
||||
require.InDelta(t, 2.0, *merged.UsagePercent, 1e-9)
|
||||
require.Equal(t, "SuperGrok", merged.Plan)
|
||||
require.InDelta(t, 15000, *merged.MonthlyLimitCents, 1e-9)
|
||||
require.Equal(t, "2026-08-01T00:00:00Z", merged.BillingPeriodEnd)
|
||||
}
|
||||
|
||||
func TestParseCentValueBareNumber(t *testing.T) {
|
||||
t.Parallel()
|
||||
raw, _ := json.Marshal(15000)
|
||||
v := parseCentValue(raw)
|
||||
require.NotNil(t, v)
|
||||
require.InDelta(t, 15000, *v, 1e-9)
|
||||
}
|
||||
|
||||
func TestBuildBillingSummaryMonthlyOnlyKeepsWeeklyUsageEmpty(t *testing.T) {
|
||||
t.Parallel()
|
||||
payload, err := ParseBillingPayload([]byte(`{"config":{"monthlyLimit":{"val":15000},"used":{"val":7500},"billingPeriodStart":"2026-07-01T00:00:00Z","billingPeriodEnd":"2026-08-01T00:00:00Z"}}`))
|
||||
require.NoError(t, err)
|
||||
|
||||
summary := BuildBillingSummary(payload.Config)
|
||||
require.NotNil(t, summary)
|
||||
require.Equal(t, "monthly", summary.PeriodType)
|
||||
require.Nil(t, summary.UsagePercent)
|
||||
require.InDelta(t, 50, *summary.UsedPercent, 1e-9)
|
||||
}
|
||||
|
||||
func TestMergeBillingProbeResultRetainsFailedWindow(t *testing.T) {
|
||||
t.Parallel()
|
||||
previous := &BillingSummary{
|
||||
PeriodType: "weekly",
|
||||
UsagePercent: floatPointer(100),
|
||||
PeriodEnd: "2026-07-16T00:00:00Z",
|
||||
MonthlyLimitCents: floatPointer(15000),
|
||||
UsedPercent: floatPointer(20),
|
||||
BillingPeriodEnd: "2026-08-01T00:00:00Z",
|
||||
WeeklyUpdatedAt: "2026-07-10T00:00:00Z",
|
||||
MonthlyUpdatedAt: "2026-07-10T00:00:00Z",
|
||||
FailedWindows: []string{"monthly"},
|
||||
}
|
||||
monthly := &BillingSummary{
|
||||
PeriodType: "monthly",
|
||||
MonthlyLimitCents: floatPointer(15000),
|
||||
UsedPercent: floatPointer(30),
|
||||
BillingPeriodEnd: "2026-08-01T00:00:00Z",
|
||||
}
|
||||
|
||||
merged := MergeBillingProbeResult(previous, nil, monthly, false, true)
|
||||
require.Equal(t, "weekly", merged.PeriodType)
|
||||
require.InDelta(t, 100, *merged.UsagePercent, 1e-9)
|
||||
require.Equal(t, previous.WeeklyUpdatedAt, merged.WeeklyUpdatedAt)
|
||||
require.InDelta(t, 30, *merged.UsedPercent, 1e-9)
|
||||
require.NotEqual(t, previous.MonthlyUpdatedAt, merged.MonthlyUpdatedAt)
|
||||
require.True(t, merged.Partial)
|
||||
require.Equal(t, []string{"weekly"}, merged.FailedWindows)
|
||||
require.Equal(t, []string{"monthly"}, previous.FailedWindows)
|
||||
}
|
||||
|
||||
func floatPointer(value float64) *float64 {
|
||||
return &value
|
||||
}
|
||||
Reference in New Issue
Block a user