fix: improve Grok OAuth, image, and usage flows

This commit is contained in:
Heatherm Huang
2026-07-09 17:54:08 +08:00
parent 668c79ea86
commit b480545c19
10 changed files with 254 additions and 6 deletions
+24 -1
View File
@@ -315,6 +315,11 @@ func (s *OpenAIGatewayService) ForwardGrokMedia(
if err != nil {
return nil, err
}
requestInfo := ParseGrokMediaRequest(contentType, body)
body, contentType, err = sanitizeGrokMediaForwardBody(endpoint, body, contentType)
if err != nil {
return nil, err
}
var bodyReader io.Reader
if endpoint.RequiresRequestBody() {
@@ -350,7 +355,6 @@ func (s *OpenAIGatewayService) ForwardGrokMedia(
defer func() { _ = resp.Body.Close() }()
requestIDHeader := firstNonEmpty(resp.Header.Get("x-request-id"), resp.Header.Get("xai-request-id"))
requestInfo := ParseGrokMediaRequest(contentType, body)
requestModel := requestInfo.Model
if resp.StatusCode >= 400 {
s.updateGrokUsageSnapshot(ctx, account.ID, xai.ParseQuotaHeaders(resp.Header, resp.StatusCode))
@@ -462,6 +466,25 @@ func normalizeGrokMediaForwardBody(endpoint GrokMediaEndpoint, body []byte, cont
return out, contentType, nil
}
func sanitizeGrokMediaForwardBody(endpoint GrokMediaEndpoint, body []byte, contentType string) ([]byte, string, error) {
if !endpoint.RequiresRequestBody() || !gjson.ValidBytes(body) {
return body, contentType, nil
}
switch endpoint {
case GrokMediaEndpointImagesGenerations, GrokMediaEndpointImagesEdits:
if !gjson.GetBytes(body, "size").Exists() {
return body, contentType, nil
}
out, err := sjson.DeleteBytes(body, "size")
if err != nil {
return nil, "", fmt.Errorf("sanitize grok media size: %w", err)
}
return out, contentType, nil
default:
return body, contentType, nil
}
}
func (r GrokMediaRequestInfo) HasInputImage() bool {
return len(r.InputImageURLs) > 0 || len(r.Uploads) > 0
}
@@ -323,6 +323,43 @@ func TestForwardGrokMediaImagesGenerationNormalizesImagineAlias(t *testing.T) {
require.Equal(t, ImageBillingSize2K, result.ImageSize)
}
func TestForwardGrokMediaImagesGenerationStripsUnsupportedSize(t *testing.T) {
t.Setenv(xai.EnvAllowUnsafeURLOverrides, "true")
gin.SetMode(gin.TestMode)
recorder := httptest.NewRecorder()
c, _ := gin.CreateTestContext(recorder)
body := []byte(`{"model":"grok-imagine-image","prompt":"draw a cat","size":"1024x1024"}`)
c.Request = httptest.NewRequest(http.MethodPost, "/v1/images/generations", bytes.NewReader(body))
c.Request.Header.Set("Content-Type", "application/json")
account := &Account{
ID: 65,
Name: "grok",
Platform: PlatformGrok,
Type: AccountTypeAPIKey,
Concurrency: 1,
Credentials: map[string]any{
"api_key": "api-key",
"base_url": "https://xai.test/v1",
},
}
upstream := &httpUpstreamRecorder{resp: &http.Response{
StatusCode: http.StatusOK,
Header: http.Header{
"Content-Type": []string{"application/json"},
},
Body: io.NopCloser(strings.NewReader(`{"data":[]}`)),
}}
svc := &OpenAIGatewayService{httpUpstream: upstream}
result, err := svc.ForwardGrokMedia(context.Background(), c, account, GrokMediaEndpointImagesGenerations, "", body, "application/json")
require.NoError(t, err)
require.JSONEq(t, `{"model":"grok-imagine-image","prompt":"draw a cat"}`, string(upstream.lastBody))
require.Equal(t, ImageBillingSize1K, result.ImageSize)
require.Equal(t, "1024x1024", result.ImageInputSize)
}
func TestForwardGrokMediaImagesEditMultipartConvertsToJSON(t *testing.T) {
t.Setenv(xai.EnvAllowUnsafeURLOverrides, "true")
gin.SetMode(gin.TestMode)
+9 -1
View File
@@ -455,7 +455,15 @@ func applyOpenAIImagesDefaults(req *OpenAIImagesRequest) {
}
func isOpenAIImageGenerationModel(model string) bool {
return strings.HasPrefix(strings.ToLower(strings.TrimSpace(model)), "gpt-image-")
model = strings.ToLower(strings.TrimSpace(model))
return strings.HasPrefix(model, "gpt-image-") || isGrokImageGenerationModel(model)
}
func isGrokImageGenerationModel(model string) bool {
model = strings.ToLower(strings.TrimSpace(model))
return model == "grok-imagine" ||
model == "grok-imagine-edit" ||
strings.HasPrefix(model, "grok-imagine-image")
}
func validateOpenAIImagesModel(model string) error {
@@ -334,6 +334,28 @@ func TestOpenAIGatewayServiceParseOpenAIImagesRequest_RejectsNonImageModel(t *te
require.ErrorContains(t, err, `images endpoint requires an image model, got "gpt-5.4"`)
}
func TestOpenAIGatewayServiceParseOpenAIImagesRequest_AllowsGrokImageModels(t *testing.T) {
gin.SetMode(gin.TestMode)
for _, model := range []string{"grok-imagine", "grok-imagine-image", "grok-imagine-image-quality", "grok-imagine-edit"} {
t.Run(model, func(t *testing.T) {
body := []byte(fmt.Sprintf(`{"model":%q,"prompt":"draw a cat","response_format":"b64_json"}`, model))
req := httptest.NewRequest(http.MethodPost, "/v1/images/generations", bytes.NewReader(body))
req.Header.Set("Content-Type", "application/json")
rec := httptest.NewRecorder()
c, _ := gin.CreateTestContext(rec)
c.Request = req
svc := &OpenAIGatewayService{}
parsed, err := svc.ParseOpenAIImagesRequest(c, body)
require.NoError(t, err)
require.NotNil(t, parsed)
require.Equal(t, model, parsed.Model)
require.Equal(t, OpenAIImagesCapabilityNative, parsed.RequiredCapability)
})
}
}
func TestOpenAIGatewayServiceParseOpenAIImagesRequest_JSONEditURLs(t *testing.T) {
gin.SetMode(gin.TestMode)
body := []byte(`{
@@ -372,6 +372,13 @@
<span class="rounded bg-gray-100 px-1.5 py-0.5 dark:bg-gray-800" :title="t('usage.accountBilled')">
A ${{ formatWindowCost(grokLocalUsage) }}
</span>
<span
v-if="grokLocalUsage.user_cost != null"
class="rounded bg-gray-100 px-1.5 py-0.5 dark:bg-gray-800"
:title="t('usage.userBilled')"
>
U ${{ formatWindowUserCost(grokLocalUsage) }}
</span>
</div>
</div>
<UsageProgressBar
@@ -1031,7 +1038,7 @@ const makeGrokQuotaBar = (quota?: { limit?: number | null; remaining?: number |
if (!quota || quota.limit == null || quota.remaining == null || quota.limit <= 0) return null
const used = Math.max(0, quota.limit - quota.remaining)
return {
utilization: Math.min(100, (used / quota.limit) * 100),
utilization: (used / quota.limit) * 100,
resetsAt: quota.reset_at || null
}
}
@@ -1087,6 +1094,7 @@ const grokRetryAfterLabel = computed(() => {
const formatWindowRequests = (stats: WindowStats) => formatCompactNumber(stats.requests, { allowBillions: false })
const formatWindowTokens = (stats: WindowStats) => formatCompactNumber(stats.tokens)
const formatWindowCost = (stats: WindowStats) => stats.cost.toFixed(2)
const formatWindowUserCost = (stats: WindowStats) => (stats.user_cost ?? 0).toFixed(2)
// 账户类型显示标签
const antigravityTierLabel = computed(() => {
@@ -566,6 +566,58 @@ describe('AccountUsageCell', () => {
expect(badges.some(node => node.attributes('title') === 'usage.userBilled')).toBe(true)
})
it('Grok OAuth 会展示本地 user billed 用量并保留超限百分比', async () => {
getUsage.mockResolvedValue({
grok_local_usage: {
requests: 4,
tokens: 1200,
cost: 0.12,
standard_cost: 0.12,
user_cost: 0.34
},
grok_request_quota: {
limit: 10,
remaining: -2,
reset_at: '2026-07-09T16:00:00Z'
},
grok_quota_snapshot_state: 'observed'
})
const wrapper = mount(AccountUsageCell, {
props: {
account: makeAccount({
id: 3861,
platform: 'grok',
type: 'oauth',
extra: {}
})
},
global: {
stubs: {
UsageProgressBar: {
props: ['label', 'utilization', 'resetsAt', 'color'],
template: '<div class="usage-bar">{{ label }}|{{ utilization }}|{{ resetsAt }}</div>'
},
AccountQuotaInfo: true,
GrokQuotaProbeCell: true
}
}
})
await flushPromises()
expect(getUsage).toHaveBeenCalledWith(3861)
expect(wrapper.text()).toContain('4 req')
expect(wrapper.text()).toContain('1.2K')
expect(wrapper.text()).toContain('A $0.12')
expect(wrapper.text()).toContain('U $0.34')
expect(wrapper.text()).toContain('admin.accounts.usageWindow.grokRequests|120|2026-07-09T16:00:00Z')
const badges = wrapper.findAll('span[title]')
expect(badges.some(node => node.attributes('title') === 'usage.accountBilled')).toBe(true)
expect(badges.some(node => node.attributes('title') === 'usage.userBilled')).toBe(true)
})
it('Key 账号在 today stats loading 时显示骨架屏', async () => {
const wrapper = mount(AccountUsageCell, {
props: {
@@ -0,0 +1,55 @@
import { describe, expect, it, vi } from 'vitest'
vi.mock('@/stores/app', () => ({
useAppStore: () => ({
showError: vi.fn()
})
}))
vi.mock('vue-i18n', () => ({
useI18n: () => ({
t: (key: string) => {
const messages: Record<string, string> = {
'admin.accounts.oauth.grok.failedToExchangeCode': 'Grok 授权码兑换失败',
'admin.accounts.oauth.grok.errors.GROK_OAUTH_INVALID_STATE':
'Grok OAuth state 与当前会话不匹配。请粘贴同一次生成的授权链接返回的回调 URL。'
}
return messages[key] ?? key
}
})
}))
vi.mock('@/api/admin', () => ({
adminAPI: {
grok: {
generateAuthUrl: vi.fn(),
exchangeCode: vi.fn(),
refreshGrokToken: vi.fn()
}
}
}))
import { useGrokOAuth } from '@/composables/useGrokOAuth'
import { adminAPI } from '@/api/admin'
describe('useGrokOAuth.exchangeAuthCode', () => {
it('shows a state mismatch recovery hint from structured backend errors', async () => {
vi.mocked(adminAPI.grok.exchangeCode).mockRejectedValueOnce({
status: 400,
reason: 'GROK_OAUTH_INVALID_STATE',
message: 'invalid oauth state'
})
const oauth = useGrokOAuth()
const tokenInfo = await oauth.exchangeAuthCode({
code: 'code',
sessionId: 'session-id',
state: 'wrong-state'
})
expect(tokenInfo).toBeNull()
expect(oauth.error.value).toBe(
'Grok OAuth state 与当前会话不匹配。请粘贴同一次生成的授权链接返回的回调 URL。'
)
})
})
+14 -3
View File
@@ -3,6 +3,7 @@ import { useI18n } from 'vue-i18n'
import { useAppStore } from '@/stores/app'
import { adminAPI } from '@/api/admin'
import type { GrokTokenInfo } from '@/api/admin/grok'
import { extractApiErrorMessage, extractI18nErrorMessage } from '@/utils/apiError'
export function useGrokOAuth() {
const appStore = useAppStore()
@@ -39,7 +40,7 @@ export function useGrokOAuth() {
state.value = response.state
return true
} catch (err: any) {
error.value = err.response?.data?.detail || t('admin.accounts.oauth.grok.failedToGenerateUrl')
error.value = extractApiErrorMessage(err, t('admin.accounts.oauth.grok.failedToGenerateUrl'))
appStore.showError(error.value)
return false
} finally {
@@ -72,7 +73,12 @@ export function useGrokOAuth() {
return await adminAPI.grok.exchangeCode(payload as any)
} catch (err: any) {
error.value = err.response?.data?.detail || t('admin.accounts.oauth.grok.failedToExchangeCode')
error.value = extractI18nErrorMessage(
err,
t,
'admin.accounts.oauth.grok.errors',
t('admin.accounts.oauth.grok.failedToExchangeCode')
)
appStore.showError(error.value)
return null
} finally {
@@ -95,7 +101,12 @@ export function useGrokOAuth() {
try {
return await adminAPI.grok.refreshGrokToken(refreshToken.trim(), proxyId)
} catch (err: any) {
error.value = err.response?.data?.detail || t('admin.accounts.oauth.grok.failedToValidateRT')
error.value = extractI18nErrorMessage(
err,
t,
'admin.accounts.oauth.grok.errors',
t('admin.accounts.oauth.grok.failedToValidateRT')
)
return null
} finally {
loading.value = false
@@ -873,6 +873,22 @@ export default {
missingExchangeParams: 'Missing authorization code, state, or OAuth session',
failedToExchangeCode: 'Failed to exchange Grok authorization code',
failedToValidateRT: 'Failed to validate Grok refresh token',
errors: {
GROK_OAUTH_SESSION_NOT_FOUND:
'Grok OAuth session was not found or has expired. Generate a new auth URL and paste the newest callback URL.',
GROK_OAUTH_INVALID_STATE:
'Grok OAuth state does not match this session. Paste the callback URL from the same generated auth link.',
GROK_OAUTH_STATE_REQUIRED:
'The callback URL is missing the OAuth state. Paste the full callback URL, not only the code.',
GROK_OAUTH_CODE_REQUIRED:
'The Grok authorization code is missing. Paste the full callback URL, query string, or code value.',
GROK_OAUTH_NO_REFRESH_TOKEN:
'The Grok response did not include a refresh token. Generate a new auth URL and approve offline access again.',
GROK_OAUTH_PROXY_NOT_AVAILABLE:
'Grok OAuth proxy lookup is unavailable. Check the selected proxy and retry.',
GROK_OAUTH_PROXY_NOT_FOUND:
'The selected proxy could not be found. Choose an available proxy and retry.'
},
oauthOnlyHint: 'Initial Grok support is OAuth subscription-backed Responses API text and reasoning traffic only.'
},
// Gemini specific
@@ -960,6 +960,22 @@ export default {
missingExchangeParams: '缺少授权码、state 或 OAuth 会话',
failedToExchangeCode: 'Grok 授权码兑换失败',
failedToValidateRT: '验证 Grok refresh token 失败',
errors: {
GROK_OAUTH_SESSION_NOT_FOUND:
'Grok OAuth 会话不存在或已过期。请重新生成授权链接,并粘贴最新的回调链接。',
GROK_OAUTH_INVALID_STATE:
'Grok OAuth state 与当前会话不匹配。请粘贴同一次生成的授权链接返回的回调 URL。',
GROK_OAUTH_STATE_REQUIRED:
'回调链接缺少 OAuth state。请粘贴完整 callback URL,不要只粘贴 code。',
GROK_OAUTH_CODE_REQUIRED:
'缺少 Grok 授权码。请粘贴完整 callback URL、查询字符串或 code 值。',
GROK_OAUTH_NO_REFRESH_TOKEN:
'Grok 响应未返回 refresh token。请重新生成授权链接,并再次确认 offline access 授权。',
GROK_OAUTH_PROXY_NOT_AVAILABLE:
'无法查询 Grok OAuth 代理配置。请检查选择的代理后重试。',
GROK_OAUTH_PROXY_NOT_FOUND:
'找不到所选代理。请选择可用代理后重试。'
},
oauthOnlyHint: '首版 Grok 支持仅包含 OAuth 订阅的 Responses API 文本/推理转发。'
},
// Gemini specific