refactor(gateway): shorten OpenAI request body retention

This commit is contained in:
name 2026-05-30 16:52:15 +08:00
parent 2caee9d884
commit 5a3e193b53
7 changed files with 137 additions and 144 deletions

View File

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

View File

@ -177,7 +177,7 @@ func TestSetClaudeCodeClientContext_FastPathAndStrictPath(t *testing.T) {
})
}
func TestSetClaudeCodeClientContext_ReuseParsedRequestAndContextCache(t *testing.T) {
func TestSetClaudeCodeClientContext_ReuseParsedRequest(t *testing.T) {
t.Run("reuse parsed request without body unmarshal", func(t *testing.T) {
c, _ := newHelperTestContext(http.MethodPost, "/v1/messages")
c.Request.Header.Set("User-Agent", "claude-cli/1.0.1")
@ -192,24 +192,6 @@ func TestSetClaudeCodeClientContext_ReuseParsedRequestAndContextCache(t *testing
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")
service.CacheOpenAIParsedRequestBody(c, []byte(`{invalid`), map[string]any{
"model": "claude-3-5-sonnet-20241022",
"system": []any{
map[string]any{"text": "You are Claude Code, Anthropic's official CLI for Claude."},
},
"metadata": map[string]any{"user_id": "user_aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa_account__session_aaaaaaaa-aaaa-aaaa-aaaa-aaaaaaaaaaaa"},
})
SetClaudeCodeClientContext(c, []byte(`{invalid`), nil)
require.True(t, service.IsClaudeCodeClient(c.Request.Context()))
})
}
func TestWaitForSlotWithPingTimeout_AccountAndUserAcquire(t *testing.T) {

View File

@ -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
}
service.CacheOpenAIParsedRequestBody(c, body, 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
}

View File

@ -36,13 +36,6 @@ import (
"go.uber.org/zap"
)
// openAIParsedRequestBodyCache 绑定 body 指纹,避免 handler 预解析的旧 body 污染后续 forwardBody。
type openAIParsedRequestBodyCache struct {
bodyHash uint64
bodyLen int
reqBody map[string]any
}
const (
// ChatGPT internal API for OAuth accounts
chatgptCodexURL = "https://chatgpt.com/backend-api/codex/responses"
@ -53,8 +46,6 @@ const (
// codex_cli_only 拒绝时单个请求头日志长度上限(字符)
codexCLIOnlyHeaderValueMaxBytes = 256
// OpenAIParsedRequestBodyKey 缓存 handler 侧已解析的请求体,避免重复解析。
OpenAIParsedRequestBodyKey = "openai_parsed_request_body"
// OpenAI WS Mode 失败后的重连次数上限(不含首次尝试)。
// 与 Codex 客户端保持一致:失败后最多重连 5 次。
openAIWSReconnectRetryLimit = 5
@ -3099,8 +3090,6 @@ func (s *OpenAIGatewayService) Forward(ctx context.Context, c *gin.Context, acco
})
s.handleFailoverSideEffects(ctx, resp, account, upstreamModel)
// reqBody 会被本次账号尝试原地修改failover 前必须释放,避免下一账号复用脏 map。
releaseOpenAIParsedRequestBody(c)
return nil, &UpstreamFailoverError{
StatusCode: resp.StatusCode,
ResponseBody: respBody,
@ -3113,7 +3102,8 @@ func (s *OpenAIGatewayService) Forward(ctx context.Context, c *gin.Context, acco
reasoningEffort := extractOpenAIReasoningEffort(reqBody, originalModel)
serviceTier := extractOpenAIServiceTier(reqBody)
releaseOpenAIParsedRequestBody(c)
// 上游接受后只保留计费需要的标量,避免响应处理期间继续保活完整 input/tools map。
reqBody = nil
// Handle normal response
var usage *OpenAIUsage
@ -6843,60 +6833,14 @@ func isEmptyBase64DataURI(raw string) bool {
return strings.TrimSpace(strings.TrimPrefix(rest, "base64,")) == ""
}
func getOpenAIRequestBodyMap(c *gin.Context, body []byte) (map[string]any, error) {
// 同一个 gin.Context 内 failover/渠道映射可能传入新 body缓存必须先校验 body 指纹。
bodyHash := xxhash.Sum64(body)
bodyLen := len(body)
if c != nil {
if cached, ok := c.Get(OpenAIParsedRequestBodyKey); ok {
if cache, ok := cached.(openAIParsedRequestBodyCache); ok && cache.reqBody != nil && cache.bodyLen == bodyLen && cache.bodyHash == bodyHash {
return cache.reqBody, nil
}
}
}
func getOpenAIRequestBodyMap(_ *gin.Context, body []byte) (map[string]any, error) {
var reqBody map[string]any
if err := json.Unmarshal(body, &reqBody); err != nil {
return nil, fmt.Errorf("parse request: %w", err)
}
if c != nil {
c.Set(OpenAIParsedRequestBodyKey, openAIParsedRequestBodyCache{bodyHash: bodyHash, bodyLen: bodyLen, reqBody: reqBody})
}
return reqBody, nil
}
// CacheOpenAIParsedRequestBody 仅缓存与当前 body 绑定的解析结果。
func CacheOpenAIParsedRequestBody(c *gin.Context, body []byte, reqBody map[string]any) {
if c == nil || reqBody == nil {
return
}
c.Set(OpenAIParsedRequestBodyKey, openAIParsedRequestBodyCache{
bodyHash: xxhash.Sum64(body),
bodyLen: len(body),
reqBody: reqBody,
})
}
// CachedOpenAIParsedRequestBody 只给同请求内不关心 body 参数的轻量识别逻辑使用。
func CachedOpenAIParsedRequestBody(c *gin.Context) map[string]any {
if c == nil {
return nil
}
if cached, ok := c.Get(OpenAIParsedRequestBodyKey); ok {
if cache, ok := cached.(openAIParsedRequestBodyCache); ok {
return cache.reqBody
}
}
return nil
}
func releaseOpenAIParsedRequestBody(c *gin.Context) {
if c == nil {
return
}
delete(c.Keys, OpenAIParsedRequestBodyKey)
}
func extractOpenAIReasoningEffort(reqBody map[string]any, requestedModel string) *string {
if value, present := getOpenAIReasoningEffortFromReqBody(reqBody); present {
if value == "" {

View File

@ -106,26 +106,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}
CacheOpenAIParsedRequestBody(c, []byte(`{invalid-json`), 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,20 +120,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"])
require.Equal(t, got, CachedOpenAIParsedRequestBody(c))
}
func TestGetOpenAIRequestBodyMap_IgnoresCacheForDifferentBody(t *testing.T) {
gin.SetMode(gin.TestMode)
rec := httptest.NewRecorder()
c, _ := gin.CreateTestContext(rec)
CacheOpenAIParsedRequestBody(c, []byte(`{"model":"cached-model"}`), map[string]any{"model": "cached-model"})
got, err := getOpenAIRequestBodyMap(c, []byte(`{"model":"forward-model"}`))
require.NoError(t, err)
require.Equal(t, "forward-model", got["model"])
require.Empty(t, c.Keys)
}
func TestSanitizeEmptyBase64InputImagesInOpenAIRequestBodyMap(t *testing.T) {

View File

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

View File

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