Merge pull request #2925 from is7Qin/feat/openai-oom

refactor(gateway): 降低大请求体内存保留
This commit is contained in:
Wesley Liddick 2026-06-01 09:40:31 +08:00 committed by GitHub
commit 0f70f84182
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
43 changed files with 2734 additions and 1900 deletions

View File

@ -154,7 +154,8 @@ func (h *GatewayHandler) Messages(c *gin.Context) {
setOpsRequestContext(c, "", false)
parsedReq, err := service.ParseGatewayRequest(body, domain.PlatformAnthropic)
bodyRef := service.NewRequestBodyRef(body)
parsedReq, err := service.ParseGatewayRequest(bodyRef, domain.PlatformAnthropic)
if err != nil {
h.errorResponse(c, http.StatusBadRequest, "invalid_request_error", "Failed to parse request body")
return
@ -509,11 +510,12 @@ func (h *GatewayHandler) Messages(c *gin.Context) {
}
// 使用量记录通过有界 worker 池提交,避免请求热路径创建无界 goroutine。
// ForceCacheBilling 提前拍成标量,避免 worker 闭包保活 failover 状态里的响应体。
forceCacheBilling := fs.ForceCacheBilling
quotaPlatform := service.QuotaPlatform(c.Request.Context(), apiKey)
h.submitUsageRecordTask(c.Request.Context(), func(ctx context.Context) {
if err := h.gatewayService.RecordUsage(ctx, &service.RecordUsageInput{
Result: result,
ParsedRequest: parsedReq,
QuotaPlatform: quotaPlatform,
APIKey: apiKey,
User: apiKey.User,
@ -524,7 +526,7 @@ func (h *GatewayHandler) Messages(c *gin.Context) {
UserAgent: userAgent,
IPAddress: clientIP,
RequestPayloadHash: requestPayloadHash,
ForceCacheBilling: fs.ForceCacheBilling,
ForceCacheBilling: forceCacheBilling,
APIKeyService: h.apiKeyService,
ChannelUsageFields: channelMapping.ToUsageFields(reqModel, result.UpstreamModel),
}); err != nil {
@ -562,6 +564,12 @@ func (h *GatewayHandler) Messages(c *gin.Context) {
retryWithFallback := false
for {
attemptParsedReq, err := parsedReq.CloneForBody(body)
if err != nil {
h.errorResponse(c, http.StatusBadRequest, "invalid_request_error", "Failed to parse request body")
return
}
// 选择支持该模型的账号
reqLog.Info("sticky.selecting_account",
zap.String("session_key", sessionKey),
@ -693,7 +701,7 @@ func (h *GatewayHandler) Messages(c *gin.Context) {
// ===== 用户消息串行队列 START =====
var queueRelease func()
umqMode := h.getUserMsgQueueMode(account, parsedReq)
umqMode := h.getUserMsgQueueMode(account, attemptParsedReq)
switch umqMode {
case config.UMQModeSerialize:
@ -740,20 +748,26 @@ func (h *GatewayHandler) Messages(c *gin.Context) {
// 用 wrapReleaseOnDone 确保 context 取消时自动释放(仅 serialize 模式有 queueRelease
queueRelease = wrapReleaseOnDone(c.Request.Context(), queueRelease)
// 注入回调到 ParsedRequest使用外层 wrapper 以便提前清理 AfterFunc
parsedReq.OnUpstreamAccepted = queueRelease
attemptParsedReq.OnUpstreamAccepted = queueRelease
// ===== 用户消息串行队列 END =====
// 应用渠道模型映射到请求
// 渠道模型映射只作用于本次账号尝试,避免 failover 后污染原始 ParsedRequest。
if channelMapping.Mapped {
parsedReq.Model = channelMapping.MappedModel
parsedReq.Body = h.gatewayService.ReplaceModelInBody(parsedReq.Body, channelMapping.MappedModel)
attemptParsedReq.Model = channelMapping.MappedModel
if err := attemptParsedReq.ReplaceBody(h.gatewayService.ReplaceModelInBody(attemptParsedReq.Body.Bytes(), channelMapping.MappedModel)); err != nil {
h.errorResponse(c, http.StatusBadRequest, "invalid_request_error", "Failed to parse request body")
return
}
}
// Bedrock CC 兼容:渠道模型映射后,清理 Anthropic API 专有字段、注入 Bedrock 必需字段
parsedReq.Body = h.gatewayService.ApplyBedrockCCCompat(c.Request.Context(), parsedReq.Body, parsedReq.Model, account, apiKey.GroupID)
body = parsedReq.Body
if err := attemptParsedReq.ReplaceBody(h.gatewayService.ApplyBedrockCCCompat(c.Request.Context(), attemptParsedReq.Body.Bytes(), attemptParsedReq.Model, account, apiKey.GroupID)); err != nil {
h.errorResponse(c, http.StatusBadRequest, "invalid_request_error", "Failed to parse request body")
return
}
attemptBody := attemptParsedReq.Body.Bytes()
// 转发请求 - 根据账号平台分流
c.Set("parsed_request", parsedReq)
c.Set("parsed_request", attemptParsedReq)
var result *service.ForwardResult
requestCtx := c.Request.Context()
if fs.SwitchCount > 0 {
@ -762,9 +776,9 @@ func (h *GatewayHandler) Messages(c *gin.Context) {
// 记录 Forward 前已写入字节数Forward 后若增加则说明 SSE 内容已发,禁止 failover
writerSizeBeforeForward := c.Writer.Size()
if account.Platform == service.PlatformAntigravity && account.Type != service.AccountTypeAPIKey {
result, err = h.antigravityGatewayService.Forward(requestCtx, c, account, body, hasBoundSession)
result, err = h.antigravityGatewayService.Forward(requestCtx, c, account, attemptBody, hasBoundSession)
} else {
result, err = h.gatewayService.Forward(requestCtx, c, account, parsedReq)
result, err = h.gatewayService.Forward(requestCtx, c, account, attemptParsedReq)
}
// 兜底释放串行锁(正常情况已通过回调提前释放)
@ -772,7 +786,7 @@ func (h *GatewayHandler) Messages(c *gin.Context) {
queueRelease()
}
// 清理回调引用,防止 failover 重试时旧回调被错误调用
parsedReq.OnUpstreamAccepted = nil
attemptParsedReq.OnUpstreamAccepted = nil
if accountReleaseFunc != nil {
accountReleaseFunc()
@ -895,20 +909,22 @@ func (h *GatewayHandler) Messages(c *gin.Context) {
// 捕获请求信息(用于异步记录,避免在 goroutine 中访问 gin.Context
userAgent := c.GetHeader("User-Agent")
clientIP := ip.GetClientIP(c)
requestPayloadHash := service.HashUsageRequestPayload(body)
// Forward 内部可能继续改写 bodyusage 去重指纹必须使用最终上游接受的当前 body。
requestPayloadHash := service.HashUsageRequestPayload(attemptParsedReq.Body.Bytes())
inboundEndpoint := GetInboundEndpoint(c)
upstreamEndpoint := GetUpstreamEndpoint(c, account.Platform)
if result.ReasoningEffort == nil {
result.ReasoningEffort = service.NormalizeClaudeOutputEffort(parsedReq.OutputEffort)
result.ReasoningEffort = service.NormalizeClaudeOutputEffort(attemptParsedReq.OutputEffort)
}
// 使用量记录通过有界 worker 池提交,避免请求热路径创建无界 goroutine。
// ForceCacheBilling 提前拍成标量,避免 worker 闭包保活 failover 状态里的响应体。
forceCacheBilling := fs.ForceCacheBilling
quotaPlatform := service.QuotaPlatform(c.Request.Context(), currentAPIKey)
h.submitUsageRecordTask(c.Request.Context(), func(ctx context.Context) {
if err := h.gatewayService.RecordUsage(ctx, &service.RecordUsageInput{
Result: result,
ParsedRequest: parsedReq,
QuotaPlatform: quotaPlatform,
APIKey: currentAPIKey,
User: currentAPIKey.User,
@ -919,7 +935,7 @@ func (h *GatewayHandler) Messages(c *gin.Context) {
UserAgent: userAgent,
IPAddress: clientIP,
RequestPayloadHash: requestPayloadHash,
ForceCacheBilling: fs.ForceCacheBilling,
ForceCacheBilling: forceCacheBilling,
APIKeyService: h.apiKeyService,
ChannelUsageFields: channelMapping.ToUsageFields(reqModel, result.UpstreamModel),
}); err != nil {
@ -1683,7 +1699,8 @@ func (h *GatewayHandler) CountTokens(c *gin.Context) {
setOpsRequestContext(c, "", false)
parsedReq, err := service.ParseGatewayRequest(body, domain.PlatformAnthropic)
bodyRef := service.NewRequestBodyRef(body)
parsedReq, err := service.ParseGatewayRequest(bodyRef, domain.PlatformAnthropic)
if err != nil {
h.errorResponse(c, http.StatusBadRequest, "invalid_request_error", "Failed to parse request body")
return

View File

@ -151,9 +151,10 @@ func (h *GatewayHandler) ChatCompletions(c *gin.Context) {
}
// Parse request for session hash
parsedReq, _ := service.ParseGatewayRequest(body, "chat_completions")
bodyRef := service.NewRequestBodyRef(body)
parsedReq, _ := service.ParseGatewayRequest(bodyRef, "chat_completions")
if parsedReq == nil {
parsedReq = &service.ParsedRequest{Model: reqModel, Stream: reqStream, Body: body}
parsedReq = &service.ParsedRequest{Model: reqModel, Stream: reqStream, Body: bodyRef}
}
parsedReq.SessionContext = &service.SessionContext{
ClientIP: ip.GetClientIP(c),

View File

@ -156,9 +156,10 @@ func (h *GatewayHandler) Responses(c *gin.Context) {
}
// Parse request for session hash
parsedReq, _ := service.ParseGatewayRequest(body, "responses")
bodyRef := service.NewRequestBodyRef(body)
parsedReq, _ := service.ParseGatewayRequest(bodyRef, "responses")
if parsedReq == nil {
parsedReq = &service.ParsedRequest{Model: reqModel, Stream: reqStream, Body: body}
parsedReq = &service.ParsedRequest{Model: reqModel, Stream: reqStream, Body: bodyRef}
}
parsedReq.SessionContext = &service.SessionContext{
ClientIP: ip.GetClientIP(c),

View File

@ -18,18 +18,12 @@ import (
// claudeCodeValidator is a singleton validator for Claude Code client detection
var claudeCodeValidator = service.NewClaudeCodeValidator()
const claudeCodeParsedRequestContextKey = "claude_code_parsed_request"
// SetClaudeCodeClientContext 检查请求是否来自 Claude Code 客户端,并设置到 context 中
// 返回更新后的 context
func SetClaudeCodeClientContext(c *gin.Context, body []byte, parsedReq *service.ParsedRequest) {
if c == nil || c.Request == nil {
return
}
if parsedReq != nil {
c.Set(claudeCodeParsedRequestContextKey, parsedReq)
}
ua := c.GetHeader("User-Agent")
// Fast path非 Claude CLI UA 直接判定 false避免热路径二次 JSON 反序列化。
if !claudeCodeValidator.ValidateUserAgent(ua) {
@ -45,9 +39,6 @@ func SetClaudeCodeClientContext(c *gin.Context, body []byte, parsedReq *service.
} else {
// 仅在确认为 Claude CLI 且 messages 路径时再做 body 解析。
bodyMap := claudeCodeBodyMapFromParsedRequest(parsedReq)
if bodyMap == nil {
bodyMap = claudeCodeBodyMapFromContextCache(c)
}
if bodyMap == nil && len(body) > 0 {
_ = json.Unmarshal(body, &bodyMap)
}
@ -74,8 +65,12 @@ func claudeCodeBodyMapFromParsedRequest(parsedReq *service.ParsedRequest) map[st
bodyMap := map[string]any{
"model": parsedReq.Model,
}
if parsedReq.System != nil || parsedReq.HasSystem {
bodyMap["system"] = parsedReq.System
if parsedReq.HasSystem {
if system, ok := parsedReq.SystemValue(); ok {
bodyMap["system"] = system
} else {
bodyMap["system"] = nil
}
}
if parsedReq.MetadataUserID != "" {
bodyMap["metadata"] = map[string]any{"user_id": parsedReq.MetadataUserID}
@ -83,26 +78,6 @@ func claudeCodeBodyMapFromParsedRequest(parsedReq *service.ParsedRequest) map[st
return bodyMap
}
func claudeCodeBodyMapFromContextCache(c *gin.Context) map[string]any {
if c == nil {
return nil
}
if cached, ok := c.Get(service.OpenAIParsedRequestBodyKey); ok {
if bodyMap, ok := cached.(map[string]any); ok {
return bodyMap
}
}
if cached, ok := c.Get(claudeCodeParsedRequestContextKey); ok {
switch v := cached.(type) {
case *service.ParsedRequest:
return claudeCodeBodyMapFromParsedRequest(v)
case service.ParsedRequest:
return claudeCodeBodyMapFromParsedRequest(&v)
}
}
return nil
}
// 并发槽位等待相关常量
//
// 性能优化说明:

View File

@ -177,7 +177,7 @@ func TestSetClaudeCodeClientContext_FastPathAndStrictPath(t *testing.T) {
})
}
func TestSetClaudeCodeClientContext_ReuseParsedRequestAndContextCache(t *testing.T) {
func TestSetClaudeCodeClientContext_ReuseParsedRequest(t *testing.T) {
t.Run("reuse parsed request without body unmarshal", func(t *testing.T) {
c, _ := newHelperTestContext(http.MethodPost, "/v1/messages")
c.Request.Header.Set("User-Agent", "claude-cli/1.0.1")
@ -185,36 +185,13 @@ func TestSetClaudeCodeClientContext_ReuseParsedRequestAndContextCache(t *testing
c.Request.Header.Set("anthropic-beta", "message-batches-2024-09-24")
c.Request.Header.Set("anthropic-version", "2023-06-01")
parsedReq := &service.ParsedRequest{
Model: "claude-3-5-sonnet-20241022",
System: []any{
map[string]any{"text": "You are Claude Code, Anthropic's official CLI for Claude."},
},
MetadataUserID: "user_aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa_account__session_aaaaaaaa-aaaa-aaaa-aaaa-aaaaaaaaaaaa",
}
parsedReq, err := service.ParseGatewayRequest(service.NewRequestBodyRef(validClaudeCodeBodyJSON()), "")
require.NoError(t, err)
// body 非法 JSON如果函数复用 parsedReq 成功则仍应判定为 Claude Code。
SetClaudeCodeClientContext(c, []byte(`{invalid`), parsedReq)
require.True(t, service.IsClaudeCodeClient(c.Request.Context()))
})
t.Run("reuse context cache without body unmarshal", func(t *testing.T) {
c, _ := newHelperTestContext(http.MethodPost, "/v1/messages")
c.Request.Header.Set("User-Agent", "claude-cli/1.0.1")
c.Request.Header.Set("X-App", "claude-code")
c.Request.Header.Set("anthropic-beta", "message-batches-2024-09-24")
c.Request.Header.Set("anthropic-version", "2023-06-01")
c.Set(service.OpenAIParsedRequestBodyKey, map[string]any{
"model": "claude-3-5-sonnet-20241022",
"system": []any{
map[string]any{"text": "You are Claude Code, Anthropic's official CLI for Claude."},
},
"metadata": map[string]any{"user_id": "user_aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa_account__session_aaaaaaaa-aaaa-aaaa-aaaa-aaaaaaaaaaaa"},
})
SetClaudeCodeClientContext(c, []byte(`{invalid`), nil)
require.True(t, service.IsClaudeCodeClient(c.Request.Context()))
})
}
func TestWaitForSlotWithPingTimeout_AccountAndUserAcquire(t *testing.T) {

View File

@ -262,7 +262,7 @@ func (h *GatewayHandler) GeminiV1BetaModels(c *gin.Context) {
sessionHash := extractGeminiCLISessionHash(c, body)
if sessionHash == "" {
// Fallback: 使用通用的会话哈希生成逻辑(适用于其他客户端)
parsedReq, _ := service.ParseGatewayRequest(body, domain.PlatformGemini)
parsedReq, _ := service.ParseGatewayRequest(service.NewRequestBodyRef(body), domain.PlatformGemini)
if parsedReq != nil {
parsedReq.SessionContext = &service.SessionContext{
ClientIP: ip.GetClientIP(c),
@ -527,6 +527,8 @@ func (h *GatewayHandler) GeminiV1BetaModels(c *gin.Context) {
requestPayloadHash := service.HashUsageRequestPayload(body)
inboundEndpoint := GetInboundEndpoint(c)
upstreamEndpoint := GetUpstreamEndpoint(c, account.Platform)
// ForceCacheBilling 提前拍成标量,避免 worker 闭包保活 failover 状态里的响应体。
forceCacheBilling := fs.ForceCacheBilling
quotaPlatform := service.QuotaPlatform(c.Request.Context(), apiKey)
h.submitUsageRecordTask(c.Request.Context(), func(ctx context.Context) {
if err := h.gatewayService.RecordUsageWithLongContext(ctx, &service.RecordUsageLongContextInput{
@ -543,7 +545,7 @@ func (h *GatewayHandler) GeminiV1BetaModels(c *gin.Context) {
RequestPayloadHash: requestPayloadHash,
LongContextThreshold: 200000, // Gemini 200K 阈值
LongContextMultiplier: 2.0, // 超出部分双倍计费
ForceCacheBilling: fs.ForceCacheBilling,
ForceCacheBilling: forceCacheBilling,
APIKeyService: h.apiKeyService,
ChannelUsageFields: channelMapping.ToUsageFields(reqModel, result.UpstreamModel),
}); err != nil {

View File

@ -948,19 +948,12 @@ func (h *OpenAIGatewayHandler) validateFunctionCallOutputRequest(c *gin.Context,
return true
}
var reqBody map[string]any
if err := json.Unmarshal(body, &reqBody); err != nil {
// 保持原有容错语义:解析失败时跳过预校验,沿用后续上游校验结果。
return true
}
c.Set(service.OpenAIParsedRequestBodyKey, reqBody)
validation := service.ValidateFunctionCallOutputContext(reqBody)
validation := service.ValidateFunctionCallOutputContextBytes(body)
if !validation.HasFunctionCallOutput {
return true
}
previousResponseID, _ := reqBody["previous_response_id"].(string)
previousResponseID := gjson.GetBytes(body, "previous_response_id").String()
if strings.TrimSpace(previousResponseID) != "" || validation.HasToolCallContext {
return true
}
@ -1379,6 +1372,7 @@ func (h *OpenAIGatewayHandler) ResponsesWebSocket(c *gin.Context) {
zap.Int("candidate_count", scheduleDecision.CandidateCount),
)
var requestPayloadHash string
hooks := &service.OpenAIWSIngressHooks{
InitialRequestModel: reqModel,
BeforeRequest: func(turn int, payload []byte, originalModel string) error {
@ -1464,7 +1458,7 @@ func (h *OpenAIGatewayHandler) ResponsesWebSocket(c *gin.Context) {
UpstreamEndpoint: upstreamEndpoint,
UserAgent: userAgent,
IPAddress: clientIP,
RequestPayloadHash: service.HashUsageRequestPayload(firstMessage),
RequestPayloadHash: requestPayloadHash,
APIKeyService: h.apiKeyService,
ChannelUsageFields: channelMappingWS.ToUsageFields(reqModel, result.UpstreamModel),
}); err != nil {
@ -1484,6 +1478,9 @@ func (h *OpenAIGatewayHandler) ResponsesWebSocket(c *gin.Context) {
wsFirstMessage = h.gatewayService.ReplaceModelInBody(firstMessage, channelMappingWS.MappedModel)
}
// WebSocket 首包可能很大hash 必须在 hooks 外算成字符串,避免 AfterTurn 闭包保活请求体。
requestPayloadHash = service.HashUsageRequestPayload(wsFirstMessage)
if err := h.gatewayService.ProxyResponsesWebSocketFromClient(ctx, c, wsConn, account, token, wsFirstMessage, hooks); err != nil {
var failoverErr *service.UpstreamFailoverError
if errors.As(err, &failoverErr) {

View File

@ -73,9 +73,10 @@ func (h *OpenAIGatewayHandler) Images(c *gin.Context) {
h.errorResponse(c, http.StatusBadRequest, "invalid_request_error", err.Error())
return
}
requestModel := parsed.Model
reqLog = reqLog.With(
zap.String("model", parsed.Model),
zap.String("model", requestModel),
zap.Bool("stream", parsed.Stream),
zap.Bool("multipart", parsed.Multipart),
zap.String("capability", string(parsed.RequiredCapability)),
@ -85,7 +86,7 @@ func (h *OpenAIGatewayHandler) Images(c *gin.Context) {
h.errorResponse(c, http.StatusForbidden, "permission_error", service.ImageGenerationPermissionMessage())
return
}
if decision := h.checkContentModeration(c, reqLog, apiKey, subject, service.ContentModerationProtocolOpenAIImages, parsed.Model, parsed.ModerationBody()); decision != nil && decision.Blocked {
if decision := h.checkContentModeration(c, reqLog, apiKey, subject, service.ContentModerationProtocolOpenAIImages, requestModel, parsed.ModerationBody()); decision != nil && decision.Blocked {
h.errorResponse(c, contentModerationStatus(decision), contentModerationErrorCode(decision), decision.Message)
return
}
@ -98,13 +99,13 @@ func (h *OpenAIGatewayHandler) Images(c *gin.Context) {
}
if parsed.Multipart {
setOpsRequestContext(c, parsed.Model, parsed.Stream)
setOpsRequestContext(c, requestModel, parsed.Stream)
} else {
setOpsRequestContext(c, parsed.Model, parsed.Stream)
setOpsRequestContext(c, requestModel, parsed.Stream)
}
setOpsEndpointContext(c, "", int16(service.RequestTypeFromLegacy(parsed.Stream, false)))
channelMapping, _ := h.gatewayService.ResolveChannelMappingAndRestrict(c.Request.Context(), apiKey.GroupID, parsed.Model)
channelMapping, _ := h.gatewayService.ResolveChannelMappingAndRestrict(c.Request.Context(), apiKey.GroupID, requestModel)
if h.errorPassthroughService != nil {
service.BindErrorPassthroughService(c, h.errorPassthroughService)
@ -147,7 +148,7 @@ func (h *OpenAIGatewayHandler) Images(c *gin.Context) {
c.Request.Context(),
apiKey.GroupID,
sessionHash,
parsed.Model,
requestModel,
failedAccountIDs,
parsed.RequiredCapability,
)
@ -324,14 +325,14 @@ func (h *OpenAIGatewayHandler) Images(c *gin.Context) {
IPAddress: clientIP,
RequestPayloadHash: requestPayloadHash,
APIKeyService: h.apiKeyService,
ChannelUsageFields: channelMapping.ToUsageFields(parsed.Model, upstreamModel),
ChannelUsageFields: channelMapping.ToUsageFields(requestModel, upstreamModel),
}); err != nil {
logger.L().With(
zap.String("component", "handler.openai_gateway.images"),
zap.Int64("user_id", subject.UserID),
zap.Int64("api_key_id", apiKey.ID),
zap.Any("group_id", apiKey.GroupID),
zap.String("model", parsed.Model),
zap.String("model", requestModel),
zap.Int64("account_id", account.ID),
).Error("openai.images.record_usage_failed", zap.Error(err))
}

View File

@ -4,6 +4,8 @@ import (
"encoding/json"
"strings"
"time"
"github.com/tidwall/gjson"
)
// Anthropic 会话 Fallback 相关常量
@ -30,30 +32,39 @@ func BuildAnthropicDigestChain(parsed *ParsedRequest) string {
var parts []string
// 1. system prompt
if parsed.System != nil {
systemData, _ := json.Marshal(parsed.System)
if len(systemData) > 0 && string(systemData) != "null" {
parts = append(parts, "s:"+shortHash(systemData))
}
if systemRaw := parsed.SystemRaw(); len(systemRaw) > 0 && string(systemRaw) != "null" {
parts = append(parts, "s:"+shortHash(canonicalAnthropicDigestJSON(systemRaw)))
}
// 2. messages
for _, msg := range parsed.Messages {
msgMap, ok := msg.(map[string]any)
if !ok {
continue
}
role, _ := msgMap["role"].(string)
prefix := rolePrefix(role)
content := msgMap["content"]
contentData, _ := json.Marshal(content)
parts = append(parts, prefix+":"+shortHash(contentData))
messages := parsed.MessagesRaw()
if len(messages) > 0 {
gjson.ParseBytes(messages).ForEach(func(_, msg gjson.Result) bool {
prefix := rolePrefix(msg.Get("role").String())
content := msg.Get("content")
parts = append(parts, prefix+":"+shortHash(canonicalAnthropicDigestJSON([]byte(content.Raw))))
return true
})
}
return strings.Join(parts, "-")
}
// canonicalAnthropicDigestJSON 保持 digest 对 JSON key 顺序和空白不敏感。
func canonicalAnthropicDigestJSON(raw []byte) []byte {
if len(raw) == 0 {
return raw
}
var value any
if err := json.Unmarshal(raw, &value); err != nil {
return raw
}
canonical, err := json.Marshal(value)
if err != nil {
return raw
}
return canonical
}
// rolePrefix 将 Anthropic 的 role 映射为单字符前缀
func rolePrefix(role string) string {
switch role {

View File

@ -1,3 +1,5 @@
//go:build unit
package service
import (
@ -5,6 +7,15 @@ import (
"testing"
)
func mustParseAnthropicDigestRequest(t *testing.T, body string) *ParsedRequest {
t.Helper()
parsed, err := ParseGatewayRequest(NewRequestBodyRef([]byte(body)), "")
if err != nil {
t.Fatalf("ParseGatewayRequest failed: %v", err)
}
return parsed
}
func TestBuildAnthropicDigestChain_NilRequest(t *testing.T) {
result := BuildAnthropicDigestChain(nil)
if result != "" {
@ -13,9 +24,7 @@ func TestBuildAnthropicDigestChain_NilRequest(t *testing.T) {
}
func TestBuildAnthropicDigestChain_EmptyMessages(t *testing.T) {
parsed := &ParsedRequest{
Messages: []any{},
}
parsed := mustParseAnthropicDigestRequest(t, `{"messages":[]}`)
result := BuildAnthropicDigestChain(parsed)
if result != "" {
t.Errorf("expected empty string for empty messages, got: %s", result)
@ -23,11 +32,7 @@ func TestBuildAnthropicDigestChain_EmptyMessages(t *testing.T) {
}
func TestBuildAnthropicDigestChain_SingleUserMessage(t *testing.T) {
parsed := &ParsedRequest{
Messages: []any{
map[string]any{"role": "user", "content": "hello"},
},
}
parsed := mustParseAnthropicDigestRequest(t, `{"messages":[{"role":"user","content":"hello"}]}`)
result := BuildAnthropicDigestChain(parsed)
parts := splitChain(result)
if len(parts) != 1 {
@ -39,12 +44,7 @@ func TestBuildAnthropicDigestChain_SingleUserMessage(t *testing.T) {
}
func TestBuildAnthropicDigestChain_UserAndAssistant(t *testing.T) {
parsed := &ParsedRequest{
Messages: []any{
map[string]any{"role": "user", "content": "hello"},
map[string]any{"role": "assistant", "content": "hi there"},
},
}
parsed := mustParseAnthropicDigestRequest(t, `{"messages":[{"role":"user","content":"hello"},{"role":"assistant","content":"hi there"}]}`)
result := BuildAnthropicDigestChain(parsed)
parts := splitChain(result)
if len(parts) != 2 {
@ -59,12 +59,7 @@ func TestBuildAnthropicDigestChain_UserAndAssistant(t *testing.T) {
}
func TestBuildAnthropicDigestChain_WithSystemString(t *testing.T) {
parsed := &ParsedRequest{
System: "You are a helpful assistant",
Messages: []any{
map[string]any{"role": "user", "content": "hello"},
},
}
parsed := mustParseAnthropicDigestRequest(t, `{"system":"You are a helpful assistant","messages":[{"role":"user","content":"hello"}]}`)
result := BuildAnthropicDigestChain(parsed)
parts := splitChain(result)
if len(parts) != 2 {
@ -79,14 +74,7 @@ func TestBuildAnthropicDigestChain_WithSystemString(t *testing.T) {
}
func TestBuildAnthropicDigestChain_WithSystemContentBlocks(t *testing.T) {
parsed := &ParsedRequest{
System: []any{
map[string]any{"type": "text", "text": "You are a helpful assistant"},
},
Messages: []any{
map[string]any{"role": "user", "content": "hello"},
},
}
parsed := mustParseAnthropicDigestRequest(t, `{"system":[{"type":"text","text":"You are a helpful assistant"}],"messages":[{"role":"user","content":"hello"}]}`)
result := BuildAnthropicDigestChain(parsed)
parts := splitChain(result)
if len(parts) != 2 {
@ -100,74 +88,33 @@ func TestBuildAnthropicDigestChain_WithSystemContentBlocks(t *testing.T) {
func TestBuildAnthropicDigestChain_ConversationPrefixRelationship(t *testing.T) {
// 核心测试:验证对话增长时链的前缀关系
// 上一轮的完整链一定是下一轮链的前缀
system := "You are a helpful assistant"
// 第 1 轮: system + user
round1 := &ParsedRequest{
System: system,
Messages: []any{
map[string]any{"role": "user", "content": "hello"},
},
}
round1 := mustParseAnthropicDigestRequest(t, `{"system":"You are a helpful assistant","messages":[{"role":"user","content":"hello"}]}`)
chain1 := BuildAnthropicDigestChain(round1)
// 第 2 轮: system + user + assistant + user
round2 := &ParsedRequest{
System: system,
Messages: []any{
map[string]any{"role": "user", "content": "hello"},
map[string]any{"role": "assistant", "content": "hi there"},
map[string]any{"role": "user", "content": "how are you?"},
},
}
round2 := mustParseAnthropicDigestRequest(t, `{"system":"You are a helpful assistant","messages":[{"role":"user","content":"hello"},{"role":"assistant","content":"hi there"},{"role":"user","content":"how are you?"}]}`)
chain2 := BuildAnthropicDigestChain(round2)
// 第 3 轮: system + user + assistant + user + assistant + user
round3 := &ParsedRequest{
System: system,
Messages: []any{
map[string]any{"role": "user", "content": "hello"},
map[string]any{"role": "assistant", "content": "hi there"},
map[string]any{"role": "user", "content": "how are you?"},
map[string]any{"role": "assistant", "content": "I'm doing well"},
map[string]any{"role": "user", "content": "great"},
},
}
round3 := mustParseAnthropicDigestRequest(t, `{"system":"You are a helpful assistant","messages":[{"role":"user","content":"hello"},{"role":"assistant","content":"hi there"},{"role":"user","content":"how are you?"},{"role":"assistant","content":"I'm doing well"},{"role":"user","content":"great"}]}`)
chain3 := BuildAnthropicDigestChain(round3)
t.Logf("Chain1: %s", chain1)
t.Logf("Chain2: %s", chain2)
t.Logf("Chain3: %s", chain3)
// chain1 是 chain2 的前缀
if !strings.HasPrefix(chain2, chain1) {
t.Errorf("chain1 should be prefix of chain2:\n chain1: %s\n chain2: %s", chain1, chain2)
}
// chain2 是 chain3 的前缀
if !strings.HasPrefix(chain3, chain2) {
t.Errorf("chain2 should be prefix of chain3:\n chain2: %s\n chain3: %s", chain2, chain3)
}
// chain1 也是 chain3 的前缀(传递性)
if !strings.HasPrefix(chain3, chain1) {
t.Errorf("chain1 should be prefix of chain3:\n chain1: %s\n chain3: %s", chain1, chain3)
}
}
func TestBuildAnthropicDigestChain_DifferentSystemProducesDifferentChain(t *testing.T) {
parsed1 := &ParsedRequest{
System: "System A",
Messages: []any{
map[string]any{"role": "user", "content": "hello"},
},
}
parsed2 := &ParsedRequest{
System: "System B",
Messages: []any{
map[string]any{"role": "user", "content": "hello"},
},
}
parsed1 := mustParseAnthropicDigestRequest(t, `{"system":"System A","messages":[{"role":"user","content":"hello"}]}`)
parsed2 := mustParseAnthropicDigestRequest(t, `{"system":"System B","messages":[{"role":"user","content":"hello"}]}`)
chain1 := BuildAnthropicDigestChain(parsed1)
chain2 := BuildAnthropicDigestChain(parsed2)
@ -176,7 +123,6 @@ func TestBuildAnthropicDigestChain_DifferentSystemProducesDifferentChain(t *test
t.Error("Different system prompts should produce different chains")
}
// 但 user 部分的 hash 应该相同
parts1 := splitChain(chain1)
parts2 := splitChain(chain2)
if parts1[1] != parts2[1] {
@ -185,20 +131,8 @@ func TestBuildAnthropicDigestChain_DifferentSystemProducesDifferentChain(t *test
}
func TestBuildAnthropicDigestChain_DifferentContentProducesDifferentChain(t *testing.T) {
parsed1 := &ParsedRequest{
Messages: []any{
map[string]any{"role": "user", "content": "hello"},
map[string]any{"role": "assistant", "content": "ORIGINAL reply"},
map[string]any{"role": "user", "content": "next"},
},
}
parsed2 := &ParsedRequest{
Messages: []any{
map[string]any{"role": "user", "content": "hello"},
map[string]any{"role": "assistant", "content": "TAMPERED reply"},
map[string]any{"role": "user", "content": "next"},
},
}
parsed1 := mustParseAnthropicDigestRequest(t, `{"messages":[{"role":"user","content":"hello"},{"role":"assistant","content":"ORIGINAL reply"},{"role":"user","content":"next"}]}`)
parsed2 := mustParseAnthropicDigestRequest(t, `{"messages":[{"role":"user","content":"hello"},{"role":"assistant","content":"TAMPERED reply"},{"role":"user","content":"next"}]}`)
chain1 := BuildAnthropicDigestChain(parsed1)
chain2 := BuildAnthropicDigestChain(parsed2)
@ -209,24 +143,16 @@ func TestBuildAnthropicDigestChain_DifferentContentProducesDifferentChain(t *tes
parts1 := splitChain(chain1)
parts2 := splitChain(chain2)
// 第一个 user message hash 应该相同
if parts1[0] != parts2[0] {
t.Error("First user message hash should be the same")
}
// assistant reply hash 应该不同
if parts1[1] == parts2[1] {
t.Error("Assistant reply hash should differ")
}
}
func TestBuildAnthropicDigestChain_Deterministic(t *testing.T) {
parsed := &ParsedRequest{
System: "test system",
Messages: []any{
map[string]any{"role": "user", "content": "hello"},
map[string]any{"role": "assistant", "content": "hi"},
},
}
parsed := mustParseAnthropicDigestRequest(t, `{"system":"test system","messages":[{"role":"user","content":"hello"},{"role":"assistant","content":"hi"}]}`)
chain1 := BuildAnthropicDigestChain(parsed)
chain2 := BuildAnthropicDigestChain(parsed)
@ -236,6 +162,18 @@ func TestBuildAnthropicDigestChain_Deterministic(t *testing.T) {
}
}
func TestBuildAnthropicDigestChain_CanonicalJSON(t *testing.T) {
parsed1 := mustParseAnthropicDigestRequest(t, `{"system":[{"type":"text","text":"system"}],"messages":[{"role":"user","content":{"type":"text","text":"hello"}}]}`)
parsed2 := mustParseAnthropicDigestRequest(t, `{"system":[{"text":"system","type":"text"}],"messages":[{"role":"user","content":{"text":"hello","type":"text"}}]}`)
chain1 := BuildAnthropicDigestChain(parsed1)
chain2 := BuildAnthropicDigestChain(parsed2)
if chain1 != chain2 {
t.Errorf("semantically equivalent JSON should produce same chain: %s vs %s", chain1, chain2)
}
}
func TestGenerateAnthropicDigestSessionKey(t *testing.T) {
tests := []struct {
name string
@ -278,7 +216,6 @@ func TestGenerateAnthropicDigestSessionKey(t *testing.T) {
})
}
// 验证不同 uuid 产生不同 sessionKey
t.Run("different uuid different key", func(t *testing.T) {
hash := "sameprefix123456"
result1 := GenerateAnthropicDigestSessionKey(hash, "uuid0001-session-a")
@ -297,18 +234,7 @@ func TestAnthropicSessionTTL(t *testing.T) {
}
func TestBuildAnthropicDigestChain_ContentBlocks(t *testing.T) {
// 测试 content 为 content blocks 数组的情况
parsed := &ParsedRequest{
Messages: []any{
map[string]any{
"role": "user",
"content": []any{
map[string]any{"type": "text", "text": "describe this image"},
map[string]any{"type": "image", "source": map[string]any{"type": "base64"}},
},
},
},
}
parsed := mustParseAnthropicDigestRequest(t, `{"messages":[{"role":"user","content":[{"type":"text","text":"describe this image"},{"type":"image","source":{"type":"base64"}}]}]}`)
result := BuildAnthropicDigestChain(parsed)
parts := splitChain(result)
if len(parts) != 1 {

View File

@ -662,7 +662,7 @@ urlFallbackLoop:
// 统一处理错误响应
if resp.StatusCode >= 400 {
respBody, _ := io.ReadAll(io.LimitReader(resp.Body, 2<<20))
respBody := s.readUpstreamErrorBody(resp)
_ = resp.Body.Close()
if overagesInjected && shouldMarkCreditsExhausted(resp, respBody, nil) {
@ -875,6 +875,22 @@ type AntigravityGatewayService struct {
internal500Cache Internal500CounterCache // INTERNAL 500 渐进惩罚计数器
}
func (s *AntigravityGatewayService) upstreamErrorBodyReadLimit() int64 {
limit := gatewayUpstreamErrorBodyReadLimit
if s != nil && s.settingService != nil && s.settingService.cfg != nil && s.settingService.cfg.Gateway.LogUpstreamErrorBody && s.settingService.cfg.Gateway.LogUpstreamErrorBodyMaxBytes > int(limit) {
limit = int64(s.settingService.cfg.Gateway.LogUpstreamErrorBodyMaxBytes)
}
return limit
}
func (s *AntigravityGatewayService) readUpstreamErrorBody(resp *http.Response) []byte {
if resp == nil || resp.Body == nil {
return nil
}
body, _ := io.ReadAll(io.LimitReader(resp.Body, s.upstreamErrorBodyReadLimit()))
return body
}
func NewAntigravityGatewayService(
accountRepo AccountRepository,
cache GatewayCache,
@ -1090,7 +1106,7 @@ func (s *AntigravityGatewayService) TestConnection(ctx context.Context, account
}
defer func() { _ = result.resp.Body.Close() }()
respBody, err := io.ReadAll(io.LimitReader(result.resp.Body, 2<<20))
respBody, err := io.ReadAll(io.LimitReader(result.resp.Body, s.upstreamErrorBodyReadLimit()))
if err != nil {
return nil, fmt.Errorf("读取响应失败: %w", err)
}
@ -1427,7 +1443,7 @@ func (s *AntigravityGatewayService) Forward(ctx context.Context, c *gin.Context,
defer func() { _ = resp.Body.Close() }()
if resp.StatusCode >= 400 {
respBody, _ := io.ReadAll(io.LimitReader(resp.Body, 2<<20))
respBody := s.readUpstreamErrorBody(resp)
// 优先检测 thinking block 的 signature 相关错误400并重试一次
// Antigravity /v1internal 链路在部分场景会对 thought/thinking signature 做严格校验,
@ -1622,7 +1638,7 @@ func (s *AntigravityGatewayService) Forward(ctx context.Context, c *gin.Context,
resp = retryResp
respBody = nil
} else {
retryBody, _ := io.ReadAll(io.LimitReader(retryResp.Body, 2<<20))
retryBody := s.readUpstreamErrorBody(retryResp)
_ = retryResp.Body.Close()
respBody = retryBody
resp = &http.Response{
@ -2189,7 +2205,7 @@ func (s *AntigravityGatewayService) ForwardGemini(ctx context.Context, c *gin.Co
// 处理错误响应
if resp.StatusCode >= 400 {
respBody, _ := io.ReadAll(io.LimitReader(resp.Body, 2<<20))
respBody := s.readUpstreamErrorBody(resp)
contentType := resp.Header.Get("Content-Type")
// 尽早关闭原始响应体,释放连接;后续逻辑仍可能需要读取 body因此用内存副本重新包装。
_ = resp.Body.Close()
@ -2270,7 +2286,7 @@ func (s *AntigravityGatewayService) ForwardGemini(ctx context.Context, c *gin.Co
if retryResp.StatusCode < 400 {
resp = retryResp
} else {
retryRespBody, _ := io.ReadAll(io.LimitReader(retryResp.Body, 2<<20))
retryRespBody := s.readUpstreamErrorBody(retryResp)
_ = retryResp.Body.Close()
retryOpsBody := retryRespBody
if retryUnwrapped, unwrapErr := s.unwrapV1InternalResponse(retryRespBody); unwrapErr == nil && len(retryUnwrapped) > 0 {
@ -4252,7 +4268,7 @@ func (s *AntigravityGatewayService) ForwardUpstream(ctx context.Context, c *gin.
// 处理错误响应
if resp.StatusCode >= 400 {
respBody, _ := io.ReadAll(io.LimitReader(resp.Body, 2<<20))
respBody := s.readUpstreamErrorBody(resp)
// 429 错误时标记账号限流
if resp.StatusCode == http.StatusTooManyRequests {

View File

@ -42,6 +42,15 @@ func newAntigravityTestService(cfg *config.Config) *AntigravityGatewayService {
}
}
func TestAntigravityUpstreamErrorBodyReadLimit_RespectsDiagnosticLimit(t *testing.T) {
svc := newAntigravityTestService(&config.Config{Gateway: config.GatewayConfig{
LogUpstreamErrorBody: true,
LogUpstreamErrorBodyMaxBytes: int(gatewayUpstreamErrorBodyReadLimit) + 1024,
}})
require.Equal(t, int64(svc.settingService.cfg.Gateway.LogUpstreamErrorBodyMaxBytes), svc.upstreamErrorBodyReadLimit())
}
func TestStripSignatureSensitiveBlocksFromClaudeRequest(t *testing.T) {
req := &antigravity.ClaudeRequest{
Model: "claude-sonnet-4-5",

View File

@ -112,7 +112,7 @@ func TestGatewayService_AnthropicAPIKeyPassthrough_ForwardStreamPreservesBodyAnd
body := []byte(`{"model":"claude-3-7-sonnet-20250219","stream":true,"system":[{"type":"text","text":"x-anthropic-billing-header keep"}],"messages":[{"role":"user","content":[{"type":"text","text":"hello"}]}]}`)
parsed := &ParsedRequest{
Body: body,
Body: NewRequestBodyRef(body),
Model: "claude-3-7-sonnet-20250219",
Stream: true,
}
@ -202,7 +202,7 @@ func TestGatewayService_AnthropicAPIKeyPassthrough_ForwardCountTokensPreservesBo
body := []byte(`{"model":"claude-3-5-sonnet-latest","messages":[{"role":"user","content":[{"type":"text","text":"hello"}]}],"thinking":{"type":"enabled"}}`)
parsed := &ParsedRequest{
Body: body,
Body: NewRequestBodyRef(body),
Model: "claude-3-5-sonnet-latest",
}
@ -344,7 +344,7 @@ func TestGatewayService_AnthropicAPIKeyPassthrough_ModelMappingEdgeCases(t *test
body := []byte(`{"model":"` + tt.model + `","messages":[{"role":"user","content":[{"type":"text","text":"hello"}]}]}`)
parsed := &ParsedRequest{
Body: body,
Body: NewRequestBodyRef(body),
Model: tt.model,
}
@ -429,7 +429,7 @@ func TestGatewayService_AnthropicAPIKeyPassthrough_ModelMappingPreservesOtherFie
// 包含复杂字段的请求体system、thinking、messages
body := []byte(`{"model":"claude-sonnet-4-20250514","system":[{"type":"text","text":"You are a helpful assistant."}],"messages":[{"role":"user","content":[{"type":"text","text":"hello world"}]}],"thinking":{"type":"enabled","budget_tokens":5000},"max_tokens":1024}`)
parsed := &ParsedRequest{
Body: body,
Body: NewRequestBodyRef(body),
Model: "claude-sonnet-4-20250514",
}
@ -485,7 +485,7 @@ func TestGatewayService_AnthropicAPIKeyPassthrough_CountTokensFiltersGenerationF
body := []byte(`{"model":"claude-sonnet-4-20250514","system":[{"type":"text","text":"sys"}],"messages":[{"role":"user","content":"hello"}],"tools":[{"name":"tool","input_schema":{"type":"object"}}],"temperature":0.7,"top_p":0.9,"top_k":40,"stream":true,"stop_sequences":["END"],"max_tokens":1024,"thinking":{"type":"enabled","budget_tokens":5000}}`)
parsed := &ParsedRequest{
Body: body,
Body: NewRequestBodyRef(body),
Model: "claude-sonnet-4-20250514",
}
@ -547,7 +547,7 @@ func TestGatewayService_AnthropicAPIKeyPassthrough_EmptyModelSkipsMapping(t *tes
body := []byte(`{"messages":[{"role":"user","content":"hello"}]}`)
parsed := &ParsedRequest{
Body: body,
Body: NewRequestBodyRef(body),
Model: "", // 空模型
}
@ -636,7 +636,7 @@ func TestGatewayService_AnthropicAPIKeyPassthrough_CountTokens404PassthroughNotE
c.Request = httptest.NewRequest(http.MethodPost, "/v1/messages/count_tokens", nil)
body := []byte(`{"model":"claude-sonnet-4-5-20250929","messages":[{"role":"user","content":"hi"}]}`)
parsed := &ParsedRequest{Body: body, Model: "claude-sonnet-4-5-20250929"}
parsed := &ParsedRequest{Body: NewRequestBodyRef(body), Model: "claude-sonnet-4-5-20250929"}
upstream := &anthropicHTTPUpstreamRecorder{
resp: &http.Response{
@ -713,7 +713,7 @@ func TestGatewayService_AnthropicAPIKeyPassthrough_BuildRequestRejectsInvalidBas
},
}
_, err := svc.buildUpstreamRequestAnthropicAPIKeyPassthrough(context.Background(), c, account, []byte(`{}`), "k")
_, _, err := svc.buildUpstreamRequestAnthropicAPIKeyPassthrough(context.Background(), c, account, []byte(`{}`), "k")
require.Error(t, err)
}
@ -738,7 +738,7 @@ func TestGatewayService_AnthropicOAuth_NotAffectedByAPIKeyPassthroughToggle(t *t
require.False(t, account.IsAnthropicAPIKeyPassthroughEnabled())
req, err := svc.buildUpstreamRequest(context.Background(), c, account, []byte(`{"model":"claude-3-7-sonnet-20250219"}`), "oauth-token", "oauth", "claude-3-7-sonnet-20250219", true, false)
req, _, err := svc.buildUpstreamRequest(context.Background(), c, account, []byte(`{"model":"claude-3-7-sonnet-20250219"}`), "oauth-token", "oauth", "claude-3-7-sonnet-20250219", true, false)
require.NoError(t, err)
require.Equal(t, "Bearer oauth-token", getHeaderRaw(req.Header, "authorization"))
require.Contains(t, getHeaderRaw(req.Header, "anthropic-beta"), claude.BetaOAuth, "OAuth 链路仍应按原逻辑补齐 oauth beta")
@ -767,7 +767,7 @@ func TestGatewayService_AnthropicOAuth_ForwardPreservesBillingHeaderSystemBlock(
c, _ := gin.CreateTestContext(rec)
c.Request = httptest.NewRequest(http.MethodPost, "/v1/messages", nil)
parsed, err := ParseGatewayRequest([]byte(tt.body), PlatformAnthropic)
parsed, err := ParseGatewayRequest(NewRequestBodyRef([]byte(tt.body)), PlatformAnthropic)
require.NoError(t, err)
upstream := &anthropicHTTPUpstreamRecorder{

View File

@ -35,7 +35,7 @@ func TestGatewayService_BuildAnthropicVertexServiceAccountRequest(t *testing.T)
body := []byte(`{"model":"claude-sonnet-4-5","stream":false,"max_tokens":32,"messages":[{"role":"user","content":"hello"}]}`)
svc := &GatewayService{}
req, err := svc.buildUpstreamRequest(
req, _, err := svc.buildUpstreamRequest(
context.Background(),
c,
account,
@ -87,7 +87,7 @@ func TestGatewayService_BuildAnthropicVertexServiceAccount_StripsContextManageme
body := []byte(`{"model":"claude-haiku-4-5","context_management":{"edits":[{"type":"clear_thinking_20251015","keep":"all"}]},"messages":[{"role":"user","content":"hi"}]}`)
svc := &GatewayService{}
req, err := svc.buildUpstreamRequest(
req, _, err := svc.buildUpstreamRequest(
context.Background(), c, account, body,
"vertex-token", "service_account", "claude-haiku-4-5@20251001", false, false,
)
@ -117,7 +117,7 @@ func TestGatewayService_BuildAnthropicVertexServiceAccount_PreservesContextManag
body := []byte(`{"model":"claude-sonnet-4-6","context_management":{"edits":[{"type":"clear_thinking_20251015"}]},"messages":[]}`)
svc := &GatewayService{}
req, err := svc.buildUpstreamRequest(
req, _, err := svc.buildUpstreamRequest(
context.Background(), c, account, body,
"vertex-token", "service_account", "claude-sonnet-4-6@20260218", false, false,
)

View File

@ -364,7 +364,7 @@ func TestBuildUpstreamRequestAnthropicAPIKeyPassthrough_StripsContextManagementW
body := []byte(`{"model":"claude-haiku-4-5","context_management":{"edits":[{"type":"clear_thinking_20251015"}]},"messages":[]}`)
svc := &GatewayService{cfg: &config.Config{}}
req, err := svc.buildUpstreamRequestAnthropicAPIKeyPassthrough(
req, _, err := svc.buildUpstreamRequestAnthropicAPIKeyPassthrough(
context.Background(), c, newAnthropicAPIKeyPassthroughAccountForBetaTest(), body, "token",
)
require.NoError(t, err)
@ -381,7 +381,7 @@ func TestBuildUpstreamRequestAnthropicAPIKeyPassthrough_PreservesContextManageme
body := []byte(`{"model":"claude-haiku-4-5","context_management":{"edits":[{"type":"clear_thinking_20251015"}]},"messages":[]}`)
svc := &GatewayService{cfg: &config.Config{}}
req, err := svc.buildUpstreamRequestAnthropicAPIKeyPassthrough(
req, _, err := svc.buildUpstreamRequestAnthropicAPIKeyPassthrough(
context.Background(), c, newAnthropicAPIKeyPassthroughAccountForBetaTest(), body, "token",
)
require.NoError(t, err)
@ -427,7 +427,7 @@ func TestBuildUpstreamRequest_OAuthMimicHaiku_StripsContextManagementEndToEnd(t
// body 必须 strip。
body := []byte(`{"model":"claude-haiku-4-5","context_management":{"edits":[{"type":"clear_thinking_20251015"}]},"messages":[]}`)
svc := &GatewayService{cfg: &config.Config{}}
req, err := svc.buildUpstreamRequest(
req, _, err := svc.buildUpstreamRequest(
context.Background(), c, account, body,
"oauth-tok", "oauth", "claude-haiku-4-5", false, true, // mimicClaudeCode=true
)
@ -457,7 +457,7 @@ func TestBuildUpstreamRequest_OAuthMimicNonHaiku_PreservesContextManagementEndTo
// body 保留。
body := []byte(`{"model":"claude-sonnet-4-6","context_management":{"edits":[{"type":"clear_thinking_20251015"}]},"messages":[]}`)
svc := &GatewayService{cfg: &config.Config{}}
req, err := svc.buildUpstreamRequest(
req, _, err := svc.buildUpstreamRequest(
context.Background(), c, account, body,
"oauth-tok", "oauth", "claude-sonnet-4-6", false, true,
)
@ -488,7 +488,7 @@ func TestBuildUpstreamRequest_OAuthTransparentHaikuWithRealCCBeta_PreservesField
}
body := []byte(`{"model":"claude-haiku-4-5","context_management":{"edits":[{"type":"clear_thinking_20251015","keep":"all"}]},"messages":[]}`)
svc := &GatewayService{cfg: &config.Config{}}
req, err := svc.buildUpstreamRequest(
req, _, err := svc.buildUpstreamRequest(
context.Background(), c, account, body,
"oauth-tok", "oauth", "claude-haiku-4-5", false, false, // mimicClaudeCode=false真 CC
)
@ -580,7 +580,7 @@ func TestBuildCountTokensRequest_OAuthMimicHaiku_PreservesContextManagementEndTo
}
body := []byte(`{"model":"claude-haiku-4-5","context_management":{"edits":[{"type":"clear_thinking_20251015"}]},"messages":[]}`)
svc := &GatewayService{cfg: &config.Config{}}
req, err := svc.buildCountTokensRequest(
req, _, err := svc.buildCountTokensRequest(
context.Background(), c, account, body,
"oauth-tok", "oauth", "claude-haiku-4-5", true, // mimicClaudeCode=true
)
@ -611,7 +611,7 @@ func TestBuildCountTokensRequest_APIKeyHaiku_StripsContextManagementEndToEnd(t *
}
body := []byte(`{"model":"claude-haiku-4-5","context_management":{"edits":[]},"messages":[]}`)
svc := &GatewayService{cfg: &config.Config{}}
req, err := svc.buildCountTokensRequest(
req, _, err := svc.buildCountTokensRequest(
context.Background(), c, account, body,
"sk-ant-xxx", "apikey", "claude-haiku-4-5", false,
)
@ -655,7 +655,7 @@ func TestBuildUpstreamRequest_APIKeyHaikuWithContextManagement_StripsField(t *te
}
body := []byte(`{"model":"claude-haiku-4-5","context_management":{"edits":[]},"messages":[]}`)
svc := &GatewayService{cfg: &config.Config{}}
req, err := svc.buildUpstreamRequest(
req, _, err := svc.buildUpstreamRequest(
context.Background(), c, account, body,
"sk-ant-xxx", "apikey", "claude-haiku-4-5", false, false,
)

View File

@ -119,7 +119,7 @@ func (s *GatewayService) ForwardAsChatCompletions(
// 10. Build upstream request
upstreamCtx, releaseUpstreamCtx := detachStreamUpstreamContext(ctx, reqStream)
upstreamReq, err := s.buildUpstreamRequest(upstreamCtx, c, account, anthropicBody, token, tokenType, mappedModel, reqStream, shouldMimicClaudeCode)
upstreamReq, _, err := s.buildUpstreamRequest(upstreamCtx, c, account, anthropicBody, token, tokenType, mappedModel, reqStream, shouldMimicClaudeCode)
releaseUpstreamCtx()
if err != nil {
return nil, fmt.Errorf("build upstream request: %w", err)
@ -148,7 +148,7 @@ func (s *GatewayService) ForwardAsChatCompletions(
// 12. Handle error response with failover
if resp.StatusCode >= 400 {
respBody, _ := io.ReadAll(io.LimitReader(resp.Body, 2<<20))
respBody, _ := s.readUpstreamErrorBody(resp)
_ = resp.Body.Close()
resp.Body = io.NopCloser(bytes.NewReader(respBody))

View File

@ -116,7 +116,7 @@ func (s *GatewayService) ForwardAsResponses(
// 10. Build upstream request
upstreamCtx, releaseUpstreamCtx := detachStreamUpstreamContext(ctx, reqStream)
upstreamReq, err := s.buildUpstreamRequest(upstreamCtx, c, account, anthropicBody, token, tokenType, mappedModel, reqStream, shouldMimicClaudeCode)
upstreamReq, _, err := s.buildUpstreamRequest(upstreamCtx, c, account, anthropicBody, token, tokenType, mappedModel, reqStream, shouldMimicClaudeCode)
releaseUpstreamCtx()
if err != nil {
return nil, fmt.Errorf("build upstream request: %w", err)
@ -145,7 +145,7 @@ func (s *GatewayService) ForwardAsResponses(
// 12. Handle error response with failover
if resp.StatusCode >= 400 {
respBody, _ := io.ReadAll(io.LimitReader(resp.Body, 2<<20))
respBody, _ := s.readUpstreamErrorBody(resp)
_ = resp.Body.Close()
resp.Body = io.NopCloser(bytes.NewReader(respBody))

View File

@ -14,8 +14,6 @@ func TestBuildOAuthMetadataUserID_FallbackWithoutAccountUUID(t *testing.T) {
Model: "claude-sonnet-4-5",
Stream: true,
MetadataUserID: "",
System: nil,
Messages: nil,
}
account := &Account{

View File

@ -51,6 +51,164 @@ type SessionContext struct {
APIKeyID int64
}
type jsonRange struct {
start int // 原始请求体中的起始偏移(闭区间)
end int // 原始请求体中的结束偏移(开区间)
kind gjson.Type // JSON 值类型,用于调用方做轻量分支
}
type RequestBodyRef struct {
data []byte
}
func NewRequestBodyRef(data []byte) *RequestBodyRef {
return &RequestBodyRef{data: data}
}
func (b *RequestBodyRef) Bytes() []byte {
if b == nil {
return nil
}
return b.data
}
func (b *RequestBodyRef) Len() int {
if b == nil {
return 0
}
return len(b.data)
}
func (b *RequestBodyRef) Replace(data []byte) {
if b == nil {
return
}
b.data = data
}
func missingJSONRange() jsonRange {
return jsonRange{start: -1, end: -1}
}
func rangeFromResult(r gjson.Result) jsonRange {
if r.Raw == "" || r.Index <= 0 {
return missingJSONRange()
}
end := r.Index + len(r.Raw)
if end < r.Index {
return missingJSONRange()
}
return jsonRange{start: r.Index, end: end, kind: r.Type}
}
func (r jsonRange) exists() bool {
return r.start >= 0 && r.end >= r.start
}
// clearGatewayRequestDerivedState 清空绑定当前 body 的轻量派生字段,防止 ReplaceBody 后读到旧值。
func clearGatewayRequestDerivedState(parsed *ParsedRequest) {
if parsed == nil {
return
}
parsed.Model = ""
parsed.Stream = false
parsed.MetadataUserID = ""
parsed.HasSystem = false
parsed.ThinkingEnabled = false
parsed.OutputEffort = ""
parsed.MaxTokens = 0
parsed.systemRange = missingJSONRange()
parsed.messagesRange = missingJSONRange()
}
func clearGatewayRequestRanges(parsed *ParsedRequest) {
if parsed == nil {
return
}
parsed.HasSystem = false
parsed.systemRange = missingJSONRange()
parsed.messagesRange = missingJSONRange()
}
func setGatewayRequestRanges(parsed *ParsedRequest, protocol string, jsonStr string) {
if parsed == nil {
return
}
switch protocol {
case domain.PlatformGemini:
if sysParts := gjson.Get(jsonStr, "systemInstruction.parts"); sysParts.Exists() && sysParts.IsArray() {
parsed.systemRange = rangeFromResult(sysParts)
}
if contents := gjson.Get(jsonStr, "contents"); contents.Exists() && contents.IsArray() {
parsed.messagesRange = rangeFromResult(contents)
}
default:
if sys := gjson.Get(jsonStr, "system"); sys.Exists() {
parsed.HasSystem = true
parsed.systemRange = rangeFromResult(sys)
}
if msgs := gjson.Get(jsonStr, "messages"); msgs.Exists() && msgs.IsArray() {
parsed.messagesRange = rangeFromResult(msgs)
}
}
}
// parseGatewayRequestCurrentBody 只做标量和 raw range 轻量解析,不恢复 system/messages 对象图。
func parseGatewayRequestCurrentBody(parsed *ParsedRequest, protocol string) error {
if parsed == nil || parsed.Body == nil {
return fmt.Errorf("empty request body")
}
bodyBytes := parsed.Body.Bytes()
if !gjson.ValidBytes(bodyBytes) {
return fmt.Errorf("invalid json")
}
// 只在当前函数内零拷贝读取 JSON 字段ReplaceBody 后必须重新进入本函数刷新派生状态。
jsonStr := *(*string)(unsafe.Pointer(&bodyBytes))
clearGatewayRequestDerivedState(parsed)
parsed.protocol = protocol
modelResult := gjson.Get(jsonStr, "model")
if modelResult.Exists() {
if modelResult.Type != gjson.String {
return fmt.Errorf("invalid model field type")
}
parsed.Model = modelResult.String()
}
streamResult := gjson.Get(jsonStr, "stream")
if streamResult.Exists() {
if streamResult.Type != gjson.True && streamResult.Type != gjson.False {
return fmt.Errorf("invalid stream field type")
}
parsed.Stream = streamResult.Bool()
}
parsed.MetadataUserID = gjson.Get(jsonStr, "metadata.user_id").String()
thinkingType := gjson.Get(jsonStr, "thinking.type").String()
parsed.ThinkingEnabled = thinkingType == "enabled" || thinkingType == "adaptive"
parsed.OutputEffort = strings.TrimSpace(gjson.Get(jsonStr, "output_config.effort").String())
maxTokensResult := gjson.Get(jsonStr, "max_tokens")
if maxTokensResult.Exists() && maxTokensResult.Type == gjson.Number {
f := maxTokensResult.Float()
if !math.IsNaN(f) && !math.IsInf(f, 0) && f == math.Trunc(f) &&
f <= float64(math.MaxInt) && f >= float64(math.MinInt) {
parsed.MaxTokens = int(f)
}
}
setGatewayRequestRanges(parsed, protocol, jsonStr)
return nil
}
func refreshGatewayRequestRanges(parsed *ParsedRequest, protocol string) error {
return parseGatewayRequestCurrentBody(parsed, protocol)
}
// ParsedRequest 保存网关请求的预解析结果
//
// 性能优化说明:
@ -64,18 +222,20 @@ type SessionContext struct {
// 2. 将解析结果 ParsedRequest 传递给 Service 层
// 3. 避免重复 json.Unmarshal减少 CPU 和内存开销
type ParsedRequest struct {
Body []byte // 原始请求体(保留用于转发)
Body *RequestBodyRef // 原始请求体引用(保留用于转发);替换内容请走 ReplaceBody
Model string // 请求的模型名称
Stream bool // 是否为流式请求
MetadataUserID string // metadata.user_id用于会话亲和
System any // system 字段内容
Messages []any // messages 数组
HasSystem bool // 是否包含 system 字段(包含 null 也视为显式传入)
ThinkingEnabled bool // 是否开启 thinking部分平台会影响最终模型名
OutputEffort string // output_config.effortClaude API 的推理强度控制)
MaxTokens int // max_tokens 值(用于探测请求拦截)
SessionContext *SessionContext // 可选请求上下文区分因子nil 时行为不变)
protocol string // 当前 Body 的协议格式,用于 Body 替换后刷新 raw range
systemRange jsonRange // system/systemInstruction.parts 的 raw JSON 范围,绑定 Body 当前内容
messagesRange jsonRange // messages/contents 的 raw JSON 范围,绑定 Body 当前内容
// GroupID 请求所属分组 ID来自 API Key
GroupID *int64
@ -130,119 +290,92 @@ func normalizeSessionUserAgentFallback(raw string) string {
// ParseGatewayRequest 解析网关请求体并返回结构化结果。
// protocol 指定请求协议格式domain.PlatformAnthropic / domain.PlatformGemini
// 不同协议使用不同的 system/messages 字段名。
func ParseGatewayRequest(body []byte, protocol string) (*ParsedRequest, error) {
// 保持与旧实现一致:请求体必须是合法 JSON。
// 注意gjson.GetBytes 对非法 JSON 不会报错,因此需要显式校验。
if !gjson.ValidBytes(body) {
return nil, fmt.Errorf("invalid json")
func ParseGatewayRequest(body *RequestBodyRef, protocol string) (*ParsedRequest, error) {
parsed := &ParsedRequest{Body: body}
if err := parseGatewayRequestCurrentBody(parsed, protocol); err != nil {
return nil, err
}
// 性能:
// - gjson.GetBytes 会把匹配的 Raw/Str 安全复制成 string对于巨大 messages 会产生额外拷贝)。
// - 这里将 body 通过 unsafe 零拷贝视为 string仅在本函数内使用且 body 不会被修改。
jsonStr := *(*string)(unsafe.Pointer(&body))
parsed := &ParsedRequest{
Body: body,
}
// --- gjson 提取简单字段(避免完整 Unmarshal ---
// model: 需要严格类型校验,非 string 返回错误
modelResult := gjson.Get(jsonStr, "model")
if modelResult.Exists() {
if modelResult.Type != gjson.String {
return nil, fmt.Errorf("invalid model field type")
}
parsed.Model = modelResult.String()
}
// stream: 需要严格类型校验,非 bool 返回错误
streamResult := gjson.Get(jsonStr, "stream")
if streamResult.Exists() {
if streamResult.Type != gjson.True && streamResult.Type != gjson.False {
return nil, fmt.Errorf("invalid stream field type")
}
parsed.Stream = streamResult.Bool()
}
// metadata.user_id: 直接路径提取,不需要严格类型校验
parsed.MetadataUserID = gjson.Get(jsonStr, "metadata.user_id").String()
// thinking.type: enabled/adaptive 都视为开启
thinkingType := gjson.Get(jsonStr, "thinking.type").String()
if thinkingType == "enabled" || thinkingType == "adaptive" {
parsed.ThinkingEnabled = true
}
// output_config.effort: Claude API 的推理强度控制参数
parsed.OutputEffort = strings.TrimSpace(gjson.Get(jsonStr, "output_config.effort").String())
// max_tokens: 仅接受整数值
maxTokensResult := gjson.Get(jsonStr, "max_tokens")
if maxTokensResult.Exists() && maxTokensResult.Type == gjson.Number {
f := maxTokensResult.Float()
if !math.IsNaN(f) && !math.IsInf(f, 0) && f == math.Trunc(f) &&
f <= float64(math.MaxInt) && f >= float64(math.MinInt) {
parsed.MaxTokens = int(f)
}
}
// --- system/messages 提取 ---
// 避免把整个 body Unmarshal 到 map会产生大量 map/接口分配)。
// 使用 gjson 抽取目标字段的 Raw再对该子树进行 Unmarshal。
switch protocol {
case domain.PlatformGemini:
// Gemini 原生格式: systemInstruction.parts / contents
if sysParts := gjson.Get(jsonStr, "systemInstruction.parts"); sysParts.Exists() && sysParts.IsArray() {
var parts []any
if err := json.Unmarshal(sliceRawFromBody(body, sysParts), &parts); err != nil {
return nil, err
}
parsed.System = parts
}
if contents := gjson.Get(jsonStr, "contents"); contents.Exists() && contents.IsArray() {
var msgs []any
if err := json.Unmarshal(sliceRawFromBody(body, contents), &msgs); err != nil {
return nil, err
}
parsed.Messages = msgs
}
default:
// Anthropic / OpenAI 格式: system / messages
// system 字段只要存在就视为显式提供(即使为 null
// 以避免客户端传 null 时被默认 system 误注入。
if sys := gjson.Get(jsonStr, "system"); sys.Exists() {
parsed.HasSystem = true
switch sys.Type {
case gjson.Null:
parsed.System = nil
case gjson.String:
// 与 encoding/json 的 Unmarshal 行为一致:返回解码后的字符串。
parsed.System = sys.String()
default:
var system any
if err := json.Unmarshal(sliceRawFromBody(body, sys), &system); err != nil {
return nil, err
}
parsed.System = system
}
}
if msgs := gjson.Get(jsonStr, "messages"); msgs.Exists() && msgs.IsArray() {
var messages []any
if err := json.Unmarshal(sliceRawFromBody(body, msgs), &messages); err != nil {
return nil, err
}
parsed.Messages = messages
}
}
return parsed, nil
}
func (p *ParsedRequest) raw(r jsonRange) []byte {
if p == nil || p.Body == nil || !r.exists() {
return nil
}
body := p.Body.Bytes()
if r.end > len(body) {
return nil
}
return body[r.start:r.end]
}
func (p *ParsedRequest) SystemRaw() []byte {
return p.raw(p.systemRange)
}
func (p *ParsedRequest) MessagesRaw() []byte {
return p.raw(p.messagesRange)
}
func (p *ParsedRequest) DecodeSystem(dst any) error {
raw := p.SystemRaw()
if len(raw) == 0 {
return nil
}
return json.Unmarshal(raw, dst)
}
func (p *ParsedRequest) DecodeMessages(dst any) error {
raw := p.MessagesRaw()
if len(raw) == 0 {
return nil
}
return json.Unmarshal(raw, dst)
}
func (p *ParsedRequest) SystemValue() (any, bool) {
raw := p.SystemRaw()
if len(raw) == 0 {
return nil, false
}
var system any
if err := json.Unmarshal(raw, &system); err != nil {
return nil, false
}
return system, true
}
// CloneForBody 为单次账号尝试创建独立 body 视图,避免 failover 复用已改写的 ParsedRequest。
func (p *ParsedRequest) CloneForBody(body []byte) (*ParsedRequest, error) {
if p == nil {
return nil, fmt.Errorf("parse request: empty request")
}
clone := *p
clone.Body = NewRequestBodyRef(body)
clone.OnUpstreamAccepted = nil
if err := refreshGatewayRequestRanges(&clone, clone.protocol); err != nil {
return nil, err
}
return &clone, nil
}
// ReplaceBody 统一刷新当前 body 和 raw range保证后续 helper 读取的是最新请求体。
func (p *ParsedRequest) ReplaceBody(data []byte) error {
if p == nil {
return fmt.Errorf("parse request: empty request")
}
if p.Body == nil {
p.Body = NewRequestBodyRef(data)
} else {
p.Body.Replace(data)
}
if err := refreshGatewayRequestRanges(p, p.protocol); err != nil {
clearGatewayRequestRanges(p)
return err
}
return nil
}
// sliceRawFromBody 返回 Result.Raw 对应的原始字节切片。
// 优先使用 Result.Index 直接从 body 切片,避免对大字段(如 messages产生额外拷贝。
// 当 Index 不可用时,退化为复制(理论上极少发生)。

View File

@ -10,24 +10,25 @@ import (
"github.com/Wei-Shaw/sub2api/internal/domain"
"github.com/stretchr/testify/require"
"github.com/tidwall/gjson"
)
func TestParseGatewayRequest(t *testing.T) {
body := []byte(`{"model":"claude-3-7-sonnet","stream":true,"metadata":{"user_id":"session_123e4567-e89b-12d3-a456-426614174000"},"system":[{"type":"text","text":"hello","cache_control":{"type":"ephemeral"}}],"messages":[{"content":"hi"}]}`)
parsed, err := ParseGatewayRequest(body, "")
parsed, err := ParseGatewayRequest(NewRequestBodyRef(body), "")
require.NoError(t, err)
require.Equal(t, "claude-3-7-sonnet", parsed.Model)
require.True(t, parsed.Stream)
require.Equal(t, "session_123e4567-e89b-12d3-a456-426614174000", parsed.MetadataUserID)
require.True(t, parsed.HasSystem)
require.NotNil(t, parsed.System)
require.Len(t, parsed.Messages, 1)
require.NotEmpty(t, parsed.SystemRaw())
require.NotEmpty(t, parsed.MessagesRaw())
require.False(t, parsed.ThinkingEnabled)
}
func TestParseGatewayRequest_ThinkingEnabled(t *testing.T) {
body := []byte(`{"model":"claude-sonnet-4-5","thinking":{"type":"enabled"},"messages":[{"content":"hi"}]}`)
parsed, err := ParseGatewayRequest(body, "")
parsed, err := ParseGatewayRequest(NewRequestBodyRef(body), "")
require.NoError(t, err)
require.Equal(t, "claude-sonnet-4-5", parsed.Model)
require.True(t, parsed.ThinkingEnabled)
@ -35,7 +36,7 @@ func TestParseGatewayRequest_ThinkingEnabled(t *testing.T) {
func TestParseGatewayRequest_ThinkingAdaptiveEnabled(t *testing.T) {
body := []byte(`{"model":"claude-sonnet-4-5","thinking":{"type":"adaptive"},"messages":[{"content":"hi"}]}`)
parsed, err := ParseGatewayRequest(body, "")
parsed, err := ParseGatewayRequest(NewRequestBodyRef(body), "")
require.NoError(t, err)
require.Equal(t, "claude-sonnet-4-5", parsed.Model)
require.True(t, parsed.ThinkingEnabled)
@ -43,36 +44,36 @@ func TestParseGatewayRequest_ThinkingAdaptiveEnabled(t *testing.T) {
func TestParseGatewayRequest_MaxTokens(t *testing.T) {
body := []byte(`{"model":"claude-haiku-4-5","max_tokens":1}`)
parsed, err := ParseGatewayRequest(body, "")
parsed, err := ParseGatewayRequest(NewRequestBodyRef(body), "")
require.NoError(t, err)
require.Equal(t, 1, parsed.MaxTokens)
}
func TestParseGatewayRequest_MaxTokensNonIntegralIgnored(t *testing.T) {
body := []byte(`{"model":"claude-haiku-4-5","max_tokens":1.5}`)
parsed, err := ParseGatewayRequest(body, "")
parsed, err := ParseGatewayRequest(NewRequestBodyRef(body), "")
require.NoError(t, err)
require.Equal(t, 0, parsed.MaxTokens)
}
func TestParseGatewayRequest_SystemNull(t *testing.T) {
body := []byte(`{"model":"claude-3","system":null}`)
parsed, err := ParseGatewayRequest(body, "")
parsed, err := ParseGatewayRequest(NewRequestBodyRef(body), "")
require.NoError(t, err)
// 显式传入 system:null 也应视为“字段已存在”,避免默认 system 被注入。
require.True(t, parsed.HasSystem)
require.Nil(t, parsed.System)
require.Equal(t, []byte("null"), parsed.SystemRaw())
}
func TestParseGatewayRequest_InvalidModelType(t *testing.T) {
body := []byte(`{"model":123}`)
_, err := ParseGatewayRequest(body, "")
_, err := ParseGatewayRequest(NewRequestBodyRef(body), "")
require.Error(t, err)
}
func TestParseGatewayRequest_InvalidStreamType(t *testing.T) {
body := []byte(`{"stream":"true"}`)
_, err := ParseGatewayRequest(body, "")
_, err := ParseGatewayRequest(NewRequestBodyRef(body), "")
require.Error(t, err)
}
@ -86,11 +87,11 @@ func TestParseGatewayRequest_GeminiContents(t *testing.T) {
{"role": "user", "parts": [{"text": "How are you?"}]}
]
}`)
parsed, err := ParseGatewayRequest(body, domain.PlatformGemini)
parsed, err := ParseGatewayRequest(NewRequestBodyRef(body), domain.PlatformGemini)
require.NoError(t, err)
require.Len(t, parsed.Messages, 3, "should parse contents as Messages")
require.Len(t, gjson.ParseBytes(parsed.MessagesRaw()).Array(), 3, "should parse contents as Messages")
require.False(t, parsed.HasSystem, "Gemini format should not set HasSystem")
require.Nil(t, parsed.System, "no systemInstruction means nil System")
require.Nil(t, parsed.SystemRaw(), "no systemInstruction means nil System")
}
func TestParseGatewayRequest_GeminiSystemInstruction(t *testing.T) {
@ -102,16 +103,13 @@ func TestParseGatewayRequest_GeminiSystemInstruction(t *testing.T) {
{"role": "user", "parts": [{"text": "Hello"}]}
]
}`)
parsed, err := ParseGatewayRequest(body, domain.PlatformGemini)
parsed, err := ParseGatewayRequest(NewRequestBodyRef(body), domain.PlatformGemini)
require.NoError(t, err)
require.NotNil(t, parsed.System, "should parse systemInstruction.parts as System")
parts, ok := parsed.System.([]any)
require.True(t, ok)
require.Len(t, parts, 1)
partMap, ok := parts[0].(map[string]any)
require.True(t, ok)
require.Equal(t, "You are a helpful assistant.", partMap["text"])
require.Len(t, parsed.Messages, 1)
system := gjson.ParseBytes(parsed.SystemRaw())
require.True(t, system.IsArray(), "should parse systemInstruction.parts as System")
require.Len(t, system.Array(), 1)
require.Equal(t, "You are a helpful assistant.", system.Get("0.text").String())
require.Len(t, gjson.ParseBytes(parsed.MessagesRaw()).Array(), 1)
}
func TestParseGatewayRequest_GeminiWithModel(t *testing.T) {
@ -119,10 +117,10 @@ func TestParseGatewayRequest_GeminiWithModel(t *testing.T) {
"model": "gemini-2.5-pro",
"contents": [{"role": "user", "parts": [{"text": "test"}]}]
}`)
parsed, err := ParseGatewayRequest(body, domain.PlatformGemini)
parsed, err := ParseGatewayRequest(NewRequestBodyRef(body), domain.PlatformGemini)
require.NoError(t, err)
require.Equal(t, "gemini-2.5-pro", parsed.Model)
require.Len(t, parsed.Messages, 1)
require.Len(t, gjson.ParseBytes(parsed.MessagesRaw()).Array(), 1)
}
func TestParseGatewayRequest_GeminiIgnoresAnthropicFields(t *testing.T) {
@ -132,25 +130,25 @@ func TestParseGatewayRequest_GeminiIgnoresAnthropicFields(t *testing.T) {
"messages": [{"role": "user", "content": "ignored"}],
"contents": [{"role": "user", "parts": [{"text": "real content"}]}]
}`)
parsed, err := ParseGatewayRequest(body, domain.PlatformGemini)
parsed, err := ParseGatewayRequest(NewRequestBodyRef(body), domain.PlatformGemini)
require.NoError(t, err)
require.False(t, parsed.HasSystem, "Gemini protocol should not parse Anthropic system field")
require.Nil(t, parsed.System, "no systemInstruction = nil System")
require.Len(t, parsed.Messages, 1, "should use contents, not messages")
require.Nil(t, parsed.SystemRaw(), "no systemInstruction = nil System")
require.Len(t, gjson.ParseBytes(parsed.MessagesRaw()).Array(), 1, "should use contents, not messages")
}
func TestParseGatewayRequest_GeminiEmptyContents(t *testing.T) {
body := []byte(`{"contents": []}`)
parsed, err := ParseGatewayRequest(body, domain.PlatformGemini)
parsed, err := ParseGatewayRequest(NewRequestBodyRef(body), domain.PlatformGemini)
require.NoError(t, err)
require.Empty(t, parsed.Messages)
require.Empty(t, gjson.ParseBytes(parsed.MessagesRaw()).Array())
}
func TestParseGatewayRequest_GeminiNoContents(t *testing.T) {
body := []byte(`{"model": "gemini-2.5-flash"}`)
parsed, err := ParseGatewayRequest(body, domain.PlatformGemini)
parsed, err := ParseGatewayRequest(NewRequestBodyRef(body), domain.PlatformGemini)
require.NoError(t, err)
require.Nil(t, parsed.Messages)
require.Nil(t, parsed.MessagesRaw())
require.Equal(t, "gemini-2.5-flash", parsed.Model)
}
@ -162,14 +160,13 @@ func TestParseGatewayRequest_AnthropicIgnoresGeminiFields(t *testing.T) {
"contents": [{"role": "user", "parts": [{"text": "ignored"}]}],
"systemInstruction": {"parts": [{"text": "ignored"}]}
}`)
parsed, err := ParseGatewayRequest(body, domain.PlatformAnthropic)
parsed, err := ParseGatewayRequest(NewRequestBodyRef(body), domain.PlatformAnthropic)
require.NoError(t, err)
require.True(t, parsed.HasSystem)
require.Equal(t, "real system", parsed.System)
require.Len(t, parsed.Messages, 1)
msg, ok := parsed.Messages[0].(map[string]any)
require.True(t, ok)
require.Equal(t, "real content", msg["content"])
require.Equal(t, "real system", gjson.ParseBytes(parsed.SystemRaw()).String())
messages := gjson.ParseBytes(parsed.MessagesRaw()).Array()
require.Len(t, messages, 1)
require.Equal(t, "real content", messages[0].Get("content").String())
}
func TestFilterThinkingBlocks(t *testing.T) {
@ -897,7 +894,7 @@ func TestParseGatewayRequest_TypeValidation(t *testing.T) {
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
_, err := ParseGatewayRequest([]byte(tt.body), "")
_, err := ParseGatewayRequest(NewRequestBodyRef([]byte(tt.body)), "")
if tt.wantErr {
require.Error(t, err)
if tt.errSubstr != "" {
@ -959,7 +956,7 @@ func TestParseGatewayRequest_OptionalFieldsMissing(t *testing.T) {
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
parsed, err := ParseGatewayRequest([]byte(tt.body), "")
parsed, err := ParseGatewayRequest(NewRequestBodyRef([]byte(tt.body)), "")
require.NoError(t, err)
require.Equal(t, tt.wantModel, parsed.Model)
@ -970,10 +967,10 @@ func TestParseGatewayRequest_OptionalFieldsMissing(t *testing.T) {
require.Equal(t, tt.wantMaxTokens, parsed.MaxTokens)
if tt.wantMessagesNil {
require.Nil(t, parsed.Messages)
require.Nil(t, parsed.MessagesRaw())
}
if tt.wantMessagesLen > 0 {
require.Len(t, parsed.Messages, tt.wantMessagesLen)
require.Len(t, gjson.ParseBytes(parsed.MessagesRaw()).Array(), tt.wantMessagesLen)
}
})
}
@ -1023,7 +1020,7 @@ func TestParseGatewayRequest_MaxTokensBoundary(t *testing.T) {
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
parsed, err := ParseGatewayRequest([]byte(tt.body), "")
parsed, err := ParseGatewayRequest(NewRequestBodyRef([]byte(tt.body)), "")
if tt.wantErr {
require.Error(t, err)
return
@ -1040,7 +1037,7 @@ func TestParseGatewayRequest_MaxTokensBoundary(t *testing.T) {
// 核心路径:先 Unmarshal 到 map[string]any再逐字段提取。
func parseGatewayRequestOld(body []byte, protocol string) (*ParsedRequest, error) {
parsed := &ParsedRequest{
Body: body,
Body: NewRequestBodyRef(body),
}
var req map[string]any
@ -1087,25 +1084,8 @@ func parseGatewayRequestOld(body []byte, protocol string) (*ParsedRequest, error
}
}
// system / messages按协议分支
switch protocol {
case domain.PlatformGemini:
if sysInst, ok := req["systemInstruction"].(map[string]any); ok {
if parts, ok := sysInst["parts"].([]any); ok {
parsed.System = parts
}
}
if contents, ok := req["contents"].([]any); ok {
parsed.Messages = contents
}
default:
if system, ok := req["system"]; ok {
parsed.HasSystem = true
parsed.System = system
}
if messages, ok := req["messages"].([]any); ok {
parsed.Messages = messages
}
if err := refreshGatewayRequestRanges(parsed, protocol); err != nil {
return nil, err
}
return parsed, nil
@ -1151,7 +1131,7 @@ func BenchmarkParseGatewayRequest_New_Small(b *testing.B) {
b.SetBytes(int64(len(data)))
b.ResetTimer()
for i := 0; i < b.N; i++ {
_, _ = ParseGatewayRequest(data, "")
_, _ = ParseGatewayRequest(NewRequestBodyRef(data), "")
}
}
@ -1203,7 +1183,7 @@ func TestParseGatewayRequest_OutputEffort(t *testing.T) {
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
parsed, err := ParseGatewayRequest([]byte(tt.body), "")
parsed, err := ParseGatewayRequest(NewRequestBodyRef([]byte(tt.body)), "")
require.NoError(t, err)
require.Equal(t, tt.wantEffort, parsed.OutputEffort)
})
@ -1245,6 +1225,6 @@ func BenchmarkParseGatewayRequest_New_Large(b *testing.B) {
b.SetBytes(int64(len(data)))
b.ResetTimer()
for i := 0; i < b.N; i++ {
_, _ = ParseGatewayRequest(data, "")
_, _ = ParseGatewayRequest(NewRequestBodyRef(data), "")
}
}

View File

@ -23,6 +23,7 @@ import (
"sync/atomic"
"syscall"
"time"
"unsafe"
"github.com/Wei-Shaw/sub2api/internal/config"
"github.com/Wei-Shaw/sub2api/internal/pkg/claude"
@ -56,6 +57,8 @@ const (
defaultModelsListCacheTTL = 15 * time.Second
postUsageBillingTimeout = 15 * time.Second
debugGatewayBodyEnv = "SUB2API_DEBUG_GATEWAY_BODY"
// 上游错误体只需要提取错误 JSON/日志摘要,默认 512KiB 避免错误风暴叠加大请求体。
gatewayUpstreamErrorBodyReadLimit int64 = 512 << 10
)
const (
@ -748,31 +751,10 @@ func (s *GatewayService) GenerateSessionHash(parsed *ParsedRequest) string {
_, _ = combined.WriteString(strconv.FormatInt(parsed.SessionContext.APIKeyID, 10))
_, _ = combined.WriteString("|")
}
if parsed.System != nil {
systemText := s.extractTextFromSystem(parsed.System)
if systemText != "" {
_, _ = combined.WriteString(systemText)
}
}
for _, msg := range parsed.Messages {
if m, ok := msg.(map[string]any); ok {
if content, exists := m["content"]; exists {
// Anthropic: messages[].content
if msgText := s.extractTextFromContent(content); msgText != "" {
_, _ = combined.WriteString(msgText)
}
} else if parts, ok := m["parts"].([]any); ok {
// Gemini: contents[].parts[].text
for _, part := range parts {
if partMap, ok := part.(map[string]any); ok {
if text, ok := partMap["text"].(string); ok {
_, _ = combined.WriteString(text)
}
}
}
}
}
if systemText := extractTextFromSystemRaw(parsed.SystemRaw()); systemText != "" {
_, _ = combined.WriteString(systemText)
}
appendMessageTextsFromRaw(&combined, parsed.MessagesRaw())
if combined.Len() > 0 {
hash := s.hashContent(combined.String())
slog.Info("sticky.hash_source",
@ -847,82 +829,135 @@ func (s *GatewayService) extractCacheableContent(parsed *ParsedRequest) string {
return ""
}
var builder strings.Builder
// 检查 system 中的 cacheable 内容
if system, ok := parsed.System.([]any); ok {
for _, part := range system {
if partMap, ok := part.(map[string]any); ok {
if cc, ok := partMap["cache_control"].(map[string]any); ok {
if cc["type"] == "ephemeral" {
if text, ok := partMap["text"].(string); ok {
_, _ = builder.WriteString(text)
}
}
}
}
}
systemText := extractCacheableTextFromSystemRaw(parsed.SystemRaw())
if messageText := extractCacheableTextFromMessagesRaw(parsed.MessagesRaw()); messageText != "" {
return messageText
}
systemText := builder.String()
// 检查 messages 中的 cacheable 内容
for _, msg := range parsed.Messages {
if msgMap, ok := msg.(map[string]any); ok {
if msgContent, ok := msgMap["content"].([]any); ok {
for _, part := range msgContent {
if partMap, ok := part.(map[string]any); ok {
if cc, ok := partMap["cache_control"].(map[string]any); ok {
if cc["type"] == "ephemeral" {
return s.extractTextFromContent(msgMap["content"])
}
}
}
}
}
}
}
return systemText
}
func (s *GatewayService) extractTextFromSystem(system any) string {
switch v := system.(type) {
case string:
return v
case []any:
var texts []string
for _, part := range v {
if partMap, ok := part.(map[string]any); ok {
if text, ok := partMap["text"].(string); ok {
texts = append(texts, text)
}
}
func parseRawJSONView(raw []byte) gjson.Result {
if len(raw) == 0 {
return gjson.Result{}
}
// 这里只做同步只读解析,避免 gjson.ParseBytes 为大 messages/contents 复制整段 raw。
return gjson.Parse(*(*string)(unsafe.Pointer(&raw)))
}
func extractTextFromSystemRaw(raw []byte) string {
system := parseRawJSONView(raw)
switch system.Type {
case gjson.String:
return system.String()
case gjson.JSON:
if !system.IsArray() {
return ""
}
return strings.Join(texts, "")
var builder strings.Builder
system.ForEach(func(_, part gjson.Result) bool {
if text := part.Get("text").String(); text != "" {
_, _ = builder.WriteString(text)
}
return true
})
return builder.String()
}
return ""
}
func (s *GatewayService) extractTextFromContent(content any) string {
switch v := content.(type) {
case string:
return v
case []any:
var texts []string
for _, part := range v {
if partMap, ok := part.(map[string]any); ok {
if partMap["type"] == "text" {
if text, ok := partMap["text"].(string); ok {
texts = append(texts, text)
}
func extractTextFromContentRaw(content gjson.Result) string {
switch content.Type {
case gjson.String:
return content.String()
case gjson.JSON:
if !content.IsArray() {
return ""
}
var builder strings.Builder
content.ForEach(func(_, part gjson.Result) bool {
if part.Get("type").String() == "text" {
if text := part.Get("text").String(); text != "" {
_, _ = builder.WriteString(text)
}
}
}
return strings.Join(texts, "")
return true
})
return builder.String()
}
return ""
}
func appendMessageTextsFromRaw(builder *strings.Builder, raw []byte) {
if builder == nil || len(raw) == 0 {
return
}
messages := parseRawJSONView(raw)
if !messages.IsArray() {
return
}
messages.ForEach(func(_, msg gjson.Result) bool {
if content := msg.Get("content"); content.Exists() {
_, _ = builder.WriteString(extractTextFromContentRaw(content))
return true
}
parts := msg.Get("parts")
if parts.IsArray() {
parts.ForEach(func(_, part gjson.Result) bool {
if text := part.Get("text").String(); text != "" {
_, _ = builder.WriteString(text)
}
return true
})
}
return true
})
}
func extractCacheableTextFromSystemRaw(raw []byte) string {
system := parseRawJSONView(raw)
if !system.IsArray() {
return ""
}
var builder strings.Builder
system.ForEach(func(_, part gjson.Result) bool {
if part.Get("cache_control.type").String() == "ephemeral" {
if text := part.Get("text").String(); text != "" {
_, _ = builder.WriteString(text)
}
}
return true
})
return builder.String()
}
func extractCacheableTextFromMessagesRaw(raw []byte) string {
messages := parseRawJSONView(raw)
if !messages.IsArray() {
return ""
}
var text string
messages.ForEach(func(_, msg gjson.Result) bool {
content := msg.Get("content")
if !content.IsArray() {
return true
}
found := false
content.ForEach(func(_, part gjson.Result) bool {
if part.Get("cache_control.type").String() == "ephemeral" {
found = true
return false
}
return true
})
if found {
text = extractTextFromContentRaw(content)
return false
}
return true
})
return text
}
func (s *GatewayService) hashContent(content string) string {
h := xxhash.Sum64String(content)
return strconv.FormatUint(h, 36)
@ -1284,7 +1319,7 @@ func (s *GatewayService) applyClaudeCodeOAuthMimicryToBody(
systemRewritten := false
if !strings.Contains(strings.ToLower(model), "haiku") {
body = rewriteSystemForNonClaudeCode(body, systemRaw)
body = rewriteSystemForNonClaudeCode(body, normalizeSystemParam(systemRaw))
systemRewritten = true
}
@ -4400,12 +4435,12 @@ func (s *GatewayService) Forward(ctx context.Context, c *gin.Context, account *A
}
// Web Search 模拟:纯 web_search 请求时,直接调用搜索 API 构造响应
if account != nil && s.shouldEmulateWebSearch(ctx, account, parsed.GroupID, parsed.Body) {
if account != nil && s.shouldEmulateWebSearch(ctx, account, parsed.GroupID, parsed.Body.Bytes()) {
return s.handleWebSearchEmulation(ctx, c, account, parsed)
}
if account != nil && account.IsAnthropicAPIKeyPassthroughEnabled() {
passthroughBody := parsed.Body
passthroughBody := parsed.Body.Bytes()
passthroughModel := parsed.Model
if passthroughModel != "" {
if mappedModel := account.GetMappedModel(passthroughModel); mappedModel != passthroughModel {
@ -4416,6 +4451,7 @@ func (s *GatewayService) Forward(ctx context.Context, c *gin.Context, account *A
}
return s.forwardAnthropicAPIKeyPassthroughWithInput(ctx, c, account, anthropicPassthroughForwardInput{
Body: passthroughBody,
Parsed: parsed,
RequestModel: passthroughModel,
OriginalModel: parsed.Model,
RequestStream: parsed.Stream,
@ -4441,7 +4477,14 @@ func (s *GatewayService) Forward(ctx context.Context, c *gin.Context, account *A
c.Set(betaPolicyFilterSetKey, filterSet)
}
body := parsed.Body
body := parsed.Body.Bytes()
replaceBody := func(next []byte) error {
if err := parsed.ReplaceBody(next); err != nil {
return fmt.Errorf("rewrite request body: %w", err)
}
body = parsed.Body.Bytes()
return nil
}
reqModel := parsed.Model
reqStream := parsed.Stream
originalModel := reqModel
@ -4474,7 +4517,10 @@ func (s *GatewayService) Forward(ctx context.Context, c *gin.Context, account *A
// Parrot 的 transform_request 从不检查客户端 system 内容,直接覆盖。
systemRewritten := false
if !strings.Contains(strings.ToLower(reqModel), "haiku") {
body = rewriteSystemForNonClaudeCode(body, parsed.System)
systemRaw, _ := parsed.SystemValue()
if err := replaceBody(rewriteSystemForNonClaudeCode(body, systemRaw)); err != nil {
return nil, err
}
systemRewritten = true
}
@ -4496,22 +4542,34 @@ func (s *GatewayService) Forward(ctx context.Context, c *gin.Context, account *A
}
}
body, reqModel = normalizeClaudeOAuthRequestBody(body, reqModel, normalizeOpts)
var normalizedBody []byte
normalizedBody, reqModel = normalizeClaudeOAuthRequestBody(body, reqModel, normalizeOpts)
if err := replaceBody(normalizedBody); err != nil {
return nil, err
}
// D/E/F: 可选 messages cache 策略 + 工具名混淆 + tools[-1] 断点
// 与 forward_as_chat_completions / forward_as_responses 路径对齐,
// 原生 /v1/messages 路径也走同一套可配置字段级改写。
body = s.rewriteMessageCacheControlIfEnabled(ctx, body)
if err := replaceBody(s.rewriteMessageCacheControlIfEnabled(ctx, body)); err != nil {
return nil, err
}
if rw := buildToolNameRewriteFromBody(body); rw != nil {
body = applyToolNameRewriteToBody(body, rw)
if err := replaceBody(applyToolNameRewriteToBody(body, rw)); err != nil {
return nil, err
}
c.Set(toolNameRewriteKey, rw)
} else {
body = applyToolsLastCacheBreakpoint(body)
if err := replaceBody(applyToolsLastCacheBreakpoint(body)); err != nil {
return nil, err
}
}
}
// 强制执行 cache_control 块数量限制(最多 4 个)
body = enforceCacheControlLimit(body)
if err := replaceBody(enforceCacheControlLimit(body)); err != nil {
return nil, err
}
// 应用模型映射:
// - APIKey 账号:使用账号级别的显式映射(如果配置),否则透传原始模型名
@ -4545,13 +4603,18 @@ func (s *GatewayService) Forward(ctx context.Context, c *gin.Context, account *A
}
if mappedModel != reqModel {
// 替换请求体中的模型名
body = s.replaceModelInBody(body, mappedModel)
if err := replaceBody(s.replaceModelInBody(body, mappedModel)); err != nil {
return nil, err
}
reqModel = mappedModel
parsed.Model = mappedModel
logger.LegacyPrintf("service.gateway", "Model mapping applied: %s -> %s (account: %s, source=%s)", originalModel, mappedModel, account.Name, mappingSource)
}
if s.shouldInjectAnthropicCacheTTL1h(ctx, account) {
body = injectAnthropicCacheControlTTL1h(body)
if err := replaceBody(injectAnthropicCacheControlTTL1h(body)); err != nil {
return nil, err
}
}
// 获取凭证
@ -4575,19 +4638,24 @@ func (s *GatewayService) Forward(ctx context.Context, c *gin.Context, account *A
logger.LegacyPrintf("service.gateway", "[Forward] Using account: ID=%d Name=%s Platform=%s Type=%s TLSFingerprint=%v Proxy=%s",
account.ID, account.Name, account.Platform, account.Type, tlsProfile, proxyURL)
// Pre-filter: strip empty text blocks (including nested in tool_result) to prevent upstream 400.
body = StripEmptyTextBlocks(body)
if err := replaceBody(StripEmptyTextBlocks(body)); err != nil {
return nil, err
}
// 重试循环
var resp *http.Response
lastWireBody := body
retryStart := time.Now()
for attempt := 1; attempt <= maxRetryAttempts; attempt++ {
// 构建上游请求(每次重试需要重新构建,因为请求体需要重新读取)
upstreamCtx, releaseUpstreamCtx := detachStreamUpstreamContext(ctx, reqStream)
upstreamReq, err := s.buildUpstreamRequest(upstreamCtx, c, account, body, token, tokenType, reqModel, reqStream, shouldMimicClaudeCode)
upstreamReq, wireBody, err := s.buildUpstreamRequest(upstreamCtx, c, account, body, token, tokenType, reqModel, reqStream, shouldMimicClaudeCode)
releaseUpstreamCtx()
if err != nil {
return nil, err
}
// 记录本次实际发送的 wire body只有请求成功后才写回 ParsedRequest避免 400 retry 基于已签名 CCH 再改写。
lastWireBody = wireBody
// 发送请求
resp, err = s.httpUpstream.DoWithTLS(upstreamReq, proxyURL, account.ID, account.Concurrency, tlsProfile)
@ -4619,7 +4687,7 @@ func (s *GatewayService) Forward(ctx context.Context, c *gin.Context, account *A
// 优先检测thinking block签名错误400并重试一次
if resp.StatusCode == 400 {
respBody, readErr := io.ReadAll(io.LimitReader(resp.Body, 2<<20))
respBody, readErr := s.readUpstreamErrorBody(resp)
if readErr == nil {
_ = resp.Body.Close()
@ -4665,18 +4733,24 @@ func (s *GatewayService) Forward(ctx context.Context, c *gin.Context, account *A
filteredBody := FilterThinkingBlocksForRetry(body)
retryCtx, releaseRetryCtx := detachStreamUpstreamContext(ctx, reqStream)
retryReq, buildErr := s.buildUpstreamRequest(retryCtx, c, account, filteredBody, token, tokenType, reqModel, reqStream, shouldMimicClaudeCode)
retryReq, retryWireBody, buildErr := s.buildUpstreamRequest(retryCtx, c, account, filteredBody, token, tokenType, reqModel, reqStream, shouldMimicClaudeCode)
releaseRetryCtx()
if buildErr == nil {
retryResp, retryErr := s.httpUpstream.DoWithTLS(retryReq, proxyURL, account.ID, account.Concurrency, tlsProfile)
if retryErr == nil {
if retryResp.StatusCode < 400 {
// 重试请求被上游接受后同步 ParsedRequest保证 usage/日志看到真实请求体。
lastWireBody = retryWireBody
if err := replaceBody(retryWireBody); err != nil {
_ = retryResp.Body.Close()
return nil, err
}
logger.LegacyPrintf("service.gateway", "Account %d: thinking block retry succeeded (blocks downgraded)", account.ID)
resp = retryResp
break
}
retryRespBody, retryReadErr := io.ReadAll(io.LimitReader(retryResp.Body, 2<<20))
retryRespBody, retryReadErr := s.readUpstreamErrorBody(retryResp)
_ = retryResp.Body.Close()
if retryReadErr == nil && retryResp.StatusCode == 400 && s.isSignatureErrorPattern(ctx, account, retryRespBody) {
appendOpsUpstreamError(c, OpsUpstreamErrorEvent{
@ -4700,11 +4774,19 @@ func (s *GatewayService) Forward(ctx context.Context, c *gin.Context, account *A
logger.LegacyPrintf("service.gateway", "Account %d: signature retry still failing and looks tool-related, retrying with tool blocks downgraded", account.ID)
filteredBody2 := FilterSignatureSensitiveBlocksForRetry(body)
retryCtx2, releaseRetryCtx2 := detachStreamUpstreamContext(ctx, reqStream)
retryReq2, buildErr2 := s.buildUpstreamRequest(retryCtx2, c, account, filteredBody2, token, tokenType, reqModel, reqStream, shouldMimicClaudeCode)
retryReq2, retryWireBody2, buildErr2 := s.buildUpstreamRequest(retryCtx2, c, account, filteredBody2, token, tokenType, reqModel, reqStream, shouldMimicClaudeCode)
releaseRetryCtx2()
if buildErr2 == nil {
retryResp2, retryErr2 := s.httpUpstream.DoWithTLS(retryReq2, proxyURL, account.ID, account.Concurrency, tlsProfile)
if retryErr2 == nil {
if retryResp2.StatusCode < 400 {
// 二阶段工具块降级成功时也必须更新当前 body。
lastWireBody = retryWireBody2
if err := replaceBody(retryWireBody2); err != nil {
_ = retryResp2.Body.Close()
return nil, err
}
}
resp = retryResp2
break
}
@ -4771,11 +4853,19 @@ func (s *GatewayService) Forward(ctx context.Context, c *gin.Context, account *A
if applied && time.Since(retryStart) < maxRetryElapsed {
logger.LegacyPrintf("service.gateway", "Account %d: detected budget_tokens constraint error, retrying with rectified budget (budget_tokens=%d, max_tokens=%d)", account.ID, BudgetRectifyBudgetTokens, BudgetRectifyMaxTokens)
budgetRetryCtx, releaseBudgetRetryCtx := detachStreamUpstreamContext(ctx, reqStream)
budgetRetryReq, buildErr := s.buildUpstreamRequest(budgetRetryCtx, c, account, rectifiedBody, token, tokenType, reqModel, reqStream, shouldMimicClaudeCode)
budgetRetryReq, budgetWireBody, buildErr := s.buildUpstreamRequest(budgetRetryCtx, c, account, rectifiedBody, token, tokenType, reqModel, reqStream, shouldMimicClaudeCode)
releaseBudgetRetryCtx()
if buildErr == nil {
budgetRetryResp, retryErr := s.httpUpstream.DoWithTLS(budgetRetryReq, proxyURL, account.ID, account.Concurrency, tlsProfile)
if retryErr == nil {
if budgetRetryResp.StatusCode < 400 {
// budget 修正请求成功后ParsedRequest 也要描述被接受的修正版。
lastWireBody = budgetWireBody
if err := replaceBody(budgetWireBody); err != nil {
_ = budgetRetryResp.Body.Close()
return nil, err
}
}
resp = budgetRetryResp
break
}
@ -4810,7 +4900,7 @@ func (s *GatewayService) Forward(ctx context.Context, c *gin.Context, account *A
break
}
respBody, _ := io.ReadAll(io.LimitReader(resp.Body, 2<<20))
respBody, _ := s.readUpstreamErrorBody(resp)
_ = resp.Body.Close()
appendOpsUpstreamError(c, OpsUpstreamErrorEvent{
Platform: account.Platform,
@ -4857,7 +4947,7 @@ func (s *GatewayService) Forward(ctx context.Context, c *gin.Context, account *A
// 处理重试耗尽的情况
if resp.StatusCode >= 400 && s.shouldRetryUpstreamError(account, resp.StatusCode) {
if s.shouldFailoverUpstreamError(resp.StatusCode) {
respBody, _ := io.ReadAll(io.LimitReader(resp.Body, 2<<20))
respBody, _ := s.readUpstreamErrorBody(resp)
_ = resp.Body.Close()
resp.Body = io.NopCloser(bytes.NewReader(respBody))
@ -4892,7 +4982,7 @@ func (s *GatewayService) Forward(ctx context.Context, c *gin.Context, account *A
// 处理可切换账号的错误
if resp.StatusCode >= 400 && s.shouldFailoverUpstreamError(resp.StatusCode) {
respBody, _ := io.ReadAll(io.LimitReader(resp.Body, 2<<20))
respBody, _ := s.readUpstreamErrorBody(resp)
_ = resp.Body.Close()
resp.Body = io.NopCloser(bytes.NewReader(respBody))
@ -4924,7 +5014,7 @@ func (s *GatewayService) Forward(ctx context.Context, c *gin.Context, account *A
if resp.StatusCode >= 400 {
// 可选:对部分 400 触发 failover默认关闭以保持语义
if resp.StatusCode == 400 && s.cfg != nil && s.cfg.Gateway.FailoverOn400 {
respBody, readErr := io.ReadAll(io.LimitReader(resp.Body, 2<<20))
respBody, readErr := s.readUpstreamErrorBody(resp)
if readErr != nil {
// ReadAll failed, fall back to normal error handling without consuming the stream
return s.handleErrorResponse(ctx, resp, c, account, reqModel)
@ -4972,6 +5062,13 @@ func (s *GatewayService) Forward(ctx context.Context, c *gin.Context, account *A
// 处理正常响应
if !bytes.Equal(lastWireBody, body) {
// 成功后再同步最终 wire body避免失败重试从已签名 CCH 的 body 继续派生。
if err := replaceBody(lastWireBody); err != nil {
return nil, err
}
}
// 触发上游接受回调(提前释放串行锁,不等流完成)
if parsed.OnUpstreamAccepted != nil {
parsed.OnUpstreamAccepted()
@ -5014,6 +5111,7 @@ func (s *GatewayService) Forward(ctx context.Context, c *gin.Context, account *A
type anthropicPassthroughForwardInput struct {
Body []byte
Parsed *ParsedRequest
RequestModel string
OriginalModel string
RequestStream bool
@ -5066,16 +5164,29 @@ func (s *GatewayService) forwardAnthropicAPIKeyPassthroughWithInput(
}
// Pre-filter: strip empty text blocks (including nested in tool_result) to prevent upstream 400.
input.Body = StripEmptyTextBlocks(input.Body)
if input.Parsed != nil {
// 透传分支也会改写实际 wire body成功 usage hash 依赖这里同步当前 body。
if err := input.Parsed.ReplaceBody(input.Body); err != nil {
return nil, err
}
}
var resp *http.Response
retryStart := time.Now()
for attempt := 1; attempt <= maxRetryAttempts; attempt++ {
upstreamCtx, releaseUpstreamCtx := detachStreamUpstreamContext(ctx, input.RequestStream)
upstreamReq, err := s.buildUpstreamRequestAnthropicAPIKeyPassthrough(upstreamCtx, c, account, input.Body, token)
upstreamReq, wireBody, err := s.buildUpstreamRequestAnthropicAPIKeyPassthrough(upstreamCtx, c, account, input.Body, token)
releaseUpstreamCtx()
if err != nil {
return nil, err
}
if input.Parsed != nil && !bytes.Equal(wireBody, input.Body) {
// build 阶段会按 beta 能力清理 body发送前同步到 ParsedRequest 当前视图。
if err := input.Parsed.ReplaceBody(wireBody); err != nil {
return nil, err
}
input.Body = input.Parsed.Body.Bytes()
}
resp, err = s.httpUpstream.DoWithTLS(upstreamReq, proxyURL, account.ID, account.Concurrency, s.tlsFPProfileService.ResolveTLSProfile(account))
if err != nil {
@ -5121,7 +5232,7 @@ func (s *GatewayService) forwardAnthropicAPIKeyPassthroughWithInput(
break
}
respBody, _ := io.ReadAll(io.LimitReader(resp.Body, 2<<20))
respBody, _ := s.readUpstreamErrorBody(resp)
_ = resp.Body.Close()
appendOpsUpstreamError(c, OpsUpstreamErrorEvent{
Platform: account.Platform,
@ -5159,7 +5270,7 @@ func (s *GatewayService) forwardAnthropicAPIKeyPassthroughWithInput(
if resp.StatusCode >= 400 && s.shouldRetryUpstreamError(account, resp.StatusCode) {
if s.shouldFailoverUpstreamError(resp.StatusCode) {
respBody, _ := io.ReadAll(io.LimitReader(resp.Body, 2<<20))
respBody, _ := s.readUpstreamErrorBody(resp)
_ = resp.Body.Close()
resp.Body = io.NopCloser(bytes.NewReader(respBody))
@ -5193,7 +5304,7 @@ func (s *GatewayService) forwardAnthropicAPIKeyPassthroughWithInput(
}
if resp.StatusCode >= 400 && s.shouldFailoverUpstreamError(resp.StatusCode) {
respBody, _ := io.ReadAll(io.LimitReader(resp.Body, 2<<20))
respBody, _ := s.readUpstreamErrorBody(resp)
_ = resp.Body.Close()
resp.Body = io.NopCloser(bytes.NewReader(respBody))
@ -5267,13 +5378,13 @@ func (s *GatewayService) buildUpstreamRequestAnthropicAPIKeyPassthrough(
account *Account,
body []byte,
token string,
) (*http.Request, error) {
) (*http.Request, []byte, error) {
targetURL := claudeAPIURL
baseURL := account.GetBaseURL()
if baseURL != "" {
validatedURL, err := s.validateUpstreamBaseURL(baseURL)
if err != nil {
return nil, err
return nil, nil, err
}
targetURL = validatedURL + "/v1/messages?beta=true"
}
@ -5291,7 +5402,7 @@ func (s *GatewayService) buildUpstreamRequestAnthropicAPIKeyPassthrough(
req, err := http.NewRequestWithContext(ctx, http.MethodPost, targetURL, bytes.NewReader(body))
if err != nil {
return nil, err
return nil, nil, err
}
if c != nil && c.Request != nil {
@ -5321,7 +5432,7 @@ func (s *GatewayService) buildUpstreamRequestAnthropicAPIKeyPassthrough(
setHeaderRaw(req.Header, "anthropic-version", "2023-06-01")
}
return req, nil
return req, body, nil
}
func (s *GatewayService) handleStreamingResponseAnthropicAPIKeyPassthrough(
@ -5735,7 +5846,7 @@ func (s *GatewayService) forwardBedrock(
) (*ForwardResult, error) {
reqModel := parsed.Model
reqStream := parsed.Stream
body := parsed.Body
body := parsed.Body.Bytes()
region := bedrockRuntimeRegion(account)
mappedModel, ok := ResolveBedrockModelID(account, reqModel)
@ -5803,6 +5914,11 @@ func (s *GatewayService) forwardBedrock(
return s.handleBedrockUpstreamErrors(ctx, resp, c, account)
}
// Bedrock 分支绕过通用 Forward 成功路径,这里保持上游接受回调语义一致。
if parsed.OnUpstreamAccepted != nil {
parsed.OnUpstreamAccepted()
}
// 响应处理
var usage *ClaudeUsage
var firstTokenMs *int
@ -5906,7 +6022,7 @@ func (s *GatewayService) executeBedrockUpstream(
break
}
respBody, _ := io.ReadAll(io.LimitReader(resp.Body, 2<<20))
respBody, _ := s.readUpstreamErrorBody(resp)
_ = resp.Body.Close()
appendOpsUpstreamError(c, OpsUpstreamErrorEvent{
Platform: account.Platform,
@ -5951,7 +6067,7 @@ func (s *GatewayService) handleBedrockUpstreamErrors(
// retry exhausted + failover
if s.shouldRetryUpstreamError(account, resp.StatusCode) {
if s.shouldFailoverUpstreamError(resp.StatusCode) {
respBody, _ := io.ReadAll(io.LimitReader(resp.Body, 2<<20))
respBody, _ := s.readUpstreamErrorBody(resp)
_ = resp.Body.Close()
resp.Body = io.NopCloser(bytes.NewReader(respBody))
@ -5978,7 +6094,7 @@ func (s *GatewayService) handleBedrockUpstreamErrors(
// non-retryable failover
if s.shouldFailoverUpstreamError(resp.StatusCode) {
respBody, _ := io.ReadAll(io.LimitReader(resp.Body, 2<<20))
respBody, _ := s.readUpstreamErrorBody(resp)
_ = resp.Body.Close()
resp.Body = io.NopCloser(bytes.NewReader(respBody))
@ -6079,9 +6195,10 @@ func (s *GatewayService) handleBedrockNonStreamingResponse(
return usage, nil
}
func (s *GatewayService) buildUpstreamRequest(ctx context.Context, c *gin.Context, account *Account, body []byte, token, tokenType, modelID string, reqStream bool, mimicClaudeCode bool) (*http.Request, error) {
func (s *GatewayService) buildUpstreamRequest(ctx context.Context, c *gin.Context, account *Account, body []byte, token, tokenType, modelID string, reqStream bool, mimicClaudeCode bool) (*http.Request, []byte, error) {
if account.Platform == PlatformAnthropic && account.Type == AccountTypeServiceAccount {
return s.buildUpstreamRequestAnthropicVertex(ctx, c, account, body, token, modelID, reqStream)
req, err := s.buildUpstreamRequestAnthropicVertex(ctx, c, account, body, token, modelID, reqStream)
return req, body, err
}
// 确定目标URL
@ -6091,18 +6208,18 @@ func (s *GatewayService) buildUpstreamRequest(ctx context.Context, c *gin.Contex
if baseURL != "" {
validatedURL, err := s.validateUpstreamBaseURL(baseURL)
if err != nil {
return nil, err
return nil, nil, err
}
targetURL = validatedURL + "/v1/messages?beta=true"
}
} else if account.IsCustomBaseURLEnabled() {
customURL := account.GetCustomBaseURL()
if customURL == "" {
return nil, fmt.Errorf("custom_base_url is enabled but not configured for account %d", account.ID)
return nil, nil, fmt.Errorf("custom_base_url is enabled but not configured for account %d", account.ID)
}
validatedURL, err := s.validateUpstreamBaseURL(customURL)
if err != nil {
return nil, err
return nil, nil, err
}
targetURL = s.buildCustomRelayURL(validatedURL, "/v1/messages", account)
}
@ -6177,7 +6294,7 @@ func (s *GatewayService) buildUpstreamRequest(ctx context.Context, c *gin.Contex
req, err := http.NewRequestWithContext(ctx, "POST", targetURL, bytes.NewReader(body))
if err != nil {
return nil, err
return nil, nil, err
}
// 设置认证头(保持原始大小写)
@ -6262,7 +6379,7 @@ func (s *GatewayService) buildUpstreamRequest(ctx context.Context, c *gin.Contex
logClaudeMimicDebug(req, body, account, tokenType, mimicClaudeCode)
}
return req, nil
return req, body, nil
}
func (s *GatewayService) buildUpstreamRequestAnthropicVertex(
@ -7120,8 +7237,19 @@ func isCountTokensUnsupported404(statusCode int, body []byte) bool {
return strings.Contains(msg, "count_tokens") && strings.Contains(msg, "not found")
}
func (s *GatewayService) readUpstreamErrorBody(resp *http.Response) ([]byte, error) {
if resp == nil || resp.Body == nil {
return nil, nil
}
limit := gatewayUpstreamErrorBodyReadLimit
if s != nil && s.cfg != nil && s.cfg.Gateway.LogUpstreamErrorBody && s.cfg.Gateway.LogUpstreamErrorBodyMaxBytes > int(limit) {
limit = int64(s.cfg.Gateway.LogUpstreamErrorBodyMaxBytes)
}
return io.ReadAll(io.LimitReader(resp.Body, limit))
}
func (s *GatewayService) handleErrorResponse(ctx context.Context, resp *http.Response, c *gin.Context, account *Account, requestedModel ...string) (*ForwardResult, error) {
body, _ := io.ReadAll(io.LimitReader(resp.Body, 2<<20))
body, _ := s.readUpstreamErrorBody(resp)
// 调试日志:打印上游错误响应
logger.LegacyPrintf("service.gateway", "[Forward] Upstream error (non-retryable): Account=%d(%s) Status=%d RequestID=%s Body=%s",
@ -7274,7 +7402,7 @@ func (s *GatewayService) handleErrorResponse(ctx context.Context, resp *http.Res
}
func (s *GatewayService) handleRetryExhaustedSideEffects(ctx context.Context, resp *http.Response, account *Account) {
body, _ := io.ReadAll(io.LimitReader(resp.Body, 2<<20))
body, _ := s.readUpstreamErrorBody(resp)
statusCode := resp.StatusCode
// OAuth/Setup Token 账号的 403标记账号异常
@ -7288,7 +7416,7 @@ func (s *GatewayService) handleRetryExhaustedSideEffects(ctx context.Context, re
}
func (s *GatewayService) handleFailoverSideEffects(ctx context.Context, resp *http.Response, account *Account, requestedModel ...string) {
body, _ := io.ReadAll(io.LimitReader(resp.Body, 2<<20))
body, _ := s.readUpstreamErrorBody(resp)
if len(requestedModel) > 0 {
s.rateLimitService.HandleUpstreamError(ctx, account, resp.StatusCode, resp.Header, body, requestedModel[0])
return
@ -7301,7 +7429,7 @@ func (s *GatewayService) handleFailoverSideEffects(ctx context.Context, resp *ht
// API Key 未配置错误码:仅返回错误,不标记账号
func (s *GatewayService) handleRetryExhaustedError(ctx context.Context, resp *http.Response, c *gin.Context, account *Account) (*ForwardResult, error) {
// Capture upstream error body before side-effects consume the stream.
respBody, _ := io.ReadAll(io.LimitReader(resp.Body, 2<<20))
respBody, _ := s.readUpstreamErrorBody(resp)
_ = resp.Body.Close()
resp.Body = io.NopCloser(bytes.NewReader(respBody))
@ -8125,10 +8253,10 @@ func (s *GatewayService) getUserGroupRateMultiplier(ctx context.Context, userID,
return resolver.Resolve(ctx, userID, groupID, groupDefaultMultiplier)
}
// RecordUsageInput 记录使用量的输入参数
// RecordUsageInput 记录使用量的输入参数。
// 异步 worker 只接收计费所需快照,不能持有 ParsedRequest/RequestBodyRef 这类大请求体引用。
type RecordUsageInput struct {
Result *ForwardResult
ParsedRequest *ParsedRequest
APIKey *APIKey
User *User
Account *Account
@ -8606,15 +8734,8 @@ func writeUsageLogBestEffort(ctx context.Context, repo UsageLogRepository, usage
}
}
// recordUsageOpts 内部选项,参数化 RecordUsage 与 RecordUsageWithLongContext 的差异点。
// recordUsageOpts 内部选项,参数化普通计费与长上下文计费的差异点。
type recordUsageOpts struct {
// Claude Max 策略所需的 ParsedRequest可选仅 Claude 路径传入)
ParsedRequest *ParsedRequest
// EnableClaudePath 启用 Claude 路径特有逻辑:
// - Claude Max 缓存计费策略
EnableClaudePath bool
// 长上下文计费(仅 Gemini 路径需要)
LongContextThreshold int
LongContextMultiplier float64
@ -8637,9 +8758,7 @@ func (s *GatewayService) RecordUsage(ctx context.Context, input *RecordUsageInpu
APIKeyService: input.APIKeyService,
QuotaPlatform: input.QuotaPlatform,
ChannelUsageFields: input.ChannelUsageFields,
}, &recordUsageOpts{
EnableClaudePath: true,
})
}, &recordUsageOpts{})
}
// RecordUsageLongContextInput 记录使用量的输入参数(支持长上下文双倍计费)
@ -8705,9 +8824,7 @@ type recordUsageCoreInput struct {
}
// recordUsageCore 是 RecordUsage 和 RecordUsageWithLongContext 的统一实现。
// opts 中的字段控制两者之间的差异行为:
// - ParsedRequest != nil → 启用 Claude Max 缓存计费策略
// - LongContextThreshold > 0 → Token 计费回退走 CalculateCostWithLongContext
// LongContextThreshold > 0 时 Token 计费回退走 CalculateCostWithLongContext。
func (s *GatewayService) recordUsageCore(ctx context.Context, input *recordUsageCoreInput, opts *recordUsageOpts) error {
result := input.Result
apiKey := input.APIKey
@ -9172,7 +9289,7 @@ func (s *GatewayService) ForwardCountTokens(ctx context.Context, c *gin.Context,
}
if account != nil && account.IsAnthropicAPIKeyPassthroughEnabled() {
passthroughBody := parsed.Body
passthroughBody := parsed.Body.Bytes()
if reqModel := parsed.Model; reqModel != "" {
if mappedModel := account.GetMappedModel(reqModel); mappedModel != reqModel {
passthroughBody = s.replaceModelInBody(passthroughBody, mappedModel)
@ -9188,24 +9305,43 @@ func (s *GatewayService) ForwardCountTokens(ctx context.Context, c *gin.Context,
return nil
}
body := parsed.Body
body := parsed.Body.Bytes()
replaceBody := func(next []byte) error {
if err := parsed.ReplaceBody(next); err != nil {
return fmt.Errorf("rewrite count_tokens body: %w", err)
}
body = parsed.Body.Bytes()
return nil
}
reqModel := parsed.Model
// Pre-filter: strip empty text blocks to prevent upstream 400.
body = StripEmptyTextBlocks(body)
if err := replaceBody(StripEmptyTextBlocks(body)); err != nil {
return err
}
isClaudeCodeCT := IsClaudeCodeClient(ctx) || isClaudeCodeClient(c.GetHeader("User-Agent"), parsed.MetadataUserID)
shouldMimicClaudeCode := account.IsOAuth() && !isClaudeCodeCT
if shouldMimicClaudeCode {
normalizeOpts := claudeOAuthNormalizeOptions{stripSystemCacheControl: true}
body, reqModel = normalizeClaudeOAuthRequestBody(body, reqModel, normalizeOpts)
var normalizedBody []byte
normalizedBody, reqModel = normalizeClaudeOAuthRequestBody(body, reqModel, normalizeOpts)
if err := replaceBody(normalizedBody); err != nil {
return err
}
body = s.rewriteMessageCacheControlIfEnabled(ctx, body)
if err := replaceBody(s.rewriteMessageCacheControlIfEnabled(ctx, body)); err != nil {
return err
}
if rw := buildToolNameRewriteFromBody(body); rw != nil {
body = applyToolNameRewriteToBody(body, rw)
if err := replaceBody(applyToolNameRewriteToBody(body, rw)); err != nil {
return err
}
} else {
body = applyToolsLastCacheBreakpoint(body)
if err := replaceBody(applyToolsLastCacheBreakpoint(body)); err != nil {
return err
}
}
}
@ -9236,9 +9372,13 @@ func (s *GatewayService) ForwardCountTokens(ctx context.Context, c *gin.Context,
}
}
if mappedModel != reqModel {
body = s.replaceModelInBody(body, mappedModel)
originalReqModel := reqModel
if err := replaceBody(s.replaceModelInBody(body, mappedModel)); err != nil {
return err
}
reqModel = mappedModel
logger.LegacyPrintf("service.gateway", "CountTokens model mapping applied: %s -> %s (account: %s, source=%s)", parsed.Model, mappedModel, account.Name, mappingSource)
parsed.Model = mappedModel
logger.LegacyPrintf("service.gateway", "CountTokens model mapping applied: %s -> %s (account: %s, source=%s)", originalReqModel, mappedModel, account.Name, mappingSource)
}
}
@ -9250,11 +9390,13 @@ func (s *GatewayService) ForwardCountTokens(ctx context.Context, c *gin.Context,
}
// 构建上游请求
upstreamReq, err := s.buildCountTokensRequest(ctx, c, account, body, token, tokenType, reqModel, shouldMimicClaudeCode)
upstreamReq, wireBody, err := s.buildCountTokensRequest(ctx, c, account, body, token, tokenType, reqModel, shouldMimicClaudeCode)
if err != nil {
s.countTokensError(c, http.StatusInternalServerError, "api_error", "Failed to build request")
return err
}
// 先记录首发 wire body如果后面进入 400 retryretry 会基于未签名的逻辑 body 重新构建。
acceptedWireBody := wireBody
// 获取代理URL自定义 base URL 模式下proxy 通过 buildCustomRelayURL 作为查询参数传递)
proxyURL := ""
@ -9290,10 +9432,14 @@ func (s *GatewayService) ForwardCountTokens(ctx context.Context, c *gin.Context,
logger.LegacyPrintf("service.gateway", "Account %d: detected thinking block signature error on count_tokens, retrying with filtered thinking blocks", account.ID)
filteredBody := FilterThinkingBlocksForRetry(body)
retryReq, buildErr := s.buildCountTokensRequest(ctx, c, account, filteredBody, token, tokenType, reqModel, shouldMimicClaudeCode)
retryReq, retryWireBody, buildErr := s.buildCountTokensRequest(ctx, c, account, filteredBody, token, tokenType, reqModel, shouldMimicClaudeCode)
if buildErr == nil {
retryResp, retryErr := s.httpUpstream.DoWithTLS(retryReq, proxyURL, account.ID, account.Concurrency, s.tlsFPProfileService.ResolveTLSProfile(account))
if retryErr == nil {
if retryResp.StatusCode < 400 {
// count_tokens 签名重试成功后记录最终 wire body错误响应仍保留原 body 便于后续处理。
acceptedWireBody = retryWireBody
}
resp = retryResp
respBody, err = ReadUpstreamResponseBody(resp.Body, s.cfg, c, countTokensTooLarge)
_ = resp.Body.Close()
@ -9307,6 +9453,13 @@ func (s *GatewayService) ForwardCountTokens(ctx context.Context, c *gin.Context,
}
}
if resp.StatusCode < 400 && !bytes.Equal(acceptedWireBody, body) {
// count_tokens 成功后再同步最终 wire body避免 retry 从已签名 body 派生。
if err := replaceBody(acceptedWireBody); err != nil {
return err
}
}
// 处理错误响应
if resp.StatusCode >= 400 {
// 标记账号状态429/529等
@ -9533,7 +9686,7 @@ func (s *GatewayService) buildCountTokensRequestAnthropicAPIKeyPassthrough(
}
// buildCountTokensRequest 构建 count_tokens 上游请求
func (s *GatewayService) buildCountTokensRequest(ctx context.Context, c *gin.Context, account *Account, body []byte, token, tokenType, modelID string, mimicClaudeCode bool) (*http.Request, error) {
func (s *GatewayService) buildCountTokensRequest(ctx context.Context, c *gin.Context, account *Account, body []byte, token, tokenType, modelID string, mimicClaudeCode bool) (*http.Request, []byte, error) {
// 确定目标 URL
targetURL := claudeAPICountTokensURL
if account.Type == AccountTypeAPIKey {
@ -9541,18 +9694,18 @@ func (s *GatewayService) buildCountTokensRequest(ctx context.Context, c *gin.Con
if baseURL != "" {
validatedURL, err := s.validateUpstreamBaseURL(baseURL)
if err != nil {
return nil, err
return nil, nil, err
}
targetURL = validatedURL + "/v1/messages/count_tokens?beta=true"
}
} else if account.IsCustomBaseURLEnabled() {
customURL := account.GetCustomBaseURL()
if customURL == "" {
return nil, fmt.Errorf("custom_base_url is enabled but not configured for account %d", account.ID)
return nil, nil, fmt.Errorf("custom_base_url is enabled but not configured for account %d", account.ID)
}
validatedURL, err := s.validateUpstreamBaseURL(customURL)
if err != nil {
return nil, err
return nil, nil, err
}
targetURL = s.buildCustomRelayURL(validatedURL, "/v1/messages/count_tokens", account)
}
@ -9608,7 +9761,7 @@ func (s *GatewayService) buildCountTokensRequest(ctx context.Context, c *gin.Con
req, err := http.NewRequestWithContext(ctx, "POST", targetURL, bytes.NewReader(body))
if err != nil {
return nil, err
return nil, nil, err
}
// 设置认证头(保持原始大小写)
@ -9672,7 +9825,7 @@ func (s *GatewayService) buildCountTokensRequest(ctx context.Context, c *gin.Con
logClaudeMimicDebug(req, body, account, tokenType, mimicClaudeCode)
}
return req, nil
return req, body, nil
}
func sanitizeCountTokensRequestBody(body []byte) []byte {

View File

@ -2,10 +2,16 @@ package service
import (
"strconv"
"strings"
"testing"
"github.com/Wei-Shaw/sub2api/internal/domain"
)
var benchmarkStringSink string
var (
benchmarkStringSink string
benchmarkIntSink int
)
// BenchmarkGenerateSessionHash_Metadata 关注 JSON 解析与正则匹配开销。
func BenchmarkGenerateSessionHash_Metadata(b *testing.B) {
@ -14,7 +20,7 @@ func BenchmarkGenerateSessionHash_Metadata(b *testing.B) {
b.ReportAllocs()
for i := 0; i < b.N; i++ {
parsed, err := ParseGatewayRequest(body, "")
parsed, err := ParseGatewayRequest(NewRequestBodyRef(body), "")
if err != nil {
b.Fatalf("解析请求失败: %v", err)
}
@ -22,6 +28,179 @@ func BenchmarkGenerateSessionHash_Metadata(b *testing.B) {
}
}
func BenchmarkParseGatewayRequest_LargeAnthropicMessages(b *testing.B) {
for _, size := range benchmarkBodySizes() {
b.Run(size.name, func(b *testing.B) {
body := buildLargeAnthropicMessagesBody(size.bytes, false)
b.SetBytes(int64(len(body)))
b.ReportAllocs()
b.ResetTimer()
for i := 0; i < b.N; i++ {
parsed, err := ParseGatewayRequest(NewRequestBodyRef(body), domain.PlatformAnthropic)
if err != nil {
b.Fatalf("解析 Anthropic 请求失败: %v", err)
}
benchmarkIntSink = len(parsed.MessagesRaw())
}
})
}
}
func BenchmarkParseGatewayRequest_LargeGeminiContents(b *testing.B) {
for _, size := range benchmarkBodySizes() {
b.Run(size.name, func(b *testing.B) {
body := buildLargeGeminiContentsBody(size.bytes)
b.SetBytes(int64(len(body)))
b.ReportAllocs()
b.ResetTimer()
for i := 0; i < b.N; i++ {
parsed, err := ParseGatewayRequest(NewRequestBodyRef(body), domain.PlatformGemini)
if err != nil {
b.Fatalf("解析 Gemini 请求失败: %v", err)
}
benchmarkIntSink = len(parsed.MessagesRaw())
}
})
}
}
func BenchmarkGenerateSessionHash_LargeAnthropicMessages(b *testing.B) {
svc := &GatewayService{}
for _, size := range benchmarkBodySizes() {
b.Run(size.name, func(b *testing.B) {
body := buildLargeAnthropicMessagesBody(size.bytes, true)
parsed, err := ParseGatewayRequest(NewRequestBodyRef(body), domain.PlatformAnthropic)
if err != nil {
b.Fatalf("解析请求失败: %v", err)
}
b.SetBytes(int64(len(body)))
b.ReportAllocs()
b.ResetTimer()
for i := 0; i < b.N; i++ {
benchmarkStringSink = svc.GenerateSessionHash(parsed)
}
})
}
}
func BenchmarkOpenAIResponses_LargeInputMeta(b *testing.B) {
for _, size := range benchmarkBodySizes() {
b.Run(size.name, func(b *testing.B) {
body := buildLargeOpenAIResponsesBody(size.bytes)
b.SetBytes(int64(len(body)))
b.ReportAllocs()
b.ResetTimer()
for i := 0; i < b.N; i++ {
model, stream, promptCacheKey := extractOpenAIRequestMetaFromBody(body)
benchmarkStringSink = model + promptCacheKey
if stream {
benchmarkIntSink++
}
}
})
}
}
func BenchmarkOpenAIResponses_LargeInputDecodeMap(b *testing.B) {
for _, size := range benchmarkBodySizes() {
b.Run(size.name, func(b *testing.B) {
body := buildLargeOpenAIResponsesBody(size.bytes)
b.SetBytes(int64(len(body)))
b.ReportAllocs()
b.ResetTimer()
for i := 0; i < b.N; i++ {
reqBody, err := getOpenAIRequestBodyMap(nil, body)
if err != nil {
b.Fatalf("解析 OpenAI 请求失败: %v", err)
}
benchmarkIntSink = len(reqBody)
}
})
}
}
func BenchmarkOpenAIResponses_LargeInputRawPatch(b *testing.B) {
for _, size := range benchmarkBodySizes() {
b.Run(size.name, func(b *testing.B) {
body := buildLargeOpenAIResponsesBody(size.bytes)
b.SetBytes(int64(len(body)))
b.ReportAllocs()
b.ResetTimer()
for i := 0; i < b.N; i++ {
view := newOpenAIRequestView(body)
view.MarkPatchSet("instructions", "You are a helpful coding assistant.")
view.MarkPatchSet("reasoning.effort", "none")
patched, err := view.ApplyPatches()
if err != nil {
b.Fatalf("应用 OpenAI raw patch 失败: %v", err)
}
benchmarkIntSink = len(patched)
}
})
}
}
func BenchmarkOpenAIResponses_LargeInputImageBillingRaw(b *testing.B) {
for _, size := range benchmarkBodySizes() {
b.Run(size.name, func(b *testing.B) {
body := buildLargeOpenAIResponsesImageToolBody(size.bytes)
b.SetBytes(int64(len(body)))
b.ReportAllocs()
b.ResetTimer()
for i := 0; i < b.N; i++ {
cfg, err := resolveOpenAIResponsesImageBillingConfigDetailedFromBody(body, "gpt-5.4")
if err != nil {
b.Fatalf("解析 OpenAI 图片计费配置失败: %v", err)
}
benchmarkStringSink = cfg.Model + cfg.SizeTier + cfg.InputSize
}
})
}
}
func BenchmarkOpenAIResponses_LargeInputEmptyBase64Guard(b *testing.B) {
for _, size := range benchmarkBodySizes() {
b.Run(size.name, func(b *testing.B) {
body := buildLargeOpenAIResponsesBody(size.bytes)
b.SetBytes(int64(len(body)))
b.ReportAllocs()
b.ResetTimer()
for i := 0; i < b.N; i++ {
if openAIRequestBodyMayContainEmptyBase64InputImage(body) {
benchmarkIntSink++
}
}
})
}
}
func BenchmarkOpenAIResponses_LargeInputFunctionCallValidation(b *testing.B) {
for _, size := range benchmarkBodySizes() {
b.Run(size.name, func(b *testing.B) {
body := buildLargeOpenAIResponsesToolContinuationBody(size.bytes)
b.SetBytes(int64(len(body)))
b.ReportAllocs()
b.ResetTimer()
for i := 0; i < b.N; i++ {
validation := ValidateFunctionCallOutputContextBytes(body)
if !validation.HasFunctionCallOutput || !validation.HasItemReferenceForAllCallIDs {
b.Fatalf("工具续链校验结果异常: %+v", validation)
}
benchmarkIntSink++
}
})
}
}
// BenchmarkExtractCacheableContent_System 关注字符串拼接路径的性能。
func BenchmarkExtractCacheableContent_System(b *testing.B) {
svc := &GatewayService{}
@ -33,18 +212,130 @@ func BenchmarkExtractCacheableContent_System(b *testing.B) {
}
}
func buildSystemCacheableRequest(parts int) *ParsedRequest {
systemParts := make([]any, 0, parts)
for i := 0; i < parts; i++ {
systemParts = append(systemParts, map[string]any{
"text": "system_part_" + strconv.Itoa(i),
"cache_control": map[string]any{
"type": "ephemeral",
},
})
}
return &ParsedRequest{
System: systemParts,
HasSystem: true,
func benchmarkBodySizes() []struct {
name string
bytes int
} {
return []struct {
name string
bytes int
}{
{name: "4MB", bytes: 4 << 20},
{name: "8MB", bytes: 8 << 20},
{name: "16MB", bytes: 16 << 20},
{name: "32MB", bytes: 32 << 20},
}
}
func buildSystemCacheableRequest(parts int) *ParsedRequest {
var builder strings.Builder
_, _ = builder.WriteString(`{"system":[`)
for i := 0; i < parts; i++ {
if i > 0 {
_ = builder.WriteByte(',')
}
_, _ = builder.WriteString(`{"text":"system_part_`)
_, _ = builder.WriteString(strconv.Itoa(i))
_, _ = builder.WriteString(`","cache_control":{"type":"ephemeral"}}`)
}
_, _ = builder.WriteString(`]}`)
parsed, err := ParseGatewayRequest(NewRequestBodyRef([]byte(builder.String())), "")
if err != nil {
panic(err)
}
return parsed
}
func buildLargeAnthropicMessagesBody(targetBytes int, includeCacheControl bool) []byte {
var builder strings.Builder
builder.Grow(targetBytes + 1024)
_, _ = builder.WriteString(`{"model":"claude-sonnet-4-5","stream":true,"system":[{"type":"text","text":"system seed"}],"messages":[`)
for i := 0; builder.Len() < targetBytes; i++ {
if i > 0 {
_ = builder.WriteByte(',')
}
_, _ = builder.WriteString(`{"role":"user","content":[{"type":"text","text":"`)
_, _ = builder.WriteString(strings.Repeat("anthropic payload ", 64))
_, _ = builder.WriteString(strconv.Itoa(i))
_ = builder.WriteByte('"')
if includeCacheControl && i%32 == 0 {
_, _ = builder.WriteString(`,"cache_control":{"type":"ephemeral"}`)
}
_, _ = builder.WriteString(`}]}`)
}
_, _ = builder.WriteString(`]}`)
return []byte(builder.String())
}
func buildLargeGeminiContentsBody(targetBytes int) []byte {
var builder strings.Builder
builder.Grow(targetBytes + 1024)
_, _ = builder.WriteString(`{"model":"gemini-2.5-pro","systemInstruction":{"parts":[{"text":"system seed"}]},"contents":[`)
for i := 0; builder.Len() < targetBytes; i++ {
if i > 0 {
_ = builder.WriteByte(',')
}
_, _ = builder.WriteString(`{"role":"user","parts":[{"text":"`)
_, _ = builder.WriteString(strings.Repeat("gemini payload ", 64))
_, _ = builder.WriteString(strconv.Itoa(i))
_, _ = builder.WriteString(`"}]}`)
}
_, _ = builder.WriteString(`]}`)
return []byte(builder.String())
}
func buildLargeOpenAIResponsesBody(targetBytes int) []byte {
var builder strings.Builder
builder.Grow(targetBytes + 1024)
_, _ = builder.WriteString(`{"model":"gpt-5.4","stream":true,"prompt_cache_key":"session-benchmark","input":[`)
for i := 0; builder.Len() < targetBytes; i++ {
if i > 0 {
_ = builder.WriteByte(',')
}
_, _ = builder.WriteString(`{"type":"message","role":"user","content":[{"type":"input_text","text":"`)
_, _ = builder.WriteString(strings.Repeat("openai responses payload ", 48))
_, _ = builder.WriteString(strconv.Itoa(i))
_, _ = builder.WriteString(`"}]}`)
}
_, _ = builder.WriteString(`],"tools":[{"type":"function","name":"lookup","parameters":{"type":"object","properties":{"query":{"type":"string"}}}}]}`)
return []byte(builder.String())
}
func buildLargeOpenAIResponsesToolContinuationBody(targetBytes int) []byte {
var builder strings.Builder
builder.Grow(targetBytes + 1024)
_, _ = builder.WriteString(`{"model":"gpt-5.4","stream":true,"previous_response_id":"resp_benchmark","input":[`)
for i := 0; builder.Len() < targetBytes; i++ {
if i > 0 {
_ = builder.WriteByte(',')
}
callID := "call_" + strconv.Itoa(i)
_, _ = builder.WriteString(`{"type":"item_reference","id":"`)
_, _ = builder.WriteString(callID)
_, _ = builder.WriteString(`"},{"type":"function_call_output","call_id":"`)
_, _ = builder.WriteString(callID)
_, _ = builder.WriteString(`","output":"`)
_, _ = builder.WriteString(strings.Repeat("tool output payload ", 48))
_, _ = builder.WriteString(strconv.Itoa(i))
_, _ = builder.WriteString(`"}`)
}
_, _ = builder.WriteString(`]}`)
return []byte(builder.String())
}
func buildLargeOpenAIResponsesImageToolBody(targetBytes int) []byte {
var builder strings.Builder
builder.Grow(targetBytes + 1024)
_, _ = builder.WriteString(`{"model":"gpt-5.4","stream":false,"tools":[{"type":"image_generation","model":"gpt-image-2","size":"2048x1152"}],"input":[`)
for i := 0; builder.Len() < targetBytes; i++ {
if i > 0 {
_ = builder.WriteByte(',')
}
_, _ = builder.WriteString(`{"type":"message","role":"user","content":[{"type":"input_text","text":"`)
_, _ = builder.WriteString(strings.Repeat("openai image billing payload ", 48))
_, _ = builder.WriteString(strconv.Itoa(i))
_, _ = builder.WriteString(`"}]}`)
}
_, _ = builder.WriteString(`]}`)
return []byte(builder.String())
}

View File

@ -150,7 +150,7 @@ func (s *GatewayService) handleWebSearchEmulation(
parsed.OnUpstreamAccepted()
}
query := extractSearchQueryFromBody(parsed.Body)
query := extractSearchQueryFromBody(parsed.Body.Bytes())
if query == "" {
return nil, fmt.Errorf("web search emulation: no query found in messages")
}

View File

@ -151,7 +151,7 @@ func (s *GeminiMessagesCompatService) forwardClaudeBodyAsChatCompletions(
}
if resp.StatusCode >= 400 && s.shouldRetryGeminiUpstreamError(account, resp.StatusCode) {
respBody, _ := io.ReadAll(io.LimitReader(resp.Body, 2<<20))
respBody := s.readUpstreamErrorBody(resp)
_ = resp.Body.Close()
if resp.StatusCode == http.StatusForbidden && isGeminiInsufficientScope(resp.Header, respBody) {
resp = &http.Response{
@ -207,7 +207,7 @@ func (s *GeminiMessagesCompatService) forwardClaudeBodyAsChatCompletions(
reasoningEffort := extractCCReasoningEffortFromBody(originalChatBody)
if resp.StatusCode >= 400 {
respBody, _ := io.ReadAll(io.LimitReader(resp.Body, 2<<20))
respBody := s.readUpstreamErrorBody(resp)
s.handleGeminiUpstreamError(ctx, account, resp.StatusCode, resp.Header, respBody)
evBody := unwrapIfNeeded(account.Type == AccountTypeOAuth, respBody)

View File

@ -56,6 +56,18 @@ type GeminiMessagesCompatService struct {
responseHeaderFilter *responseheaders.CompiledHeaderFilter
}
func (s *GeminiMessagesCompatService) readUpstreamErrorBody(resp *http.Response) []byte {
if resp == nil || resp.Body == nil {
return nil
}
limit := gatewayUpstreamErrorBodyReadLimit
if s != nil && s.cfg != nil && s.cfg.Gateway.LogUpstreamErrorBody && s.cfg.Gateway.LogUpstreamErrorBodyMaxBytes > int(limit) {
limit = int64(s.cfg.Gateway.LogUpstreamErrorBodyMaxBytes)
}
body, _ := io.ReadAll(io.LimitReader(resp.Body, limit))
return body
}
func NewGeminiMessagesCompatService(
accountRepo AccountRepository,
groupRepo GroupRepository,
@ -789,7 +801,7 @@ func (s *GeminiMessagesCompatService) Forward(ctx context.Context, c *gin.Contex
// Special-case: signature/thought_signature validation errors are not transient, but may be fixed by
// downgrading Claude thinking/tool history to plain text (conservative two-stage retry).
if resp.StatusCode == http.StatusBadRequest && signatureRetryStage < 2 {
respBody, _ := io.ReadAll(io.LimitReader(resp.Body, 2<<20))
respBody := s.readUpstreamErrorBody(resp)
_ = resp.Body.Close()
if isGeminiSignatureRelatedError(respBody) {
@ -860,7 +872,7 @@ func (s *GeminiMessagesCompatService) Forward(ctx context.Context, c *gin.Contex
}
if resp.StatusCode >= 400 && s.shouldRetryGeminiUpstreamError(account, resp.StatusCode) {
respBody, _ := io.ReadAll(io.LimitReader(resp.Body, 2<<20))
respBody := s.readUpstreamErrorBody(resp)
_ = resp.Body.Close()
// Don't treat insufficient-scope as transient.
if resp.StatusCode == 403 && isGeminiInsufficientScope(resp.Header, respBody) {
@ -919,7 +931,7 @@ func (s *GeminiMessagesCompatService) Forward(ctx context.Context, c *gin.Contex
defer func() { _ = resp.Body.Close() }()
if resp.StatusCode >= 400 {
respBody, _ := io.ReadAll(io.LimitReader(resp.Body, 2<<20))
respBody := s.readUpstreamErrorBody(resp)
// 统一错误策略:自定义错误码 + 临时不可调度
if s.rateLimitService != nil {
switch s.rateLimitService.CheckErrorPolicy(ctx, account, resp.StatusCode, respBody) {
@ -1329,7 +1341,7 @@ func (s *GeminiMessagesCompatService) ForwardNative(ctx context.Context, c *gin.
}
if resp.StatusCode >= 400 && s.shouldRetryGeminiUpstreamError(account, resp.StatusCode) {
respBody, _ := io.ReadAll(io.LimitReader(resp.Body, 2<<20))
respBody := s.readUpstreamErrorBody(resp)
_ = resp.Body.Close()
// Don't treat insufficient-scope as transient.
if resp.StatusCode == 403 && isGeminiInsufficientScope(resp.Header, respBody) {
@ -1410,7 +1422,7 @@ func (s *GeminiMessagesCompatService) ForwardNative(ctx context.Context, c *gin.
isOAuth := account.Type == AccountTypeOAuth
if resp.StatusCode >= 400 {
respBody, _ := io.ReadAll(io.LimitReader(resp.Body, 2<<20))
respBody := s.readUpstreamErrorBody(resp)
// Best-effort fallback for OAuth tokens missing AI Studio scopes when calling countTokens.
// This avoids Gemini SDKs failing hard during preflight token counting.
// Checked before error policy so it always works regardless of custom error codes.
@ -1619,7 +1631,7 @@ func (s *GeminiMessagesCompatService) checkErrorPolicyInLoop(
if resp.StatusCode < 400 || s.rateLimitService == nil {
return false, resp
}
body, _ := io.ReadAll(io.LimitReader(resp.Body, 2<<20))
body := s.readUpstreamErrorBody(resp)
_ = resp.Body.Close()
rebuilt = &http.Response{
StatusCode: resp.StatusCode,

File diff suppressed because it is too large Load Diff

View File

@ -1,7 +1,6 @@
package service
import (
"encoding/json"
"strings"
"github.com/tidwall/gjson"
@ -91,7 +90,7 @@ func openAIJSONToolsContainImageGeneration(tools gjson.Result) bool {
}
found := false
tools.ForEach(func(_, item gjson.Result) bool {
if strings.TrimSpace(item.Get("type").String()) == "image_generation" {
if openAIJSONString(item.Get("type")) == "image_generation" {
found = true
return false
}
@ -100,6 +99,36 @@ func openAIJSONToolsContainImageGeneration(tools gjson.Result) bool {
return found
}
func openAIRequestBodyHasImageGenerationTool(body []byte) bool {
if len(body) == 0 || !gjson.ValidBytes(body) {
return false
}
return openAIJSONToolsContainImageGeneration(gjson.GetBytes(body, "tools"))
}
func openAIRequestBodyImageGenerationToolNeedsNormalization(body []byte) bool {
if len(body) == 0 || !gjson.ValidBytes(body) {
return false
}
tools := gjson.GetBytes(body, "tools")
if !tools.IsArray() {
return false
}
needsNormalization := false
tools.ForEach(func(_, item gjson.Result) bool {
if openAIJSONString(item.Get("type")) != "image_generation" {
return true
}
// 只有旧字段需要迁移时才进入 map 修改,纯计费读取保持 raw 路径。
if item.Get("format").Exists() || item.Get("compression").Exists() {
needsNormalization = true
return false
}
return true
})
return needsNormalization
}
func openAIJSONToolChoiceSelectsImageGeneration(choice gjson.Result) bool {
if !choice.Exists() {
return false
@ -159,17 +188,6 @@ func apiKeyGroup(apiKey *APIKey) *Group {
return apiKey.Group
}
func cloneRequestMapForImageIntent(body []byte) map[string]any {
if len(body) == 0 {
return nil
}
var out map[string]any
if err := json.Unmarshal(body, &out); err != nil {
return nil
}
return out
}
type OpenAIResponsesImageBillingConfig struct {
Model string
SizeTier string
@ -225,8 +243,43 @@ func resolveOpenAIResponsesImageBillingConfigFromBody(body []byte, fallbackModel
}
func resolveOpenAIResponsesImageBillingConfigDetailedFromBody(body []byte, fallbackModel string) (OpenAIResponsesImageBillingConfig, error) {
reqBody := cloneRequestMapForImageIntent(body)
return resolveOpenAIResponsesImageBillingConfigDetailed(reqBody, fallbackModel)
imageModel := ""
imageSize := ""
hasImageTool := false
if len(body) > 0 && gjson.ValidBytes(body) {
tools := gjson.GetBytes(body, "tools")
if tools.IsArray() {
tools.ForEach(func(_, item gjson.Result) bool {
if openAIJSONString(item.Get("type")) != "image_generation" {
return true
}
hasImageTool = true
imageModel = openAIJSONString(item.Get("model"))
imageSize = openAIJSONString(item.Get("size"))
return false
})
}
if imageSize == "" {
imageSize = openAIJSONString(gjson.GetBytes(body, "size"))
}
if imageModel == "" {
bodyModel := openAIJSONString(gjson.GetBytes(body, "model"))
if isOpenAIImageBillingModelAlias(bodyModel) || !hasImageTool {
imageModel = bodyModel
}
}
}
if imageModel == "" && hasImageTool {
imageModel = "gpt-image-2"
}
if imageModel == "" {
imageModel = strings.TrimSpace(fallbackModel)
}
return OpenAIResponsesImageBillingConfig{
Model: imageModel,
SizeTier: normalizeOpenAIImageSizeTier(imageSize),
InputSize: imageSize,
}, nil
}
func isOpenAIImageBillingModelAlias(model string) bool {
@ -236,3 +289,10 @@ func isOpenAIImageBillingModelAlias(model string) bool {
}
return isOpenAIImageGenerationModel(normalized) || strings.Contains(normalized, "image")
}
func openAIJSONString(value gjson.Result) string {
if value.Type != gjson.String {
return ""
}
return strings.TrimSpace(value.String())
}

View File

@ -84,6 +84,17 @@ func TestResolveOpenAIResponsesImageBillingConfigToolModelWins(t *testing.T) {
require.Equal(t, "2K", imageSize)
}
func TestResolveOpenAIResponsesImageBillingConfigFromBodyIgnoresUnrelatedLargeInput(t *testing.T) {
cfg, err := resolveOpenAIResponsesImageBillingConfigDetailedFromBody(
[]byte(`{"model":"mapped-text-model","tools":[{"type":"image_generation","model":"gpt-image-2","size":"2048x1152"}],"input":[{"type":"message","content":[{"type":"input_text","text":"hi","nonce":1e1000000}]}]}`),
"requested-model",
)
require.NoError(t, err)
require.Equal(t, "gpt-image-2", cfg.Model)
require.Equal(t, "2K", cfg.SizeTier)
require.Equal(t, "2048x1152", cfg.InputSize)
}
func TestResolveOpenAIResponsesImageBillingConfigSupportsOfficialAndCustomSizes(t *testing.T) {
tests := []struct {
name string

View File

@ -104,7 +104,7 @@ func (s *OpenAIGatewayService) ForwardEmbeddings(
defer func() { _ = resp.Body.Close() }()
if resp.StatusCode >= 400 {
respBody, _ := io.ReadAll(io.LimitReader(resp.Body, 2<<20))
respBody := s.readUpstreamErrorBody(resp)
_ = resp.Body.Close()
resp.Body = io.NopCloser(bytes.NewReader(respBody))

View File

@ -241,7 +241,7 @@ func (s *OpenAIGatewayService) ForwardAsChatCompletions(
// 8. Handle error response with failover
if resp.StatusCode >= 400 {
respBody, _ := io.ReadAll(io.LimitReader(resp.Body, 2<<20))
respBody := s.readUpstreamErrorBody(resp)
_ = resp.Body.Close()
resp.Body = io.NopCloser(bytes.NewReader(respBody))

View File

@ -183,7 +183,7 @@ func (s *OpenAIGatewayService) forwardAsRawChatCompletions(
// 7. Handle error response with failover
if resp.StatusCode >= 400 {
respBody, _ := io.ReadAll(io.LimitReader(resp.Body, 2<<20))
respBody := s.readUpstreamErrorBody(resp)
_ = resp.Body.Close()
resp.Body = io.NopCloser(bytes.NewReader(respBody))

View File

@ -300,7 +300,7 @@ func (s *OpenAIGatewayService) ForwardAsAnthropic(
// 8. Handle error response with failover
if resp.StatusCode >= 400 {
respBody, _ := io.ReadAll(io.LimitReader(resp.Body, 2<<20))
respBody := s.readUpstreamErrorBody(resp)
_ = resp.Body.Close()
resp.Body = io.NopCloser(bytes.NewReader(respBody))

View File

@ -163,7 +163,7 @@ func (s *OpenAIGatewayService) forwardResponsesViaRawChatCompletions(
defer func() { _ = resp.Body.Close() }()
if resp.StatusCode >= 400 {
respBody, _ := io.ReadAll(io.LimitReader(resp.Body, 2<<20))
respBody := s.readUpstreamErrorBody(resp)
_ = resp.Body.Close()
resp.Body = io.NopCloser(bytes.NewReader(respBody))

File diff suppressed because it is too large Load Diff

View File

@ -1,14 +1,533 @@
package service
import (
"context"
"encoding/json"
"io"
"net/http"
"net/http/httptest"
"strings"
"testing"
"github.com/Wei-Shaw/sub2api/internal/config"
"github.com/gin-gonic/gin"
"github.com/stretchr/testify/require"
"github.com/tidwall/gjson"
)
func TestOpenAIRequestView_ExtractsRawScalars(t *testing.T) {
view := newOpenAIRequestView([]byte(`{"model":" gpt-5 ","stream":true,"prompt_cache_key":" ses-1 ","previous_response_id":" resp-1 ","service_tier":" fast ","reasoning":{"effort":" medium "}}`))
require.Equal(t, "gpt-5", view.Model)
require.True(t, view.Stream)
require.Equal(t, "ses-1", view.PromptCacheKey)
require.Equal(t, "resp-1", view.PreviousResponseID)
require.Equal(t, "fast", view.ServiceTier)
require.Equal(t, "medium", view.ReasoningEffort)
}
func TestOpenAIRequestView_DecodeKeepsFullMapBehavior(t *testing.T) {
view := newOpenAIRequestView([]byte(`{"model":"gpt-5","stream":true,"input":[{"type":"message","content":"hi"}]}`))
reqBody, err := view.Decode(nil)
require.NoError(t, err)
require.Equal(t, "gpt-5", reqBody["model"])
require.IsType(t, []any{}, reqBody["input"])
}
func TestOpenAIRequestView_ApplyPatches(t *testing.T) {
view := newOpenAIRequestView([]byte(`{"model":"gpt-5","previous_response_id":"resp_1","reasoning":{"effort":"minimal"},"input":[{"type":"message","content":"hi"}]}`))
view.MarkPatchSet("model", "gpt-5.1")
view.MarkPatchDelete("previous_response_id")
view.MarkPatchSet("reasoning.effort", "none")
patched, err := view.ApplyPatches()
require.NoError(t, err)
require.JSONEq(t, `{"model":"gpt-5.1","reasoning":{"effort":"none"},"input":[{"type":"message","content":"hi"}]}`, string(patched))
}
func TestOpenAIRequestView_RejectsEscapedPatchPath(t *testing.T) {
view := newOpenAIRequestView([]byte(`{"metadata":{"user.id":"old"}}`))
view.MarkPatchSet(`metadata.user\.id`, "new")
require.False(t, view.HasPatches())
_, err := view.ApplyPatches()
require.Error(t, err)
}
func TestOpenAIRequestView_ApplyPatchesDisabled(t *testing.T) {
view := newOpenAIRequestView([]byte(`{"model":"gpt-5"}`))
view.MarkPatchSet("model", "gpt-5.1")
view.DisablePatches()
_, err := view.ApplyPatches()
require.Error(t, err)
}
func TestOpenAIRequestView_HasPatches(t *testing.T) {
view := newOpenAIRequestView([]byte(`{"model":"gpt-5"}`))
require.False(t, view.HasPatches())
view.MarkPatchSet("model", "gpt-5.1")
require.True(t, view.HasPatches())
view.DisablePatches()
require.False(t, view.HasPatches())
}
func TestOpenAIGatewayService_Forward_HTTPPatchPathKeepsLargeInputRaw(t *testing.T) {
gin.SetMode(gin.TestMode)
upstream := &httpUpstreamRecorder{
resp: &http.Response{
StatusCode: http.StatusOK,
Header: http.Header{"Content-Type": []string{"application/json"}},
Body: io.NopCloser(strings.NewReader(
`{"usage":{"input_tokens":1,"output_tokens":2,"input_tokens_details":{"cached_tokens":0}}}`,
)),
},
}
cfg := &config.Config{}
cfg.Security.URLAllowlist.Enabled = false
svc := &OpenAIGatewayService{cfg: cfg, httpUpstream: upstream}
account := &Account{
ID: 1,
Name: "openai-apikey",
Platform: PlatformOpenAI,
Type: AccountTypeAPIKey,
Concurrency: 1,
Credentials: map[string]any{
"api_key": "sk-test",
"base_url": "https://example.com",
},
Extra: map[string]any{"use_responses_api": true},
}
rec := httptest.NewRecorder()
c, _ := gin.CreateTestContext(rec)
c.Request = httptest.NewRequest(http.MethodPost, "/openai/v1/responses", nil)
SetOpenAIClientTransport(c, OpenAIClientTransportHTTP)
body := []byte(`{"model":"gpt-5","stream":false,"reasoning":{"effort":"minimal"},"input":[{"type":"message","content":[{"type":"input_text","text":"hi","nonce":9007199254740993}]}]}`)
result, err := svc.Forward(context.Background(), c, account, body)
require.NoError(t, err)
require.NotNil(t, result)
require.NotNil(t, upstream.lastReq)
require.JSONEq(t, `{"model":"gpt-5","stream":false,"reasoning":{"effort":"none"},"instructions":"You are a helpful coding assistant.","input":[{"type":"message","content":[{"type":"input_text","text":"hi","nonce":9007199254740993}]}]}`, string(upstream.lastBody))
require.Equal(t, "9007199254740993", gjson.GetBytes(upstream.lastBody, "input.0.content.0.nonce").Raw)
}
func TestOpenAIGatewayService_Forward_DecodedMutationKeepsLaterFieldDeletes(t *testing.T) {
gin.SetMode(gin.TestMode)
upstream := &httpUpstreamRecorder{
resp: &http.Response{
StatusCode: http.StatusOK,
Header: http.Header{"Content-Type": []string{"application/json"}},
Body: io.NopCloser(strings.NewReader(`{"usage":{"input_tokens":1,"output_tokens":2}}`)),
},
}
cfg := &config.Config{}
cfg.Security.URLAllowlist.Enabled = false
svc := &OpenAIGatewayService{cfg: cfg, httpUpstream: upstream}
account := &Account{
ID: 2,
Name: "openai-apikey",
Platform: PlatformOpenAI,
Type: AccountTypeAPIKey,
Concurrency: 1,
Credentials: map[string]any{
"api_key": "sk-test",
"base_url": "https://example.com",
},
Extra: map[string]any{"use_responses_api": true},
}
rec := httptest.NewRecorder()
c, _ := gin.CreateTestContext(rec)
c.Request = httptest.NewRequest(http.MethodPost, "/openai/v1/responses", nil)
SetOpenAIClientTransport(c, OpenAIClientTransportHTTP)
body := []byte(`{"model":"gpt-5.4","stream":false,"max_completion_tokens":12,"tools":[{"type":"image_generation","format":"png"}],"input":[{"type":"message","content":"draw"}]}`)
result, err := svc.Forward(context.Background(), c, account, body)
require.NoError(t, err)
require.NotNil(t, result)
require.False(t, gjson.GetBytes(upstream.lastBody, "max_completion_tokens").Exists())
require.False(t, gjson.GetBytes(upstream.lastBody, "tools.0.format").Exists())
require.Equal(t, "png", gjson.GetBytes(upstream.lastBody, "tools.0.output_format").String())
}
func TestOpenAIGatewayService_Forward_MappedImageModelUsesImageGate(t *testing.T) {
gin.SetMode(gin.TestMode)
upstream := &httpUpstreamRecorder{
resp: &http.Response{
StatusCode: http.StatusOK,
Header: http.Header{"Content-Type": []string{"application/json"}},
Body: io.NopCloser(strings.NewReader(`{"usage":{"input_tokens":1,"output_tokens":2}}`)),
},
}
cfg := &config.Config{}
cfg.Security.URLAllowlist.Enabled = false
svc := &OpenAIGatewayService{cfg: cfg, httpUpstream: upstream}
account := &Account{
ID: 3,
Name: "openai-apikey",
Platform: PlatformOpenAI,
Type: AccountTypeAPIKey,
Concurrency: 1,
Credentials: map[string]any{
"api_key": "sk-test",
"base_url": "https://example.com",
"model_mapping": map[string]any{"draw-alias": "gpt-image-2"},
},
Extra: map[string]any{"use_responses_api": true},
}
rec := httptest.NewRecorder()
c, _ := gin.CreateTestContext(rec)
c.Request = httptest.NewRequest(http.MethodPost, "/openai/v1/responses", nil)
c.Set("api_key", &APIKey{Group: &Group{AllowImageGeneration: false}})
SetOpenAIClientTransport(c, OpenAIClientTransportHTTP)
body := []byte(`{"model":"draw-alias","stream":false,"input":"draw"}`)
result, err := svc.Forward(context.Background(), c, account, body)
require.Error(t, err)
require.Nil(t, result)
require.Nil(t, upstream.lastReq)
require.Equal(t, http.StatusForbidden, rec.Code)
}
func TestOpenAIGatewayService_Forward_TextDataImageDoesNotForceMapMarshal(t *testing.T) {
gin.SetMode(gin.TestMode)
upstream := &httpUpstreamRecorder{
resp: &http.Response{
StatusCode: http.StatusOK,
Header: http.Header{"Content-Type": []string{"application/json"}},
Body: io.NopCloser(strings.NewReader(`{"usage":{"input_tokens":1,"output_tokens":2}}`)),
},
}
cfg := &config.Config{}
cfg.Security.URLAllowlist.Enabled = false
svc := &OpenAIGatewayService{cfg: cfg, httpUpstream: upstream}
account := &Account{
ID: 4,
Name: "openai-apikey",
Platform: PlatformOpenAI,
Type: AccountTypeAPIKey,
Concurrency: 1,
Credentials: map[string]any{
"api_key": "sk-test",
"base_url": "https://example.com",
},
Extra: map[string]any{"use_responses_api": true},
}
rec := httptest.NewRecorder()
c, _ := gin.CreateTestContext(rec)
c.Request = httptest.NewRequest(http.MethodPost, "/openai/v1/responses", nil)
SetOpenAIClientTransport(c, OpenAIClientTransportHTTP)
body := []byte(`{"model":"gpt-5","stream":false,"input":[{"type":"message","content":[{"type":"input_text","text":"literal data:image/png;base64, only","nonce":1e1000000}]}]}`)
result, err := svc.Forward(context.Background(), c, account, body)
require.NoError(t, err)
require.NotNil(t, result)
require.Equal(t, "1e1000000", gjson.GetBytes(upstream.lastBody, "input.0.content.0.nonce").Raw)
}
func TestOpenAIGatewayService_Forward_ImageToolBillingDoesNotForceFullDecode(t *testing.T) {
gin.SetMode(gin.TestMode)
upstream := &httpUpstreamRecorder{
resp: &http.Response{
StatusCode: http.StatusOK,
Header: http.Header{"Content-Type": []string{"application/json"}},
Body: io.NopCloser(strings.NewReader(
`{"output":[{"id":"ig_1","type":"image_generation_call","result":"final-image"}],"usage":{"input_tokens":1,"output_tokens":2}}`,
)),
},
}
cfg := &config.Config{}
cfg.Security.URLAllowlist.Enabled = false
svc := &OpenAIGatewayService{cfg: cfg, httpUpstream: upstream}
account := &Account{
ID: 9,
Name: "openai-apikey",
Platform: PlatformOpenAI,
Type: AccountTypeAPIKey,
Concurrency: 1,
Credentials: map[string]any{
"api_key": "sk-test",
"base_url": "https://example.com",
},
Extra: map[string]any{"use_responses_api": true},
}
rec := httptest.NewRecorder()
c, _ := gin.CreateTestContext(rec)
c.Request = httptest.NewRequest(http.MethodPost, "/openai/v1/responses", nil)
SetOpenAIClientTransport(c, OpenAIClientTransportHTTP)
body := []byte(`{"model":"gpt-5","stream":false,"tools":[{"type":"image_generation","model":"gpt-image-2","size":"2048x1152"}],"input":[{"type":"message","content":[{"type":"input_text","text":"draw","nonce":1e1000000}]}]}`)
result, err := svc.Forward(context.Background(), c, account, body)
require.NoError(t, err)
require.NotNil(t, result)
require.Equal(t, "1e1000000", gjson.GetBytes(upstream.lastBody, "input.0.content.0.nonce").Raw)
require.Equal(t, 1, result.ImageCount)
require.Equal(t, "2K", result.ImageSize)
require.Equal(t, "gpt-image-2", result.BillingModel)
}
func TestOpenAIGatewayService_Forward_ImageToolWithImageOnlyModelIsNormalized(t *testing.T) {
gin.SetMode(gin.TestMode)
upstream := &httpUpstreamRecorder{
resp: &http.Response{
StatusCode: http.StatusOK,
Header: http.Header{"Content-Type": []string{"application/json"}},
Body: io.NopCloser(strings.NewReader(`{"usage":{"input_tokens":1,"output_tokens":2}}`)),
},
}
cfg := &config.Config{}
cfg.Security.URLAllowlist.Enabled = false
svc := &OpenAIGatewayService{cfg: cfg, httpUpstream: upstream}
account := &Account{
ID: 11,
Name: "openai-apikey",
Platform: PlatformOpenAI,
Type: AccountTypeAPIKey,
Concurrency: 1,
Credentials: map[string]any{
"api_key": "sk-test",
"base_url": "https://example.com",
},
Extra: map[string]any{"use_responses_api": true},
}
rec := httptest.NewRecorder()
c, _ := gin.CreateTestContext(rec)
c.Request = httptest.NewRequest(http.MethodPost, "/openai/v1/responses", nil)
SetOpenAIClientTransport(c, OpenAIClientTransportHTTP)
body := []byte(`{"model":"gpt-image-2","stream":false,"tools":[{"type":"image_generation","model":"gpt-image-2"}],"input":"draw"}`)
result, err := svc.Forward(context.Background(), c, account, body)
require.NoError(t, err)
require.NotNil(t, result)
require.Equal(t, openAIImagesResponsesMainModel, gjson.GetBytes(upstream.lastBody, "model").String())
}
func TestOpenAIGatewayService_Forward_HTTPRetryRecoveryDoesNotDecodeBeforeError(t *testing.T) {
gin.SetMode(gin.TestMode)
upstream := &httpUpstreamRecorder{
responses: []*http.Response{
{
StatusCode: http.StatusBadRequest,
Header: http.Header{"Content-Type": []string{"application/json"}},
Body: io.NopCloser(strings.NewReader(`{"error":{"code":"invalid_encrypted_content","type":"invalid_request_error","message":"bad encrypted content"}}`)),
},
{
StatusCode: http.StatusOK,
Header: http.Header{"Content-Type": []string{"application/json"}},
Body: io.NopCloser(strings.NewReader(`{"usage":{"input_tokens":1,"output_tokens":2}}`)),
},
},
}
cfg := &config.Config{}
cfg.Security.URLAllowlist.Enabled = false
svc := &OpenAIGatewayService{cfg: cfg, httpUpstream: upstream}
account := &Account{
ID: 10,
Name: "openai-apikey",
Platform: PlatformOpenAI,
Type: AccountTypeAPIKey,
Concurrency: 1,
Credentials: map[string]any{
"api_key": "sk-test",
"base_url": "https://example.com",
},
Extra: map[string]any{"use_responses_api": true},
}
rec := httptest.NewRecorder()
c, _ := gin.CreateTestContext(rec)
c.Request = httptest.NewRequest(http.MethodPost, "/openai/v1/responses", nil)
SetOpenAIClientTransport(c, OpenAIClientTransportHTTP)
body := []byte(`{"model":"gpt-5","stream":false,"input":[{"type":"reasoning","encrypted_content":"gAAA","summary":[{"type":"summary_text","text":"keep me"}]},{"type":"message","content":[{"type":"input_text","text":"hi","nonce":9007199254740993}]}]}`)
result, err := svc.Forward(context.Background(), c, account, body)
require.NoError(t, err)
require.NotNil(t, result)
require.Len(t, upstream.bodies, 2)
require.Equal(t, "gAAA", gjson.GetBytes(upstream.bodies[0], "input.0.encrypted_content").String())
require.Equal(t, "9007199254740993", gjson.GetBytes(upstream.bodies[0], "input.1.content.0.nonce").Raw)
require.False(t, gjson.GetBytes(upstream.bodies[1], "input.0.encrypted_content").Exists())
require.Equal(t, "summary_text", gjson.GetBytes(upstream.bodies[1], "input.0.summary.0.type").String())
}
func TestOpenAIGatewayService_Forward_CodexSparkRejectsEscapedInputImage(t *testing.T) {
gin.SetMode(gin.TestMode)
upstream := &httpUpstreamRecorder{
resp: &http.Response{
StatusCode: http.StatusOK,
Header: http.Header{"Content-Type": []string{"application/json"}},
Body: io.NopCloser(strings.NewReader(`{"usage":{"input_tokens":1,"output_tokens":2}}`)),
},
}
cfg := &config.Config{}
cfg.Security.URLAllowlist.Enabled = false
svc := &OpenAIGatewayService{cfg: cfg, httpUpstream: upstream}
account := &Account{
ID: 5,
Name: "openai-apikey",
Platform: PlatformOpenAI,
Type: AccountTypeAPIKey,
Concurrency: 1,
Credentials: map[string]any{
"api_key": "sk-test",
"base_url": "https://example.com",
},
Extra: map[string]any{"use_responses_api": true},
}
rec := httptest.NewRecorder()
c, _ := gin.CreateTestContext(rec)
c.Request = httptest.NewRequest(http.MethodPost, "/openai/v1/responses", nil)
SetOpenAIClientTransport(c, OpenAIClientTransportHTTP)
body := []byte(`{"model":"gpt-5.3-codex-spark","stream":false,"input":[{"type":"input_` + "\\u0069" + `mage","file_id":"file_1"}]}`)
result, err := svc.Forward(context.Background(), c, account, body)
require.Error(t, err)
require.Nil(t, result)
require.Nil(t, upstream.lastReq)
require.Equal(t, http.StatusBadRequest, rec.Code)
}
func TestOpenAIGatewayService_Forward_CodexBridgeInjectionSetsImageBilling(t *testing.T) {
gin.SetMode(gin.TestMode)
upstream := &httpUpstreamRecorder{
resp: &http.Response{
StatusCode: http.StatusOK,
Header: http.Header{"Content-Type": []string{"application/json"}},
Body: io.NopCloser(strings.NewReader(
`{"output":[{"id":"ig_1","type":"image_generation_call","result":"final-image","size":"1024x1024"}],"usage":{"input_tokens":1,"output_tokens":2}}`,
)),
},
}
cfg := &config.Config{}
cfg.Security.URLAllowlist.Enabled = false
cfg.Gateway.ForceCodexCLI = true
cfg.Gateway.CodexImageGenerationBridgeEnabled = true
svc := &OpenAIGatewayService{cfg: cfg, httpUpstream: upstream}
account := &Account{
ID: 7,
Name: "openai-apikey",
Platform: PlatformOpenAI,
Type: AccountTypeAPIKey,
Concurrency: 1,
Credentials: map[string]any{
"api_key": "sk-test",
"base_url": "https://example.com",
},
Extra: map[string]any{"use_responses_api": true},
}
rec := httptest.NewRecorder()
c, _ := gin.CreateTestContext(rec)
c.Request = httptest.NewRequest(http.MethodPost, "/openai/v1/responses", nil)
c.Set("api_key", &APIKey{Group: &Group{AllowImageGeneration: true}})
SetOpenAIClientTransport(c, OpenAIClientTransportHTTP)
body := []byte(`{"model":"gpt-5","stream":false,"input":"draw if needed"}`)
result, err := svc.Forward(context.Background(), c, account, body)
require.NoError(t, err)
require.NotNil(t, result)
require.Equal(t, 1, result.ImageCount)
require.Equal(t, "2K", result.ImageSize)
require.Equal(t, "gpt-image-2", result.BillingModel)
}
func TestOpenAIGatewayService_Forward_HTTPDeletesPreviousResponseIDWhenPresent(t *testing.T) {
gin.SetMode(gin.TestMode)
cfg := &config.Config{}
cfg.Security.URLAllowlist.Enabled = false
account := &Account{
ID: 8,
Name: "openai-apikey",
Platform: PlatformOpenAI,
Type: AccountTypeAPIKey,
Concurrency: 1,
Credentials: map[string]any{
"api_key": "sk-test",
"base_url": "https://example.com",
},
Extra: map[string]any{"use_responses_api": true},
}
for _, body := range [][]byte{
[]byte(`{"model":"gpt-5","stream":false,"previous_response_id":"","input":"hi"}`),
[]byte(`{"model":"gpt-5","stream":false,"previous_response_id":null,"input":"hi"}`),
} {
upstream := &httpUpstreamRecorder{
resp: &http.Response{
StatusCode: http.StatusOK,
Header: http.Header{"Content-Type": []string{"application/json"}},
Body: io.NopCloser(strings.NewReader(`{"usage":{"input_tokens":1,"output_tokens":2}}`)),
},
}
svc := &OpenAIGatewayService{cfg: cfg, httpUpstream: upstream}
rec := httptest.NewRecorder()
c, _ := gin.CreateTestContext(rec)
c.Request = httptest.NewRequest(http.MethodPost, "/openai/v1/responses", nil)
SetOpenAIClientTransport(c, OpenAIClientTransportHTTP)
result, err := svc.Forward(context.Background(), c, account, body)
require.NoError(t, err)
require.NotNil(t, result)
require.False(t, gjson.GetBytes(upstream.lastBody, "previous_response_id").Exists())
}
}
func TestOpenAIRequestBodyMayContainEmptyBase64InputImageSeesEscapedJSON(t *testing.T) {
body := []byte(`{"input":[{"type":"message","content":[{"type":"input_image","image_` + "\\u0075" + `rl":"data:image/png;base64` + "\\u002c" + ` "}]}]}`)
require.True(t, openAIRequestBodyMayContainEmptyBase64InputImage(body))
}
func TestOpenAIRequestBodyMayContainEmptyBase64InputImageSeesEscapedImageType(t *testing.T) {
body := []byte(`{"input":[{"type":"message","content":[{"type":"input_` + "\\u0069" + `mage","image_url":"data:image/png;base64, "}]}]}`)
require.True(t, openAIRequestBodyMayContainEmptyBase64InputImage(body))
}
func TestOpenAIRequestBodyMayContainEmptyBase64InputImageSeesEscapedInputPrefix(t *testing.T) {
body := []byte(`{"input":[{"type":"message","content":[{"type":"inp` + "\\u0075" + `t_image","image_url":"data:image/png;base64, "}]}]}`)
require.True(t, openAIRequestBodyMayContainEmptyBase64InputImage(body))
}
func TestOpenAIGatewayService_Forward_ImageOnlyModelKeepsSupportedVerbosity(t *testing.T) {
gin.SetMode(gin.TestMode)
upstream := &httpUpstreamRecorder{
resp: &http.Response{
StatusCode: http.StatusOK,
Header: http.Header{"Content-Type": []string{"application/json"}},
Body: io.NopCloser(strings.NewReader(`{"usage":{"input_tokens":1,"output_tokens":2}}`)),
},
}
cfg := &config.Config{}
cfg.Security.URLAllowlist.Enabled = false
svc := &OpenAIGatewayService{cfg: cfg, httpUpstream: upstream}
account := &Account{
ID: 6,
Name: "openai-apikey",
Platform: PlatformOpenAI,
Type: AccountTypeAPIKey,
Concurrency: 1,
Credentials: map[string]any{
"api_key": "sk-test",
"base_url": "https://example.com",
},
Extra: map[string]any{"use_responses_api": true},
}
rec := httptest.NewRecorder()
c, _ := gin.CreateTestContext(rec)
c.Request = httptest.NewRequest(http.MethodPost, "/openai/v1/responses", nil)
SetOpenAIClientTransport(c, OpenAIClientTransportHTTP)
body := []byte(`{"model":"gpt-image-2","stream":false,"text":{"verbosity":"low"},"input":"draw"}`)
result, err := svc.Forward(context.Background(), c, account, body)
require.NoError(t, err)
require.NotNil(t, result)
require.Equal(t, "low", gjson.GetBytes(upstream.lastBody, "text.verbosity").String())
require.Equal(t, openAIImagesResponsesMainModel, gjson.GetBytes(upstream.lastBody, "model").String())
}
func TestExtractOpenAIRequestMetaFromBody(t *testing.T) {
tests := []struct {
name string
@ -106,26 +625,13 @@ func TestExtractOpenAIReasoningEffortFromBody(t *testing.T) {
}
}
func TestGetOpenAIRequestBodyMap_UsesContextCache(t *testing.T) {
gin.SetMode(gin.TestMode)
rec := httptest.NewRecorder()
c, _ := gin.CreateTestContext(rec)
cached := map[string]any{"model": "cached-model", "stream": true}
c.Set(OpenAIParsedRequestBodyKey, cached)
got, err := getOpenAIRequestBodyMap(c, []byte(`{invalid-json`))
require.NoError(t, err)
require.Equal(t, cached, got)
}
func TestGetOpenAIRequestBodyMap_ParseErrorWithoutCache(t *testing.T) {
func TestGetOpenAIRequestBodyMap_ParseError(t *testing.T) {
_, err := getOpenAIRequestBodyMap(nil, []byte(`{invalid-json`))
require.Error(t, err)
require.Contains(t, err.Error(), "parse request")
}
func TestGetOpenAIRequestBodyMap_WriteBackContextCache(t *testing.T) {
func TestGetOpenAIRequestBodyMap_DoesNotWriteContextCache(t *testing.T) {
gin.SetMode(gin.TestMode)
rec := httptest.NewRecorder()
c, _ := gin.CreateTestContext(rec)
@ -133,12 +639,7 @@ func TestGetOpenAIRequestBodyMap_WriteBackContextCache(t *testing.T) {
got, err := getOpenAIRequestBodyMap(c, []byte(`{"model":"gpt-5","stream":true}`))
require.NoError(t, err)
require.Equal(t, "gpt-5", got["model"])
cached, ok := c.Get(OpenAIParsedRequestBodyKey)
require.True(t, ok)
cachedMap, ok := cached.(map[string]any)
require.True(t, ok)
require.Equal(t, got, cachedMap)
require.Empty(t, c.Keys)
}
func TestSanitizeEmptyBase64InputImagesInOpenAIRequestBodyMap(t *testing.T) {

View File

@ -622,7 +622,7 @@ func (s *OpenAIGatewayService) forwardOpenAIImagesAPIKey(
return nil, fmt.Errorf("upstream request failed: %s", safeErr)
}
if resp.StatusCode >= 400 {
respBody, _ := io.ReadAll(io.LimitReader(resp.Body, 2<<20))
respBody := s.readUpstreamErrorBody(resp)
_ = resp.Body.Close()
resp.Body = io.NopCloser(bytes.NewReader(respBody))
upstreamMsg := strings.TrimSpace(extractUpstreamErrorMessage(respBody))
@ -1276,21 +1276,22 @@ func resolveOpenAIImageBytes(
headers http.Header,
conversationID string,
pointer openAIImagePointerInfo,
errorBodyReadLimit int64,
) ([]byte, error) {
if normalized := normalizeOpenAIImageBase64(pointer.B64JSON); normalized != "" {
return base64.StdEncoding.DecodeString(normalized)
}
if downloadURL := strings.TrimSpace(pointer.DownloadURL); downloadURL != "" {
return downloadOpenAIImageBytes(ctx, client, headers, downloadURL)
return downloadOpenAIImageBytes(ctx, client, headers, downloadURL, errorBodyReadLimit)
}
if strings.TrimSpace(pointer.Pointer) == "" {
return nil, fmt.Errorf("image asset is missing pointer, url, and base64 data")
}
downloadURL, err := fetchOpenAIImageDownloadURL(ctx, client, headers, conversationID, pointer.Pointer)
downloadURL, err := fetchOpenAIImageDownloadURL(ctx, client, headers, conversationID, pointer.Pointer, errorBodyReadLimit)
if err != nil {
return nil, err
}
return downloadOpenAIImageBytes(ctx, client, headers, downloadURL)
return downloadOpenAIImageBytes(ctx, client, headers, downloadURL, errorBodyReadLimit)
}
func normalizeOpenAIImageBase64(raw string) string {
@ -1395,6 +1396,7 @@ func fetchOpenAIImageDownloadURL(
headers http.Header,
conversationID string,
pointer string,
errorBodyReadLimit int64,
) (string, error) {
url := ""
allowConversationRetry := false
@ -1425,7 +1427,7 @@ func fetchOpenAIImageDownloadURL(
} else if resp.IsSuccessState() && strings.TrimSpace(result.DownloadURL) != "" {
return strings.TrimSpace(result.DownloadURL), nil
} else {
statusErr := newOpenAIImageStatusError(resp, "fetch image download url failed")
statusErr := newOpenAIImageStatusError(resp, "fetch image download url failed", errorBodyReadLimit)
if !allowConversationRetry || !isOpenAIImageTransientConversationNotFoundError(statusErr) {
return "", statusErr
}
@ -1450,7 +1452,7 @@ func fetchOpenAIImageDownloadURL(
return "", lastErr
}
func downloadOpenAIImageBytes(ctx context.Context, client *req.Client, headers http.Header, downloadURL string) ([]byte, error) {
func downloadOpenAIImageBytes(ctx context.Context, client *req.Client, headers http.Header, downloadURL string, errorBodyReadLimit int64) ([]byte, error) {
request := client.R().
SetContext(ctx).
DisableAutoReadResponse()
@ -1478,7 +1480,7 @@ func downloadOpenAIImageBytes(ctx context.Context, client *req.Client, headers h
}
}()
if resp.StatusCode < 200 || resp.StatusCode >= 300 {
return nil, newOpenAIImageStatusError(resp, "download image bytes failed")
return nil, newOpenAIImageStatusError(resp, "download image bytes failed", errorBodyReadLimit)
}
return io.ReadAll(io.LimitReader(resp.Body, openAIImageMaxDownloadBytes))
}
@ -1505,7 +1507,7 @@ func (e *openAIImageStatusError) Error() string {
return "openai image backend request failed"
}
func newOpenAIImageStatusError(resp *req.Response, fallback string) error {
func newOpenAIImageStatusError(resp *req.Response, fallback string, errorBodyReadLimit int64) error {
if resp == nil {
if strings.TrimSpace(fallback) == "" {
fallback = "openai image backend request failed"
@ -1526,7 +1528,10 @@ func newOpenAIImageStatusError(resp *req.Response, fallback string) error {
requestURL = resp.Request.URL.String()
}
if resp.Body != nil {
body, _ = io.ReadAll(io.LimitReader(resp.Body, 2<<20))
if errorBodyReadLimit <= 0 {
errorBodyReadLimit = openAIUpstreamErrorBodyReadLimit
}
body, _ = io.ReadAll(io.LimitReader(resp.Body, errorBodyReadLimit))
_ = resp.Body.Close()
}
}

View File

@ -1172,7 +1172,7 @@ func (s *OpenAIGatewayService) forwardOpenAIImagesOAuth(
return nil, fmt.Errorf("upstream request failed: %s", safeErr)
}
if resp.StatusCode >= 400 {
respBody, _ := io.ReadAll(io.LimitReader(resp.Body, 2<<20))
respBody := s.readUpstreamErrorBody(resp)
_ = resp.Body.Close()
resp.Body = io.NopCloser(bytes.NewReader(respBody))
upstreamMsg := strings.TrimSpace(extractUpstreamErrorMessage(respBody))

View File

@ -4,6 +4,7 @@ import (
"bytes"
"context"
"errors"
"fmt"
"io"
"mime/multipart"
"net/http"
@ -14,6 +15,7 @@ import (
"github.com/Wei-Shaw/sub2api/internal/config"
"github.com/gin-gonic/gin"
"github.com/imroc/req/v3"
"github.com/stretchr/testify/require"
"github.com/tidwall/gjson"
)
@ -398,11 +400,38 @@ func TestCollectOpenAIImagePointers_RecognizesDirectAssets(t *testing.T) {
func TestResolveOpenAIImageBytes_PrefersInlineBase64(t *testing.T) {
data, err := resolveOpenAIImageBytes(context.Background(), nil, nil, "", openAIImagePointerInfo{
B64JSON: "data:image/png;base64,QUJD",
})
}, openAIUpstreamErrorBodyReadLimit)
require.NoError(t, err)
require.Equal(t, []byte("ABC"), data)
}
func TestNewOpenAIImageStatusError_UsesProvidedReadLimit(t *testing.T) {
padding := strings.Repeat("x", int(openAIUpstreamErrorBodyReadLimit)+1024)
body := fmt.Sprintf(`{"error":{"padding":"%s","message":"diagnostic-marker"}}`, padding)
resp := &req.Response{Response: &http.Response{
StatusCode: http.StatusBadGateway,
Header: http.Header{},
Body: io.NopCloser(strings.NewReader(body)),
}}
err := newOpenAIImageStatusError(resp, "download image bytes failed", int64(len(body)))
require.Error(t, err)
require.Equal(t, "diagnostic-marker", err.Error())
var statusErr *openAIImageStatusError
require.ErrorAs(t, err, &statusErr)
require.Len(t, statusErr.ResponseBody, len(body))
}
func TestOpenAIUpstreamErrorBodyReadLimitForConfig_RespectsDiagnosticLimit(t *testing.T) {
cfg := &config.Config{Gateway: config.GatewayConfig{
LogUpstreamErrorBody: true,
LogUpstreamErrorBodyMaxBytes: int(openAIUpstreamErrorBodyReadLimit) + 1024,
}}
require.Equal(t, int64(cfg.Gateway.LogUpstreamErrorBodyMaxBytes), openAIUpstreamErrorBodyReadLimitForConfig(cfg))
}
func TestAccountSupportsOpenAIImageCapability_OAuthSupportsNative(t *testing.T) {
account := &Account{
Platform: PlatformOpenAI,

View File

@ -14,6 +14,7 @@ import (
"github.com/Wei-Shaw/sub2api/internal/config"
"github.com/Wei-Shaw/sub2api/internal/pkg/logger"
"github.com/Wei-Shaw/sub2api/internal/pkg/openai_compat"
"github.com/Wei-Shaw/sub2api/internal/pkg/tlsfingerprint"
"github.com/gin-gonic/gin"
"github.com/stretchr/testify/require"
@ -101,6 +102,57 @@ func TestOpenAIGatewayService_ResponsesUnknownModelDoesNotFallbackToGPT54(t *tes
require.True(t, rec.Code >= http.StatusBadRequest)
}
func TestOpenAIGatewayService_NativeResponsesBodyModificationPreservesHTMLChars(t *testing.T) {
gin.SetMode(gin.TestMode)
payloadText := strings.Repeat(`<tag>&value</tag>`, 128)
originalBody := []byte(fmt.Sprintf(`{"model":"gpt-5.5","stream":false,"max_output_tokens":100,"previous_response_id":"resp_prev","input":[{"type":"message","role":"user","content":[{"type":"input_text","text":%q}]}]}`, payloadText))
rec := httptest.NewRecorder()
c, _ := gin.CreateTestContext(rec)
c.Request = httptest.NewRequest(http.MethodPost, "/v1/responses", bytes.NewReader(originalBody))
c.Request.Header.Set("Content-Type", "application/json")
upstream := &httpUpstreamRecorder{resp: &http.Response{
StatusCode: http.StatusBadRequest,
Header: http.Header{"Content-Type": []string{"application/json"}, "x-request-id": []string{"rid_native_reencode"}},
Body: io.NopCloser(strings.NewReader(`{"error":{"type":"invalid_request_error","message":"stop after capture"}}`)),
}}
svc := &OpenAIGatewayService{
cfg: &config.Config{Security: config.SecurityConfig{URLAllowlist: config.URLAllowlistConfig{
Enabled: false,
AllowInsecureHTTP: true,
}}},
httpUpstream: upstream,
}
account := &Account{
ID: 456,
Name: "openai-apikey",
Platform: PlatformOpenAI,
Type: AccountTypeAPIKey,
Concurrency: 1,
Credentials: map[string]any{
"api_key": "sk-test",
"base_url": "http://upstream.example",
},
Extra: map[string]any{
openai_compat.ExtraKeyResponsesMode: string(openai_compat.ResponsesSupportModeAuto),
openai_compat.ExtraKeyResponsesSupported: true,
},
Status: StatusActive,
Schedulable: true,
}
result, err := svc.Forward(context.Background(), c, account, originalBody)
require.Error(t, err)
require.Nil(t, result)
require.NotNil(t, upstream.lastReq)
require.Equal(t, "http://upstream.example/v1/responses", upstream.lastReq.URL.String())
require.Contains(t, string(upstream.lastBody), payloadText)
require.NotContains(t, string(upstream.lastBody), `\\u003c`)
require.NotContains(t, string(upstream.lastBody), `\\u003e`)
require.NotContains(t, string(upstream.lastBody), `\\u0026`)
}
func TestOpenAIGatewayService_OAuthMessagesBridgeDoesNotInjectDefaultInstructions(t *testing.T) {
gin.SetMode(gin.TestMode)

View File

@ -1,6 +1,10 @@
package service
import "strings"
import (
"strings"
"github.com/tidwall/gjson"
)
// ToolContinuationSignals 聚合工具续链相关信号,避免重复遍历 input。
type ToolContinuationSignals struct {
@ -150,6 +154,67 @@ func AnalyzeToolContinuationSignals(reqBody map[string]any) ToolContinuationSign
return signals
}
// ValidateFunctionCallOutputContextBytes 基于 raw JSON 校验工具输出续链,避免 handler 预校验阶段全量解码大 input。
func ValidateFunctionCallOutputContextBytes(body []byte) FunctionCallOutputValidation {
result := FunctionCallOutputValidation{}
if len(body) == 0 {
return result
}
// handler 热路径只读扫描 input避免 GetBytes 为大 Responses body 复制整段 JSON。
input := parseRawJSONView(body).Get("input")
if !input.IsArray() {
return result
}
var callIDs map[string]struct{}
var referenceIDs map[string]struct{}
input.ForEach(func(_, item gjson.Result) bool {
if !item.IsObject() {
return true
}
itemType := item.Get("type").String()
switch {
case isCodexToolCallOutputItemType(itemType):
result.HasFunctionCallOutput = true
callID := strings.TrimSpace(item.Get("call_id").String())
if callID == "" {
result.HasFunctionCallOutputMissingCallID = true
return true
}
if callIDs == nil {
callIDs = make(map[string]struct{})
}
callIDs[callID] = struct{}{}
case isCodexToolCallContextItemType(itemType):
if strings.TrimSpace(item.Get("call_id").String()) != "" {
result.HasToolCallContext = true
}
case itemType == "item_reference":
idValue := strings.TrimSpace(item.Get("id").String())
if idValue == "" {
return true
}
if referenceIDs == nil {
referenceIDs = make(map[string]struct{})
}
referenceIDs[idValue] = struct{}{}
}
return !result.HasFunctionCallOutput || !result.HasToolCallContext
})
if !result.HasFunctionCallOutput || result.HasToolCallContext || len(callIDs) == 0 || len(referenceIDs) == 0 {
return result
}
allReferenced := true
for callID := range callIDs {
if _, ok := referenceIDs[callID]; !ok {
allReferenced = false
break
}
}
result.HasItemReferenceForAllCallIDs = allReferenced
return result
}
// ValidateFunctionCallOutputContext 为 handler 提供低开销校验结果:
// 1) 无工具输出直接返回
// 2) 若已存在工具调用上下文则提前返回

View File

@ -1,6 +1,7 @@
package service
import (
"encoding/json"
"testing"
"github.com/stretchr/testify/require"
@ -118,3 +119,68 @@ func TestHasItemReferenceForCallIDs(t *testing.T) {
require.True(t, HasItemReferenceForCallIDs(req, []string{"call_1", "call_2"}))
require.False(t, HasItemReferenceForCallIDs(req, []string{"call_1", "call_3"}))
}
func TestValidateFunctionCallOutputContextBytesMatchesMapValidation(t *testing.T) {
// handler 预校验走 raw JSON 扫描,语义必须与 service 内部 map 校验保持一致。
cases := []struct {
name string
body map[string]any
}{
{
name: "no_input",
body: map[string]any{"model": "gpt-5.4"},
},
{
name: "missing_call_id",
body: map[string]any{"input": []any{map[string]any{"type": "function_call_output"}}},
},
{
name: "call_id_without_reference",
body: map[string]any{"input": []any{map[string]any{"type": "function_call_output", "call_id": "call_1"}}},
},
{
name: "matching_reference",
body: map[string]any{"input": []any{
map[string]any{"type": "function_call_output", "call_id": "call_1"},
map[string]any{"type": "item_reference", "id": "call_1"},
}},
},
{
name: "partial_reference",
body: map[string]any{"input": []any{
map[string]any{"type": "function_call_output", "call_id": "call_1"},
map[string]any{"type": "tool_search_output", "call_id": "call_2"},
map[string]any{"type": "item_reference", "id": "call_1"},
}},
},
{
name: "tool_context",
body: map[string]any{"input": []any{
map[string]any{"type": "function_call_output", "call_id": "call_1"},
map[string]any{"type": "function_call", "call_id": "call_1"},
}},
},
{
name: "all_codex_tool_outputs",
body: map[string]any{"input": []any{
map[string]any{"type": "function_call_output", "call_id": "call_function"},
map[string]any{"type": "tool_search_output", "call_id": "call_search"},
map[string]any{"type": "custom_tool_call_output", "call_id": "call_custom"},
map[string]any{"type": "mcp_tool_call_output", "call_id": "call_mcp"},
map[string]any{"type": "item_reference", "id": "call_function"},
map[string]any{"type": "item_reference", "id": "call_search"},
map[string]any{"type": "item_reference", "id": "call_custom"},
map[string]any{"type": "item_reference", "id": "call_mcp"},
}},
},
}
for _, tt := range cases {
t.Run(tt.name, func(t *testing.T) {
bodyBytes, err := json.Marshal(tt.body)
require.NoError(t, err)
require.Equal(t, ValidateFunctionCallOutputContext(tt.body), ValidateFunctionCallOutputContextBytes(bodyBytes))
})
}
}

View File

@ -171,6 +171,91 @@ func TestOpenAIGatewayService_Forward_WSv2_SuccessAndBindSticky(t *testing.T) {
require.Equal(t, "resp_new_1", gjson.GetBytes(responseBody, "id").String())
}
func TestOpenAIGatewayService_Forward_WSv2_UsesPatchedBodyAfterValidationDecode(t *testing.T) {
gin.SetMode(gin.TestMode)
type receivedPayload struct {
MaxCompletionTokensExists bool
}
receivedCh := make(chan receivedPayload, 1)
upgrader := websocket.Upgrader{CheckOrigin: func(r *http.Request) bool { return true }}
wsServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
conn, err := upgrader.Upgrade(w, r, nil)
if err != nil {
t.Errorf("upgrade websocket failed: %v", err)
return
}
defer func() { _ = conn.Close() }()
var request map[string]any
if err := conn.ReadJSON(&request); err != nil {
t.Errorf("read ws request failed: %v", err)
return
}
requestJSON := requestToJSONString(request)
receivedCh <- receivedPayload{MaxCompletionTokensExists: gjson.Get(requestJSON, "max_completion_tokens").Exists()}
if err := conn.WriteJSON(map[string]any{
"type": "response.completed",
"response": map[string]any{
"id": "resp_patched_ws_1",
"model": "gpt-5.3-codex-spark",
"usage": map[string]any{"input_tokens": 1, "output_tokens": 1},
},
}); err != nil {
t.Errorf("write response.completed failed: %v", err)
return
}
}))
defer wsServer.Close()
rec := httptest.NewRecorder()
c, _ := gin.CreateTestContext(rec)
c.Request = httptest.NewRequest(http.MethodPost, "/openai/v1/responses", nil)
c.Request.Header.Set("User-Agent", "unit-test-agent/1.0")
cfg := &config.Config{}
cfg.Security.URLAllowlist.Enabled = false
cfg.Security.URLAllowlist.AllowInsecureHTTP = true
cfg.Gateway.OpenAIWS.Enabled = true
cfg.Gateway.OpenAIWS.APIKeyEnabled = true
cfg.Gateway.OpenAIWS.ResponsesWebsocketsV2 = true
cfg.Gateway.OpenAIWS.DialTimeoutSeconds = 3
cfg.Gateway.OpenAIWS.ReadTimeoutSeconds = 30
cfg.Gateway.OpenAIWS.WriteTimeoutSeconds = 10
svc := &OpenAIGatewayService{
cfg: cfg,
openaiWSResolver: NewOpenAIWSProtocolResolver(cfg),
toolCorrector: NewCodexToolCorrector(),
}
account := &Account{
ID: 10,
Name: "openai-ws",
Platform: PlatformOpenAI,
Type: AccountTypeAPIKey,
Status: StatusActive,
Schedulable: true,
Concurrency: 1,
Credentials: map[string]any{
"api_key": "sk-test",
"base_url": wsServer.URL,
},
Extra: map[string]any{"responses_websockets_v2_enabled": true},
}
body := []byte(`{"model":"gpt-5.4","stream":false,"max_completion_tokens":12,"tools":[{"type":"image_generation"}],"input":[{"type":"input_text","text":"hello"}]}`)
result, err := svc.Forward(context.Background(), c, account, body)
require.NoError(t, err)
require.NotNil(t, result)
require.True(t, result.OpenAIWSMode)
received := <-receivedCh
require.False(t, received.MaxCompletionTokensExists)
}
func TestOpenAIGatewayService_Forward_WSv2_ImageGenerationCountsOutputs(t *testing.T) {
gin.SetMode(gin.TestMode)

View File

@ -12,6 +12,7 @@ import (
"github.com/Wei-Shaw/sub2api/internal/config"
"github.com/Wei-Shaw/sub2api/internal/pkg/logger"
"github.com/tidwall/gjson"
)
// UserMsgQueueCache 用户消息串行队列 Redis 缓存接口
@ -62,43 +63,48 @@ func NewUserMessageQueueService(cache UserMsgQueueCache, rpmCache RPMCache, cfg
// 2. 最后一条消息 role == "user"
// 3. 最后一条消息 content如果是数组中不含 type:"tool_result" / "tool_use_result"
func IsRealUserMessage(parsed *ParsedRequest) bool {
if parsed == nil || len(parsed.Messages) == 0 {
if parsed == nil {
return false
}
messagesRaw := parsed.MessagesRaw()
if len(messagesRaw) == 0 {
return false
}
lastMsg := parsed.Messages[len(parsed.Messages)-1]
msgMap, ok := lastMsg.(map[string]any)
if !ok {
messages := gjson.ParseBytes(messagesRaw)
if !messages.IsArray() {
return false
}
lastMsg := gjson.Result{}
messages.ForEach(func(_, msg gjson.Result) bool {
lastMsg = msg
return true
})
if !lastMsg.Exists() || !lastMsg.IsObject() {
return false
}
if lastMsg.Get("role").String() != "user" {
return false
}
role, _ := msgMap["role"].(string)
if role != "user" {
return false
content := lastMsg.Get("content")
if !content.Exists() {
return true
}
if !content.IsArray() {
return true
}
// 检查 content 是否包含 tool_result 类型
content, ok := msgMap["content"]
if !ok {
return true // 没有 content 字段,视为普通用户消息
}
contentArr, ok := content.([]any)
if !ok {
return true // content 不是数组(可能是 string视为普通用户消息
}
for _, item := range contentArr {
itemMap, ok := item.(map[string]any)
if !ok {
continue
}
itemType, _ := itemMap["type"].(string)
isReal := true
content.ForEach(func(_, item gjson.Result) bool {
itemType := item.Get("type").String()
if itemType == "tool_result" || itemType == "tool_use_result" {
isReal = false
return false
}
}
return true
return true
})
return isReal
}
// TryAcquire 尝试立即获取串行锁