Files
new-api/relay/compatible_handler.go
T
LucasandClaude Fable 5 d6b5ce99de fix(relay): set Request.GetBody so the HTTP/2 transport can transparently retry after an upstream stream reset (#6249)
* fix(relay): set Request.GetBody so the HTTP/2 transport can transparently retry after an upstream stream reset

The outbound request body is a type-erased io.Reader over BodyStorage, so
net/http cannot derive Request.GetBody (it only does so for *bytes.Reader,
*bytes.Buffer and *strings.Reader). With GetBody nil, the HTTP/2 transport
cannot transparently retry a request once the body has been written and the
upstream resets the stream with a retryable error (REFUSED_STREAM, or a
connection-level GOAWAY); the relay request then fails with:

    http2: Transport: cannot retry err [...] after Request.Body was written;
    define Request.GetBody to avoid this error

This affects every relay path that goes through DoApiRequest (chat, claude,
gemini, responses, embedding, image, rerank).

BodyStorage (memory and disk) already implements io.Seeker, so replay support
only needed wiring:

- NewOutboundJSONBody additionally returns a getBody that rewinds the storage
  and hands out a fresh non-closing reader. The transport only calls GetBody
  after the previous attempt's body has been abandoned, so the rewind cannot
  race an in-flight read.
- RelayInfo carries it in the new UpstreamRequestGetBody field, set alongside
  UpstreamRequestBodySize by the handlers that build storage-backed bodies.
- applyUpstreamGetBody (symmetric with applyUpstreamContentLength) wires it
  into DoApiRequest/DoFormRequest/DoTaskApiRequest, only when req.GetBody is
  still nil.

Also remove the hand-rolled GetBody override in DoTaskApiRequest: it returned
the same already-consumed reader, so any transport-level replay would have
silently sent an empty body, and it clobbered the correct snapshot-based
GetBody that net/http derives from the *bytes.Reader bodies the task adaptors
pass in. For non-replayable bodies GetBody now stays nil, so a retry fails
loudly instead of corrupting the request.

Covered by unit tests plus an end-to-end raw-frame HTTP/2 test that resets
the first stream with REFUSED_STREAM after the body is written and asserts
the transport transparently retries with the complete body.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

* fix(relay): hand out independent readers from GetBody (address review)

Per the http.Request.GetBody contract ("returns a new copy of Body"),
each call must yield a reader with its own cursor. The previous
implementation rewound and reused the shared BodyStorage, so two
consecutive GetBody readers would interfere with each other, and a
replay could disturb the primary body's offset under extreme transport
timing (e.g. attempt N's body write not yet fully abandoned when the
transport builds attempt N+1).

Instead of snapshotting the payload (an extra copy), add
BodyStorage.NewReader, which returns an independent zero-copy reader:

- memory mode: a fresh bytes.Reader over the same immutable backing
  array;
- disk mode: a separate file descriptor over the cache file, so the
  transport closing a replayed body only closes that descriptor.

NewOutboundJSONBody's getBody now simply hands out storage.NewReader,
and once the handler releases the storage, GetBody fails with
ErrStorageClosed instead of replaying stale data.

Tests: interleaved reads across two replay readers and the primary
body each observe exactly their own byte stream, for both the memory
and the disk-backed storage; the existing GetBody and HTTP/2 retry
suites still pass (h2 e2e tests flake-free with -count=20).

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

* fix(relay): bind replayable metadata on pass-through requests

* fix(relay): reset upstream body metadata between channels

* test(relay): cover replay across retries and channel attempts

* fix(relay): stop following upstream redirects

---------

Co-authored-by: Claude Fable 5 <noreply@anthropic.com>
2026-08-06 15:56:38 +08:00

227 lines
8.0 KiB
Go

package relay
import (
"fmt"
"io"
"net/http"
"strings"
"github.com/QuantumNous/new-api/common"
"github.com/QuantumNous/new-api/constant"
"github.com/QuantumNous/new-api/logger"
relaycommon "github.com/QuantumNous/new-api/relay/common"
relayconstant "github.com/QuantumNous/new-api/relay/constant"
"github.com/QuantumNous/new-api/relay/helper"
"github.com/QuantumNous/new-api/relaykit/dto"
"github.com/QuantumNous/new-api/relaykit/types"
"github.com/QuantumNous/new-api/service"
"github.com/QuantumNous/new-api/setting/model_setting"
"github.com/QuantumNous/new-api/setting/ratio_setting"
"github.com/samber/lo"
"github.com/gin-gonic/gin"
)
func TextHelper(c *gin.Context, info *relaycommon.RelayInfo) (newAPIError *types.NewAPIError) {
info.InitChannelMeta(c)
textReq, ok := info.Request.(*dto.GeneralOpenAIRequest)
if !ok {
return types.NewErrorWithStatusCode(fmt.Errorf("invalid request type, expected dto.GeneralOpenAIRequest, got %T", info.Request), types.ErrorCodeInvalidRequest, http.StatusBadRequest, types.ErrOptionWithSkipRetry())
}
request, err := common.DeepCopy(textReq)
if err != nil {
return types.NewError(fmt.Errorf("failed to copy request to GeneralOpenAIRequest: %w", err), types.ErrorCodeInvalidRequest, types.ErrOptionWithSkipRetry())
}
if request.WebSearchOptions != nil {
c.Set("chat_completion_web_search_context_size", request.WebSearchOptions.SearchContextSize)
}
err = helper.ModelMappedHelper(c, info, request)
if err != nil {
return types.NewError(err, types.ErrorCodeChannelModelMappedError, types.ErrOptionWithSkipRetry())
}
includeUsage := true
// 判断用户是否需要返回使用情况
if request.StreamOptions != nil {
includeUsage = request.StreamOptions.IncludeUsage
}
// 如果不支持StreamOptions,将StreamOptions设置为nil
if !info.SupportStreamOptions || !lo.FromPtrOr(request.Stream, false) {
request.StreamOptions = nil
} else {
// 如果支持StreamOptions,且请求中没有设置StreamOptions,根据配置文件设置StreamOptions
if constant.ForceStreamOption {
request.StreamOptions = &dto.StreamOptions{
IncludeUsage: true,
}
}
}
info.ShouldIncludeUsage = includeUsage
adaptor := GetAdaptor(info.ApiType)
if adaptor == nil {
return types.NewError(fmt.Errorf("invalid api type: %d", info.ApiType), types.ErrorCodeInvalidApiType, types.ErrOptionWithSkipRetry())
}
adaptor.Init(info)
passThroughGlobal := model_setting.GetGlobalSettings().PassThroughRequestEnabled
if info.RelayMode == relayconstant.RelayModeChatCompletions &&
!passThroughGlobal &&
!info.ChannelSetting.PassThroughBodyEnabled &&
service.ShouldChatCompletionsUseResponsesGlobal(info.ChannelId, info.ChannelType, info.OriginModelName) {
applySystemPromptIfNeeded(c, info, request)
usage, newApiErr := chatCompletionsViaResponses(c, info, adaptor, request)
if newApiErr != nil {
return newApiErr
}
var containAudioTokens = usage.CompletionTokenDetails.AudioTokens > 0 || usage.PromptTokensDetails.AudioTokens > 0
var containsAudioRatios = ratio_setting.ContainsAudioRatio(info.OriginModelName) || ratio_setting.ContainsAudioCompletionRatio(info.OriginModelName)
if containAudioTokens && containsAudioRatios {
service.PostAudioConsumeQuota(c, info, usage, "")
} else {
service.PostTextConsumeQuota(c, info, usage, nil)
}
return nil
}
var requestBody io.Reader
if passThroughGlobal || info.ChannelSetting.PassThroughBodyEnabled {
storage, err := common.GetBodyStorage(c)
if err != nil {
return types.NewErrorWithStatusCode(err, types.ErrorCodeReadRequestBodyFailed, http.StatusBadRequest, types.ErrOptionWithSkipRetry())
}
if common.DebugEnabled {
if debugBytes, bErr := storage.Bytes(); bErr == nil {
logger.LogDebug(c, "requestBody: %s", debugBytes)
}
}
info.UpstreamRequestBodySize = storage.Size()
info.UpstreamRequestGetBody = storage.NewReader
requestBody = common.ReaderOnly(storage)
} else {
convertedRequest, err := adaptor.ConvertOpenAIRequest(c, info, request)
if err != nil {
return types.NewError(err, types.ErrorCodeConvertRequestFailed, types.ErrOptionWithSkipRetry())
}
relaycommon.AppendRequestConversionFromRequest(info, convertedRequest)
if info.ChannelSetting.SystemPrompt != "" {
// 如果有系统提示,则将其添加到请求中
request, ok := convertedRequest.(*dto.GeneralOpenAIRequest)
if ok {
containSystemPrompt := false
for _, message := range request.Messages {
if message.Role == request.GetSystemRoleName() {
containSystemPrompt = true
break
}
}
if !containSystemPrompt {
// 如果没有系统提示,则添加系统提示
systemMessage := dto.Message{
Role: request.GetSystemRoleName(),
Content: info.ChannelSetting.SystemPrompt,
}
request.Messages = append([]dto.Message{systemMessage}, request.Messages...)
} else if info.ChannelSetting.SystemPromptOverride {
common.SetContextKey(c, constant.ContextKeySystemPromptOverride, true)
// 如果有系统提示,且允许覆盖,则拼接到前面
for i, message := range request.Messages {
if message.Role == request.GetSystemRoleName() {
if message.IsStringContent() {
request.Messages[i].SetStringContent(info.ChannelSetting.SystemPrompt + "\n" + message.StringContent())
} else {
contents := message.ParseContent()
contents = append([]dto.MediaContent{
{
Type: dto.ContentTypeText,
Text: info.ChannelSetting.SystemPrompt,
},
}, contents...)
request.Messages[i].Content = contents
}
break
}
}
}
}
}
jsonData, err := common.Marshal(convertedRequest)
if err != nil {
return types.NewError(err, types.ErrorCodeJsonMarshalFailed, types.ErrOptionWithSkipRetry())
}
// remove disabled fields for OpenAI API
jsonData, err = relaycommon.RemoveDisabledFields(jsonData, info.ChannelOtherSettings, info.ChannelSetting.PassThroughBodyEnabled)
if err != nil {
return types.NewError(err, types.ErrorCodeConvertRequestFailed, types.ErrOptionWithSkipRetry())
}
// apply param override
if len(info.ParamOverride) > 0 {
jsonData, err = relaycommon.ApplyParamOverrideWithRelayInfo(jsonData, info)
if err != nil {
return newAPIErrorFromParamOverride(err)
}
}
logger.LogDebug(c, "text request body: %s", jsonData)
body, size, getBody, closer, err := relaycommon.NewOutboundJSONBody(jsonData)
if err != nil {
return types.NewError(err, types.ErrorCodeConvertRequestFailed, types.ErrOptionWithSkipRetry())
}
defer closer.Close()
jsonData = nil
info.UpstreamRequestBodySize = size
info.UpstreamRequestGetBody = getBody
requestBody = body
}
var httpResp *http.Response
resp, err := adaptor.DoRequest(c, info, requestBody)
if err != nil {
return types.NewOpenAIError(err, types.ErrorCodeDoRequestFailed, http.StatusInternalServerError)
}
statusCodeMappingStr := c.GetString("status_code_mapping")
if resp != nil {
httpResp = resp.(*http.Response)
info.IsStream = info.IsStream || strings.HasPrefix(httpResp.Header.Get("Content-Type"), "text/event-stream")
if httpResp.StatusCode != http.StatusOK {
newApiErr := service.RelayErrorHandler(c.Request.Context(), httpResp, false)
// reset status code 重置状态码
service.ResetStatusCode(newApiErr, statusCodeMappingStr)
return newApiErr
}
}
usage, newApiErr := adaptor.DoResponse(c, httpResp, info)
if newApiErr != nil {
// reset status code 重置状态码
service.ResetStatusCode(newApiErr, statusCodeMappingStr)
return newApiErr
}
var containAudioTokens = usage.(*dto.Usage).CompletionTokenDetails.AudioTokens > 0 || usage.(*dto.Usage).PromptTokensDetails.AudioTokens > 0
var containsAudioRatios = ratio_setting.ContainsAudioRatio(info.OriginModelName) || ratio_setting.ContainsAudioCompletionRatio(info.OriginModelName)
if containAudioTokens && containsAudioRatios {
service.PostAudioConsumeQuota(c, info, usage.(*dto.Usage), "")
} else {
service.PostTextConsumeQuota(c, info, usage.(*dto.Usage), nil)
}
return nil
}