mirror of
https://github.com/Wei-Shaw/sub2api.git
synced 2026-09-24 16:05:44 +08:00
fix: stabilize Grok quota probe model
This commit is contained in:
@@ -51,6 +51,7 @@ func DefaultModelMapping() map[string]string {
|
||||
mapping["grok-build"] = "grok-build-0.1"
|
||||
mapping["grok-build-latest"] = "grok-4.5"
|
||||
mapping["grok-composer"] = "grok-composer-2.5-fast"
|
||||
mapping["composer-2.5"] = "grok-composer-2.5-fast"
|
||||
mapping["grok-4.20-reasoning"] = "grok-4.20-0309-reasoning"
|
||||
mapping["grok-4.20-non-reasoning"] = "grok-4.20-0309-non-reasoning"
|
||||
return mapping
|
||||
|
||||
@@ -214,6 +214,7 @@ func TestDefaultModelMappingIncludesGrokAliases(t *testing.T) {
|
||||
require.Equal(t, "grok-build-0.1", mapping["grok-build"])
|
||||
require.Equal(t, "grok-4.5", mapping["grok-build-latest"])
|
||||
require.Equal(t, "grok-composer-2.5-fast", mapping["grok-composer"])
|
||||
require.Equal(t, "grok-composer-2.5-fast", mapping["composer-2.5"])
|
||||
require.Equal(t, "grok-4.20-0309-reasoning", mapping["grok-4.20-reasoning"])
|
||||
require.Equal(t, "grok-4.20-0309-non-reasoning", mapping["grok-4.20-non-reasoning"])
|
||||
require.Equal(t, "grok-4.20-multi-agent-0309", mapping["grok-4.20-multi-agent-0309"])
|
||||
|
||||
@@ -22,6 +22,7 @@ const (
|
||||
|
||||
type GrokQuotaProbeResult struct {
|
||||
Source string `json:"source"`
|
||||
Model string `json:"model"`
|
||||
Snapshot *xai.QuotaSnapshot `json:"snapshot,omitempty"`
|
||||
StatusCode int `json:"status_code,omitempty"`
|
||||
HeadersObserved bool `json:"headers_observed"`
|
||||
@@ -62,7 +63,8 @@ func (s *GrokQuotaService) ProbeUsage(ctx context.Context, accountID int64) (*Gr
|
||||
return nil, err
|
||||
}
|
||||
|
||||
body, err := buildGrokQuotaProbeBody(account)
|
||||
probeModel := grokQuotaProbeModel()
|
||||
body, err := buildGrokQuotaProbeBody(probeModel)
|
||||
if err != nil {
|
||||
return nil, infraerrors.Newf(http.StatusBadRequest, "GROK_QUOTA_PROBE_BODY_ERROR", "failed to build probe body: %v", err)
|
||||
}
|
||||
@@ -95,6 +97,7 @@ func (s *GrokQuotaService) ProbeUsage(ctx context.Context, accountID int64) (*Gr
|
||||
|
||||
result := &GrokQuotaProbeResult{
|
||||
Source: "active_probe",
|
||||
Model: probeModel,
|
||||
Snapshot: snapshot,
|
||||
StatusCode: resp.StatusCode,
|
||||
HeadersObserved: snapshot.HeadersObserved,
|
||||
@@ -107,8 +110,8 @@ func (s *GrokQuotaService) ProbeUsage(ctx context.Context, accountID int64) (*Gr
|
||||
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)
|
||||
slog.Warn("grok_quota_probe_failed", "account_id", account.ID, "model", probeModel, "status", resp.StatusCode, "body", bodyText)
|
||||
return nil, infraerrors.Newf(mapUpstreamStatus(resp.StatusCode), "GROK_QUOTA_PROBE_UPSTREAM_ERROR", "upstream returned %d for probe model %q: %s", resp.StatusCode, probeModel, bodyText)
|
||||
}
|
||||
return result, nil
|
||||
}
|
||||
@@ -175,12 +178,14 @@ func (s *GrokQuotaService) loadGrokOAuthAccount(ctx context.Context, accountID i
|
||||
return account, nil
|
||||
}
|
||||
|
||||
func buildGrokQuotaProbeBody(account *Account) ([]byte, error) {
|
||||
model := grokQuotaDefaultModel
|
||||
if account != nil {
|
||||
if mapped := strings.TrimSpace(account.GetMappedModel("grok")); mapped != "" {
|
||||
model = mapped
|
||||
}
|
||||
func grokQuotaProbeModel() string {
|
||||
return grokQuotaDefaultModel
|
||||
}
|
||||
|
||||
func buildGrokQuotaProbeBody(model string) ([]byte, error) {
|
||||
model = strings.TrimSpace(model)
|
||||
if model == "" {
|
||||
model = grokQuotaDefaultModel
|
||||
}
|
||||
return json.Marshal(map[string]any{
|
||||
"model": model,
|
||||
|
||||
@@ -13,6 +13,7 @@ import (
|
||||
infraerrors "github.com/Wei-Shaw/sub2api/internal/pkg/errors"
|
||||
"github.com/Wei-Shaw/sub2api/internal/pkg/xai"
|
||||
"github.com/stretchr/testify/require"
|
||||
"github.com/tidwall/gjson"
|
||||
)
|
||||
|
||||
type grokQuotaAccountRepo struct {
|
||||
@@ -85,6 +86,7 @@ func TestGrokQuotaServiceProbeUsageStoresHeaders(t *testing.T) {
|
||||
result, err := svc.ProbeUsage(context.Background(), 42)
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, http.StatusOK, result.StatusCode)
|
||||
require.Equal(t, "grok-4.3", result.Model)
|
||||
require.True(t, result.HeadersObserved)
|
||||
require.NotNil(t, result.Snapshot)
|
||||
require.True(t, result.Snapshot.HeadersObserved)
|
||||
@@ -96,11 +98,79 @@ func TestGrokQuotaServiceProbeUsageStoresHeaders(t *testing.T) {
|
||||
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.Equal(t, "grok-4.3", gjson.GetBytes(upstream.lastBody, "model").String())
|
||||
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 TestGrokQuotaServiceProbeUsageIgnoresAccountGrokMapping(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
account := &Account{
|
||||
ID: 47,
|
||||
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),
|
||||
"model_mapping": map[string]any{
|
||||
"grok": "grok-composer",
|
||||
"grok-composer": "grok-composer-2.5-fast",
|
||||
},
|
||||
},
|
||||
}
|
||||
repo := &grokQuotaAccountRepo{
|
||||
mockAccountRepoForPlatform: &mockAccountRepoForPlatform{
|
||||
accountsByID: map[int64]*Account{47: account},
|
||||
},
|
||||
}
|
||||
upstream := &httpUpstreamRecorder{resp: &http.Response{
|
||||
StatusCode: http.StatusOK,
|
||||
Header: http.Header{},
|
||||
Body: io.NopCloser(strings.NewReader(`{"id":"resp_probe"}`)),
|
||||
}}
|
||||
svc := NewGrokQuotaService(repo, nil, NewGrokTokenProvider(repo, nil), upstream)
|
||||
|
||||
result, err := svc.ProbeUsage(context.Background(), 47)
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, "grok-4.3", result.Model)
|
||||
require.Equal(t, "grok-4.3", gjson.GetBytes(upstream.lastBody, "model").String())
|
||||
require.NotContains(t, string(upstream.lastBody), "grok-composer")
|
||||
}
|
||||
|
||||
func TestGrokQuotaServiceProbeUsageReportsProbeModelOnUpstreamError(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
account := &Account{
|
||||
ID: 48,
|
||||
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{48: account},
|
||||
},
|
||||
}
|
||||
upstream := &httpUpstreamRecorder{resp: &http.Response{
|
||||
StatusCode: http.StatusBadRequest,
|
||||
Header: http.Header{},
|
||||
Body: io.NopCloser(strings.NewReader(`{"code":"invalid-argument","error":"Model not found"}`)),
|
||||
}}
|
||||
svc := NewGrokQuotaService(repo, nil, NewGrokTokenProvider(repo, nil), upstream)
|
||||
|
||||
_, err := svc.ProbeUsage(context.Background(), 48)
|
||||
require.Error(t, err)
|
||||
require.Equal(t, "GROK_QUOTA_PROBE_UPSTREAM_ERROR", infraerrors.Reason(err))
|
||||
require.Contains(t, infraerrors.Message(err), `probe model "grok-4.3"`)
|
||||
}
|
||||
|
||||
func TestGrokQuotaServiceProbeUsageLoadsProxyWhenAccountEdgeMissing(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
|
||||
@@ -63,6 +63,7 @@ export interface GrokQuotaSnapshot {
|
||||
|
||||
export interface GrokQuotaProbeResult {
|
||||
source: 'active_probe'
|
||||
model: string
|
||||
snapshot?: GrokQuotaSnapshot | null
|
||||
status_code?: number
|
||||
headers_observed: boolean
|
||||
|
||||
@@ -69,6 +69,14 @@ describe('useModelWhitelist', () => {
|
||||
})
|
||||
})
|
||||
|
||||
it('grok 模型列表包含 Composer 默认项和兼容别名', () => {
|
||||
const models = getModelsByPlatform('grok')
|
||||
|
||||
expect(models).toContain('grok-composer-2.5-fast')
|
||||
expect(models).toContain('grok-composer')
|
||||
expect(models).toContain('composer-2.5')
|
||||
})
|
||||
|
||||
it('gemini 模型列表包含原生生图模型', () => {
|
||||
const models = getModelsByPlatform('gemini')
|
||||
|
||||
|
||||
@@ -148,6 +148,7 @@ const xaiModels = [
|
||||
'grok-build',
|
||||
'grok-build-latest',
|
||||
'grok-composer',
|
||||
'composer-2.5',
|
||||
'grok-4.20-reasoning',
|
||||
'grok-4.20-non-reasoning',
|
||||
'grok-imagine',
|
||||
@@ -306,6 +307,7 @@ const grokPresetMappings = [
|
||||
{ label: 'Build 0.1', from: 'grok-build', to: 'grok-build-0.1', color: 'bg-cyan-100 text-cyan-700 hover:bg-cyan-200 dark:bg-cyan-900/30 dark:text-cyan-400' },
|
||||
{ label: 'Build Latest', from: 'grok-build-latest', to: 'grok-4.5', color: 'bg-teal-100 text-teal-700 hover:bg-teal-200 dark:bg-teal-900/30 dark:text-teal-400' },
|
||||
{ label: 'Composer 2.5', from: 'grok-composer', to: 'grok-composer-2.5-fast', color: 'bg-teal-100 text-teal-700 hover:bg-teal-200 dark:bg-teal-900/30 dark:text-teal-400' },
|
||||
{ label: 'Composer legacy', from: 'composer-2.5', to: 'grok-composer-2.5-fast', color: 'bg-teal-100 text-teal-700 hover:bg-teal-200 dark:bg-teal-900/30 dark:text-teal-400' },
|
||||
{ label: '4.20 Reasoning', from: 'grok-4.20-reasoning', to: 'grok-4.20-0309-reasoning', color: 'bg-indigo-100 text-indigo-700 hover:bg-indigo-200 dark:bg-indigo-900/30 dark:text-indigo-400' },
|
||||
{ label: '4.20 Non Reasoning', from: 'grok-4.20-non-reasoning', to: 'grok-4.20-0309-non-reasoning', color: 'bg-violet-100 text-violet-700 hover:bg-violet-200 dark:bg-violet-900/30 dark:text-violet-400' },
|
||||
{ label: 'Imagine Image', from: 'grok-imagine', to: 'grok-imagine-image-quality', color: 'bg-sky-100 text-sky-700 hover:bg-sky-200 dark:bg-sky-900/30 dark:text-sky-400' },
|
||||
|
||||
Reference in New Issue
Block a user