seewe/sub2api
1
0
Fork 0
mirror of https://github.com/Wei-Shaw/sub2api.git synced 2026-10-14 12:02:27 +08:00
Code Issues Packages Projects Releases Wiki Activity

fix(grok): align Free probes with health checks

Browse Source
This commit is contained in:
superman2003
2026-07-18 21:05:50 +08:00
parent 4e8ea7d568
commit dd7a2b22f0
9 changed files with 366 additions and 17 deletions
Download Patch File Download Diff File Expand all files Collapse all files
backend/internal/handler/admin/grok_oauth_handler_test.go
+5 -1
View File
@@ -136,8 +136,12 @@ func TestGrokOAuthHandlerQueryQuotaProbesUpstream(t *testing.T) {
for i, upstreamReq := range requests {
require.Equal(t, "Bearer access-token", upstreamReq.Header.Get("Authorization"))
if upstreamReq.URL.String() == xai.DefaultCLIBaseURL+"/responses" {
require.Equal(t, "application/json, text/event-stream", upstreamReq.Header.Get("Accept"))
require.Contains(t, string(bodies[i]), `"model":"grok-4.5"`)
require.Contains(t, string(bodies[i]), `"store":false`)
require.Contains(t, string(bodies[i]), `"input":"hi"`)
require.Contains(t, string(bodies[i]), `"stream":true`)
require.NotContains(t, string(bodies[i]), `"max_output_tokens"`)
require.NotContains(t, string(bodies[i]), `"store"`)
}
}
require.NotNil(t, repo.updates[42])
backend/internal/service/account_test_service.go
+1 -5
View File
@@ -735,11 +735,7 @@ func (s *AccountTestService) testGrokAccountConnection(c *gin.Context, account *
c.Writer.Header().Set("X-Accel-Buffering", "no")
c.Writer.Flush()
payloadBytes, err := json.Marshal(map[string]any{
"model": testModelID,
"input": "hi",
"stream": true,
})
payloadBytes, err := buildGrokQuotaProbeBody(testModelID)
if err != nil {
return s.sendErrorAndEnd(c, "Failed to create Grok test payload")
}
backend/internal/service/account_test_service_grok_test.go
+5
View File
@@ -75,7 +75,12 @@ func TestAccountTestService_TestAccountConnection_GrokUsesXAIResponses(t *testin
require.Equal(t, "https://cli-chat-proxy.grok.com/v1/responses", upstream.lastReq.URL.String())
require.Equal(t, "Bearer grok-access-token", upstream.lastReq.Header.Get("Authorization"))
require.Equal(t, grokCLIVersion, upstream.lastReq.Header.Get("X-Grok-Client-Version"))
require.Equal(t, "application/json, text/event-stream", upstream.lastReq.Header.Get("Accept"))
require.Equal(t, "grok-4.3", gjson.GetBytes(upstream.lastBody, "model").String())
require.Equal(t, grokQuotaProbeInput, gjson.GetBytes(upstream.lastBody, "input").String())
require.True(t, gjson.GetBytes(upstream.lastBody, "stream").Bool())
require.False(t, gjson.GetBytes(upstream.lastBody, "max_output_tokens").Exists())
require.False(t, gjson.GetBytes(upstream.lastBody, "store").Exists())
require.NotContains(t, rec.Body.String(), "claude")
require.Contains(t, rec.Body.String(), `"model":"grok-4.3"`)
require.Contains(t, rec.Body.String(), `"type":"test_complete"`)
backend/internal/service/grok_quota_fetcher.go
+35 -1
View File
@@ -32,6 +32,7 @@ func (f *GrokQuotaFetcher) BuildUsageInfo(account *Account) *UsageInfo {
billing, _ := grokBillingSnapshotFromExtra(account.Extra)
snapshot, err := grokQuotaSnapshotFromExtra(account.Extra)
activeProbeClearsForbidden := newerSuccessfulGrokActiveProbeClearsBillingForbidden(billing, snapshot)
if billing != nil {
usage.GrokBilling = billing
if billing.Plan != "" {
@@ -87,7 +88,13 @@ func (f *GrokQuotaFetcher) BuildUsageInfo(account *Account) *UsageInfo {
usage.GrokLastQuotaProbeAt = snapshot.LastProbeAt
}
usage.GrokLastHeadersSeenAt = snapshot.LastHeadersSeenAt
if snapshot.StatusCode >= http.StatusBadRequest || usage.GrokLastStatusCode == 0 {
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() {
@@ -117,9 +124,36 @@ func (f *GrokQuotaFetcher) BuildUsageInfo(account *Account) *UsageInfo {
}
}
applyGrokCredentialUsageFallback(usage, account)
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) {
if usage == nil || account == nil {
return
backend/internal/service/grok_quota_fetcher_test.go
+125
View File
@@ -121,6 +121,131 @@ func TestGrokQuotaFetcherSnapshotErrorOverridesSuccessfulBillingStatus(t *testin
require.Equal(t, http.StatusTooManyRequests, usage.GrokLastStatusCode)
}
func TestGrokQuotaFetcherNewerSuccessfulActiveProbeClearsBillingForbidden(t *testing.T) {
t.Parallel()
billingAt := "2030-01-01T00:00:00Z"
probeAt := "2030-01-01T00:05:00Z"
account := &Account{
Platform: PlatformGrok,
Type: AccountTypeOAuth,
Credentials: map[string]any{
"entitlement_status": "forbidden",
},
Extra: map[string]any{
grokBillingExtraKey: &xai.BillingSummary{
StatusCode: http.StatusForbidden,
UpdatedAt: billingAt,
},
grokQuotaSnapshotExtraKey: &xai.QuotaSnapshot{
StatusCode: http.StatusOK,
ObservationSource: "active_probe",
LastProbeAt: probeAt,
UpdatedAt: probeAt,
},
},
}
usage := NewGrokQuotaFetcher().BuildUsageInfo(account)
require.False(t, usage.IsForbidden)
require.Empty(t, usage.ForbiddenType)
require.Empty(t, usage.ErrorCode)
require.Empty(t, usage.GrokEntitlementStatus)
require.Equal(t, http.StatusOK, usage.GrokLastStatusCode)
require.Equal(t, probeAt, usage.GrokLastQuotaProbeAt)
require.NotNil(t, usage.UpdatedAt)
require.True(t, usage.UpdatedAt.Equal(time.Date(2030, 1, 1, 0, 5, 0, 0, time.UTC)))
}
func TestGrokQuotaFetcherSameSecondSuccessfulActiveProbeClearsBillingForbidden(t *testing.T) {
t.Parallel()
observedAt := "2030-01-01T00:05:00Z"
account := &Account{
Platform: PlatformGrok,
Type: AccountTypeOAuth,
Extra: map[string]any{
grokBillingExtraKey: &xai.BillingSummary{
StatusCode: http.StatusForbidden,
UpdatedAt: observedAt,
},
grokQuotaSnapshotExtraKey: &xai.QuotaSnapshot{
StatusCode: http.StatusOK,
ObservationSource: "active_probe",
LastProbeAt: observedAt,
UpdatedAt: observedAt,
},
},
}
usage := NewGrokQuotaFetcher().BuildUsageInfo(account)
require.False(t, usage.IsForbidden)
require.Empty(t, usage.ForbiddenType)
require.Empty(t, usage.ErrorCode)
require.Equal(t, http.StatusOK, usage.GrokLastStatusCode)
}
func TestGrokQuotaFetcherDoesNotClearBillingForbiddenWithoutNewerSuccessfulActiveProbe(t *testing.T) {
t.Parallel()
tests := []struct {
name string
snapshot xai.QuotaSnapshot
}{
{
name: "older active probe",
snapshot: xai.QuotaSnapshot{
StatusCode: http.StatusOK,
ObservationSource: "active_probe",
LastProbeAt: "2030-01-01T00:04:59Z",
UpdatedAt: "2030-01-01T00:04:59Z",
},
},
{
name: "newer passive response",
snapshot: xai.QuotaSnapshot{
StatusCode: http.StatusOK,
ObservationSource: "upstream_response",
UpdatedAt: "2030-01-01T00:05:01Z",
},
},
{
name: "newer failed active probe",
snapshot: xai.QuotaSnapshot{
StatusCode: http.StatusTooManyRequests,
ObservationSource: "active_probe",
LastProbeAt: "2030-01-01T00:05:01Z",
UpdatedAt: "2030-01-01T00:05:01Z",
},
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
t.Parallel()
account := &Account{
Platform: PlatformGrok,
Type: AccountTypeOAuth,
Extra: map[string]any{
grokBillingExtraKey: &xai.BillingSummary{
StatusCode: http.StatusForbidden,
UpdatedAt: "2030-01-01T00:05:00Z",
},
grokQuotaSnapshotExtraKey: tt.snapshot,
},
}
usage := NewGrokQuotaFetcher().BuildUsageInfo(account)
require.True(t, usage.IsForbidden)
require.Equal(t, "forbidden", usage.ForbiddenType)
require.Equal(t, "forbidden", usage.ErrorCode)
})
}
}
func TestGrokQuotaFetcherBuildUsageInfoFromNoHeadersProbe(t *testing.T) {
t.Parallel()
backend/internal/service/grok_quota_service.go
+5 -6
View File
@@ -20,7 +20,7 @@ import (
const (
grokQuotaUpstreamTimeout = 20 * time.Second
grokQuotaProbeInput = "."
grokQuotaProbeInput = "hi"
grokQuotaDefaultModel = grokDefaultResponsesModel
grokBillingExtraKey = "grok_billing_snapshot"
)
@@ -152,7 +152,7 @@ func (s *GrokQuotaService) probeUsage(ctx context.Context, accountID int64) (*Gr
}
req.Header.Set("Authorization", "Bearer "+token)
req.Header.Set("Content-Type", "application/json")
req.Header.Set("Accept", "application/json")
req.Header.Set("Accept", "application/json, text/event-stream")
if account.IsGrokOAuth() {
applyGrokCLIHeaders(req.Header)
}
@@ -502,10 +502,9 @@ func buildGrokQuotaProbeBody(model string) ([]byte, error) {
model = grokQuotaDefaultModel
}
return json.Marshal(map[string]any{
"model": model,
"input": grokQuotaProbeInput,
"max_output_tokens": 1,
"store": false,
"model": model,
"input": grokQuotaProbeInput,
"stream": true,
})
}
backend/internal/service/grok_quota_service_test.go
+10 -3
View File
@@ -266,9 +266,12 @@ func TestGrokQuotaServiceProbeUsageStoresHeaders(t *testing.T) {
require.Equal(t, "https://cli-chat-proxy.grok.com/v1/responses", upstream.lastReq.URL.String())
require.Equal(t, "Bearer access-token", upstream.lastReq.Header.Get("Authorization"))
require.Equal(t, grokCLIVersion, upstream.lastReq.Header.Get("X-Grok-Client-Version"))
require.Equal(t, "application/json, text/event-stream", upstream.lastReq.Header.Get("Accept"))
require.Equal(t, "grok-4.5", 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.Equal(t, grokQuotaProbeInput, gjson.GetBytes(upstream.lastBody, "input").String())
require.True(t, gjson.GetBytes(upstream.lastBody, "stream").Bool())
require.False(t, gjson.GetBytes(upstream.lastBody, "max_output_tokens").Exists())
require.False(t, gjson.GetBytes(upstream.lastBody, "store").Exists())
require.NotNil(t, repo.updates[42][grokQuotaSnapshotExtraKey])
}
@@ -498,8 +501,12 @@ func TestGrokQuotaServiceQueryQuotaFreeFallsBackToGrok45(t *testing.T) {
}
responseCalls++
require.Equal(t, http.MethodPost, req.Method)
require.Equal(t, "application/json, text/event-stream", req.Header.Get("Accept"))
require.Equal(t, "grok-4.5", gjson.GetBytes(bodies[i], "model").String())
require.EqualValues(t, 1, gjson.GetBytes(bodies[i], "max_output_tokens").Int())
require.Equal(t, grokQuotaProbeInput, gjson.GetBytes(bodies[i], "input").String())
require.True(t, gjson.GetBytes(bodies[i], "stream").Bool())
require.False(t, gjson.GetBytes(bodies[i], "max_output_tokens").Exists())
require.False(t, gjson.GetBytes(bodies[i], "store").Exists())
}
require.Equal(t, 1, responseCalls)
}
frontend/src/components/account/AccountUsageCell.vue
+22 -1
View File
@@ -1354,6 +1354,21 @@ const handleGrokProbed = (result: GrokQuotaProbeResult) => {
const current = usageInfo.value
if (!current) return
const snapshot = result.snapshot
const statusCode = snapshot?.status_code ?? result.status_code
const hasActiveProbeSnapshot = snapshot != null && (
result.source === 'active_probe' ||
result.source === 'hybrid_probe' ||
snapshot.observation_source === 'active_probe'
)
const probeSucceeded = hasActiveProbeSnapshot &&
statusCode != null && statusCode >= 200 && statusCode < 300
const snapshotEntitlement = snapshot?.entitlement_status?.trim()
const currentEntitlement = current.grok_entitlement_status?.trim()
const entitlementStatus = snapshotEntitlement || (
probeSucceeded && currentEntitlement?.toLowerCase() === 'forbidden'
? undefined
: current.grok_entitlement_status
)
const merged: AccountUsageInfo = {
...current,
grok_billing: result.billing ?? current.grok_billing,
@@ -1363,7 +1378,7 @@ const handleGrokProbed = (result: GrokQuotaProbeResult) => {
grok_request_quota: snapshot?.requests ?? current.grok_request_quota,
grok_token_quota: snapshot?.tokens ?? current.grok_token_quota,
grok_retry_after_seconds: snapshot?.retry_after_seconds ?? current.grok_retry_after_seconds,
grok_entitlement_status: snapshot?.entitlement_status || current.grok_entitlement_status,
grok_entitlement_status: entitlementStatus,
grok_quota_snapshot_state: result.billing
? 'billing_observed'
: snapshot?.headers_observed
@@ -1372,6 +1387,12 @@ const handleGrokProbed = (result: GrokQuotaProbeResult) => {
grok_last_quota_probe_at: result.billing?.fetched_at ?? snapshot?.last_probe_at ?? current.grok_last_quota_probe_at,
grok_last_headers_seen_at: snapshot?.last_headers_seen_at ?? current.grok_last_headers_seen_at,
grok_last_status_code: result.status_code ?? snapshot?.status_code ?? current.grok_last_status_code,
is_forbidden: probeSucceeded ? false : current.is_forbidden,
forbidden_reason: probeSucceeded ? undefined : current.forbidden_reason,
forbidden_type: probeSucceeded ? undefined : current.forbidden_type,
validation_url: probeSucceeded ? undefined : current.validation_url,
needs_verify: probeSucceeded ? false : current.needs_verify,
is_banned: probeSucceeded ? false : current.is_banned,
error: result.billing || snapshot ? undefined : current.error,
error_code: result.billing || snapshot ? undefined : current.error_code
}
frontend/src/components/account/__tests__/AccountUsageCell.spec.ts
+158
View File
@@ -1005,6 +1005,164 @@ describe('AccountUsageCell', () => {
expect(wrapper.text()).not.toContain('stale error')
})
it('Grok successful probes immediately clear stale forbidden state', async () => {
getUsage.mockResolvedValue({
is_forbidden: true,
forbidden_reason: 'stale forbidden response',
forbidden_type: 'validation',
validation_url: 'https://example.com/verify',
needs_verify: true,
is_banned: true,
grok_entitlement_status: 'forbidden',
grok_quota_snapshot_state: 'no_headers',
error: 'stale forbidden response',
error_code: 'forbidden'
})
const wrapper = mount(AccountUsageCell, {
props: {
account: makeAccount({ id: 4503, platform: 'grok', type: 'oauth', extra: {} })
},
global: {
stubs: {
UsageProgressBar: true,
AccountQuotaInfo: true,
GrokQuotaProbeCell: true
}
}
})
await flushPromises()
expect(wrapper.text()).toContain('forbidden')
const setupState = wrapper.vm.$.setupState as {
handleGrokProbed: (result: Record<string, unknown>) => void
usageInfo: Record<string, unknown> | null
}
setupState.handleGrokProbed({
source: 'active_probe',
snapshot: {
headers_observed: false,
updated_at: '2026-07-18T00:00:00Z',
status_code: 200
},
status_code: 200,
headers_observed: false,
reset_supported: false,
fetched_at: 1
})
await wrapper.vm.$nextTick()
expect(setupState.usageInfo).toMatchObject({
is_forbidden: false,
needs_verify: false,
is_banned: false,
grok_last_status_code: 200
})
expect(setupState.usageInfo?.forbidden_reason).toBeUndefined()
expect(setupState.usageInfo?.forbidden_type).toBeUndefined()
expect(setupState.usageInfo?.validation_url).toBeUndefined()
expect(setupState.usageInfo?.grok_entitlement_status).toBeUndefined()
expect(wrapper.text()).not.toContain('admin.accounts.forbidden')
})
it('Grok successful probes preserve the entitlement reported by the latest snapshot', async () => {
getUsage.mockResolvedValue({
is_forbidden: true,
grok_entitlement_status: 'forbidden'
})
const wrapper = mount(AccountUsageCell, {
props: {
account: makeAccount({ id: 4504, platform: 'grok', type: 'oauth', extra: {} })
},
global: {
stubs: {
UsageProgressBar: true,
AccountQuotaInfo: true,
GrokQuotaProbeCell: true
}
}
})
await flushPromises()
const setupState = wrapper.vm.$.setupState as {
handleGrokProbed: (result: Record<string, unknown>) => void
usageInfo: Record<string, unknown> | null
}
setupState.handleGrokProbed({
source: 'active_probe',
snapshot: {
headers_observed: true,
updated_at: '2026-07-18T00:00:00Z',
entitlement_status: 'ACTIVE',
status_code: 200
},
status_code: 200,
headers_observed: true,
reset_supported: false,
fetched_at: 1
})
await wrapper.vm.$nextTick()
expect(setupState.usageInfo?.grok_entitlement_status).toBe('ACTIVE')
expect(wrapper.text()).toContain('ACTIVE')
expect(wrapper.text()).not.toContain('admin.accounts.forbidden')
})
it('Grok billing-only success does not clear an active-probe forbidden state', async () => {
getUsage.mockResolvedValue({
is_forbidden: true,
forbidden_type: 'forbidden',
needs_verify: true,
is_banned: true,
grok_entitlement_status: 'forbidden'
})
const wrapper = mount(AccountUsageCell, {
props: {
account: makeAccount({ id: 4505, platform: 'grok', type: 'oauth', extra: {} })
},
global: {
stubs: {
UsageProgressBar: true,
AccountQuotaInfo: true,
GrokQuotaProbeCell: true
}
}
})
await flushPromises()
const setupState = wrapper.vm.$.setupState as {
handleGrokProbed: (result: Record<string, unknown>) => void
usageInfo: Record<string, unknown> | null
}
setupState.handleGrokProbed({
source: 'billing_probe',
billing: {
period_type: 'weekly',
usage_percent: 10,
plan: 'SuperGrok'
},
status_code: 200,
headers_observed: false,
reset_supported: false,
fetched_at: 1
})
await wrapper.vm.$nextTick()
expect(setupState.usageInfo).toMatchObject({
is_forbidden: true,
forbidden_type: 'forbidden',
needs_verify: true,
is_banned: true,
grok_entitlement_status: 'forbidden'
})
expect(wrapper.text()).toContain('forbidden')
})
it('Grok Free manual probes merge rolling 24h usage', async () => {
getUsage.mockResolvedValue({
subscription_tier: 'FREE',