mirror of
https://github.com/Wei-Shaw/sub2api.git
synced 2026-09-24 16:05:44 +08:00
fix(grok): route cacheable chat requests via responses
This commit is contained in:
@@ -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)
|
||||
|
||||
@@ -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"`
|
||||
|
||||
Reference in New Issue
Block a user