mirror of
https://github.com/Wei-Shaw/sub2api.git
synced 2026-09-24 16:05:44 +08:00
feat: add grok quota probe parity
This commit is contained in:
Heatherm Huang
@@ -199,7 +199,8 @@ func initializeApplication(buildInfo handler.BuildInfo) (*Application, error) {
|
||||
openAIOAuthHandler := admin.NewOpenAIOAuthHandler(openAIOAuthService, adminService, openAIQuotaService)
|
||||
geminiOAuthHandler := admin.NewGeminiOAuthHandler(geminiOAuthService)
|
||||
antigravityOAuthHandler := admin.NewAntigravityOAuthHandler(antigravityOAuthService)
|
||||
grokOAuthHandler := admin.NewGrokOAuthHandler(grokOAuthService, adminService)
|
||||
grokQuotaService := service.ProvideGrokQuotaService(accountRepository, grokTokenProvider, httpUpstream)
|
||||
grokOAuthHandler := admin.NewGrokOAuthHandler(grokOAuthService, adminService, grokQuotaService)
|
||||
proxyHandler := admin.NewProxyHandler(adminService)
|
||||
adminRedeemHandler := admin.NewRedeemHandler(adminService, redeemService)
|
||||
promoHandler := admin.NewPromoHandler(promoService)
|
||||
|
||||
@@ -13,12 +13,18 @@ import (
|
||||
type GrokOAuthHandler struct {
|
||||
grokOAuthService *service.GrokOAuthService
|
||||
adminService service.AdminService
|
||||
quotaService *service.GrokQuotaService
|
||||
}
|
||||
|
||||
func NewGrokOAuthHandler(grokOAuthService *service.GrokOAuthService, adminService service.AdminService) *GrokOAuthHandler {
|
||||
func NewGrokOAuthHandler(
|
||||
grokOAuthService *service.GrokOAuthService,
|
||||
adminService service.AdminService,
|
||||
quotaService *service.GrokQuotaService,
|
||||
) *GrokOAuthHandler {
|
||||
return &GrokOAuthHandler{
|
||||
grokOAuthService: grokOAuthService,
|
||||
adminService: adminService,
|
||||
quotaService: quotaService,
|
||||
}
|
||||
}
|
||||
|
||||
@@ -197,3 +203,39 @@ func (h *GrokOAuthHandler) CreateAccountFromOAuth(c *gin.Context) {
|
||||
}
|
||||
response.Success(c, dto.AccountFromService(account))
|
||||
}
|
||||
|
||||
func (h *GrokOAuthHandler) QueryQuota(c *gin.Context) {
|
||||
accountID, err := strconv.ParseInt(c.Param("id"), 10, 64)
|
||||
if err != nil {
|
||||
response.BadRequest(c, "Invalid account ID")
|
||||
return
|
||||
}
|
||||
if h.quotaService == nil {
|
||||
response.BadRequest(c, "grok quota service is not enabled")
|
||||
return
|
||||
}
|
||||
result, err := h.quotaService.ProbeUsage(c.Request.Context(), accountID)
|
||||
if err != nil {
|
||||
response.ErrorFrom(c, err)
|
||||
return
|
||||
}
|
||||
response.Success(c, result)
|
||||
}
|
||||
|
||||
func (h *GrokOAuthHandler) ResetQuota(c *gin.Context) {
|
||||
accountID, err := strconv.ParseInt(c.Param("id"), 10, 64)
|
||||
if err != nil {
|
||||
response.BadRequest(c, "Invalid account ID")
|
||||
return
|
||||
}
|
||||
if h.quotaService == nil {
|
||||
response.BadRequest(c, "grok quota service is not enabled")
|
||||
return
|
||||
}
|
||||
result, err := h.quotaService.ResetQuota(c.Request.Context(), accountID)
|
||||
if err != nil {
|
||||
response.ErrorFrom(c, err)
|
||||
return
|
||||
}
|
||||
response.Success(c, result)
|
||||
}
|
||||
|
||||
@@ -0,0 +1,127 @@
|
||||
//go:build unit
|
||||
|
||||
package admin
|
||||
|
||||
import (
|
||||
"context"
|
||||
"io"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"strings"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/stretchr/testify/require"
|
||||
|
||||
"github.com/Wei-Shaw/sub2api/internal/pkg/tlsfingerprint"
|
||||
"github.com/Wei-Shaw/sub2api/internal/pkg/xai"
|
||||
"github.com/Wei-Shaw/sub2api/internal/service"
|
||||
)
|
||||
|
||||
type grokQuotaHandlerAccountRepo struct {
|
||||
service.AccountRepository
|
||||
account *service.Account
|
||||
updates map[int64]map[string]any
|
||||
}
|
||||
|
||||
func (r *grokQuotaHandlerAccountRepo) GetByID(_ context.Context, id int64) (*service.Account, error) {
|
||||
if r.account != nil && r.account.ID == id {
|
||||
return r.account, nil
|
||||
}
|
||||
return nil, service.ErrAccountNotFound
|
||||
}
|
||||
|
||||
func (r *grokQuotaHandlerAccountRepo) UpdateExtra(_ context.Context, id int64, updates map[string]any) error {
|
||||
if r.updates == nil {
|
||||
r.updates = make(map[int64]map[string]any)
|
||||
}
|
||||
r.updates[id] = updates
|
||||
return nil
|
||||
}
|
||||
|
||||
type grokQuotaHandlerUpstream struct {
|
||||
resp *http.Response
|
||||
lastReq *http.Request
|
||||
lastBody []byte
|
||||
}
|
||||
|
||||
func (u *grokQuotaHandlerUpstream) Do(req *http.Request, _ string, _ int64, _ int) (*http.Response, error) {
|
||||
u.lastReq = req
|
||||
if req.Body != nil {
|
||||
u.lastBody, _ = io.ReadAll(req.Body)
|
||||
}
|
||||
return u.resp, nil
|
||||
}
|
||||
|
||||
func (u *grokQuotaHandlerUpstream) DoWithTLS(
|
||||
req *http.Request,
|
||||
proxyURL string,
|
||||
accountID int64,
|
||||
accountConcurrency int,
|
||||
_ *tlsfingerprint.Profile,
|
||||
) (*http.Response, error) {
|
||||
return u.Do(req, proxyURL, accountID, accountConcurrency)
|
||||
}
|
||||
|
||||
func TestGrokOAuthHandlerQueryQuotaProbesUpstream(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
|
||||
repo := &grokQuotaHandlerAccountRepo{account: &service.Account{
|
||||
ID: 42,
|
||||
Platform: service.PlatformGrok,
|
||||
Type: service.AccountTypeOAuth,
|
||||
Concurrency: 1,
|
||||
Credentials: map[string]any{
|
||||
"access_token": "access-token",
|
||||
"expires_at": time.Now().Add(time.Hour).UTC().Format(time.RFC3339),
|
||||
},
|
||||
}}
|
||||
upstream := &grokQuotaHandlerUpstream{resp: &http.Response{
|
||||
StatusCode: http.StatusOK,
|
||||
Header: http.Header{
|
||||
"X-Ratelimit-Limit-Requests": []string{"10"},
|
||||
"X-Ratelimit-Remaining-Requests": []string{"8"},
|
||||
},
|
||||
Body: io.NopCloser(strings.NewReader(`{"id":"resp_probe"}`)),
|
||||
}}
|
||||
quotaService := service.NewGrokQuotaService(repo, service.NewGrokTokenProvider(repo, nil, nil), upstream)
|
||||
handler := NewGrokOAuthHandler(nil, nil, quotaService)
|
||||
|
||||
router := gin.New()
|
||||
router.GET("/api/v1/admin/grok/accounts/:id/quota", handler.QueryQuota)
|
||||
rec := httptest.NewRecorder()
|
||||
req := httptest.NewRequest(http.MethodGet, "/api/v1/admin/grok/accounts/42/quota", nil)
|
||||
router.ServeHTTP(rec, req)
|
||||
|
||||
require.Equal(t, http.StatusOK, rec.Code)
|
||||
require.Contains(t, rec.Body.String(), `"source":"active_probe"`)
|
||||
require.Contains(t, rec.Body.String(), `"headers_observed":true`)
|
||||
require.NotContains(t, rec.Body.String(), "access-token")
|
||||
require.Equal(t, xai.DefaultBaseURL+"/responses", upstream.lastReq.URL.String())
|
||||
require.Equal(t, "Bearer access-token", upstream.lastReq.Header.Get("Authorization"))
|
||||
require.Contains(t, string(upstream.lastBody), `"store":false`)
|
||||
require.NotNil(t, repo.updates[42])
|
||||
}
|
||||
|
||||
func TestGrokOAuthHandlerResetQuotaReturnsUnsupported(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
|
||||
repo := &grokQuotaHandlerAccountRepo{account: &service.Account{
|
||||
ID: 43,
|
||||
Platform: service.PlatformGrok,
|
||||
Type: service.AccountTypeOAuth,
|
||||
}}
|
||||
quotaService := service.NewGrokQuotaService(repo, nil, nil)
|
||||
handler := NewGrokOAuthHandler(nil, nil, quotaService)
|
||||
|
||||
router := gin.New()
|
||||
router.POST("/api/v1/admin/grok/accounts/:id/reset-quota", handler.ResetQuota)
|
||||
rec := httptest.NewRecorder()
|
||||
req := httptest.NewRequest(http.MethodPost, "/api/v1/admin/grok/accounts/43/reset-quota", nil)
|
||||
router.ServeHTTP(rec, req)
|
||||
|
||||
require.Equal(t, http.StatusNotImplemented, rec.Code)
|
||||
require.Contains(t, rec.Body.String(), `"reason":"GROK_QUOTA_RESET_UNSUPPORTED"`)
|
||||
require.NotContains(t, rec.Body.String(), "access-token")
|
||||
}
|
||||
@@ -260,6 +260,14 @@ func BuildResponsesURL(baseURL string) string {
|
||||
return EffectiveBaseURL(baseURL) + "/responses"
|
||||
}
|
||||
|
||||
func BuildChatCompletionsURL(baseURL string) (string, error) {
|
||||
validatedBaseURL, err := ValidatedBaseURL(baseURL)
|
||||
if err != nil {
|
||||
return "", fmt.Errorf("invalid base url: %w", err)
|
||||
}
|
||||
return validatedBaseURL + "/chat/completions", nil
|
||||
}
|
||||
|
||||
// TokenResponse represents xAI OAuth token responses.
|
||||
type TokenResponse struct {
|
||||
AccessToken string `json:"access_token"`
|
||||
|
||||
@@ -88,6 +88,58 @@ func TestBuildAuthorizationURLIncludesHermesCompatibleParameters(t *testing.T) {
|
||||
require.Equal(t, "sub2api", values.Get("referrer"))
|
||||
}
|
||||
|
||||
func TestValidateXAIURLsAllowOfficialOAuthAndGatewayHosts(t *testing.T) {
|
||||
authorizeURL, err := ValidateOAuthEndpointURL(DefaultAuthorizeURL)
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, DefaultAuthorizeURL, authorizeURL)
|
||||
|
||||
tokenURL, err := ValidateOAuthEndpointURL(DefaultTokenURL)
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, DefaultTokenURL, tokenURL)
|
||||
|
||||
baseURL, err := ValidateBaseURL(DefaultBaseURL)
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, DefaultBaseURL, baseURL)
|
||||
|
||||
cliBaseURL, err := ValidateBaseURL(DefaultCLIBaseURL)
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, DefaultCLIBaseURL, cliBaseURL)
|
||||
|
||||
baseURLNoPath, err := ValidateBaseURL("https://api.x.ai")
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, DefaultBaseURL, baseURLNoPath)
|
||||
|
||||
chatURL, err := BuildChatCompletionsURL(DefaultCLIBaseURL + "/")
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, DefaultCLIBaseURL+"/chat/completions", chatURL)
|
||||
}
|
||||
|
||||
func TestValidateXAIURLsRejectArbitraryHostsByDefault(t *testing.T) {
|
||||
_, err := ValidateOAuthEndpointURL("https://auth.example.test/oauth2/token")
|
||||
require.Error(t, err)
|
||||
|
||||
_, err = ValidateBaseURL("https://xai.test/v1")
|
||||
require.Error(t, err)
|
||||
|
||||
_, err = ValidateBaseURL("http://127.0.0.1:8080/v1")
|
||||
require.Error(t, err)
|
||||
|
||||
_, err = ValidateBaseURL("https://api.x.ai/custom")
|
||||
require.Error(t, err)
|
||||
}
|
||||
|
||||
func TestValidateXAIURLsAllowUnsafeDevOverride(t *testing.T) {
|
||||
t.Setenv(EnvAllowUnsafeURLOverrides, "true")
|
||||
|
||||
tokenURL, err := ValidateOAuthEndpointURL("http://127.0.0.1:8080/oauth2/token")
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, "http://127.0.0.1:8080/oauth2/token", tokenURL)
|
||||
|
||||
baseURL, err := ValidateBaseURL("http://127.0.0.1:8080/v1/")
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, "http://127.0.0.1:8080/v1", baseURL)
|
||||
}
|
||||
|
||||
func TestDefaultModelMappingIncludesGrokAliases(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
|
||||
@@ -396,6 +396,8 @@ func registerGrokOAuthRoutes(admin *gin.RouterGroup, h *handler.Handlers) {
|
||||
grok.POST("/oauth/refresh-token", h.Admin.GrokOAuth.RefreshToken)
|
||||
grok.POST("/oauth/create-from-oauth", h.Admin.GrokOAuth.CreateAccountFromOAuth)
|
||||
grok.POST("/accounts/:id/refresh", h.Admin.GrokOAuth.RefreshAccountToken)
|
||||
grok.GET("/accounts/:id/quota", h.Admin.GrokOAuth.QueryQuota)
|
||||
grok.POST("/accounts/:id/reset-quota", h.Admin.GrokOAuth.ResetQuota)
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -0,0 +1,186 @@
|
||||
package service
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"encoding/json"
|
||||
"io"
|
||||
"log/slog"
|
||||
"net/http"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
infraerrors "github.com/Wei-Shaw/sub2api/internal/pkg/errors"
|
||||
"github.com/Wei-Shaw/sub2api/internal/pkg/xai"
|
||||
)
|
||||
|
||||
const (
|
||||
grokQuotaUpstreamTimeout = 20 * time.Second
|
||||
grokQuotaProbeInput = "."
|
||||
grokQuotaDefaultModel = "grok-4.3"
|
||||
)
|
||||
|
||||
type GrokQuotaProbeResult struct {
|
||||
Source string `json:"source"`
|
||||
Snapshot *xai.QuotaSnapshot `json:"snapshot,omitempty"`
|
||||
StatusCode int `json:"status_code,omitempty"`
|
||||
HeadersObserved bool `json:"headers_observed"`
|
||||
ResetSupported bool `json:"reset_supported"`
|
||||
FetchedAt int64 `json:"fetched_at"`
|
||||
}
|
||||
|
||||
type GrokQuotaResetResult struct {
|
||||
Supported bool `json:"supported"`
|
||||
Code string `json:"code"`
|
||||
Message string `json:"message"`
|
||||
}
|
||||
|
||||
type GrokQuotaService struct {
|
||||
accountRepo AccountRepository
|
||||
tokenProvider *GrokTokenProvider
|
||||
httpUpstream HTTPUpstream
|
||||
}
|
||||
|
||||
func NewGrokQuotaService(
|
||||
accountRepo AccountRepository,
|
||||
tokenProvider *GrokTokenProvider,
|
||||
httpUpstream HTTPUpstream,
|
||||
) *GrokQuotaService {
|
||||
return &GrokQuotaService{
|
||||
accountRepo: accountRepo,
|
||||
tokenProvider: tokenProvider,
|
||||
httpUpstream: httpUpstream,
|
||||
}
|
||||
}
|
||||
|
||||
func (s *GrokQuotaService) ProbeUsage(ctx context.Context, accountID int64) (*GrokQuotaProbeResult, error) {
|
||||
account, token, proxyURL, err := s.prepareProbe(ctx, accountID)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
body, err := buildGrokQuotaProbeBody(account)
|
||||
if err != nil {
|
||||
return nil, infraerrors.Newf(http.StatusBadRequest, "GROK_QUOTA_PROBE_BODY_ERROR", "failed to build probe body: %v", err)
|
||||
}
|
||||
targetURL, err := xai.BuildResponsesURL(account.GetGrokBaseURL())
|
||||
if err != nil {
|
||||
return nil, infraerrors.Newf(http.StatusBadRequest, "GROK_QUOTA_BASE_URL_INVALID", "invalid Grok base_url: %v", err)
|
||||
}
|
||||
|
||||
callCtx, cancel := context.WithTimeout(ctx, grokQuotaUpstreamTimeout)
|
||||
defer cancel()
|
||||
req, err := http.NewRequestWithContext(callCtx, http.MethodPost, targetURL, bytes.NewReader(body))
|
||||
if err != nil {
|
||||
return nil, infraerrors.Newf(http.StatusInternalServerError, "GROK_QUOTA_PROBE_REQUEST_BUILD_FAILED", "failed to build upstream request: %v", err)
|
||||
}
|
||||
req.Header.Set("Authorization", "Bearer "+token)
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
req.Header.Set("Accept", "application/json")
|
||||
req.Header.Set("User-Agent", "sub2api-grok-quota-probe/1.0")
|
||||
|
||||
resp, err := s.httpUpstream.Do(req, proxyURL, account.ID, maxInt(account.Concurrency, 1))
|
||||
if err != nil {
|
||||
return nil, infraerrors.Newf(http.StatusBadGateway, "GROK_QUOTA_PROBE_REQUEST_FAILED", "upstream probe failed: %v", err)
|
||||
}
|
||||
defer func() { _ = resp.Body.Close() }()
|
||||
|
||||
snapshot := xai.ParseQuotaHeaders(resp.Header, resp.StatusCode)
|
||||
if snapshot != nil {
|
||||
_ = s.accountRepo.UpdateExtra(ctx, account.ID, map[string]any{
|
||||
grokQuotaSnapshotExtraKey: snapshot,
|
||||
})
|
||||
}
|
||||
|
||||
result := &GrokQuotaProbeResult{
|
||||
Source: "active_probe",
|
||||
Snapshot: snapshot,
|
||||
StatusCode: resp.StatusCode,
|
||||
HeadersObserved: snapshot != nil,
|
||||
ResetSupported: false,
|
||||
FetchedAt: time.Now().Unix(),
|
||||
}
|
||||
if resp.StatusCode == http.StatusTooManyRequests {
|
||||
return result, nil
|
||||
}
|
||||
if resp.StatusCode >= 400 {
|
||||
bodyBytes, _ := io.ReadAll(io.LimitReader(resp.Body, 240))
|
||||
bodyText := truncate(strings.TrimSpace(string(bodyBytes)), 240)
|
||||
slog.Warn("grok_quota_probe_failed", "account_id", account.ID, "status", resp.StatusCode, "body", bodyText)
|
||||
return nil, infraerrors.Newf(mapUpstreamStatus(resp.StatusCode), "GROK_QUOTA_PROBE_UPSTREAM_ERROR", "upstream returned %d: %s", resp.StatusCode, bodyText)
|
||||
}
|
||||
return result, nil
|
||||
}
|
||||
|
||||
func (s *GrokQuotaService) ResetQuota(ctx context.Context, accountID int64) (*GrokQuotaResetResult, error) {
|
||||
if _, err := s.loadGrokOAuthAccount(ctx, accountID); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return nil, infraerrors.New(http.StatusNotImplemented, "GROK_QUOTA_RESET_UNSUPPORTED", "xAI does not expose a Grok subscription quota reset endpoint for OAuth accounts")
|
||||
}
|
||||
|
||||
func (s *GrokQuotaService) prepareProbe(ctx context.Context, accountID int64) (*Account, string, string, error) {
|
||||
if s == nil || s.tokenProvider == nil || s.httpUpstream == nil {
|
||||
return nil, "", "", infraerrors.New(http.StatusInternalServerError, "GROK_QUOTA_NOT_CONFIGURED", "grok quota service is not configured")
|
||||
}
|
||||
account, err := s.loadGrokOAuthAccount(ctx, accountID)
|
||||
if err != nil {
|
||||
return nil, "", "", err
|
||||
}
|
||||
|
||||
token, err := s.tokenProvider.GetAccessToken(ctx, account)
|
||||
if err != nil {
|
||||
return nil, "", "", infraerrors.Newf(http.StatusBadGateway, "GROK_QUOTA_TOKEN_UNAVAILABLE", "failed to acquire access token: %v", err)
|
||||
}
|
||||
if strings.TrimSpace(token) == "" {
|
||||
return nil, "", "", infraerrors.New(http.StatusBadGateway, "GROK_QUOTA_TOKEN_UNAVAILABLE", "access token is empty")
|
||||
}
|
||||
|
||||
proxyURL := ""
|
||||
if account.ProxyID != nil && account.Proxy != nil {
|
||||
proxyURL = account.Proxy.URL()
|
||||
}
|
||||
return account, token, proxyURL, nil
|
||||
}
|
||||
|
||||
func (s *GrokQuotaService) loadGrokOAuthAccount(ctx context.Context, accountID int64) (*Account, error) {
|
||||
if s == nil || s.accountRepo == nil {
|
||||
return nil, infraerrors.New(http.StatusInternalServerError, "GROK_QUOTA_NOT_CONFIGURED", "grok quota service is not configured")
|
||||
}
|
||||
account, err := s.accountRepo.GetByID(ctx, accountID)
|
||||
if err != nil {
|
||||
return nil, infraerrors.Newf(http.StatusNotFound, "GROK_QUOTA_ACCOUNT_NOT_FOUND", "account not found: %v", err)
|
||||
}
|
||||
if account == nil {
|
||||
return nil, infraerrors.New(http.StatusNotFound, "GROK_QUOTA_ACCOUNT_NOT_FOUND", "account not found")
|
||||
}
|
||||
if account.Platform != PlatformGrok {
|
||||
return nil, infraerrors.New(http.StatusBadRequest, "GROK_QUOTA_INVALID_PLATFORM", "account is not a Grok account")
|
||||
}
|
||||
if account.Type != AccountTypeOAuth {
|
||||
return nil, infraerrors.New(http.StatusBadRequest, "GROK_QUOTA_INVALID_TYPE", "account is not an OAuth account")
|
||||
}
|
||||
return account, nil
|
||||
}
|
||||
|
||||
func buildGrokQuotaProbeBody(account *Account) ([]byte, error) {
|
||||
model := grokQuotaDefaultModel
|
||||
if account != nil {
|
||||
if mapped := strings.TrimSpace(account.GetMappedModel("grok")); mapped != "" {
|
||||
model = mapped
|
||||
}
|
||||
}
|
||||
return json.Marshal(map[string]any{
|
||||
"model": model,
|
||||
"input": grokQuotaProbeInput,
|
||||
"max_output_tokens": 1,
|
||||
"store": false,
|
||||
})
|
||||
}
|
||||
|
||||
func maxInt(a, b int) int {
|
||||
if a > b {
|
||||
return a
|
||||
}
|
||||
return b
|
||||
}
|
||||
@@ -0,0 +1,183 @@
|
||||
//go:build unit
|
||||
|
||||
package service
|
||||
|
||||
import (
|
||||
"context"
|
||||
"io"
|
||||
"net/http"
|
||||
"strings"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
infraerrors "github.com/Wei-Shaw/sub2api/internal/pkg/errors"
|
||||
"github.com/Wei-Shaw/sub2api/internal/pkg/xai"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
type grokQuotaAccountRepo struct {
|
||||
*mockAccountRepoForPlatform
|
||||
updates map[int64]map[string]any
|
||||
}
|
||||
|
||||
func (r *grokQuotaAccountRepo) UpdateExtra(_ context.Context, id int64, updates map[string]any) error {
|
||||
if r.updates == nil {
|
||||
r.updates = make(map[int64]map[string]any)
|
||||
}
|
||||
r.updates[id] = updates
|
||||
return nil
|
||||
}
|
||||
|
||||
func TestGrokQuotaServiceProbeUsageStoresHeaders(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
account := &Account{
|
||||
ID: 42,
|
||||
Platform: PlatformGrok,
|
||||
Type: AccountTypeOAuth,
|
||||
Concurrency: 1,
|
||||
Credentials: map[string]any{
|
||||
"access_token": "access-token",
|
||||
"expires_at": time.Now().Add(time.Hour).UTC().Format(time.RFC3339),
|
||||
},
|
||||
}
|
||||
repo := &grokQuotaAccountRepo{
|
||||
mockAccountRepoForPlatform: &mockAccountRepoForPlatform{
|
||||
accountsByID: map[int64]*Account{42: account},
|
||||
},
|
||||
}
|
||||
upstream := &httpUpstreamRecorder{resp: &http.Response{
|
||||
StatusCode: http.StatusOK,
|
||||
Header: http.Header{
|
||||
"X-Ratelimit-Limit-Requests": []string{"10"},
|
||||
"X-Ratelimit-Remaining-Requests": []string{"7"},
|
||||
"X-Ratelimit-Reset-Requests": []string{"2000000000"},
|
||||
"X-Ratelimit-Limit-Tokens": []string{"1000"},
|
||||
"X-Ratelimit-Remaining-Tokens": []string{"900"},
|
||||
},
|
||||
Body: io.NopCloser(strings.NewReader(`{"id":"resp_probe"}`)),
|
||||
}}
|
||||
svc := NewGrokQuotaService(repo, NewGrokTokenProvider(repo, nil, nil), upstream)
|
||||
|
||||
result, err := svc.ProbeUsage(context.Background(), 42)
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, http.StatusOK, result.StatusCode)
|
||||
require.True(t, result.HeadersObserved)
|
||||
require.NotNil(t, result.Snapshot)
|
||||
require.NotNil(t, result.Snapshot.Requests)
|
||||
require.EqualValues(t, 10, *result.Snapshot.Requests.Limit)
|
||||
require.EqualValues(t, 7, *result.Snapshot.Requests.Remaining)
|
||||
require.Equal(t, "https://api.x.ai/v1/responses", upstream.lastReq.URL.String())
|
||||
require.Equal(t, "Bearer access-token", upstream.lastReq.Header.Get("Authorization"))
|
||||
require.Contains(t, string(upstream.lastBody), `"max_output_tokens":1`)
|
||||
require.Contains(t, string(upstream.lastBody), `"store":false`)
|
||||
require.NotNil(t, repo.updates[42][grokQuotaSnapshotExtraKey])
|
||||
}
|
||||
|
||||
func TestGrokQuotaServiceProbeUsageReturnsRateLimitedSnapshot(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
account := &Account{
|
||||
ID: 43,
|
||||
Platform: PlatformGrok,
|
||||
Type: AccountTypeOAuth,
|
||||
Credentials: map[string]any{
|
||||
"access_token": "access-token",
|
||||
"expires_at": time.Now().Add(time.Hour).UTC().Format(time.RFC3339),
|
||||
},
|
||||
}
|
||||
repo := &grokQuotaAccountRepo{
|
||||
mockAccountRepoForPlatform: &mockAccountRepoForPlatform{
|
||||
accountsByID: map[int64]*Account{43: account},
|
||||
},
|
||||
}
|
||||
upstream := &httpUpstreamRecorder{resp: &http.Response{
|
||||
StatusCode: http.StatusTooManyRequests,
|
||||
Header: http.Header{"Retry-After": []string{"45"}},
|
||||
Body: io.NopCloser(strings.NewReader(`{"error":{"message":"rate limited"}}`)),
|
||||
}}
|
||||
svc := NewGrokQuotaService(repo, NewGrokTokenProvider(repo, nil, nil), upstream)
|
||||
|
||||
result, err := svc.ProbeUsage(context.Background(), 43)
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, http.StatusTooManyRequests, result.StatusCode)
|
||||
require.NotNil(t, result.Snapshot)
|
||||
require.NotNil(t, result.Snapshot.RetryAfterSeconds)
|
||||
require.Equal(t, 45, *result.Snapshot.RetryAfterSeconds)
|
||||
}
|
||||
|
||||
func TestGrokQuotaServiceResetQuotaUnsupported(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
account := &Account{
|
||||
ID: 44,
|
||||
Platform: PlatformGrok,
|
||||
Type: AccountTypeOAuth,
|
||||
}
|
||||
repo := &grokQuotaAccountRepo{
|
||||
mockAccountRepoForPlatform: &mockAccountRepoForPlatform{
|
||||
accountsByID: map[int64]*Account{44: account},
|
||||
},
|
||||
}
|
||||
svc := NewGrokQuotaService(repo, nil, nil)
|
||||
|
||||
_, err := svc.ResetQuota(context.Background(), 44)
|
||||
require.Error(t, err)
|
||||
require.Equal(t, http.StatusNotImplemented, infraerrors.Code(err))
|
||||
require.Equal(t, "GROK_QUOTA_RESET_UNSUPPORTED", infraerrors.Reason(err))
|
||||
}
|
||||
|
||||
func TestShouldAutoPauseGrokAccountByQuota(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
zero := int64(0)
|
||||
limit := int64(10)
|
||||
resetFuture := time.Now().Add(time.Minute).Unix()
|
||||
retryAfter := 30
|
||||
tests := []struct {
|
||||
name string
|
||||
snapshot xai.QuotaSnapshot
|
||||
want bool
|
||||
}{
|
||||
{
|
||||
name: "remaining requests exhausted",
|
||||
snapshot: xai.QuotaSnapshot{
|
||||
Requests: &xai.QuotaWindow{Limit: &limit, Remaining: &zero, ResetUnix: &resetFuture},
|
||||
UpdatedAt: time.Now().UTC().Format(time.RFC3339),
|
||||
},
|
||||
want: true,
|
||||
},
|
||||
{
|
||||
name: "retry after active",
|
||||
snapshot: xai.QuotaSnapshot{
|
||||
RetryAfterSeconds: &retryAfter,
|
||||
UpdatedAt: time.Now().UTC().Format(time.RFC3339),
|
||||
},
|
||||
want: true,
|
||||
},
|
||||
{
|
||||
name: "stale snapshot ignored",
|
||||
snapshot: xai.QuotaSnapshot{
|
||||
Requests: &xai.QuotaWindow{Limit: &limit, Remaining: &zero, ResetUnix: &resetFuture},
|
||||
UpdatedAt: time.Now().Add(-3 * time.Hour).UTC().Format(time.RFC3339),
|
||||
},
|
||||
want: false,
|
||||
},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
tt := tt
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
t.Parallel()
|
||||
account := &Account{
|
||||
Platform: PlatformGrok,
|
||||
Type: AccountTypeOAuth,
|
||||
Extra: map[string]any{
|
||||
grokQuotaSnapshotExtraKey: tt.snapshot,
|
||||
},
|
||||
}
|
||||
got, _ := shouldAutoPauseGrokAccountByQuota(account)
|
||||
require.Equal(t, tt.want, got)
|
||||
})
|
||||
}
|
||||
}
|
||||
@@ -74,6 +74,10 @@ func (s *OpenAIGatewayService) ForwardAsChatCompletions(
|
||||
return nil, errors.New("codex_cli_only restriction: only codex official clients are allowed")
|
||||
}
|
||||
|
||||
if account.Platform == PlatformGrok {
|
||||
return s.forwardAsRawChatCompletions(ctx, c, account, body, defaultMappedModel)
|
||||
}
|
||||
|
||||
// 入口分流：APIKey 账号 + 强制或已探测确认上游不支持 Responses，走 CC 直转。
|
||||
// 自动模式下标记缺失（未探测）按"现状即证据"原则继续走下方原 Responses 转换路径。
|
||||
if account.Type == AccountTypeAPIKey && !openai_compat.ShouldUseResponsesAPI(account.Extra) {
|
||||
|
||||
@@ -14,6 +14,7 @@ import (
|
||||
|
||||
"github.com/Wei-Shaw/sub2api/internal/pkg/apicompat"
|
||||
"github.com/Wei-Shaw/sub2api/internal/pkg/logger"
|
||||
"github.com/Wei-Shaw/sub2api/internal/pkg/xai"
|
||||
"github.com/Wei-Shaw/sub2api/internal/util/responseheaders"
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/tidwall/gjson"
|
||||
@@ -121,19 +122,18 @@ func (s *OpenAIGatewayService) forwardAsRawChatCompletions(
|
||||
)
|
||||
|
||||
// 5. Build upstream request
|
||||
apiKey := account.GetOpenAIApiKey()
|
||||
if apiKey == "" {
|
||||
return nil, fmt.Errorf("account %d missing api_key", account.ID)
|
||||
}
|
||||
baseURL := account.GetOpenAIBaseURL()
|
||||
if baseURL == "" {
|
||||
baseURL = "https://api.openai.com"
|
||||
}
|
||||
validatedURL, err := s.validateUpstreamBaseURL(baseURL)
|
||||
token, tokenKind, err := s.GetAccessToken(ctx, account)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("invalid base_url: %w", err)
|
||||
return nil, err
|
||||
}
|
||||
if strings.TrimSpace(token) == "" {
|
||||
return nil, fmt.Errorf("account %d missing %s credential", account.ID, tokenKind)
|
||||
}
|
||||
|
||||
targetURL, err := s.rawChatCompletionsURL(account)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
targetURL := buildOpenAIChatCompletionsURL(validatedURL)
|
||||
|
||||
upstreamCtx, releaseUpstreamCtx := detachUpstreamContext(ctx)
|
||||
upstreamReq, err := http.NewRequestWithContext(upstreamCtx, http.MethodPost, targetURL, bytes.NewReader(upstreamBody))
|
||||
@@ -143,7 +143,7 @@ func (s *OpenAIGatewayService) forwardAsRawChatCompletions(
|
||||
}
|
||||
upstreamReq = upstreamReq.WithContext(WithHTTPUpstreamProfile(upstreamReq.Context(), HTTPUpstreamProfileOpenAI))
|
||||
upstreamReq.Header.Set("Content-Type", "application/json")
|
||||
upstreamReq.Header.Set("Authorization", "Bearer "+apiKey)
|
||||
upstreamReq.Header.Set("Authorization", "Bearer "+token)
|
||||
if clientStream {
|
||||
upstreamReq.Header.Set("Accept", "text/event-stream")
|
||||
} else {
|
||||
@@ -162,6 +162,8 @@ func (s *OpenAIGatewayService) forwardAsRawChatCompletions(
|
||||
customUA := account.GetOpenAIUserAgent()
|
||||
if customUA != "" {
|
||||
upstreamReq.Header.Set("user-agent", customUA)
|
||||
} else if account.Platform == PlatformGrok {
|
||||
upstreamReq.Header.Set("user-agent", "sub2api-grok/1.0")
|
||||
}
|
||||
|
||||
// 6. Send request
|
||||
@@ -180,9 +182,32 @@ func (s *OpenAIGatewayService) forwardAsRawChatCompletions(
|
||||
respBody := s.readUpstreamErrorBody(resp)
|
||||
_ = resp.Body.Close()
|
||||
resp.Body = io.NopCloser(bytes.NewReader(respBody))
|
||||
if account.Platform == PlatformGrok {
|
||||
s.updateGrokUsageSnapshot(ctx, account.ID, xai.ParseQuotaHeaders(resp.Header, resp.StatusCode))
|
||||
}
|
||||
|
||||
upstreamMsg := strings.TrimSpace(extractUpstreamErrorMessage(respBody))
|
||||
upstreamMsg = sanitizeUpstreamErrorMessage(upstreamMsg)
|
||||
if account.Platform == PlatformGrok {
|
||||
appendOpsUpstreamError(c, OpsUpstreamErrorEvent{
|
||||
Platform: account.Platform,
|
||||
AccountID: account.ID,
|
||||
AccountName: account.Name,
|
||||
UpstreamStatusCode: resp.StatusCode,
|
||||
UpstreamRequestID: firstNonEmpty(resp.Header.Get("x-request-id"), resp.Header.Get("xai-request-id")),
|
||||
Kind: "failover",
|
||||
Message: upstreamMsg,
|
||||
})
|
||||
s.handleGrokAccountUpstreamError(ctx, account, resp.StatusCode, resp.Header, respBody)
|
||||
if s.shouldFailoverUpstreamError(resp.StatusCode) {
|
||||
return nil, &UpstreamFailoverError{
|
||||
StatusCode: resp.StatusCode,
|
||||
ResponseBody: respBody,
|
||||
RetryableOnSameAccount: account.IsPoolMode() && account.IsPoolModeRetryableStatus(resp.StatusCode),
|
||||
}
|
||||
}
|
||||
return s.handleChatCompletionsErrorResponse(resp, c, account, billingModel)
|
||||
}
|
||||
if s.shouldFailoverOpenAIUpstreamResponse(resp.StatusCode, upstreamMsg, respBody) {
|
||||
upstreamDetail := ""
|
||||
if s.cfg != nil && s.cfg.Gateway.LogUpstreamErrorBody {
|
||||
@@ -212,6 +237,10 @@ func (s *OpenAIGatewayService) forwardAsRawChatCompletions(
|
||||
return s.handleChatCompletionsErrorResponse(resp, c, account, billingModel)
|
||||
}
|
||||
|
||||
if account.Platform == PlatformGrok {
|
||||
s.updateGrokUsageSnapshot(ctx, account.ID, xai.ParseQuotaHeaders(resp.Header, resp.StatusCode))
|
||||
}
|
||||
|
||||
// 8. Forward response
|
||||
if clientStream {
|
||||
return s.streamRawChatCompletions(c, resp, account, originalModel, billingModel, upstreamModel, reasoningEffort, serviceTier, startTime, len(body))
|
||||
@@ -219,6 +248,26 @@ func (s *OpenAIGatewayService) forwardAsRawChatCompletions(
|
||||
return s.bufferRawChatCompletions(c, resp, originalModel, billingModel, upstreamModel, reasoningEffort, serviceTier, startTime)
|
||||
}
|
||||
|
||||
func (s *OpenAIGatewayService) rawChatCompletionsURL(account *Account) (string, error) {
|
||||
if account.Platform == PlatformGrok {
|
||||
targetURL, err := xai.BuildChatCompletionsURL(account.GetGrokBaseURL())
|
||||
if err != nil {
|
||||
return "", fmt.Errorf("invalid grok base_url: %w", err)
|
||||
}
|
||||
return targetURL, nil
|
||||
}
|
||||
|
||||
baseURL := account.GetOpenAIBaseURL()
|
||||
if baseURL == "" {
|
||||
baseURL = "https://api.openai.com"
|
||||
}
|
||||
validatedURL, err := s.validateUpstreamBaseURL(baseURL)
|
||||
if err != nil {
|
||||
return "", fmt.Errorf("invalid base_url: %w", err)
|
||||
}
|
||||
return buildOpenAIChatCompletionsURL(validatedURL), nil
|
||||
}
|
||||
|
||||
// streamRawChatCompletions 透传上游 CC SSE 流到客户端，并提取 usage（包括
|
||||
// 末尾 [DONE] 之前的 chunk 中的 usage 字段，按 OpenAI CC 协议）。
|
||||
//
|
||||
|
||||
@@ -3,13 +3,18 @@
|
||||
package service
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"encoding/json"
|
||||
"io"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"strings"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/Wei-Shaw/sub2api/internal/pkg/xai"
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/stretchr/testify/require"
|
||||
"github.com/tidwall/gjson"
|
||||
)
|
||||
@@ -57,3 +62,75 @@ func TestBuildGrokResponsesRequestUsesAccountBaseURLAndBearerToken(t *testing.T)
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, `{"model":"grok-4.3"}`, strings.TrimSpace(string(data)))
|
||||
}
|
||||
|
||||
func TestBuildGrokResponsesRequestRejectsUnsafeAccountBaseURL(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
account := &Account{
|
||||
Platform: PlatformGrok,
|
||||
Type: AccountTypeOAuth,
|
||||
Credentials: map[string]any{
|
||||
"base_url": "https://xai.test/v1",
|
||||
},
|
||||
}
|
||||
|
||||
_, err := buildGrokResponsesRequest(context.Background(), nil, account, []byte(`{"model":"grok-4.3"}`), "access-token")
|
||||
require.Error(t, err)
|
||||
require.Contains(t, err.Error(), "invalid base url")
|
||||
}
|
||||
|
||||
func TestForwardAsChatCompletionsForGrokUsesXAIChatCompletionsAndSnapshots(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
|
||||
recorder := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(recorder)
|
||||
body := []byte(`{"model":"grok","messages":[{"role":"user","content":"hi"}],"stream":false}`)
|
||||
c.Request = httptest.NewRequest(http.MethodPost, "/v1/chat/completions", bytes.NewReader(body))
|
||||
|
||||
account := &Account{
|
||||
ID: 51,
|
||||
Name: "grok",
|
||||
Platform: PlatformGrok,
|
||||
Type: AccountTypeOAuth,
|
||||
Concurrency: 1,
|
||||
Credentials: map[string]any{
|
||||
"access_token": "access-token",
|
||||
"expires_at": time.Now().Add(time.Hour).UTC().Format(time.RFC3339),
|
||||
"base_url": xai.DefaultCLIBaseURL,
|
||||
},
|
||||
}
|
||||
repo := &grokQuotaAccountRepo{
|
||||
mockAccountRepoForPlatform: &mockAccountRepoForPlatform{
|
||||
accountsByID: map[int64]*Account{51: account},
|
||||
},
|
||||
}
|
||||
upstream := &httpUpstreamRecorder{resp: &http.Response{
|
||||
StatusCode: http.StatusOK,
|
||||
Header: http.Header{
|
||||
"Content-Type": []string{"application/json"},
|
||||
"Xai-Request-Id": []string{"xai-req"},
|
||||
"X-Ratelimit-Limit-Requests": []string{"10"},
|
||||
"X-Ratelimit-Remaining-Requests": []string{"9"},
|
||||
"X-Ratelimit-Limit-Tokens": []string{"1000"},
|
||||
"X-Ratelimit-Remaining-Tokens": []string{"990"},
|
||||
},
|
||||
Body: io.NopCloser(strings.NewReader(`{"id":"chatcmpl","object":"chat.completion","model":"grok-4.3","choices":[],"usage":{"prompt_tokens":1,"completion_tokens":2}}`)),
|
||||
}}
|
||||
svc := &OpenAIGatewayService{
|
||||
httpUpstream: upstream,
|
||||
grokTokenProvider: NewGrokTokenProvider(repo, nil, nil),
|
||||
accountRepo: repo,
|
||||
}
|
||||
|
||||
result, err := svc.ForwardAsChatCompletions(context.Background(), c, account, body, "", "")
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, xai.DefaultCLIBaseURL+"/chat/completions", upstream.lastReq.URL.String())
|
||||
require.Equal(t, "Bearer access-token", upstream.lastReq.Header.Get("Authorization"))
|
||||
require.Equal(t, "grok-4.3", gjson.GetBytes(upstream.lastBody, "model").String())
|
||||
require.Equal(t, "grok", result.Model)
|
||||
require.Equal(t, "grok-4.3", result.UpstreamModel)
|
||||
require.Equal(t, 1, result.Usage.InputTokens)
|
||||
require.Equal(t, 2, result.Usage.OutputTokens)
|
||||
require.NotNil(t, repo.updates[51][grokQuotaSnapshotExtraKey])
|
||||
require.Equal(t, http.StatusOK, recorder.Code)
|
||||
}
|
||||
|
||||
@@ -26,6 +26,7 @@ import (
|
||||
"github.com/Wei-Shaw/sub2api/internal/pkg/logger"
|
||||
"github.com/Wei-Shaw/sub2api/internal/pkg/openai"
|
||||
"github.com/Wei-Shaw/sub2api/internal/pkg/openai_compat"
|
||||
"github.com/Wei-Shaw/sub2api/internal/pkg/xai"
|
||||
"github.com/Wei-Shaw/sub2api/internal/util/responseheaders"
|
||||
"github.com/Wei-Shaw/sub2api/internal/util/urlvalidator"
|
||||
"github.com/cespare/xxhash/v2"
|
||||
@@ -1376,6 +1377,17 @@ func isOpenAICompatibleAccountEligibleForRequest(ctx context.Context, account *A
|
||||
return false
|
||||
}
|
||||
}
|
||||
if account.IsGrok() {
|
||||
if paused, reason := shouldAutoPauseGrokAccountByQuota(account); paused {
|
||||
slog.Debug("grok_account_auto_paused_by_quota",
|
||||
"account_id", account.ID,
|
||||
"window", reason.window,
|
||||
"threshold", reason.threshold,
|
||||
"utilization", reason.utilization,
|
||||
)
|
||||
return false
|
||||
}
|
||||
}
|
||||
if requestedModel != "" && !account.IsModelSupported(requestedModel) {
|
||||
return false
|
||||
}
|
||||
@@ -1394,6 +1406,55 @@ type openAIQuotaAutoPauseDecision struct {
|
||||
utilization float64
|
||||
}
|
||||
|
||||
func shouldAutoPauseGrokAccountByQuota(account *Account) (bool, openAIQuotaAutoPauseDecision) {
|
||||
if account == nil || !account.IsGrok() || account.Type != AccountTypeOAuth {
|
||||
return false, openAIQuotaAutoPauseDecision{}
|
||||
}
|
||||
snapshot, err := grokQuotaSnapshotFromExtra(account.Extra)
|
||||
if err != nil || snapshot == nil {
|
||||
return false, openAIQuotaAutoPauseDecision{}
|
||||
}
|
||||
now := time.Now()
|
||||
if grokQuotaSnapshotStaleForPause(snapshot, now) {
|
||||
return false, openAIQuotaAutoPauseDecision{}
|
||||
}
|
||||
if snapshot.RetryAfterSeconds != nil && *snapshot.RetryAfterSeconds > 0 {
|
||||
return true, openAIQuotaAutoPauseDecision{window: "retry_after", threshold: 1, utilization: 1}
|
||||
}
|
||||
if paused, decision := shouldAutoPauseGrokQuotaWindow("requests", snapshot.Requests, now); paused {
|
||||
return true, decision
|
||||
}
|
||||
if paused, decision := shouldAutoPauseGrokQuotaWindow("tokens", snapshot.Tokens, now); paused {
|
||||
return true, decision
|
||||
}
|
||||
return false, openAIQuotaAutoPauseDecision{}
|
||||
}
|
||||
|
||||
func shouldAutoPauseGrokQuotaWindow(name string, window *xai.QuotaWindow, now time.Time) (bool, openAIQuotaAutoPauseDecision) {
|
||||
if window == nil || window.Limit == nil || window.Remaining == nil || *window.Limit <= 0 {
|
||||
return false, openAIQuotaAutoPauseDecision{}
|
||||
}
|
||||
if window.ResetUnix != nil && *window.ResetUnix > 0 && !now.Before(time.Unix(*window.ResetUnix, 0)) {
|
||||
return false, openAIQuotaAutoPauseDecision{}
|
||||
}
|
||||
utilization := float64(*window.Limit-*window.Remaining) / float64(*window.Limit)
|
||||
if *window.Remaining <= 0 || utilization >= 1 {
|
||||
return true, openAIQuotaAutoPauseDecision{window: name, threshold: 1, utilization: utilization}
|
||||
}
|
||||
return false, openAIQuotaAutoPauseDecision{}
|
||||
}
|
||||
|
||||
func grokQuotaSnapshotStaleForPause(snapshot *xai.QuotaSnapshot, now time.Time) bool {
|
||||
if snapshot == nil || strings.TrimSpace(snapshot.UpdatedAt) == "" {
|
||||
return false
|
||||
}
|
||||
updatedAt, err := parseTime(snapshot.UpdatedAt)
|
||||
if err != nil {
|
||||
return false
|
||||
}
|
||||
return now.Sub(updatedAt) >= openAICodexAutoPauseStaleAfter
|
||||
}
|
||||
|
||||
func shouldAutoPauseOpenAIAccountByQuota(ctx context.Context, account *Account) (bool, openAIQuotaAutoPauseDecision) {
|
||||
if account == nil || !account.IsOpenAI() {
|
||||
return false, openAIQuotaAutoPauseDecision{}
|
||||
|
||||
@@ -125,6 +125,14 @@ func ProvideOpenAIQuotaService(
|
||||
return NewOpenAIQuotaService(accountRepo, proxyRepo, tokenProvider, privacyClientFactory)
|
||||
}
|
||||
|
||||
func ProvideGrokQuotaService(
|
||||
accountRepo AccountRepository,
|
||||
tokenProvider *GrokTokenProvider,
|
||||
httpUpstream HTTPUpstream,
|
||||
) *GrokQuotaService {
|
||||
return NewGrokQuotaService(accountRepo, tokenProvider, httpUpstream)
|
||||
}
|
||||
|
||||
// ProvideGeminiTokenProvider creates GeminiTokenProvider with OAuthRefreshAPI injection
|
||||
func ProvideGeminiTokenProvider(
|
||||
accountRepo AccountRepository,
|
||||
@@ -565,6 +573,7 @@ var ProviderSet = wire.NewSet(
|
||||
ProvideGrokTokenProvider,
|
||||
ProvideOpenAITokenProvider,
|
||||
ProvideOpenAIQuotaService,
|
||||
ProvideGrokQuotaService,
|
||||
ProvideClaudeTokenProvider,
|
||||
NewAntigravityGatewayService,
|
||||
ProvideRateLimitService,
|
||||
|
||||
@@ -39,6 +39,39 @@ export interface GrokTokenInfo {
|
||||
[key: string]: unknown
|
||||
}
|
||||
|
||||
export interface GrokQuotaWindow {
|
||||
limit?: number | null
|
||||
remaining?: number | null
|
||||
reset_unix?: number | null
|
||||
reset_at?: string | null
|
||||
}
|
||||
|
||||
export interface GrokQuotaSnapshot {
|
||||
requests?: GrokQuotaWindow | null
|
||||
tokens?: GrokQuotaWindow | null
|
||||
retry_after_seconds?: number | null
|
||||
subscription_tier?: string
|
||||
entitlement_status?: string
|
||||
status_code?: number
|
||||
headers?: Record<string, string>
|
||||
updated_at: string
|
||||
}
|
||||
|
||||
export interface GrokQuotaProbeResult {
|
||||
source: 'active_probe'
|
||||
snapshot?: GrokQuotaSnapshot | null
|
||||
status_code?: number
|
||||
headers_observed: boolean
|
||||
reset_supported: boolean
|
||||
fetched_at: number
|
||||
}
|
||||
|
||||
export interface GrokQuotaResetResult {
|
||||
supported: boolean
|
||||
code: string
|
||||
message: string
|
||||
}
|
||||
|
||||
export async function generateAuthUrl(
|
||||
payload: GrokAuthUrlRequest
|
||||
): Promise<GrokAuthUrlResponse> {
|
||||
@@ -71,4 +104,14 @@ export async function refreshGrokToken(
|
||||
return data
|
||||
}
|
||||
|
||||
export default { generateAuthUrl, exchangeCode, refreshGrokToken }
|
||||
export async function queryQuota(id: number): Promise<GrokQuotaProbeResult> {
|
||||
const { data } = await apiClient.get<GrokQuotaProbeResult>(`/admin/grok/accounts/${id}/quota`)
|
||||
return data
|
||||
}
|
||||
|
||||
export async function resetQuota(id: number): Promise<GrokQuotaResetResult> {
|
||||
const { data } = await apiClient.post<GrokQuotaResetResult>(`/admin/grok/accounts/${id}/reset-quota`)
|
||||
return data
|
||||
}
|
||||
|
||||
export default { generateAuthUrl, exchangeCode, refreshGrokToken, queryQuota, resetQuota }
|
||||
|
||||
@@ -384,6 +384,7 @@
|
||||
<div v-else-if="usageInfo.error" class="truncate text-xs text-amber-600 dark:text-amber-400 max-w-[200px]" :title="usageInfo.error">
|
||||
{{ usageErrorLabel }}
|
||||
</div>
|
||||
<GrokQuotaProbeCell :account="account" />
|
||||
</div>
|
||||
<div v-else class="text-xs text-gray-400">-</div>
|
||||
</template>
|
||||
@@ -583,6 +584,7 @@ import { formatCompactNumber } from '@/utils/format'
|
||||
import UsageProgressBar from './UsageProgressBar.vue'
|
||||
import AccountQuotaInfo from './AccountQuotaInfo.vue'
|
||||
import OpenAIQuotaResetCell from './OpenAIQuotaResetCell.vue'
|
||||
import GrokQuotaProbeCell from './GrokQuotaProbeCell.vue'
|
||||
|
||||
// Module-level cache shared across all AccountUsageCell instances
|
||||
const _usageCache = new Map<number, { data: AccountUsageInfo; ts: number }>()
|
||||
|
||||
@@ -0,0 +1,135 @@
|
||||
<template>
|
||||
<div v-if="visible" class="space-y-1">
|
||||
<div class="flex flex-wrap items-center gap-1.5">
|
||||
<button
|
||||
type="button"
|
||||
class="inline-flex items-center gap-0.5 rounded px-1.5 py-0.5 text-[10px] font-medium text-cyan-700 transition-colors hover:bg-cyan-50 disabled:cursor-not-allowed disabled:opacity-50 dark:text-cyan-300 dark:hover:bg-cyan-900/30"
|
||||
:disabled="loading"
|
||||
:title="t('admin.accounts.usageWindow.grokProbeTooltip')"
|
||||
@click="handleProbe"
|
||||
>
|
||||
<svg
|
||||
class="h-2.5 w-2.5"
|
||||
:class="{ 'animate-spin': loading }"
|
||||
fill="none"
|
||||
stroke="currentColor"
|
||||
viewBox="0 0 24 24"
|
||||
>
|
||||
<path
|
||||
stroke-linecap="round"
|
||||
stroke-linejoin="round"
|
||||
stroke-width="2"
|
||||
d="M4 4v5h.582m15.356 2A8.001 8.001 0 004.582 9m0 0H9m11 11v-5h-.581m0 0a8.003 8.003 0 01-15.357-2m15.357 2H15"
|
||||
/>
|
||||
</svg>
|
||||
{{ t('admin.accounts.usageWindow.grokProbe') }}
|
||||
</button>
|
||||
|
||||
<button
|
||||
type="button"
|
||||
class="inline-flex cursor-not-allowed items-center gap-0.5 rounded px-1.5 py-0.5 text-[10px] font-medium text-gray-400 opacity-70 dark:text-gray-500"
|
||||
disabled
|
||||
:title="t('admin.accounts.usageWindow.grokResetUnsupportedTooltip')"
|
||||
>
|
||||
{{ t('admin.accounts.usageWindow.grokResetUnsupported') }}
|
||||
</button>
|
||||
</div>
|
||||
|
||||
<div v-if="summary" class="text-[10px] text-gray-600 dark:text-gray-300">
|
||||
{{ summary }}
|
||||
</div>
|
||||
<div v-if="error" class="truncate text-[10px] text-red-600 dark:text-red-400" :title="error">
|
||||
{{ truncatedError }}
|
||||
</div>
|
||||
</div>
|
||||
</template>
|
||||
|
||||
<script setup lang="ts">
|
||||
import { computed, ref, watch } from 'vue'
|
||||
import { useI18n } from 'vue-i18n'
|
||||
import { adminAPI } from '@/api/admin'
|
||||
import type { GrokQuotaProbeResult, GrokQuotaWindow } from '@/api/admin/grok'
|
||||
import type { Account } from '@/types'
|
||||
|
||||
const props = defineProps<{
|
||||
account: Account
|
||||
}>()
|
||||
|
||||
const { t } = useI18n()
|
||||
|
||||
const visible = computed(() => props.account.platform === 'grok' && props.account.type === 'oauth')
|
||||
const loading = ref(false)
|
||||
const error = ref<string | null>(null)
|
||||
const data = ref<GrokQuotaProbeResult | null>(null)
|
||||
|
||||
const extractErrorMessage = (e: unknown): string => {
|
||||
const err = e as {
|
||||
message?: string
|
||||
reason?: string
|
||||
response?: { data?: { message?: string; error?: string } }
|
||||
}
|
||||
return (
|
||||
err?.message ||
|
||||
err?.reason ||
|
||||
err?.response?.data?.message ||
|
||||
err?.response?.data?.error ||
|
||||
t('common.error')
|
||||
)
|
||||
}
|
||||
|
||||
const formatWindow = (label: string, window?: GrokQuotaWindow | null): string | null => {
|
||||
if (!window || window.limit == null || window.remaining == null) return null
|
||||
return `${label} ${window.remaining}/${window.limit}`
|
||||
}
|
||||
|
||||
const retryAfterLabel = computed(() => {
|
||||
const seconds = data.value?.snapshot?.retry_after_seconds
|
||||
if (seconds == null || seconds <= 0) return null
|
||||
if (seconds < 60) return `${seconds}s`
|
||||
return `${Math.ceil(seconds / 60)}m`
|
||||
})
|
||||
|
||||
const summary = computed(() => {
|
||||
const snapshot = data.value?.snapshot
|
||||
if (!data.value) return ''
|
||||
if (!snapshot) return t('admin.accounts.usageWindow.grokNoHeaders')
|
||||
const parts = [
|
||||
formatWindow(t('admin.accounts.usageWindow.grokRequests'), snapshot.requests),
|
||||
formatWindow(t('admin.accounts.usageWindow.grokTokens'), snapshot.tokens)
|
||||
].filter(Boolean)
|
||||
if (retryAfterLabel.value) {
|
||||
parts.push(t('admin.accounts.usageWindow.grokRetryAfter', { time: retryAfterLabel.value }))
|
||||
}
|
||||
if (snapshot.entitlement_status) {
|
||||
parts.push(snapshot.entitlement_status)
|
||||
}
|
||||
return parts.length > 0 ? parts.join(' | ') : t('admin.accounts.usageWindow.grokNoHeaders')
|
||||
})
|
||||
|
||||
const truncatedError = computed(() => {
|
||||
if (!error.value) return ''
|
||||
return error.value.length > 80 ? `${error.value.slice(0, 80)}...` : error.value
|
||||
})
|
||||
|
||||
const handleProbe = async () => {
|
||||
if (loading.value) return
|
||||
loading.value = true
|
||||
error.value = null
|
||||
try {
|
||||
data.value = await adminAPI.grok.queryQuota(props.account.id)
|
||||
} catch (e) {
|
||||
error.value = extractErrorMessage(e)
|
||||
} finally {
|
||||
loading.value = false
|
||||
}
|
||||
}
|
||||
|
||||
watch(
|
||||
() => props.account.id,
|
||||
() => {
|
||||
data.value = null
|
||||
error.value = null
|
||||
loading.value = false
|
||||
}
|
||||
)
|
||||
</script>
|
||||
@@ -4154,6 +4154,11 @@ export default {
|
||||
grokTokens: 'Tok',
|
||||
grokUnknown: 'Grok quota is unknown until the first upstream response includes xAI rate-limit headers.',
|
||||
grokRetryAfter: 'Retry after {time}',
|
||||
grokProbe: 'Probe',
|
||||
grokProbeTooltip: 'Send a minimal xAI Responses probe and read quota headers',
|
||||
grokResetUnsupported: 'Reset unsupported',
|
||||
grokResetUnsupportedTooltip: 'xAI does not expose reset credits for Grok OAuth accounts',
|
||||
grokNoHeaders: 'No quota headers observed',
|
||||
passiveSampled: 'Passive',
|
||||
activeQuery: 'Query'
|
||||
},
|
||||
|
||||
@@ -3412,6 +3412,11 @@ export default {
|
||||
grokTokens: 'Token',
|
||||
grokUnknown: 'Grok 配额需等待首次上游响应返回 xAI rate-limit 头后显示。',
|
||||
grokRetryAfter: '{time} 后重试',
|
||||
grokProbe: '探测',
|
||||
grokProbeTooltip: '发送最小 xAI Responses 探测请求并读取配额响应头',
|
||||
grokResetUnsupported: '不支持重置',
|
||||
grokResetUnsupportedTooltip: 'xAI 未向 Grok OAuth 账号开放重置额度接口',
|
||||
grokNoHeaders: '未观察到配额响应头',
|
||||
passiveSampled: '被动采样',
|
||||
activeQuery: '查询'
|
||||
},
|
||||
|
||||
Reference in New Issue
Block a user