Merge pull request #2925 from is7Qin/feat/openai-oom
refactor(gateway): 降低大请求体内存保留
This commit is contained in:
commit
0f70f84182
@ -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 内部可能继续改写 body,usage 去重指纹必须使用最终上游接受的当前 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
|
||||
|
||||
@ -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),
|
||||
|
||||
@ -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),
|
||||
|
||||
@ -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
|
||||
}
|
||||
|
||||
// 并发槽位等待相关常量
|
||||
//
|
||||
// 性能优化说明:
|
||||
|
||||
@ -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) {
|
||||
|
||||
@ -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 {
|
||||
|
||||
@ -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) {
|
||||
|
||||
@ -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))
|
||||
}
|
||||
|
||||
@ -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 {
|
||||
|
||||
@ -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 {
|
||||
|
||||
@ -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 {
|
||||
|
||||
@ -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",
|
||||
|
||||
@ -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{
|
||||
|
||||
@ -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,
|
||||
)
|
||||
|
||||
@ -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,
|
||||
)
|
||||
|
||||
@ -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))
|
||||
|
||||
|
||||
@ -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))
|
||||
|
||||
|
||||
@ -14,8 +14,6 @@ func TestBuildOAuthMetadataUserID_FallbackWithoutAccountUUID(t *testing.T) {
|
||||
Model: "claude-sonnet-4-5",
|
||||
Stream: true,
|
||||
MetadataUserID: "",
|
||||
System: nil,
|
||||
Messages: nil,
|
||||
}
|
||||
|
||||
account := &Account{
|
||||
|
||||
@ -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.effort(Claude 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 不可用时,退化为复制(理论上极少发生)。
|
||||
|
||||
@ -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), "")
|
||||
}
|
||||
}
|
||||
|
||||
@ -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 retry,retry 会基于未签名的逻辑 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 {
|
||||
|
||||
@ -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())
|
||||
}
|
||||
|
||||
@ -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")
|
||||
}
|
||||
|
||||
@ -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)
|
||||
|
||||
|
||||
@ -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
@ -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())
|
||||
}
|
||||
|
||||
@ -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
|
||||
|
||||
@ -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))
|
||||
|
||||
|
||||
@ -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))
|
||||
|
||||
|
||||
@ -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))
|
||||
|
||||
|
||||
@ -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))
|
||||
|
||||
|
||||
@ -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
@ -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) {
|
||||
|
||||
@ -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()
|
||||
}
|
||||
}
|
||||
|
||||
@ -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))
|
||||
|
||||
@ -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,
|
||||
|
||||
@ -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)
|
||||
|
||||
|
||||
@ -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) 若已存在工具调用上下文则提前返回
|
||||
|
||||
@ -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))
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
@ -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)
|
||||
|
||||
|
||||
@ -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 尝试立即获取串行锁
|
||||
|
||||
Loading…
x
Reference in New Issue
Block a user