refactor(relay): move replay metadata onto request bodies

This commit is contained in:
CaIon
2026-08-06 17:33:10 +08:00
parent d6b5ce99de
commit ea4f021012
20 changed files with 211 additions and 267 deletions
+29 -4
View File
@@ -29,6 +29,14 @@ type BodyStorage interface {
NewReader() (io.ReadCloser, error)
}
// ReplayableBody is an outbound request body that can report its byte size and
// create independent readers for transport-level retries.
type ReplayableBody interface {
io.Reader
Size() int64
NewReader() (io.ReadCloser, error)
}
// ErrStorageClosed 存储已关闭错误
var ErrStorageClosed = fmt.Errorf("body storage is closed")
@@ -339,10 +347,27 @@ func CreateBodyStorageFromReader(reader io.Reader, contentLength int64, maxBytes
return storage, nil
}
// ReaderOnly wraps an io.Reader to hide io.Closer, preventing http.NewRequest
// from type-asserting io.ReadCloser and closing the underlying BodyStorage.
func ReaderOnly(r io.Reader) io.Reader {
return struct{ io.Reader }{r}
type replayableBodyReader struct {
storage BodyStorage
}
func (r replayableBodyReader) Read(p []byte) (int, error) {
return r.storage.Read(p)
}
func (r replayableBodyReader) Size() int64 {
return r.storage.Size()
}
func (r replayableBodyReader) NewReader() (io.ReadCloser, error) {
return r.storage.NewReader()
}
// NewReplayableBodyReader exposes the replay capabilities of storage without
// exposing io.Closer. This keeps ownership of the storage lifecycle with the
// caller instead of allowing net/http to close it as the request body.
func NewReplayableBodyReader(storage BodyStorage) ReplayableBody {
return replayableBodyReader{storage: storage}
}
// CleanupOldCacheFiles 清理旧的缓存文件(用于启动时清理残留)
+37
View File
@@ -0,0 +1,37 @@
package common
import (
"io"
"net/http"
"testing"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
)
func TestNewReplayableBodyReaderKeepsStorageLifecycleWithCaller(t *testing.T) {
payload := []byte(`{"model":"test-model","input":"hello"}`)
storage, err := CreateBodyStorage(payload)
require.NoError(t, err)
defer storage.Close()
body := NewReplayableBodyReader(storage)
assert.EqualValues(t, len(payload), body.Size())
_, exposesCloser := any(body).(io.Closer)
assert.False(t, exposesCloser, "the request body must not expose the storage closer")
req, err := http.NewRequest(http.MethodPost, "https://example.com", body)
require.NoError(t, err)
require.NoError(t, req.Body.Close())
replayBody, err := body.NewReader()
require.NoError(t, err, "closing the HTTP request body must not close the storage")
replay, err := io.ReadAll(replayBody)
require.NoError(t, err)
require.NoError(t, replayBody.Close())
assert.Equal(t, payload, replay)
require.NoError(t, storage.Close())
_, err = body.NewReader()
require.ErrorIs(t, err, ErrStorageClosed)
}
+1 -3
View File
@@ -62,13 +62,11 @@ func AlphaSearchHelper(c *gin.Context, info *relaycommon.RelayInfo) (newAPIError
}
logger.LogDebug(c, "requestBody: %s", jsonData)
body, size, getBody, closer, err := relaycommon.NewOutboundJSONBody(jsonData)
body, closer, err := relaycommon.NewOutboundJSONBody(jsonData)
if err != nil {
return types.NewError(err, types.ErrorCodeConvertRequestFailed, types.ErrOptionWithSkipRetry())
}
defer closer.Close()
info.UpstreamRequestBodySize = size
info.UpstreamRequestGetBody = getBody
adaptor := GetAdaptor(info.ApiType)
if adaptor == nil {
+20 -41
View File
@@ -25,50 +25,29 @@ import (
"github.com/gorilla/websocket"
)
// applyUpstreamContentLength populates req.ContentLength when the upstream
// body is wrapped in a BodyStorage (see relay/common/outbound_body.go).
//
// net/http.NewRequest only auto-detects ContentLength for *bytes.Reader,
// *bytes.Buffer and *strings.Reader. When the body is a type-erased io.Reader
// (which is the case for ReaderOnly(BodyStorage)), the Content-Length header
// would otherwise be omitted, forcing chunked transfer encoding and breaking
// some upstreams that require an explicit Content-Length.
func applyUpstreamContentLength(req *http.Request, info *common.RelayInfo) {
if info == nil {
// ApplyUpstreamBodyMetadata restores metadata that net/http cannot infer from
// a ReplayableBody. Callers must pass the original body because NewRequest
// hides its dynamic type behind req.Body's io.ReadCloser wrapper.
func ApplyUpstreamBodyMetadata(req *http.Request, body io.Reader) {
replayable, ok := body.(common2.ReplayableBody)
if !ok {
return
}
if info.UpstreamRequestBodySize > 0 && req.ContentLength <= 0 {
req.ContentLength = info.UpstreamRequestBodySize
}
}
// applyUpstreamGetBody populates req.GetBody when the upstream body is wrapped
// in a BodyStorage (see relay/common/outbound_body.go).
//
// net/http.NewRequest only auto-populates GetBody for *bytes.Reader,
// *bytes.Buffer and *strings.Reader. When the body is a type-erased io.Reader
// (which is the case for ReaderOnly(BodyStorage)), GetBody would otherwise stay
// nil, and the HTTP/2 transport cannot transparently retry the request once the
// upstream resets the stream after the body was already written; the request
// then fails with "http2: Transport: cannot retry err ... after Request.Body
// was written; define Request.GetBody to avoid this error".
func applyUpstreamGetBody(req *http.Request, info *common.RelayInfo) {
if info == nil || info.UpstreamRequestGetBody == nil {
return
// BodyStorage structurally satisfies ReplayableBody, but it also exposes
// io.Closer. If a caller passes the storage directly instead of using
// NewReplayableBodyReader, hide Close before the transport takes ownership
// of req.Body so the shared replay source remains available to GetBody.
if _, rawStorage := body.(common2.BodyStorage); rawStorage {
req.Body = io.NopCloser(body)
}
req.ContentLength = replayable.Size()
if req.GetBody == nil {
req.GetBody = info.UpstreamRequestGetBody
req.GetBody = replayable.NewReader
}
}
// ApplyUpstreamBodyMetadata restores metadata that net/http cannot infer when
// a BodyStorage is exposed through a type-erased reader. Provider adaptors
// that construct requests directly should call this before sending them.
func ApplyUpstreamBodyMetadata(req *http.Request, info *common.RelayInfo) {
applyUpstreamContentLength(req, info)
applyUpstreamGetBody(req, info)
}
func SetupApiRequestHeader(info *common.RelayInfo, c *gin.Context, req *http.Header) {
if info.RelayMode == constant.RelayModeAudioTranscription || info.RelayMode == constant.RelayModeAudioTranslation {
// multipart/form-data
@@ -341,7 +320,7 @@ func DoApiRequest(a Adaptor, c *gin.Context, info *common.RelayInfo, requestBody
if err != nil {
return nil, fmt.Errorf("new request failed: %w", err)
}
ApplyUpstreamBodyMetadata(req, info)
ApplyUpstreamBodyMetadata(req, requestBody)
headers := req.Header
err = a.SetupRequestHeader(c, &headers, info)
if err != nil {
@@ -371,7 +350,7 @@ func DoFormRequest(a Adaptor, c *gin.Context, info *common.RelayInfo, requestBod
if err != nil {
return nil, fmt.Errorf("new request failed: %w", err)
}
ApplyUpstreamBodyMetadata(req, info)
ApplyUpstreamBodyMetadata(req, requestBody)
// set form data
req.Header.Set("Content-Type", c.Request.Header.Get("Content-Type"))
headers := req.Header
@@ -588,13 +567,13 @@ func DoTaskApiRequest(a TaskAdaptor, c *gin.Context, info *common.RelayInfo, req
if err != nil {
return nil, fmt.Errorf("new request failed: %w", err)
}
ApplyUpstreamBodyMetadata(req, info)
ApplyUpstreamBodyMetadata(req, requestBody)
// Do NOT wrap requestBody in a GetBody closure here: returning the same
// (already consumed) reader would make any transport-level retry silently
// replay an empty body. http.NewRequest already derives a correct,
// snapshot-based GetBody for *bytes.Reader/Buffer/strings.Reader bodies
// (which most task adaptors pass in); for type-erased readers,
// ApplyUpstreamBodyMetadata wires a replayable body when one is available.
// (which most task adaptors pass in); ApplyUpstreamBodyMetadata wires the
// same contract for bodies that explicitly implement ReplayableBody.
// Otherwise GetBody stays nil so the transport fails the retry instead of
// sending a corrupted request.
+82 -91
View File
@@ -23,27 +23,29 @@ import (
"golang.org/x/net/http2/hpack"
)
func TestApplyUpstreamGetBody_SetsReplayableGetBody(t *testing.T) {
func TestApplyUpstreamBodyMetadataSetsReplayableMetadata(t *testing.T) {
t.Parallel()
payload := []byte(`{"model":"test-model","messages":[{"role":"user","content":"hi"}]}`)
body, size, getBody, closer, err := relaycommon.NewOutboundJSONBody(payload)
body, closer, err := relaycommon.NewOutboundJSONBody(payload)
require.NoError(t, err)
defer closer.Close()
// Mirror DoApiRequest: a type-erased io.Reader gives net/http neither
// ContentLength nor GetBody.
// NewRequest hides the body's dynamic type behind req.Body, so metadata
// extraction must use the original body passed to NewRequest.
req, err := http.NewRequest(http.MethodPost, "https://example.com/v1/chat/completions", body)
require.NoError(t, err)
assert.Nil(t, req.GetBody)
assert.Zero(t, req.ContentLength)
_, requestBodyIsReplayable := req.Body.(common.ReplayableBody)
assert.False(t, requestBodyIsReplayable)
info := &relaycommon.RelayInfo{
UpstreamRequestBodySize: size,
UpstreamRequestGetBody: getBody,
}
ApplyUpstreamBodyMetadata(req, info)
ApplyUpstreamBodyMetadata(req, req.Body)
assert.Nil(t, req.GetBody)
assert.Zero(t, req.ContentLength)
ApplyUpstreamBodyMetadata(req, body)
assert.EqualValues(t, len(payload), req.ContentLength)
require.NotNil(t, req.GetBody)
@@ -64,7 +66,40 @@ func TestApplyUpstreamGetBody_SetsReplayableGetBody(t *testing.T) {
}
}
func TestApplyUpstreamGetBody_KeepsExistingGetBody(t *testing.T) {
func TestApplyUpstreamBodyMetadataHidesRawBodyStorageCloser(t *testing.T) {
t.Parallel()
payload := []byte(`{"model":"test-model","input":"raw storage"}`)
storage, err := common.CreateBodyStorage(payload)
require.NoError(t, err)
defer storage.Close()
req, err := http.NewRequest(http.MethodPost, "https://example.com/v1/chat/completions", storage)
require.NoError(t, err)
_, exposesStorageBeforeApply := req.Body.(common.BodyStorage)
require.True(t, exposesStorageBeforeApply)
ApplyUpstreamBodyMetadata(req, storage)
_, exposesStorageAfterApply := req.Body.(common.BodyStorage)
assert.False(t, exposesStorageAfterApply)
assert.EqualValues(t, len(payload), req.ContentLength)
require.NotNil(t, req.GetBody)
sent, err := io.ReadAll(req.Body)
require.NoError(t, err)
assert.Equal(t, payload, sent)
require.NoError(t, req.Body.Close())
replayBody, err := req.GetBody()
require.NoError(t, err, "closing the HTTP request body must not close the shared storage")
replay, err := io.ReadAll(replayBody)
require.NoError(t, err)
require.NoError(t, replayBody.Close())
assert.Equal(t, payload, replay)
}
func TestApplyUpstreamBodyMetadataKeepsNativeMetadataForNonReplayableBody(t *testing.T) {
tests := []struct {
name string
body func() io.Reader
@@ -79,17 +114,12 @@ func TestApplyUpstreamGetBody_KeepsExistingGetBody(t *testing.T) {
t.Run(test.name, func(t *testing.T) {
t.Parallel()
req, err := http.NewRequest(http.MethodPost, "https://example.com/v1/chat/completions", test.body())
body := test.body()
req, err := http.NewRequest(http.MethodPost, "https://example.com/v1/chat/completions", body)
require.NoError(t, err)
require.NotNil(t, req.GetBody, "net/http must derive GetBody for the concrete reader")
info := &relaycommon.RelayInfo{
UpstreamRequestBodySize: 99,
UpstreamRequestGetBody: func() (io.ReadCloser, error) {
return io.NopCloser(bytes.NewReader([]byte("override"))), nil
},
}
ApplyUpstreamBodyMetadata(req, info)
ApplyUpstreamBodyMetadata(req, body)
rc, err := req.GetBody()
require.NoError(t, err)
@@ -102,36 +132,43 @@ func TestApplyUpstreamGetBody_KeepsExistingGetBody(t *testing.T) {
}
}
func TestApplyUpstreamGetBody_NoopWithoutReplaySource(t *testing.T) {
func TestApplyUpstreamBodyMetadataKeepsExistingGetBody(t *testing.T) {
t.Parallel()
storageBody, _, _, closer, err := relaycommon.NewOutboundJSONBody([]byte(`{}`))
payload := []byte(`{"model":"test-model"}`)
body, closer, err := relaycommon.NewOutboundJSONBody(payload)
require.NoError(t, err)
defer closer.Close()
req, err := http.NewRequest(http.MethodPost, "https://example.com/v1/chat/completions", storageBody)
req, err := http.NewRequest(http.MethodPost, "https://example.com/v1/chat/completions", body)
require.NoError(t, err)
applyUpstreamGetBody(req, nil)
assert.Nil(t, req.GetBody)
applyUpstreamGetBody(req, &relaycommon.RelayInfo{})
assert.Nil(t, req.GetBody)
req.ContentLength = 99
req.GetBody = func() (io.ReadCloser, error) {
return io.NopCloser(bytes.NewReader([]byte("existing"))), nil
}
func TestApplyUpstreamBodyMetadata_EmptyStorageRemainsReplayable(t *testing.T) {
ApplyUpstreamBodyMetadata(req, body)
assert.EqualValues(t, len(payload), req.ContentLength)
rc, err := req.GetBody()
require.NoError(t, err)
got, err := io.ReadAll(rc)
require.NoError(t, err)
require.NoError(t, rc.Close())
assert.Equal(t, "existing", string(got))
}
func TestApplyUpstreamBodyMetadataEmptyStorageRemainsReplayable(t *testing.T) {
t.Parallel()
storage, err := common.CreateBodyStorage(nil)
require.NoError(t, err)
defer storage.Close()
body := common.NewReplayableBodyReader(storage)
req, err := http.NewRequest(http.MethodPost, "https://example.com/v1/chat/completions", common.ReaderOnly(storage))
req, err := http.NewRequest(http.MethodPost, "https://example.com/v1/chat/completions", body)
require.NoError(t, err)
ApplyUpstreamBodyMetadata(req, &relaycommon.RelayInfo{
UpstreamRequestBodySize: storage.Size(),
UpstreamRequestGetBody: storage.NewReader,
})
ApplyUpstreamBodyMetadata(req, body)
assert.Zero(t, req.ContentLength)
require.NotNil(t, req.GetBody)
@@ -143,45 +180,6 @@ func TestApplyUpstreamBodyMetadata_EmptyStorageRemainsReplayable(t *testing.T) {
assert.Empty(t, replay)
}
func TestUpstreamBodyMetadataIsReboundAcrossChannelAttempts(t *testing.T) {
firstPayload := []byte(`{"attempt":"first-with-longer-body"}`)
_, firstSize, firstGetBody, firstCloser, err := relaycommon.NewOutboundJSONBody(firstPayload)
require.NoError(t, err)
info := &relaycommon.RelayInfo{
UpstreamRequestBodySize: firstSize,
UpstreamRequestGetBody: firstGetBody,
}
require.NoError(t, firstCloser.Close())
_, err = firstGetBody()
require.ErrorIs(t, err, common.ErrStorageClosed)
c, _ := gin.CreateTestContext(httptest.NewRecorder())
info.InitChannelMeta(c)
assert.Zero(t, info.UpstreamRequestBodySize)
assert.Nil(t, info.UpstreamRequestGetBody)
secondPayload := []byte(`{"attempt":"second"}`)
secondStorage, err := common.CreateBodyStorage(secondPayload)
require.NoError(t, err)
defer secondStorage.Close()
info.UpstreamRequestBodySize = secondStorage.Size()
info.UpstreamRequestGetBody = secondStorage.NewReader
req, err := http.NewRequest(http.MethodPost, "https://example.com/v1/chat/completions", common.ReaderOnly(secondStorage))
require.NoError(t, err)
ApplyUpstreamBodyMetadata(req, info)
assert.EqualValues(t, len(secondPayload), req.ContentLength)
require.NotNil(t, req.GetBody)
rc, err := req.GetBody()
require.NoError(t, err)
replay, err := io.ReadAll(rc)
require.NoError(t, err)
require.NoError(t, rc.Close())
assert.Equal(t, secondPayload, replay)
}
// stubTaskAdaptor implements just enough of TaskAdaptor for DoTaskApiRequest.
type stubTaskAdaptor struct {
TaskAdaptor
@@ -449,14 +447,11 @@ func newH2PriorKnowledgeClient(ln net.Listener) (*http.Client, *http2.Transport)
return &http.Client{Transport: transport, Timeout: 15 * time.Second}, transport
}
func newPassThroughBody(t *testing.T, payload []byte) (io.Reader, *relaycommon.RelayInfo, common.BodyStorage) {
func newPassThroughBody(t *testing.T, payload []byte) (common.ReplayableBody, common.BodyStorage) {
t.Helper()
storage, err := common.CreateBodyStorage(payload)
require.NoError(t, err)
return common.ReaderOnly(storage), &relaycommon.RelayInfo{
UpstreamRequestBodySize: storage.Size(),
UpstreamRequestGetBody: storage.NewReader,
}, storage
return common.NewReplayableBodyReader(storage), storage
}
// TestUpstreamGetBody_HTTP2RetryAfterUpstreamStreamReset exercises the actual
@@ -475,19 +470,15 @@ func TestUpstreamGetBody_HTTP2RetryAfterUpstreamStreamReset(t *testing.T) {
client, transport := newH2PriorKnowledgeClient(ln)
defer transport.CloseIdleConnections()
// Build the upstream request exactly the way DoApiRequest does: a
// type-erased BodyStorage reader plus the applyUpstream* helpers.
body, size, getBody, closer, err := relaycommon.NewOutboundJSONBody(payload)
// Build the upstream request exactly the way DoApiRequest does: pass the
// original replayable body to the metadata helper after NewRequest.
body, closer, err := relaycommon.NewOutboundJSONBody(payload)
require.NoError(t, err)
defer closer.Close()
req, err := http.NewRequest(http.MethodPost, "http://upstream.test/v1/chat/completions", body)
require.NoError(t, err)
info := &relaycommon.RelayInfo{
UpstreamRequestBodySize: size,
UpstreamRequestGetBody: getBody,
}
ApplyUpstreamBodyMetadata(req, info)
ApplyUpstreamBodyMetadata(req, body)
require.NotNil(t, req.GetBody)
resp, err := client.Do(req)
@@ -514,11 +505,11 @@ func TestUpstreamGetBody_HTTP2RetryAfterUpstreamStreamReset_PassThrough(t *testi
client, transport := newH2PriorKnowledgeClient(ln)
defer transport.CloseIdleConnections()
body, info, storage := newPassThroughBody(t, payload)
body, storage := newPassThroughBody(t, payload)
defer storage.Close()
req, err := http.NewRequest(http.MethodPost, "http://upstream.test/v1/chat/completions", body)
require.NoError(t, err)
ApplyUpstreamBodyMetadata(req, info)
ApplyUpstreamBodyMetadata(req, body)
require.NotNil(t, req.GetBody)
assert.EqualValues(t, len(payload), req.ContentLength)
@@ -546,11 +537,11 @@ func TestUpstreamGetBody_HTTP2RetryAfterGracefulGoAway_PassThrough(t *testing.T)
client, transport := newH2PriorKnowledgeClient(ln)
defer transport.CloseIdleConnections()
body, info, storage := newPassThroughBody(t, payload)
body, storage := newPassThroughBody(t, payload)
defer storage.Close()
req, err := http.NewRequest(http.MethodPost, "http://upstream.test/v1/chat/completions", body)
require.NoError(t, err)
ApplyUpstreamBodyMetadata(req, info)
ApplyUpstreamBodyMetadata(req, body)
require.NotNil(t, req.GetBody)
resp, err := client.Do(req)
@@ -580,13 +571,13 @@ func TestUpstreamGetBody_HTTP2CannotRetryWithoutGetBody(t *testing.T) {
client, transport := newH2PriorKnowledgeClient(ln)
defer transport.CloseIdleConnections()
body, size, _, closer, err := relaycommon.NewOutboundJSONBody(payload)
body, closer, err := relaycommon.NewOutboundJSONBody(payload)
require.NoError(t, err)
defer closer.Close()
req, err := http.NewRequest(http.MethodPost, "http://upstream.test/v1/chat/completions", body)
require.NoError(t, err)
applyUpstreamContentLength(req, &relaycommon.RelayInfo{UpstreamRequestBodySize: size})
req.ContentLength = body.Size()
assert.Nil(t, req.GetBody)
resp, err := client.Do(req) //nolint:bodyclose // Do fails, no body to close
+1 -1
View File
@@ -112,7 +112,7 @@ func (a *Adaptor) DoRequest(c *gin.Context, info *relaycommon.RelayInfo, request
if err != nil {
return nil, fmt.Errorf("new request failed: %w", err)
}
channel.ApplyUpstreamBodyMetadata(req, info)
channel.ApplyUpstreamBodyMetadata(req, requestBody)
err = Sign(c, req, info.ApiKey)
if err != nil {
return nil, fmt.Errorf("setup request header failed: %w", err)
+1 -3
View File
@@ -216,9 +216,7 @@ func (a *TaskAdaptor) BuildRequestBody(c *gin.Context, info *relaycommon.RelayIn
return &buf, nil
}
info.UpstreamRequestBodySize = storage.Size()
info.UpstreamRequestGetBody = storage.NewReader
return common.ReaderOnly(storage), nil
return common.NewReplayableBodyReader(storage), nil
}
// DoRequest delegates to common helper.
+5 -4
View File
@@ -14,7 +14,7 @@ import (
"github.com/stretchr/testify/require"
)
func TestSoraBuildRequestBodyBindsReplayMetadataForPassThrough(t *testing.T) {
func TestSoraBuildRequestBodyReturnsReplayablePassThroughBody(t *testing.T) {
payload := []byte("opaque-sora-request-body")
c, _ := gin.CreateTestContext(httptest.NewRecorder())
c.Request = httptest.NewRequest(http.MethodPost, "/v1/videos", bytes.NewReader(payload))
@@ -24,14 +24,15 @@ func TestSoraBuildRequestBodyBindsReplayMetadataForPassThrough(t *testing.T) {
info := &relaycommon.RelayInfo{}
body, err := (&TaskAdaptor{}).BuildRequestBody(c, info)
require.NoError(t, err)
replayable, ok := body.(common.ReplayableBody)
require.True(t, ok)
sent, err := io.ReadAll(body)
require.NoError(t, err)
assert.Equal(t, payload, sent)
assert.EqualValues(t, len(payload), info.UpstreamRequestBodySize)
require.NotNil(t, info.UpstreamRequestGetBody)
assert.EqualValues(t, len(payload), replayable.Size())
replayBody, err := info.UpstreamRequestGetBody()
replayBody, err := replayable.NewReader()
require.NoError(t, err)
replay, err := io.ReadAll(replayBody)
require.NoError(t, err)
+1 -3
View File
@@ -128,14 +128,12 @@ func chatCompletionsViaResponses(c *gin.Context, info *relaycommon.RelayInfo, ad
return nil, types.NewError(err, types.ErrorCodeConvertRequestFailed, types.ErrOptionWithSkipRetry())
}
body, size, getBody, closer, err := relaycommon.NewOutboundJSONBody(jsonData)
body, closer, err := relaycommon.NewOutboundJSONBody(jsonData)
if err != nil {
return nil, types.NewError(err, types.ErrorCodeConvertRequestFailed, types.ErrOptionWithSkipRetry())
}
defer closer.Close()
jsonData = nil
info.UpstreamRequestBodySize = size
info.UpstreamRequestGetBody = getBody
var requestBody io.Reader = body
var httpResp *http.Response
+2 -6
View File
@@ -159,9 +159,7 @@ func ClaudeHelper(c *gin.Context, info *relaycommon.RelayInfo) (newAPIError *typ
if err != nil {
return types.NewErrorWithStatusCode(err, types.ErrorCodeReadRequestBodyFailed, http.StatusBadRequest, types.ErrOptionWithSkipRetry())
}
info.UpstreamRequestBodySize = storage.Size()
info.UpstreamRequestGetBody = storage.NewReader
requestBody = common.ReaderOnly(storage)
requestBody = common.NewReplayableBodyReader(storage)
} else {
convertedRequest, err := adaptor.ConvertClaudeRequest(c, info, request)
if err != nil {
@@ -188,14 +186,12 @@ func ClaudeHelper(c *gin.Context, info *relaycommon.RelayInfo) (newAPIError *typ
}
logger.LogDebug(c, "requestBody: %s", jsonData)
body, size, getBody, closer, err := relaycommon.NewOutboundJSONBody(jsonData)
body, 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
}
+7 -19
View File
@@ -18,26 +18,14 @@ import (
// The caller MUST invoke closer.Close() once the upstream call has finished
// (typically via defer) to release the disk file / memory accounting.
//
// The returned reader is wrapped with common.ReaderOnly to prevent the HTTP
// transport from prematurely closing the underlying BodyStorage. The returned
// size is meant to be propagated to http.Request.ContentLength because the
// type-erased io.Reader prevents net/http from auto-detecting it.
//
// The returned getBody hands out a new, independent reader over the full body
// on every call, per the http.Request.GetBody contract of returning a fresh
// copy of the body. It is meant to be propagated to http.Request.GetBody
// (which net/http likewise cannot derive from a type-erased io.Reader) so the
// HTTP/2 transport can transparently retry the request when the upstream
// resets the stream after the body was already written ("http2: Transport:
// cannot retry err ... after Request.Body was written"). Each reader has its
// own cursor — in memory mode a fresh bytes.Reader over the shared immutable
// backing array, in disk mode a separate file descriptor — so replays never
// share seek state with the primary body or with each other, and closing a
// replayed reader never releases the underlying storage.
func NewOutboundJSONBody(data []byte) (body io.Reader, size int64, getBody func() (io.ReadCloser, error), closer io.Closer, err error) {
// The returned body exposes its size and replay capability without exposing
// io.Closer. Request construction uses that metadata to populate ContentLength
// and GetBody, while the caller retains ownership of the underlying storage
// through the separately returned closer.
func NewOutboundJSONBody(data []byte) (body common.ReplayableBody, closer io.Closer, err error) {
storage, err := common.CreateBodyStorage(data)
if err != nil {
return nil, 0, nil, nil, err
return nil, nil, err
}
return common.ReaderOnly(storage), storage.Size(), storage.NewReader, storage, nil
return common.NewReplayableBodyReader(storage), storage, nil
}
+15 -16
View File
@@ -14,12 +14,11 @@ func TestNewOutboundJSONBody_GetBodyReplaysFullBody(t *testing.T) {
payload := []byte(`{"model":"test-model","messages":[{"role":"user","content":"hello"}]}`)
body, size, getBody, closer, err := NewOutboundJSONBody(payload)
body, closer, err := NewOutboundJSONBody(payload)
require.NoError(t, err)
defer closer.Close()
assert.EqualValues(t, len(payload), size)
require.NotNil(t, getBody)
assert.EqualValues(t, len(payload), body.Size())
// Consume the primary body, as the HTTP transport does on the first attempt.
first, err := io.ReadAll(body)
@@ -29,7 +28,7 @@ func TestNewOutboundJSONBody_GetBodyReplaysFullBody(t *testing.T) {
// GetBody must hand out the complete body again — and repeatedly, since the
// transport may need more than one retry.
for i := 0; i < 2; i++ {
rc, err := getBody()
rc, err := body.NewReader()
require.NoError(t, err)
replay, err := io.ReadAll(rc)
require.NoError(t, err)
@@ -43,7 +42,7 @@ func TestNewOutboundJSONBody_GetBodyAfterPartialRead(t *testing.T) {
payload := []byte(`{"model":"test-model","input":"0123456789"}`)
body, _, getBody, closer, err := NewOutboundJSONBody(payload)
body, closer, err := NewOutboundJSONBody(payload)
require.NoError(t, err)
defer closer.Close()
@@ -52,7 +51,7 @@ func TestNewOutboundJSONBody_GetBodyAfterPartialRead(t *testing.T) {
_, err = io.ReadFull(body, partial)
require.NoError(t, err)
rc, err := getBody()
rc, err := body.NewReader()
require.NoError(t, err)
replay, err := io.ReadAll(rc)
require.NoError(t, err)
@@ -61,7 +60,7 @@ func TestNewOutboundJSONBody_GetBodyAfterPartialRead(t *testing.T) {
// Closing the replayed body must not close the underlying storage: the
// handler owns the storage lifetime via the returned closer.
require.NoError(t, rc.Close())
rc2, err := getBody()
rc2, err := body.NewReader()
require.NoError(t, err)
replay2, err := io.ReadAll(rc2)
require.NoError(t, err)
@@ -73,7 +72,7 @@ func TestNewOutboundJSONBody_GetBodyAfterPartialRead(t *testing.T) {
// independent cursors, per the http.Request.GetBody contract of returning a
// new copy of the body: interleaved reads across two replay readers and the
// primary body each observe exactly their own byte stream.
func assertIndependentReplayReaders(t *testing.T, payload []byte, body io.Reader, getBody func() (io.ReadCloser, error)) {
func assertIndependentReplayReaders(t *testing.T, payload []byte, body common.ReplayableBody) {
t.Helper()
half := len(payload) / 2
@@ -87,9 +86,9 @@ func assertIndependentReplayReaders(t *testing.T, payload []byte, body io.Reader
// Interleave two replay readers: A reads half, B reads everything, then A
// reads the rest.
a, err := getBody()
a, err := body.NewReader()
require.NoError(t, err)
b, err := getBody()
b, err := body.NewReader()
require.NoError(t, err)
aHead := make([]byte, half)
@@ -118,16 +117,16 @@ func TestNewOutboundJSONBody_GetBodyReadersAreIndependent(t *testing.T) {
payload := []byte(`{"model":"test-model","input":"abcdefghijklmnopqrstuvwxyz"}`)
body, _, getBody, closer, err := NewOutboundJSONBody(payload)
body, closer, err := NewOutboundJSONBody(payload)
require.NoError(t, err)
defer closer.Close()
assertIndependentReplayReaders(t, payload, body, getBody)
assertIndependentReplayReaders(t, payload, body)
// Once the handler releases the storage, GetBody must fail loudly instead
// of replaying stale data.
require.NoError(t, closer.Close())
_, err = getBody()
_, err = body.NewReader()
require.ErrorIs(t, err, common.ErrStorageClosed)
}
@@ -147,7 +146,7 @@ func TestNewOutboundJSONBody_GetBodyReadersAreIndependent_DiskStorage(t *testing
payload := []byte(`{"model":"test-model","input":"abcdefghijklmnopqrstuvwxyz"}`)
body, _, getBody, closer, err := NewOutboundJSONBody(payload)
body, closer, err := NewOutboundJSONBody(payload)
require.NoError(t, err)
defer closer.Close()
@@ -155,9 +154,9 @@ func TestNewOutboundJSONBody_GetBodyReadersAreIndependent_DiskStorage(t *testing
require.True(t, ok)
assert.True(t, storage.IsDisk(), "the payload must have taken the diskStorage path")
assertIndependentReplayReaders(t, payload, body, getBody)
assertIndependentReplayReaders(t, payload, body)
require.NoError(t, closer.Close())
_, err = getBody()
_, err = body.NewReader()
require.ErrorIs(t, err, common.ErrStorageClosed)
}
-24
View File
@@ -4,7 +4,6 @@ import (
"encoding/json"
"errors"
"fmt"
"io"
"strconv"
"strings"
"time"
@@ -150,23 +149,6 @@ type RelayInfo struct {
UseRuntimeHeadersOverride bool
ParamOverrideAudit []string
// UpstreamRequestBodySize is the byte size of the marshaled upstream request
// body. It is set when the body is wrapped in a BodyStorage (see
// relay/common/outbound_body.go), so that DoApiRequest can populate
// http.Request.ContentLength manually (net/http only auto-detects it for
// *bytes.Reader/Buffer/strings.Reader). 0 means "let net/http decide".
UpstreamRequestBodySize int64
// UpstreamRequestGetBody returns a fresh reader over the full marshaled
// upstream request body. It is set alongside UpstreamRequestBodySize when
// the body is wrapped in a BodyStorage (see relay/common/outbound_body.go),
// so that DoApiRequest can populate http.Request.GetBody manually (net/http
// only auto-populates it for *bytes.Reader/Buffer/strings.Reader). Without
// GetBody the HTTP/2 transport cannot transparently retry a request whose
// stream was reset by the upstream after the body was already written.
// nil means "no safe replay available".
UpstreamRequestGetBody func() (io.ReadCloser, error)
PriceData hosttypes.PriceData
// QuotaClamp is set (non-nil) when a quota conversion saturated at the
@@ -204,12 +186,6 @@ type RelayInfo struct {
}
func (info *RelayInfo) InitChannelMeta(c *gin.Context) {
// RelayInfo is reused across channel attempts. Body metadata belongs to the
// current attempt and may reference storage that its handler has closed, so
// discard it before the next channel binds its outbound body.
info.UpstreamRequestBodySize = 0
info.UpstreamRequestGetBody = nil
channelType := common.GetContextKeyInt(c, constant.ContextKeyChannelType)
paramOverride := common.GetContextKeyStringMap(c, constant.ContextKeyChannelParamOverride)
headerOverride := common.GetContextKeyStringMap(c, constant.ContextKeyChannelHeaderOverride)
-18
View File
@@ -1,32 +1,14 @@
package common
import (
"io"
"net/http/httptest"
"testing"
"github.com/QuantumNous/new-api/relaykit/relayconvert/convmeta"
"github.com/QuantumNous/new-api/relaykit/types"
"github.com/gin-gonic/gin"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
)
func TestInitChannelMetaClearsUpstreamBodyMetadata(t *testing.T) {
c, _ := gin.CreateTestContext(httptest.NewRecorder())
info := &RelayInfo{
UpstreamRequestBodySize: 37,
UpstreamRequestGetBody: func() (io.ReadCloser, error) {
return nil, nil
},
}
info.InitChannelMeta(c)
assert.Zero(t, info.UpstreamRequestBodySize)
assert.Nil(t, info.UpstreamRequestGetBody)
}
func TestRelayInfoGetFinalRequestRelayFormatPrefersExplicitFinal(t *testing.T) {
info := &RelayInfo{
RelayFormat: types.RelayFormatOpenAI,
+2 -6
View File
@@ -104,9 +104,7 @@ func TextHelper(c *gin.Context, info *relaycommon.RelayInfo) (newAPIError *types
logger.LogDebug(c, "requestBody: %s", debugBytes)
}
}
info.UpstreamRequestBodySize = storage.Size()
info.UpstreamRequestGetBody = storage.NewReader
requestBody = common.ReaderOnly(storage)
requestBody = common.NewReplayableBodyReader(storage)
} else {
convertedRequest, err := adaptor.ConvertOpenAIRequest(c, info, request)
if err != nil {
@@ -177,14 +175,12 @@ func TextHelper(c *gin.Context, info *relaycommon.RelayInfo) (newAPIError *types
logger.LogDebug(c, "text request body: %s", jsonData)
body, size, getBody, closer, err := relaycommon.NewOutboundJSONBody(jsonData)
body, 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
}
+1 -3
View File
@@ -58,14 +58,12 @@ func EmbeddingHelper(c *gin.Context, info *relaycommon.RelayInfo) (newAPIError *
}
logger.LogDebug(c, "converted embedding request body: %s", jsonData)
body, size, getBody, closer, err := relaycommon.NewOutboundJSONBody(jsonData)
body, 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
var requestBody io.Reader = body
statusCodeMappingStr := c.GetString("status_code_mapping")
resp, err := adaptor.DoRequest(c, info, requestBody)
+3 -9
View File
@@ -141,9 +141,7 @@ func GeminiHelper(c *gin.Context, info *relaycommon.RelayInfo) (newAPIError *typ
if err != nil {
return types.NewErrorWithStatusCode(err, types.ErrorCodeReadRequestBodyFailed, http.StatusBadRequest, types.ErrOptionWithSkipRetry())
}
info.UpstreamRequestBodySize = storage.Size()
info.UpstreamRequestGetBody = storage.NewReader
requestBody = common.ReaderOnly(storage)
requestBody = common.NewReplayableBodyReader(storage)
} else {
// 使用 ConvertGeminiRequest 转换请求格式
convertedRequest, err := adaptor.ConvertGeminiRequest(c, info, request)
@@ -166,14 +164,12 @@ func GeminiHelper(c *gin.Context, info *relaycommon.RelayInfo) (newAPIError *typ
logger.LogDebug(c, "Gemini request body: %s", jsonData)
body, size, getBody, closer, err := relaycommon.NewOutboundJSONBody(jsonData)
body, 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
}
@@ -272,14 +268,12 @@ func GeminiEmbeddingHandler(c *gin.Context, info *relaycommon.RelayInfo) (newAPI
}
}
logger.LogDebug(c, "Gemini embedding request body: %s", jsonData)
body, size, getBody, closer, err := relaycommon.NewOutboundJSONBody(jsonData)
body, 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
resp, err := adaptor.DoRequest(c, info, requestBody)
+2 -6
View File
@@ -51,9 +51,7 @@ func ImageHelper(c *gin.Context, info *relaycommon.RelayInfo) (newAPIError *type
if err != nil {
return types.NewErrorWithStatusCode(err, types.ErrorCodeReadRequestBodyFailed, http.StatusBadRequest, types.ErrOptionWithSkipRetry())
}
info.UpstreamRequestBodySize = storage.Size()
info.UpstreamRequestGetBody = storage.NewReader
requestBody = common.ReaderOnly(storage)
requestBody = common.NewReplayableBodyReader(storage)
} else {
convertedRequest, err := adaptor.ConvertImageRequest(c, info, *request)
if err != nil {
@@ -79,14 +77,12 @@ func ImageHelper(c *gin.Context, info *relaycommon.RelayInfo) (newAPIError *type
}
logger.LogDebug(c, "image request body: %s", jsonData)
body, size, getBody, closer, err := relaycommon.NewOutboundJSONBody(jsonData)
body, 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
}
}
+2 -6
View File
@@ -47,9 +47,7 @@ func RerankHelper(c *gin.Context, info *relaycommon.RelayInfo) (newAPIError *typ
if err != nil {
return types.NewErrorWithStatusCode(err, types.ErrorCodeReadRequestBodyFailed, http.StatusBadRequest, types.ErrOptionWithSkipRetry())
}
info.UpstreamRequestBodySize = storage.Size()
info.UpstreamRequestGetBody = storage.NewReader
requestBody = common.ReaderOnly(storage)
requestBody = common.NewReplayableBodyReader(storage)
} else {
convertedRequest, err := adaptor.ConvertRerankRequest(c, info.RelayMode, *request)
if err != nil {
@@ -70,14 +68,12 @@ func RerankHelper(c *gin.Context, info *relaycommon.RelayInfo) (newAPIError *typ
}
logger.LogDebug(c, "Rerank request body: %s", jsonData)
body, size, getBody, closer, err := relaycommon.NewOutboundJSONBody(jsonData)
body, 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
}
+2 -6
View File
@@ -82,9 +82,7 @@ func ResponsesHelper(c *gin.Context, info *relaycommon.RelayInfo) (newAPIError *
if err != nil {
return types.NewError(err, types.ErrorCodeReadRequestBodyFailed, types.ErrOptionWithSkipRetry())
}
info.UpstreamRequestBodySize = storage.Size()
info.UpstreamRequestGetBody = storage.NewReader
requestBody = common.ReaderOnly(storage)
requestBody = common.NewReplayableBodyReader(storage)
} else {
convertedRequest, err := adaptor.ConvertOpenAIResponsesRequest(c, info, *request)
if err != nil {
@@ -111,14 +109,12 @@ func ResponsesHelper(c *gin.Context, info *relaycommon.RelayInfo) (newAPIError *
}
logger.LogDebug(c, "requestBody: %s", jsonData)
body, size, getBody, closer, err := relaycommon.NewOutboundJSONBody(jsonData)
body, 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
}