fix: stabilize Grok quota probe model

This commit is contained in:
Heatherm Huang
2026-07-09 17:58:31 +08:00
parent b480545c19
commit 1a0a6ea9a8
7 changed files with 97 additions and 9 deletions
+1
View File
@@ -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
+1
View File
@@ -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"])
+14 -9
View File
@@ -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()
+1
View File
@@ -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' },