fix(grok): route cacheable chat requests via responses

This commit is contained in:
superman2003
2026-07-12 11:50:56 +08:00
parent 1dedb2097d
commit 7050070aa3
10 changed files with 803 additions and 17 deletions
+4 -2
View File
@@ -155,8 +155,10 @@ func isBareOrSubpathOf(path, root string) bool {
// account platform and the normalized inbound endpoint.
//
// Platform-specific rules:
// - OpenAI always forwards to /v1/responses (with optional subpath
// such as /v1/responses/compact preserved from the raw URL).
// - OpenAI and Grok default to /v1/responses (with optional subpath
// such as /v1/responses/compact preserved from the raw URL). Grok raw Chat
// requests override this through the forwarding result consumed by
// resolveOpenAIUpstreamEndpoint.
// - Anthropic → /v1/messages
// - Gemini → /v1beta/models
// - Antigravity → /v1/messages (Claude) or gemini (Gemini)
+55
View File
@@ -121,6 +121,8 @@ func TestDeriveUpstreamEndpoint(t *testing.T) {
{"openai embeddings", EndpointEmbeddings, "/v1/embeddings", service.PlatformOpenAI, EndpointEmbeddings},
{"openai image generations", EndpointImagesGenerations, "/v1/images/generations", service.PlatformOpenAI, EndpointImagesGenerations},
{"openai image edits", EndpointImagesEdits, "/openai/v1/images/edits", service.PlatformOpenAI, EndpointImagesEdits},
{"grok chat defaults to responses without runtime result", EndpointChatCompletions, "/v1/chat/completions", service.PlatformGrok, EndpointResponses},
{"grok responses", EndpointResponses, "/v1/responses", service.PlatformGrok, EndpointResponses},
{"grok video generations", EndpointVideosGenerations, "/v1/videos/generations", service.PlatformGrok, EndpointVideosGenerations},
{"grok video status", EndpointVideos, "/videos/req_123", service.PlatformGrok, EndpointVideos},
@@ -138,6 +140,59 @@ func TestDeriveUpstreamEndpoint(t *testing.T) {
}
}
func TestResolveOpenAIUpstreamEndpointPrefersForwardResult(t *testing.T) {
tests := []struct {
name string
account *service.Account
result *service.OpenAIForwardResult
runtimeEndpoint string
want string
}{
{
name: "grok raw chat result overrides stale context",
account: &service.Account{Platform: service.PlatformGrok, Type: service.AccountTypeOAuth},
result: &service.OpenAIForwardResult{UpstreamEndpoint: EndpointChatCompletions},
runtimeEndpoint: EndpointResponses,
want: EndpointChatCompletions,
},
{
name: "grok chat bridged to responses",
account: &service.Account{Platform: service.PlatformGrok, Type: service.AccountTypeOAuth},
result: &service.OpenAIForwardResult{UpstreamEndpoint: EndpointResponses},
want: EndpointResponses,
},
{
name: "grok empty result keeps responses default",
account: &service.Account{Platform: service.PlatformGrok, Type: service.AccountTypeOAuth},
result: &service.OpenAIForwardResult{},
want: EndpointResponses,
},
{
name: "grok raw error uses runtime endpoint",
account: &service.Account{Platform: service.PlatformGrok, Type: service.AccountTypeOAuth},
runtimeEndpoint: EndpointChatCompletions,
want: EndpointChatCompletions,
},
{
name: "openai behavior remains responses",
account: &service.Account{Platform: service.PlatformOpenAI, Type: service.AccountTypeOAuth},
result: &service.OpenAIForwardResult{},
want: EndpointResponses,
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
rec := httptest.NewRecorder()
c, _ := gin.CreateTestContext(rec)
c.Request = httptest.NewRequest(http.MethodPost, EndpointChatCompletions, nil)
c.Set(ctxKeyInboundEndpoint, EndpointChatCompletions)
service.SetActualOpenAIUpstreamEndpoint(c, tt.runtimeEndpoint)
require.Equal(t, tt.want, resolveOpenAIUpstreamEndpoint(c, tt.account, tt.result))
})
}
}
// ──────────────────────────────────────────────────────────
// responsesSubpathSuffix
// ──────────────────────────────────────────────────────────
@@ -5,6 +5,7 @@ import (
"errors"
"net/http"
"strconv"
"strings"
"time"
"github.com/Wei-Shaw/sub2api/internal/pkg/ip"
@@ -298,7 +299,7 @@ func (h *OpenAIGatewayHandler) ChatCompletions(c *gin.Context) {
userAgent := c.GetHeader("User-Agent")
clientIP := ip.GetClientIP(c)
inboundEndpoint := GetInboundEndpoint(c)
upstreamEndpoint := resolveOpenAIUpstreamEndpoint(c, account)
upstreamEndpoint := resolveOpenAIUpstreamEndpoint(c, account, result)
quotaPlatform := service.QuotaPlatform(c.Request.Context(), apiKey)
cyberBlocked := service.GetOpsCyberPolicy(c) != nil
@@ -337,14 +338,22 @@ func (h *OpenAIGatewayHandler) ChatCompletions(c *gin.Context) {
}
// resolveOpenAIUpstreamEndpoint returns the actual upstream endpoint for an
// OpenAI account, used by every OpenAI usage-recording site. APIKey accounts
// whose upstream is forced or probed to not support the Responses API are
// served directly via /v1/chat/completions (the raw chat path) regardless of
// the inbound endpoint; everything else goes through the Responses API.
func resolveOpenAIUpstreamEndpoint(c *gin.Context, account *service.Account) string {
// OpenAI-compatible account. A forwarding result is authoritative because a
// single inbound route may choose raw Chat or a Responses bridge at runtime.
// The account-based derivation remains as a fallback for existing callers and
// forwarding paths that do not report their endpoint yet.
func resolveOpenAIUpstreamEndpoint(c *gin.Context, account *service.Account, result *service.OpenAIForwardResult) string {
if result != nil {
if endpoint := strings.TrimSpace(result.UpstreamEndpoint); endpoint != "" {
return endpoint
}
}
if endpoint := service.GetActualOpenAIUpstreamEndpoint(c); endpoint != "" {
return endpoint
}
if account != nil && account.Type == service.AccountTypeAPIKey &&
!openai_compat.ShouldUseResponsesAPI(account.Extra) {
return "/v1/chat/completions"
return EndpointChatCompletions
}
return GetUpstreamEndpoint(c, account.Platform)
}
@@ -522,7 +522,7 @@ func (h *OpenAIGatewayHandler) Responses(c *gin.Context) {
clientIP := ip.GetClientIP(c)
requestPayloadHash := service.HashUsageRequestPayload(body)
inboundEndpoint := GetInboundEndpoint(c)
upstreamEndpoint := resolveOpenAIUpstreamEndpoint(c, account)
upstreamEndpoint := resolveOpenAIUpstreamEndpoint(c, account, result)
quotaPlatform := service.QuotaPlatform(c.Request.Context(), apiKey)
// 使用量记录通过有界 worker 池提交,避免请求热路径创建无界 goroutine。
@@ -994,7 +994,7 @@ func (h *OpenAIGatewayHandler) Messages(c *gin.Context) {
clientIP := ip.GetClientIP(c)
requestPayloadHash := service.HashUsageRequestPayload(body)
inboundEndpoint := GetInboundEndpoint(c)
upstreamEndpoint := resolveOpenAIUpstreamEndpoint(c, account)
upstreamEndpoint := resolveOpenAIUpstreamEndpoint(c, account, result)
quotaPlatform := service.QuotaPlatform(c.Request.Context(), apiKey)
cyberBlocked := service.GetOpsCyberPolicy(c) != nil
@@ -1601,7 +1601,7 @@ func (h *OpenAIGatewayHandler) ResponsesWebSocket(c *gin.Context) {
}
h.gatewayService.ReportOpenAIAccountScheduleResult(account.ID, true, result.FirstTokenMs)
inboundEndpoint := GetInboundEndpoint(c)
upstreamEndpoint := resolveOpenAIUpstreamEndpoint(c, account)
upstreamEndpoint := resolveOpenAIUpstreamEndpoint(c, account, result)
quotaPlatform := service.QuotaPlatform(c.Request.Context(), apiKey)
cyberBlocked := service.GetOpsCyberPolicy(c) != nil
h.submitOpenAIUsageRecordTask(ctx, result, func(taskCtx context.Context) {
@@ -2441,7 +2441,7 @@ func (h *OpenAIGatewayHandler) recordCyberPolicyIfMarked(c *gin.Context, apiKey
var accountID int64
if account != nil {
accountID = account.ID
upstreamEndpoint = resolveOpenAIUpstreamEndpoint(c, account)
upstreamEndpoint = resolveOpenAIUpstreamEndpoint(c, account, nil)
}
stream := false
if v, ok := c.Get(opsStreamKey); ok {
@@ -72,6 +72,16 @@ func (s *OpenAIGatewayService) ForwardAsChatCompletions(
}
if account.Platform == PlatformGrok {
if account.IsGrokOAuth() {
if eligible, reason := grokChatResponsesBridgeEligibility(body); eligible {
return s.forwardGrokChatCompletionsViaResponses(ctx, c, account, body, promptCacheKey, defaultMappedModel)
} else {
logger.L().Debug("grok chat_completions: using raw fallback",
zap.Int64("account_id", account.ID),
zap.String("reason", reason),
)
}
}
return s.forwardAsRawChatCompletions(ctx, c, account, body, defaultMappedModel)
}
@@ -160,6 +160,7 @@ func (s *OpenAIGatewayService) forwardAsRawChatCompletions(
if err != nil {
return nil, err
}
SetActualOpenAIUpstreamEndpoint(c, grokChatRawEndpoint)
customUA := account.GetOpenAIUserAgent()
if customUA == "" && account.Platform == PlatformGrok {
customUA = "sub2api-grok/1.0"
@@ -213,6 +214,7 @@ func (s *OpenAIGatewayService) forwardAsRawChatCompletions(
}
if result != nil {
addOpenAIUsage(&result.Usage, bridgeUsage)
result.UpstreamEndpoint = grokChatRawEndpoint
}
return result, forwardErr
}
@@ -0,0 +1,320 @@
package service
import (
"context"
"encoding/json"
"errors"
"fmt"
"net/http"
"strings"
"time"
"github.com/Wei-Shaw/sub2api/internal/pkg/apicompat"
"github.com/Wei-Shaw/sub2api/internal/pkg/xai"
"github.com/gin-gonic/gin"
)
const (
grokChatResponsesEndpoint = "/v1/responses"
grokChatRawEndpoint = "/v1/chat/completions"
)
var grokChatResponsesBridgeTopLevelFields = map[string]struct{}{
"model": {},
"messages": {},
"stream": {},
"stream_options": {},
"max_tokens": {},
"max_completion_tokens": {},
"temperature": {},
"top_p": {},
"prompt_cache_key": {},
"tools": {},
"tool_choice": {},
"functions": {},
"function_call": {},
}
// grokChatResponsesBridgeEligibility deliberately accepts only request shapes
// whose Chat Completions semantics are preserved by the Responses bridge.
// Everything else stays on raw Chat Completions rather than being silently
// dropped or rewritten.
func grokChatResponsesBridgeEligibility(body []byte) (bool, string) {
var root map[string]json.RawMessage
if err := json.Unmarshal(body, &root); err != nil || root == nil {
return false, "invalid_json"
}
for _, field := range []string{"stop", "reasoning_effort"} {
if _, exists := root[field]; exists {
return false, "unsupported_" + field
}
}
for _, field := range []string{"tools", "functions"} {
if raw, exists := root[field]; exists && !grokChatNullOrEmptyArray(raw) {
return false, "unsupported_" + field
}
}
if raw, exists := root["tool_choice"]; exists && !grokChatNullOrNone(raw) {
return false, "unsupported_tool_choice"
}
if raw, exists := root["function_call"]; exists && !grokChatNullOrNone(raw) {
return false, "unsupported_function_call"
}
for field := range root {
if _, supported := grokChatResponsesBridgeTopLevelFields[field]; !supported {
return false, "unknown_field_" + field
}
}
var model string
if raw, ok := root["model"]; !ok || json.Unmarshal(raw, &model) != nil || strings.TrimSpace(model) == "" {
return false, "invalid_model"
}
if raw, ok := root["stream"]; ok {
var stream *bool
if json.Unmarshal(raw, &stream) != nil || stream == nil {
return false, "invalid_stream"
}
}
if raw, ok := root["stream_options"]; ok {
var options map[string]json.RawMessage
if json.Unmarshal(raw, &options) != nil || options == nil {
return false, "invalid_stream_options"
}
for field, value := range options {
if field != "include_usage" {
return false, "unknown_stream_option_" + field
}
var includeUsage *bool
if json.Unmarshal(value, &includeUsage) != nil || includeUsage == nil {
return false, "invalid_stream_include_usage"
}
}
}
for _, field := range []string{"max_tokens", "max_completion_tokens"} {
if raw, ok := root[field]; ok {
var value *int
if json.Unmarshal(raw, &value) != nil || value == nil || *value < 128 {
return false, "unsafe_" + field
}
}
}
if _, hasMaxTokens := root["max_tokens"]; hasMaxTokens {
if _, hasMaxCompletionTokens := root["max_completion_tokens"]; hasMaxCompletionTokens {
return false, "conflicting_max_tokens"
}
}
for _, field := range []string{"temperature", "top_p"} {
if raw, ok := root[field]; ok {
var value *float64
if json.Unmarshal(raw, &value) != nil || value == nil {
return false, "invalid_" + field
}
}
}
if raw, ok := root["prompt_cache_key"]; ok {
var key string
if json.Unmarshal(raw, &key) != nil {
return false, "invalid_prompt_cache_key"
}
}
var messages []map[string]json.RawMessage
rawMessages, ok := root["messages"]
if !ok || json.Unmarshal(rawMessages, &messages) != nil || len(messages) == 0 {
return false, "invalid_messages"
}
for _, message := range messages {
for field := range message {
if field != "role" && field != "content" {
return false, "unsafe_message_field_" + field
}
}
var role string
if raw, exists := message["role"]; !exists || json.Unmarshal(raw, &role) != nil {
return false, "invalid_message_role"
}
switch role {
case "system", "user", "assistant":
default:
return false, "unsupported_message_role_" + role
}
var content string
if raw, exists := message["content"]; !exists || json.Unmarshal(raw, &content) != nil {
// Structured content includes image_url and other parts whose exact
// behavior is not guaranteed by this bridge.
return false, "non_text_message_content"
}
if strings.TrimSpace(content) == "" {
return false, "empty_message_content"
}
}
return true, ""
}
func grokChatNullOrEmptyArray(raw json.RawMessage) bool {
if strings.TrimSpace(string(raw)) == "null" {
return true
}
var values []json.RawMessage
return json.Unmarshal(raw, &values) == nil && len(values) == 0
}
func grokChatNullOrNone(raw json.RawMessage) bool {
if strings.TrimSpace(string(raw)) == "null" {
return true
}
var value string
return json.Unmarshal(raw, &value) == nil && strings.EqualFold(strings.TrimSpace(value), "none")
}
func grokChatCacheIntentBody(body []byte) ([]byte, error) {
var root map[string]json.RawMessage
if err := json.Unmarshal(body, &root); err != nil {
return nil, err
}
for _, field := range []string{"tools", "tool_choice", "functions", "function_call"} {
delete(root, field)
}
return json.Marshal(root)
}
func grokChatResponsesRuntimeEligible(upstreamModel, cacheIdentity string) bool {
return strings.TrimSpace(upstreamModel) == "grok-4.5" && strings.TrimSpace(cacheIdentity) != ""
}
// forwardGrokChatCompletionsViaResponses converts a strictly compatible Chat
// request into xAI Responses format and reuses the established Responses-to-
// Chat response translators. It intentionally does not run the Codex OAuth
// transform because Grok CLI is a separate upstream protocol.
func (s *OpenAIGatewayService) forwardGrokChatCompletionsViaResponses(
ctx context.Context,
c *gin.Context,
account *Account,
body []byte,
promptCacheKey string,
defaultMappedModel string,
) (*OpenAIForwardResult, error) {
startTime := time.Now()
var chatReq apicompat.ChatCompletionsRequest
if err := json.Unmarshal(body, &chatReq); err != nil {
return nil, fmt.Errorf("parse grok chat completions request: %w", err)
}
originalModel := chatReq.Model
clientStream := chatReq.Stream
billingModel := resolveOpenAIForwardModel(account, originalModel, defaultMappedModel)
upstreamModel := normalizeOpenAIModelForUpstream(account, billingModel)
cacheIdentity := resolveGrokCacheIdentity(c, body, promptCacheKey, upstreamModel)
if !grokChatResponsesRuntimeEligible(upstreamModel, cacheIdentity) {
return s.forwardAsRawChatCompletions(ctx, c, account, body, defaultMappedModel)
}
responsesReq, err := apicompat.ChatCompletionsToResponses(&chatReq)
if err != nil {
return nil, fmt.Errorf("convert grok chat completions to responses: %w", err)
}
responsesReq.Model = upstreamModel
responsesReq.Stream = true
// These fields are useful to Codex but are not needed by the Grok CLI
// protocol. Keep the bridge request as close as possible to native Grok.
responsesReq.Include = nil
responsesReq.Store = nil
responsesBody, err := json.Marshal(responsesReq)
if err != nil {
return nil, fmt.Errorf("marshal grok responses bridge request: %w", err)
}
responsesBody, err = patchGrokResponsesBody(responsesBody, upstreamModel)
if err != nil {
return nil, fmt.Errorf("patch grok responses bridge request: %w", err)
}
intentBody, err := grokChatCacheIntentBody(body)
if err != nil {
return nil, fmt.Errorf("normalize grok responses bridge tool intent: %w", err)
}
responsesBody, err = applyGrokResponsesCacheIdentity(responsesBody, intentBody, cacheIdentity, true)
if err != nil {
return nil, fmt.Errorf("apply grok responses bridge cache identity: %w", err)
}
updatedBody, policyErr := s.applyOpenAIFastPolicyToBody(ctx, account, upstreamModel, responsesBody)
if policyErr != nil {
var blocked *OpenAIFastBlockedError
if errors.As(policyErr, &blocked) {
MarkOpsClientBusinessLimited(c, OpsClientBusinessLimitedReasonLocalPolicyDenied)
writeChatCompletionsError(c, http.StatusForbidden, "permission_error", blocked.Message)
}
return nil, policyErr
}
responsesBody = updatedBody
token, _, err := s.GetAccessToken(ctx, account)
if err != nil {
return nil, fmt.Errorf("get grok access token: %w", err)
}
upstreamCtx, releaseUpstreamCtx := detachUpstreamContext(ctx)
upstreamReq, err := buildGrokResponsesRequest(upstreamCtx, c, account, responsesBody, token, cacheIdentity)
releaseUpstreamCtx()
if err != nil {
return nil, fmt.Errorf("build grok responses bridge request: %w", err)
}
SetActualOpenAIUpstreamEndpoint(c, grokChatResponsesEndpoint)
proxyURL := ""
if account.ProxyID != nil && account.Proxy != nil {
proxyURL = account.Proxy.URL()
}
resp, err := s.httpUpstream.Do(upstreamReq, proxyURL, account.ID, account.Concurrency)
if err != nil {
return nil, s.handleOpenAIUpstreamTransportError(ctx, c, account, err, false)
}
defer func() { _ = resp.Body.Close() }()
if resp.StatusCode >= http.StatusBadRequest {
respBody, upstreamMsg := s.readOpenAIUpstreamError(resp)
if upstreamMsg == "" {
upstreamMsg = fmt.Sprintf("xAI upstream returned status %d", resp.StatusCode)
}
appendOpsUpstreamError(c, OpsUpstreamErrorEvent{
Platform: account.Platform,
AccountID: account.ID,
AccountName: account.Name,
UpstreamStatusCode: resp.StatusCode,
UpstreamRequestID: firstNonEmpty(resp.Header.Get("x-request-id"), resp.Header.Get("xai-request-id")),
Kind: "failover",
Message: upstreamMsg,
})
s.handleGrokAccountUpstreamError(ctx, account, resp.StatusCode, resp.Header, respBody)
if s.shouldFailoverUpstreamError(resp.StatusCode) {
return nil, &UpstreamFailoverError{
StatusCode: resp.StatusCode,
ResponseBody: respBody,
RetryableOnSameAccount: account.IsPoolMode() && account.IsPoolModeRetryableStatus(resp.StatusCode),
}
}
return s.handleChatCompletionsErrorResponse(resp, c, account, billingModel)
}
s.updateGrokUsageSnapshot(ctx, account, xai.ParseQuotaHeaders(resp.Header, resp.StatusCode))
var result *OpenAIForwardResult
if clientStream {
result, err = s.handleChatStreamingResponse(resp, c, account, originalModel, billingModel, upstreamModel, startTime, len(body))
} else {
result, err = s.handleChatBufferedStreamingResponse(resp, c, account, originalModel, billingModel, upstreamModel, startTime)
}
if result != nil {
result.UpstreamEndpoint = grokChatResponsesEndpoint
result.ResponseHeaders = resp.Header.Clone()
if result.RequestID == "" {
result.RequestID = firstNonEmpty(resp.Header.Get("x-request-id"), resp.Header.Get("xai-request-id"))
}
result.ReasoningEffort = extractOpenAIReasoningEffortFromBody(body, upstreamModel, billingModel, originalModel)
}
return result, err
}
@@ -0,0 +1,358 @@
//go:build unit
package service
import (
"bytes"
"context"
"errors"
"io"
"net/http"
"net/http/httptest"
"strconv"
"strings"
"testing"
"time"
"github.com/Wei-Shaw/sub2api/internal/pkg/xai"
"github.com/gin-gonic/gin"
"github.com/stretchr/testify/require"
"github.com/tidwall/gjson"
)
func TestGrokChatResponsesBridgeEligibility(t *testing.T) {
t.Parallel()
tests := []struct {
name string
body string
want bool
reason string
}{
{
name: "plain text chat",
body: `{"model":"grok","messages":[{"role":"system","content":"concise"},{"role":"user","content":"hi"}],"stream":false}`,
want: true,
},
{
name: "safe generation options",
body: `{"model":"grok","messages":[{"role":"user","content":"hi"}],"stream":true,"stream_options":{"include_usage":true},"max_completion_tokens":256,"temperature":0.2,"top_p":0.9,"prompt_cache_key":"session","tools":[],"functions":null,"tool_choice":"none"}`,
want: true,
},
{
name: "stop falls back",
body: `{"model":"grok","messages":[{"role":"user","content":"hi"}],"stop":"done"}`,
reason: "unsupported_stop",
},
{
name: "developer role falls back",
body: `{"model":"grok","messages":[{"role":"developer","content":"rules"},{"role":"user","content":"hi"}]}`,
reason: "unsupported_message_role_developer",
},
{
name: "image content falls back",
body: `{"model":"grok","messages":[{"role":"user","content":[{"type":"image_url","image_url":{"url":"data:image/png;base64,QQ=="}}]}]}`,
reason: "non_text_message_content",
},
{
name: "function tools fall back",
body: `{"model":"grok","messages":[{"role":"user","content":"hi"}],"tools":[{"type":"function","function":{"name":"lookup"}}]}`,
reason: "unsupported_tools",
},
{
name: "automatic tool choice falls back",
body: `{"model":"grok","messages":[{"role":"user","content":"hi"}],"tools":[],"tool_choice":"auto"}`,
reason: "unsupported_tool_choice",
},
{
name: "reasoning effort falls back because conversion adds summary",
body: `{"model":"grok","messages":[{"role":"user","content":"hi"}],"reasoning_effort":"high"}`,
reason: "unsupported_reasoning_effort",
},
{
name: "both token limits fall back",
body: `{"model":"grok","messages":[{"role":"user","content":"hi"}],"max_tokens":256,"max_completion_tokens":256}`,
reason: "conflicting_max_tokens",
},
{
name: "empty message falls back",
body: `{"model":"grok","messages":[{"role":"assistant","content":""},{"role":"user","content":"hi"}]}`,
reason: "empty_message_content",
},
{
name: "tool history falls back",
body: `{"model":"grok","messages":[{"role":"assistant","content":"","tool_calls":[]}]}`,
reason: "unsafe_message_field_tool_calls",
},
{
name: "unknown field falls back",
body: `{"model":"grok","messages":[{"role":"user","content":"hi"}],"seed":7}`,
reason: "unknown_field_seed",
},
{
name: "small max tokens falls back because conversion clamps it",
body: `{"model":"grok","messages":[{"role":"user","content":"hi"}],"max_tokens":32}`,
reason: "unsafe_max_tokens",
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
got, reason := grokChatResponsesBridgeEligibility([]byte(tt.body))
require.Equal(t, tt.want, got)
require.Equal(t, tt.reason, reason)
})
}
}
func TestGrokChatResponsesRuntimeEligibility(t *testing.T) {
t.Parallel()
require.True(t, grokChatResponsesRuntimeEligible("grok-4.5", "isolated-id"))
require.False(t, grokChatResponsesRuntimeEligible("grok-4.3", "isolated-id"))
require.False(t, grokChatResponsesRuntimeEligible("grok-4.5-build-free", "isolated-id"))
require.False(t, grokChatResponsesRuntimeEligible("grok-4.5", ""))
}
func TestForwardGrokChatViaResponsesNonStreamingCachesAndReturnsChat(t *testing.T) {
gin.SetMode(gin.TestMode)
body := []byte(`{"model":"grok","messages":[{"role":"system","content":"be concise"},{"role":"user","content":"hi"}],"stream":false,"prompt_cache_key":"stable-session","tools":[],"functions":null,"tool_choice":"none"}`)
recorder := httptest.NewRecorder()
c, _ := gin.CreateTestContext(recorder)
c.Request = httptest.NewRequest(http.MethodPost, grokChatRawEndpoint, bytes.NewReader(body))
c.Set("api_key", &APIKey{ID: 7101})
account := grokChatBridgeTestAccount(71)
repo := &grokQuotaAccountRepo{mockAccountRepoForPlatform: &mockAccountRepoForPlatform{
accountsByID: map[int64]*Account{account.ID: account},
}}
upstream := &httpUpstreamRecorder{resp: grokChatBridgeCompletedResponse("resp_grok_chat_cache", 9856)}
svc := &OpenAIGatewayService{
httpUpstream: upstream,
grokTokenProvider: NewGrokTokenProvider(repo, nil),
accountRepo: repo,
}
result, err := svc.ForwardAsChatCompletions(context.Background(), c, account, body, "", "")
require.NoError(t, err)
require.NotNil(t, result)
require.Equal(t, xai.DefaultCLIBaseURL+"/responses", upstream.lastReq.URL.String())
require.Equal(t, grokChatResponsesEndpoint, result.UpstreamEndpoint)
require.Equal(t, "grok-4.5", result.UpstreamModel)
require.Equal(t, 9908, result.Usage.InputTokens)
require.Equal(t, 12, result.Usage.OutputTokens)
require.Equal(t, 9856, result.Usage.CacheReadInputTokens)
identity := gjson.GetBytes(upstream.lastBody, "prompt_cache_key").String()
require.NotEmpty(t, identity)
require.NotEqual(t, "stable-session", identity)
require.Equal(t, identity, upstream.lastReq.Header.Get(grokConversationIDHeader))
require.Equal(t, "web_search", gjson.GetBytes(upstream.lastBody, "tools.0.type").String())
require.Equal(t, "x_search", gjson.GetBytes(upstream.lastBody, "tools.1.type").String())
require.Equal(t, grokFreeCacheDisabledToolChoice, gjson.GetBytes(upstream.lastBody, "tool_choice").String())
require.True(t, gjson.GetBytes(upstream.lastBody, "stream").Bool())
require.Equal(t, "system", gjson.GetBytes(upstream.lastBody, "input.0.role").String())
require.Equal(t, "user", gjson.GetBytes(upstream.lastBody, "input.1.role").String())
require.False(t, gjson.GetBytes(upstream.lastBody, "instructions").Exists())
require.False(t, gjson.GetBytes(upstream.lastBody, "include").Exists())
require.False(t, gjson.GetBytes(upstream.lastBody, "store").Exists())
require.Equal(t, http.StatusOK, recorder.Code)
require.Equal(t, "cached ok", gjson.Get(recorder.Body.String(), "choices.0.message.content").String())
require.Equal(t, int64(9856), gjson.Get(recorder.Body.String(), "usage.prompt_tokens_details.cached_tokens").Int())
require.NotNil(t, repo.updates[account.ID][grokQuotaSnapshotExtraKey])
}
func TestForwardGrokChatViaResponsesStreamingPropagatesCachedUsage(t *testing.T) {
gin.SetMode(gin.TestMode)
body := []byte(`{"model":"grok","messages":[{"role":"user","content":"hi"}],"stream":true}`)
recorder := httptest.NewRecorder()
c, _ := gin.CreateTestContext(recorder)
c.Request = httptest.NewRequest(http.MethodPost, grokChatRawEndpoint, bytes.NewReader(body))
c.Set("api_key", &APIKey{ID: 7201})
account := grokChatBridgeTestAccount(72)
repo := &grokQuotaAccountRepo{mockAccountRepoForPlatform: &mockAccountRepoForPlatform{
accountsByID: map[int64]*Account{account.ID: account},
}}
upstream := &httpUpstreamRecorder{resp: grokChatBridgeCompletedResponse("resp_grok_chat_stream", 4096)}
svc := &OpenAIGatewayService{
httpUpstream: upstream,
grokTokenProvider: NewGrokTokenProvider(repo, nil),
accountRepo: repo,
}
result, err := svc.ForwardAsChatCompletions(context.Background(), c, account, body, "", "")
require.NoError(t, err)
require.NotNil(t, result)
require.True(t, result.Stream)
require.Equal(t, grokChatResponsesEndpoint, result.UpstreamEndpoint)
require.Equal(t, 4096, result.Usage.CacheReadInputTokens)
require.Contains(t, recorder.Header().Get("Content-Type"), "text/event-stream")
require.Contains(t, recorder.Body.String(), `"content":"cached ok"`)
require.Contains(t, recorder.Body.String(), `"cached_tokens":4096`)
require.Contains(t, recorder.Body.String(), "data: [DONE]")
}
func TestForwardGrokChatRuntimeGateFallsBackToRaw(t *testing.T) {
gin.SetMode(gin.TestMode)
tests := []struct {
name string
setAPIKey bool
mappedModel string
wantUpstream string
}{
{name: "missing cache identity", wantUpstream: "grok-4.5"},
{name: "non cache capable mapped model", setAPIKey: true, mappedModel: "grok-4.3", wantUpstream: "grok-4.3"},
}
for index, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
body := []byte(`{"model":"grok","messages":[{"role":"user","content":"hi"}],"stream":false}`)
recorder := httptest.NewRecorder()
c, _ := gin.CreateTestContext(recorder)
c.Request = httptest.NewRequest(http.MethodPost, grokChatRawEndpoint, bytes.NewReader(body))
if tt.setAPIKey {
c.Set("api_key", &APIKey{ID: int64(7301 + index)})
}
account := grokChatBridgeTestAccount(int64(73 + index))
if tt.mappedModel != "" {
account.Credentials["model_mapping"] = map[string]any{"grok": tt.mappedModel}
}
repo := &grokQuotaAccountRepo{mockAccountRepoForPlatform: &mockAccountRepoForPlatform{
accountsByID: map[int64]*Account{account.ID: account},
}}
upstream := &httpUpstreamRecorder{resp: &http.Response{
StatusCode: http.StatusOK,
Header: http.Header{"Content-Type": []string{"application/json"}},
Body: io.NopCloser(strings.NewReader(
`{"id":"chat_raw","object":"chat.completion","model":"` + tt.wantUpstream + `","choices":[{"index":0,"message":{"role":"assistant","content":"raw ok"},"finish_reason":"stop"}],"usage":{"prompt_tokens":2,"completion_tokens":1,"total_tokens":3}}`,
)),
}}
svc := &OpenAIGatewayService{
httpUpstream: upstream,
grokTokenProvider: NewGrokTokenProvider(repo, nil),
accountRepo: repo,
}
result, err := svc.ForwardAsChatCompletions(context.Background(), c, account, body, "", "")
require.NoError(t, err)
require.NotNil(t, result)
require.Equal(t, xai.DefaultCLIBaseURL+"/chat/completions", upstream.lastReq.URL.String())
require.Equal(t, grokChatRawEndpoint, result.UpstreamEndpoint)
require.Equal(t, tt.wantUpstream, result.UpstreamModel)
require.False(t, gjson.GetBytes(upstream.lastBody, "tools").Exists())
require.Equal(t, "raw ok", gjson.Get(recorder.Body.String(), "choices.0.message.content").String())
})
}
}
func TestForwardGrokChatViaResponses429UsesGrokRateLimitPolicy(t *testing.T) {
gin.SetMode(gin.TestMode)
body := []byte(`{"model":"grok","messages":[{"role":"user","content":"hi"}],"stream":false}`)
recorder := httptest.NewRecorder()
c, _ := gin.CreateTestContext(recorder)
c.Request = httptest.NewRequest(http.MethodPost, grokChatRawEndpoint, bytes.NewReader(body))
c.Set("api_key", &APIKey{ID: 7501})
account := grokChatBridgeTestAccount(75)
repo := &grokQuotaAccountRepo{mockAccountRepoForPlatform: &mockAccountRepoForPlatform{
accountsByID: map[int64]*Account{account.ID: account},
}}
upstream := &httpUpstreamRecorder{resp: &http.Response{
StatusCode: http.StatusTooManyRequests,
Header: http.Header{
"Content-Type": []string{"application/json"},
"Retry-After": []string{"45"},
},
Body: io.NopCloser(strings.NewReader(`{"error":{"message":"rate limited"}}`)),
}}
svc := &OpenAIGatewayService{
httpUpstream: upstream,
grokTokenProvider: NewGrokTokenProvider(repo, nil),
accountRepo: repo,
}
before := time.Now()
result, err := svc.ForwardAsChatCompletions(context.Background(), c, account, body, "", "")
require.Error(t, err)
require.Nil(t, result)
var failoverErr *UpstreamFailoverError
require.True(t, errors.As(err, &failoverErr))
require.Equal(t, http.StatusTooManyRequests, failoverErr.StatusCode)
require.Equal(t, xai.DefaultCLIBaseURL+"/responses", upstream.lastReq.URL.String())
require.Equal(t, grokChatResponsesEndpoint, GetActualOpenAIUpstreamEndpoint(c))
require.Equal(t, 1, repo.rateLimitedCalls)
require.Zero(t, repo.tempUnschedCalls)
require.WithinDuration(t, before.Add(45*time.Second), repo.lastRateLimitResetAt, time.Second)
require.True(t, svc.isOpenAIAccountRuntimeBlocked(account))
}
func TestForwardGrokRawChatErrorRecordsActualEndpoint(t *testing.T) {
gin.SetMode(gin.TestMode)
body := []byte(`{"model":"grok","messages":[{"role":"user","content":"hi"}],"stream":false,"stop":"done"}`)
recorder := httptest.NewRecorder()
c, _ := gin.CreateTestContext(recorder)
c.Request = httptest.NewRequest(http.MethodPost, grokChatRawEndpoint, bytes.NewReader(body))
c.Set("api_key", &APIKey{ID: 7601})
account := grokChatBridgeTestAccount(76)
repo := &grokQuotaAccountRepo{mockAccountRepoForPlatform: &mockAccountRepoForPlatform{
accountsByID: map[int64]*Account{account.ID: account},
}}
upstream := &httpUpstreamRecorder{resp: &http.Response{
StatusCode: http.StatusBadRequest,
Header: http.Header{"Content-Type": []string{"application/json"}},
Body: io.NopCloser(strings.NewReader(`{"error":{"message":"bad request"}}`)),
}}
svc := &OpenAIGatewayService{
httpUpstream: upstream,
grokTokenProvider: NewGrokTokenProvider(repo, nil),
accountRepo: repo,
}
result, err := svc.ForwardAsChatCompletions(context.Background(), c, account, body, "", "")
require.Error(t, err)
require.Nil(t, result)
require.Equal(t, xai.DefaultCLIBaseURL+"/chat/completions", upstream.lastReq.URL.String())
require.Equal(t, grokChatRawEndpoint, GetActualOpenAIUpstreamEndpoint(c))
}
func grokChatBridgeTestAccount(id int64) *Account {
return &Account{
ID: id,
Name: "grok-cache-bridge",
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),
"base_url": xai.DefaultCLIBaseURL,
},
}
}
func grokChatBridgeCompletedResponse(responseID string, cachedTokens int) *http.Response {
body := strings.Join([]string{
`data: {"type":"response.output_text.delta","sequence_number":0,"delta":"cached ok"}`,
"",
`data: {"type":"response.completed","sequence_number":1,"response":{"id":"` + responseID + `","object":"response","model":"grok-4.5","status":"completed","output":[{"type":"message","id":"msg_1","role":"assistant","status":"completed","content":[{"type":"output_text","text":"cached ok"}]}],"usage":{"input_tokens":9908,"output_tokens":12,"total_tokens":9920,"input_tokens_details":{"cached_tokens":` + strconv.Itoa(cachedTokens) + `}}}}`,
"",
}, "\n")
return &http.Response{
StatusCode: http.StatusOK,
Header: http.Header{
"Content-Type": []string{"text/event-stream"},
"Xai-Request-Id": []string{responseID + "-request"},
"X-Ratelimit-Limit-Requests": []string{"10"},
"X-Ratelimit-Remaining-Requests": []string{"9"},
},
Body: io.NopCloser(strings.NewReader(body)),
}
}
@@ -615,12 +615,12 @@ func TestForwardGrokMedia429ReconcilesRateLimitBeforeCustomErrorBypass(t *testin
require.True(t, svc.isOpenAIAccountRuntimeBlocked(account))
}
func TestForwardAsChatCompletionsForGrokUsesXAIChatCompletionsAndSnapshots(t *testing.T) {
func TestForwardAsChatCompletionsForGrokStopFallsBackToXAIChatCompletions(t *testing.T) {
gin.SetMode(gin.TestMode)
recorder := httptest.NewRecorder()
c, _ := gin.CreateTestContext(recorder)
body := []byte(`{"model":"grok","messages":[{"role":"user","content":"hi"}],"stream":false,"prompt_cache_key":"raw-client-cache-key"}`)
body := []byte(`{"model":"grok","messages":[{"role":"user","content":"hi"}],"stream":false,"stop":"done","prompt_cache_key":"raw-client-cache-key"}`)
c.Request = httptest.NewRequest(http.MethodPost, "/v1/chat/completions", bytes.NewReader(body))
c.Set("api_key", &APIKey{ID: 5101})
@@ -881,12 +881,12 @@ func TestForwardGrokResponsesFailoverKeepsCacheIdentityAcrossAccounts(t *testing
require.Equal(t, "Bearer access-token-b", upstream.requests[1].Header.Get("Authorization"))
}
func TestForwardAsChatCompletionsForGrokStreamingUsesRawXAIChatCompletions(t *testing.T) {
func TestForwardAsChatCompletionsForGrokStreamingStopFallsBackToRawXAIChatCompletions(t *testing.T) {
gin.SetMode(gin.TestMode)
recorder := httptest.NewRecorder()
c, _ := gin.CreateTestContext(recorder)
body := []byte(`{"model":"grok","messages":[{"role":"user","content":"hi"}],"stream":true}`)
body := []byte(`{"model":"grok","messages":[{"role":"user","content":"hi"}],"stream":true,"stop":"done"}`)
c.Request = httptest.NewRequest(http.MethodPost, "/v1/chat/completions", bytes.NewReader(body))
c.Request.Header.Set("Content-Type", "application/json")
c.Request.Header.Set(grokConversationIDHeader, "native-client-conversation")
@@ -49,6 +49,7 @@ const (
openAIWSRetryBackoffMaxDefault = 2 * time.Second
openAIWSRetryJitterRatioDefault = 0.2
openAICompactSessionSeedKey = "openai_compact_session_seed"
openAIUpstreamEndpointContextKey = "openai_actual_upstream_endpoint"
codexCLIVersion = "0.144.1"
// Codex 限额快照仅用于后台展示/诊断,不需要每个成功请求都立即落库。
openAICodexSnapshotPersistMinInterval = 30 * time.Second
@@ -223,6 +224,9 @@ type OpenAIForwardResult struct {
// UpstreamModel is the actual model sent to the upstream provider after mapping.
// Empty when no mapping was applied (requested model was used as-is).
UpstreamModel string
// UpstreamEndpoint is the actual upstream API path used for this request.
// It avoids guessing when one downstream protocol can use multiple upstream endpoints.
UpstreamEndpoint string
// ServiceTier records the OpenAI Responses API service tier, e.g. "priority" / "flex".
// Nil means the request did not specify a recognized tier.
ServiceTier *string
@@ -251,6 +255,32 @@ type OpenAIForwardResult struct {
wsReplayInputExists bool
}
// SetActualOpenAIUpstreamEndpoint records the endpoint selected by the current
// forwarding attempt. It covers error paths where no OpenAIForwardResult is
// available for usage and operations logging.
func SetActualOpenAIUpstreamEndpoint(c *gin.Context, endpoint string) {
if c == nil {
return
}
if endpoint = strings.TrimSpace(endpoint); endpoint != "" {
c.Set(openAIUpstreamEndpointContextKey, endpoint)
}
}
// GetActualOpenAIUpstreamEndpoint returns the endpoint recorded by the latest
// forwarding attempt in this request.
func GetActualOpenAIUpstreamEndpoint(c *gin.Context) string {
if c == nil {
return ""
}
value, exists := c.Get(openAIUpstreamEndpointContextKey)
if !exists {
return ""
}
endpoint, _ := value.(string)
return strings.TrimSpace(endpoint)
}
type OpenAIWSRetryMetricsSnapshot struct {
RetryAttemptsTotal int64 `json:"retry_attempts_total"`
RetryBackoffMsTotal int64 `json:"retry_backoff_ms_total"`