From ea4f021012cddc52126123ab4ed8ced3df260b85 Mon Sep 17 00:00:00 2001 From: CaIon Date: Thu, 6 Aug 2026 17:33:10 +0800 Subject: [PATCH] refactor(relay): move replay metadata onto request bodies --- common/body_storage.go | 33 ++++- common/body_storage_test.go | 37 +++++ relay/alpha_search_handler.go | 4 +- relay/channel/api_request.go | 61 +++----- relay/channel/api_request_getbody_test.go | 169 ++++++++++------------ relay/channel/jimeng/adaptor.go | 2 +- relay/channel/task/sora/adaptor.go | 4 +- relay/channel/task/sora/adaptor_test.go | 9 +- relay/chat_completions_via_responses.go | 4 +- relay/claude_handler.go | 8 +- relay/common/outbound_body.go | 26 +--- relay/common/outbound_body_test.go | 31 ++-- relay/common/relay_info.go | 24 --- relay/common/relay_info_test.go | 18 --- relay/compatible_handler.go | 8 +- relay/embedding_handler.go | 4 +- relay/gemini_handler.go | 12 +- relay/image_handler.go | 8 +- relay/rerank_handler.go | 8 +- relay/responses_handler.go | 8 +- 20 files changed, 211 insertions(+), 267 deletions(-) create mode 100644 common/body_storage_test.go diff --git a/common/body_storage.go b/common/body_storage.go index 0c4a849d..515d749e 100644 --- a/common/body_storage.go +++ b/common/body_storage.go @@ -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 清理旧的缓存文件（用于启动时清理残留） diff --git a/common/body_storage_test.go b/common/body_storage_test.go new file mode 100644 index 00000000..c3f877a6 --- /dev/null +++ b/common/body_storage_test.go @@ -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) +} diff --git a/relay/alpha_search_handler.go b/relay/alpha_search_handler.go index c66cb1de..246eebe8 100644 --- a/relay/alpha_search_handler.go +++ b/relay/alpha_search_handler.go @@ -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 { diff --git a/relay/channel/api_request.go b/relay/channel/api_request.go index 2ed9f69f..48241b14 100644 --- a/relay/channel/api_request.go +++ b/relay/channel/api_request.go @@ -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. diff --git a/relay/channel/api_request_getbody_test.go b/relay/channel/api_request_getbody_test.go index 847468e1..9a2de733 100644 --- a/relay/channel/api_request_getbody_test.go +++ b/relay/channel/api_request_getbody_test.go @@ -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) + req.ContentLength = 99 + req.GetBody = func() (io.ReadCloser, error) { + return io.NopCloser(bytes.NewReader([]byte("existing"))), nil + } - applyUpstreamGetBody(req, nil) - assert.Nil(t, req.GetBody) + ApplyUpstreamBodyMetadata(req, body) - applyUpstreamGetBody(req, &relaycommon.RelayInfo{}) - assert.Nil(t, req.GetBody) + 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 TestApplyUpstreamBodyMetadata_EmptyStorageRemainsReplayable(t *testing.T) { +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 diff --git a/relay/channel/jimeng/adaptor.go b/relay/channel/jimeng/adaptor.go index d2f9c389..776c1829 100644 --- a/relay/channel/jimeng/adaptor.go +++ b/relay/channel/jimeng/adaptor.go @@ -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) diff --git a/relay/channel/task/sora/adaptor.go b/relay/channel/task/sora/adaptor.go index 59a7a258..7f81e533 100644 --- a/relay/channel/task/sora/adaptor.go +++ b/relay/channel/task/sora/adaptor.go @@ -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. diff --git a/relay/channel/task/sora/adaptor_test.go b/relay/channel/task/sora/adaptor_test.go index ef514952..7021f2a9 100644 --- a/relay/channel/task/sora/adaptor_test.go +++ b/relay/channel/task/sora/adaptor_test.go @@ -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) diff --git a/relay/chat_completions_via_responses.go b/relay/chat_completions_via_responses.go index 660cc889..b8a6fc87 100644 --- a/relay/chat_completions_via_responses.go +++ b/relay/chat_completions_via_responses.go @@ -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 diff --git a/relay/claude_handler.go b/relay/claude_handler.go index bcbe9908..c8b01f7a 100644 --- a/relay/claude_handler.go +++ b/relay/claude_handler.go @@ -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 } diff --git a/relay/common/outbound_body.go b/relay/common/outbound_body.go index 0316cf74..ac9e110e 100644 --- a/relay/common/outbound_body.go +++ b/relay/common/outbound_body.go @@ -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 } diff --git a/relay/common/outbound_body_test.go b/relay/common/outbound_body_test.go index 23e2ff95..28c3aa69 100644 --- a/relay/common/outbound_body_test.go +++ b/relay/common/outbound_body_test.go @@ -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) } diff --git a/relay/common/relay_info.go b/relay/common/relay_info.go index 52c74115..594e2640 100644 --- a/relay/common/relay_info.go +++ b/relay/common/relay_info.go @@ -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) diff --git a/relay/common/relay_info_test.go b/relay/common/relay_info_test.go index edf5fdbf..9e58f3f9 100644 --- a/relay/common/relay_info_test.go +++ b/relay/common/relay_info_test.go @@ -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, diff --git a/relay/compatible_handler.go b/relay/compatible_handler.go index 24f355f9..8edb3362 100644 --- a/relay/compatible_handler.go +++ b/relay/compatible_handler.go @@ -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 } diff --git a/relay/embedding_handler.go b/relay/embedding_handler.go index d3b617b6..44d3de48 100644 --- a/relay/embedding_handler.go +++ b/relay/embedding_handler.go @@ -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) diff --git a/relay/gemini_handler.go b/relay/gemini_handler.go index 38ac0d87..57010d87 100644 --- a/relay/gemini_handler.go +++ b/relay/gemini_handler.go @@ -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) diff --git a/relay/image_handler.go b/relay/image_handler.go index 2d99a63c..690f229f 100644 --- a/relay/image_handler.go +++ b/relay/image_handler.go @@ -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 } } diff --git a/relay/rerank_handler.go b/relay/rerank_handler.go index e2546dda..460fe305 100644 --- a/relay/rerank_handler.go +++ b/relay/rerank_handler.go @@ -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 } diff --git a/relay/responses_handler.go b/relay/responses_handler.go index 2531b3a6..f8b4dec2 100644 --- a/relay/responses_handler.go +++ b/relay/responses_handler.go @@ -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 }