From dd7a2b22f051ab5cfd77de85473eb5cd4bbf1289 Mon Sep 17 00:00:00 2001 From: superman2003 <2112076433zcr@gmail.com> Date: Sat, 18 Jul 2026 20:46:02 +0800 Subject: [PATCH 1/2] fix(grok): align Free probes with health checks --- .../handler/admin/grok_oauth_handler_test.go | 6 +- .../internal/service/account_test_service.go | 6 +- .../service/account_test_service_grok_test.go | 5 + .../internal/service/grok_quota_fetcher.go | 36 +++- .../service/grok_quota_fetcher_test.go | 125 ++++++++++++++ .../internal/service/grok_quota_service.go | 11 +- .../service/grok_quota_service_test.go | 13 +- .../components/account/AccountUsageCell.vue | 23 ++- .../__tests__/AccountUsageCell.spec.ts | 158 ++++++++++++++++++ 9 files changed, 366 insertions(+), 17 deletions(-) diff --git a/backend/internal/handler/admin/grok_oauth_handler_test.go b/backend/internal/handler/admin/grok_oauth_handler_test.go index 9d87fc0a98..6214394882 100644 --- a/backend/internal/handler/admin/grok_oauth_handler_test.go +++ b/backend/internal/handler/admin/grok_oauth_handler_test.go @@ -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]) diff --git a/backend/internal/service/account_test_service.go b/backend/internal/service/account_test_service.go index 1d8966007f..11c3853221 100644 --- a/backend/internal/service/account_test_service.go +++ b/backend/internal/service/account_test_service.go @@ -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") } diff --git a/backend/internal/service/account_test_service_grok_test.go b/backend/internal/service/account_test_service_grok_test.go index 72f995a0c7..fcfe1257b6 100644 --- a/backend/internal/service/account_test_service_grok_test.go +++ b/backend/internal/service/account_test_service_grok_test.go @@ -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"`) diff --git a/backend/internal/service/grok_quota_fetcher.go b/backend/internal/service/grok_quota_fetcher.go index f220fe33b9..1f2d6374ad 100644 --- a/backend/internal/service/grok_quota_fetcher.go +++ b/backend/internal/service/grok_quota_fetcher.go @@ -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 diff --git a/backend/internal/service/grok_quota_fetcher_test.go b/backend/internal/service/grok_quota_fetcher_test.go index 1de9b51c9e..6bea2814f2 100644 --- a/backend/internal/service/grok_quota_fetcher_test.go +++ b/backend/internal/service/grok_quota_fetcher_test.go @@ -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() diff --git a/backend/internal/service/grok_quota_service.go b/backend/internal/service/grok_quota_service.go index db11fe6cc8..cf28f6c6ec 100644 --- a/backend/internal/service/grok_quota_service.go +++ b/backend/internal/service/grok_quota_service.go @@ -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, }) } diff --git a/backend/internal/service/grok_quota_service_test.go b/backend/internal/service/grok_quota_service_test.go index 728277e62d..ae3dacc4a7 100644 --- a/backend/internal/service/grok_quota_service_test.go +++ b/backend/internal/service/grok_quota_service_test.go @@ -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) } diff --git a/frontend/src/components/account/AccountUsageCell.vue b/frontend/src/components/account/AccountUsageCell.vue index 301b512c25..9a3d6a86e9 100644 --- a/frontend/src/components/account/AccountUsageCell.vue +++ b/frontend/src/components/account/AccountUsageCell.vue @@ -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 } diff --git a/frontend/src/components/account/__tests__/AccountUsageCell.spec.ts b/frontend/src/components/account/__tests__/AccountUsageCell.spec.ts index 7c2fe48a56..5f61494c5f 100644 --- a/frontend/src/components/account/__tests__/AccountUsageCell.spec.ts +++ b/frontend/src/components/account/__tests__/AccountUsageCell.spec.ts @@ -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) => void + usageInfo: Record | 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) => void + usageInfo: Record | 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) => void + usageInfo: Record | 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', From e14fb2b6ff3b93b42311471c05eb04bd5dd1518f Mon Sep 17 00:00:00 2001 From: superman2003 <2112076433zcr@gmail.com> Date: Sat, 18 Jul 2026 20:46:03 +0800 Subject: [PATCH 2/2] fix(grok): recover invalid encrypted content once --- .../internal/service/openai_gateway_grok.go | 98 ++++++++- .../service/openai_gateway_grok_test.go | 195 ++++++++++++++++++ 2 files changed, 285 insertions(+), 8 deletions(-) diff --git a/backend/internal/service/openai_gateway_grok.go b/backend/internal/service/openai_gateway_grok.go index c093029332..198055f20e 100644 --- a/backend/internal/service/openai_gateway_grok.go +++ b/backend/internal/service/openai_gateway_grok.go @@ -75,10 +75,6 @@ func (s *OpenAIGatewayService) forwardGrokResponses( upstreamCtx, releaseUpstreamCtx := detachUpstreamContext(ctx) defer releaseUpstreamCtx() - upstreamReq, err := buildGrokResponsesRequest(upstreamCtx, c, account, patchedBody, token, cacheIdentity, s.cfg) - if err != nil { - return nil, err - } proxyURL := "" if account.ProxyID != nil && account.Proxy != nil { @@ -86,10 +82,45 @@ func (s *OpenAIGatewayService) forwardGrokResponses( } upstreamStart := time.Now() - resp, err := s.httpUpstream.Do(upstreamReq, proxyURL, account.ID, account.Concurrency) - SetOpsLatencyMs(c, OpsUpstreamLatencyMsKey, time.Since(upstreamStart).Milliseconds()) - if err != nil { - return nil, s.handleOpenAIUpstreamTransportError(ctx, c, account, err, false) + var resp *http.Response + for attempt := 0; ; attempt++ { + upstreamReq, buildErr := buildGrokResponsesRequest(upstreamCtx, c, account, patchedBody, token, cacheIdentity, s.cfg) + if buildErr != nil { + return nil, buildErr + } + + resp, err = s.httpUpstream.Do(upstreamReq, proxyURL, account.ID, account.Concurrency) + SetOpsLatencyMs(c, OpsUpstreamLatencyMsKey, time.Since(upstreamStart).Milliseconds()) + if err != nil { + return nil, s.handleOpenAIUpstreamTransportError(ctx, c, account, err, false) + } + + // xAI can reject encrypted reasoning copied from a response produced under + // another account or cache identity. Retry once with the same routing and + // credential after removing only the rejected encrypted reasoning payload. + if attempt > 0 || resp.StatusCode != http.StatusBadRequest { + break + } + respBody := s.readUpstreamErrorBody(resp) + if resp.Body != nil { + _ = resp.Body.Close() + } + if !isGrokInvalidEncryptedContentResponse(resp.StatusCode, respBody) { + resp.Body = io.NopCloser(bytes.NewReader(respBody)) + break + } + + retryBody, changed, trimErr := trimGrokInvalidEncryptedContentRetryBody(patchedBody) + if trimErr != nil { + return nil, fmt.Errorf("prepare Grok invalid encrypted_content retry: %w", trimErr) + } + if !changed { + resp.Body = io.NopCloser(bytes.NewReader(respBody)) + break + } + + patchedBody = retryBody + slog.Info("grok_invalid_encrypted_content_retry", "account_id", account.ID, "cache_identity_present", cacheIdentity != "") } defer func() { _ = resp.Body.Close() }() @@ -162,6 +193,57 @@ func (s *OpenAIGatewayService) forwardGrokResponses( }, nil } +func isGrokInvalidEncryptedContentResponse(statusCode int, body []byte) bool { + if statusCode != http.StatusBadRequest { + return false + } + + code := gjson.GetBytes(body, "code") + message := gjson.GetBytes(body, "error") + if code.Type != gjson.String || message.Type != gjson.String || + !strings.EqualFold(strings.TrimSpace(code.String()), "invalid-argument") { + return false + } + + normalizedMessage := strings.ToLower(message.String()) + return strings.Contains(normalizedMessage, "decrypt") && strings.Contains(normalizedMessage, "encrypted_content") +} + +func trimGrokInvalidEncryptedContentRetryBody(body []byte) ([]byte, bool, error) { + input := gjson.GetBytes(body, "input") + items := input.Array() + if input.IsObject() { + items = []gjson.Result{input} + } + + hasEncryptedReasoning := false + for _, item := range items { + if strings.TrimSpace(item.Get("type").String()) == "reasoning" && item.Get("encrypted_content").Exists() { + hasEncryptedReasoning = true + break + } + } + if !hasEncryptedReasoning { + return body, false, nil + } + + var requestBody map[string]any + decoder := json.NewDecoder(bytes.NewReader(body)) + decoder.UseNumber() + if err := decoder.Decode(&requestBody); err != nil { + return nil, false, err + } + if !trimOpenAIEncryptedReasoningItems(requestBody) { + return body, false, nil + } + + retryBody, err := marshalOpenAIUpstreamJSON(requestBody) + if err != nil { + return nil, false, err + } + return retryBody, true, nil +} + func patchGrokResponsesBody(body []byte, upstreamModel string) ([]byte, error) { if !json.Valid(body) { return nil, fmt.Errorf("invalid json request body") diff --git a/backend/internal/service/openai_gateway_grok_test.go b/backend/internal/service/openai_gateway_grok_test.go index 8995c3b1de..806b9d5f19 100644 --- a/backend/internal/service/openai_gateway_grok_test.go +++ b/backend/internal/service/openai_gateway_grok_test.go @@ -1186,6 +1186,201 @@ func TestForwardGrokResponsesAPIKeyUsesXAIResponses(t *testing.T) { require.Equal(t, 1, result.Usage.OutputTokens) } +func TestForwardGrokResponsesRetriesInvalidEncryptedContentOnce(t *testing.T) { + gin.SetMode(gin.TestMode) + + recorder := httptest.NewRecorder() + c, _ := gin.CreateTestContext(recorder) + body := []byte(`{ + "model":"grok", + "input":[ + {"type":"reasoning","summary":[{"type":"summary_text","text":"keep this summary"}],"encrypted_content":"encrypted-reasoning"}, + {"type":"message","role":"user","content":[{"type":"input_text","text":"hi"}]} + ], + "metadata":{"large_id":9007199254740993}, + "stream":false + }`) + c.Request = httptest.NewRequest(http.MethodPost, "/v1/responses", bytes.NewReader(body)) + c.Set("api_key", &APIKey{ID: 4535}) + + account := &Account{ + ID: 4535, + Name: "grok-api-key", + Platform: PlatformGrok, + Type: AccountTypeAPIKey, + Status: StatusActive, + Schedulable: true, + Concurrency: 2, + Credentials: map[string]any{ + "api_key": "same-token", + "base_url": "https://api.x.ai/v1", + }, + } + upstream := &httpUpstreamRecorder{responses: []*http.Response{ + { + StatusCode: http.StatusBadRequest, + Header: http.Header{ + "Content-Type": []string{"application/json"}, + "Xai-Request-Id": []string{"recoverable-first"}, + }, + Body: io.NopCloser(strings.NewReader(`{"code":"invalid-argument","error":"Could not decrypt the provided encrypted_content. Ensure the value is unmodified."}`)), + }, + { + StatusCode: http.StatusOK, + Header: http.Header{ + "Content-Type": []string{"application/json"}, + "Xai-Request-Id": []string{"recovered-second"}, + }, + Body: io.NopCloser(strings.NewReader(`{"id":"resp_recovered","object":"response","model":"grok-4.5","status":"completed","output":[],"usage":{"input_tokens":2,"output_tokens":1}}`)), + }, + }} + svc := &OpenAIGatewayService{httpUpstream: upstream} + + result, err := svc.forwardGrokResponses(context.Background(), c, account, body, "grok", false, time.Now()) + require.NoError(t, err) + require.NotNil(t, result) + require.Equal(t, "resp_recovered", result.ResponseID) + require.Equal(t, "recovered-second", result.RequestID) + require.Len(t, upstream.requests, 2) + require.Len(t, upstream.bodies, 2) + + require.Equal(t, "reasoning", gjson.GetBytes(upstream.bodies[0], "input.0.type").String()) + require.Equal(t, "encrypted-reasoning", gjson.GetBytes(upstream.bodies[0], "input.0.encrypted_content").String()) + require.Equal(t, "reasoning", gjson.GetBytes(upstream.bodies[1], "input.0.type").String()) + require.False(t, gjson.GetBytes(upstream.bodies[1], "input.0.encrypted_content").Exists()) + require.Equal(t, "keep this summary", gjson.GetBytes(upstream.bodies[1], "input.0.summary.0.text").String()) + require.Equal(t, "message", gjson.GetBytes(upstream.bodies[1], "input.1.type").String()) + require.Equal(t, "9007199254740993", gjson.GetBytes(upstream.bodies[0], "metadata.large_id").Raw) + require.Equal(t, "9007199254740993", gjson.GetBytes(upstream.bodies[1], "metadata.large_id").Raw) + + firstIdentity := gjson.GetBytes(upstream.bodies[0], "prompt_cache_key").String() + secondIdentity := gjson.GetBytes(upstream.bodies[1], "prompt_cache_key").String() + require.NotEmpty(t, firstIdentity) + require.Equal(t, firstIdentity, secondIdentity) + for _, req := range upstream.requests { + require.Equal(t, "Bearer same-token", req.Header.Get("Authorization")) + require.Equal(t, firstIdentity, req.Header.Get(grokConversationIDHeader)) + } + require.Equal(t, StatusActive, account.Status) + _, hasUpstreamErrors := c.Get(OpsUpstreamErrorsKey) + require.False(t, hasUpstreamErrors) + _, hasTerminalStatus := c.Get(OpsUpstreamStatusCodeKey) + require.False(t, hasTerminalStatus) +} + +func TestForwardGrokResponsesInvalidEncryptedContentRecoveryDoesNotOvermatch(t *testing.T) { + gin.SetMode(gin.TestMode) + + matchingError := `{"code":"invalid-argument","error":"Could not decrypt the provided encrypted_content."}` + tests := []struct { + name string + requestBody string + responseBody string + }{ + { + name: "different top-level code", + requestBody: `{"model":"grok","input":[{"type":"reasoning","encrypted_content":"cipher"}],"stream":false}`, + responseBody: `{"code":"bad-request","error":"Could not decrypt the provided encrypted_content."}`, + }, + { + name: "message does not mention decryption", + requestBody: `{"model":"grok","input":[{"type":"reasoning","encrypted_content":"cipher"}],"stream":false}`, + responseBody: `{"code":"invalid-argument","error":"The provided encrypted_content is invalid."}`, + }, + { + name: "nested OpenAI error shape", + requestBody: `{"model":"grok","input":[{"type":"reasoning","encrypted_content":"cipher"}],"stream":false}`, + responseBody: `{"code":"invalid-argument","error":{"message":"Could not decrypt the provided encrypted_content."}}`, + }, + { + name: "request has no encrypted reasoning", + requestBody: `{"model":"grok","input":[{"type":"message","role":"user","content":"hi"}],"stream":false}`, + responseBody: matchingError, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + recorder := httptest.NewRecorder() + c, _ := gin.CreateTestContext(recorder) + body := []byte(tt.requestBody) + c.Request = httptest.NewRequest(http.MethodPost, "/v1/responses", bytes.NewReader(body)) + + account := &Account{ + ID: 4536, + Name: "grok-api-key", + Platform: PlatformGrok, + Type: AccountTypeAPIKey, + Concurrency: 1, + Credentials: map[string]any{"api_key": "token", "base_url": "https://api.x.ai/v1"}, + } + upstream := &httpUpstreamRecorder{responses: []*http.Response{{ + StatusCode: http.StatusBadRequest, + Header: http.Header{"Content-Type": []string{"application/json"}}, + Body: io.NopCloser(strings.NewReader(tt.responseBody)), + }}} + svc := &OpenAIGatewayService{httpUpstream: upstream} + + result, err := svc.forwardGrokResponses(context.Background(), c, account, body, "grok", false, time.Now()) + require.Nil(t, result) + require.Error(t, err) + require.Len(t, upstream.requests, 1) + require.Len(t, upstream.bodies, 1) + }) + } +} + +func TestForwardGrokResponsesInvalidEncryptedContentRetryFailureIsTerminal(t *testing.T) { + gin.SetMode(gin.TestMode) + + recorder := httptest.NewRecorder() + c, _ := gin.CreateTestContext(recorder) + body := []byte(`{"model":"grok","input":[{"type":"reasoning","encrypted_content":"cipher"},{"type":"message","role":"user","content":"hi"}],"stream":false}`) + c.Request = httptest.NewRequest(http.MethodPost, "/v1/responses", bytes.NewReader(body)) + + account := &Account{ + ID: 4537, + Name: "grok-api-key", + Platform: PlatformGrok, + Type: AccountTypeAPIKey, + Concurrency: 1, + Credentials: map[string]any{"api_key": "same-token", "base_url": "https://api.x.ai/v1"}, + } + newInvalidEncryptedResponse := func(requestID string) *http.Response { + return &http.Response{ + StatusCode: http.StatusBadRequest, + Header: http.Header{ + "Content-Type": []string{"application/json"}, + "Xai-Request-Id": []string{requestID}, + }, + Body: io.NopCloser(strings.NewReader(`{"code":"invalid-argument","error":"Could not decrypt the provided encrypted_content."}`)), + } + } + upstream := &httpUpstreamRecorder{responses: []*http.Response{ + newInvalidEncryptedResponse("recoverable-first"), + newInvalidEncryptedResponse("terminal-second"), + }} + svc := &OpenAIGatewayService{httpUpstream: upstream} + + result, err := svc.forwardGrokResponses(context.Background(), c, account, body, "grok", false, time.Now()) + require.Nil(t, result) + require.Error(t, err) + require.Len(t, upstream.requests, 2) + require.Len(t, upstream.bodies, 2) + require.True(t, gjson.GetBytes(upstream.bodies[0], "input.0.encrypted_content").Exists()) + require.False(t, gjson.GetBytes(upstream.bodies[1], `input.#(type=="reasoning")`).Exists()) + + rawEvents, ok := c.Get(OpsUpstreamErrorsKey) + require.True(t, ok) + events, ok := rawEvents.([]*OpsUpstreamErrorEvent) + require.True(t, ok) + require.NotEmpty(t, events) + for _, event := range events { + require.NotEqual(t, "recoverable-first", event.UpstreamRequestID) + } + require.Equal(t, http.StatusBadRequest, c.GetInt(OpsUpstreamStatusCodeKey)) +} + func TestForwardAsChatCompletionsForGrokAPIKeyUsesConfiguredRawEndpointWithoutOAuthIdentity(t *testing.T) { gin.SetMode(gin.TestMode) recorder := httptest.NewRecorder()