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) 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 存储已关闭错误 // ErrStorageClosed 存储已关闭错误
var ErrStorageClosed = fmt.Errorf("body storage is closed") var ErrStorageClosed = fmt.Errorf("body storage is closed")
@@ -339,10 +347,27 @@ func CreateBodyStorageFromReader(reader io.Reader, contentLength int64, maxBytes
return storage, nil return storage, nil
} }
// ReaderOnly wraps an io.Reader to hide io.Closer, preventing http.NewRequest type replayableBodyReader struct {
// from type-asserting io.ReadCloser and closing the underlying BodyStorage. storage BodyStorage
func ReaderOnly(r io.Reader) io.Reader { }
return struct{ io.Reader }{r}
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 清理旧的缓存文件(用于启动时清理残留) // 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) logger.LogDebug(c, "requestBody: %s", jsonData)
body, size, getBody, closer, err := relaycommon.NewOutboundJSONBody(jsonData) body, closer, err := relaycommon.NewOutboundJSONBody(jsonData)
if err != nil { if err != nil {
return types.NewError(err, types.ErrorCodeConvertRequestFailed, types.ErrOptionWithSkipRetry()) return types.NewError(err, types.ErrorCodeConvertRequestFailed, types.ErrOptionWithSkipRetry())
} }
defer closer.Close() defer closer.Close()
info.UpstreamRequestBodySize = size
info.UpstreamRequestGetBody = getBody
adaptor := GetAdaptor(info.ApiType) adaptor := GetAdaptor(info.ApiType)
if adaptor == nil { if adaptor == nil {
+20 -41
View File
@@ -25,50 +25,29 @@ import (
"github.com/gorilla/websocket" "github.com/gorilla/websocket"
) )
// applyUpstreamContentLength populates req.ContentLength when the upstream // ApplyUpstreamBodyMetadata restores metadata that net/http cannot infer from
// body is wrapped in a BodyStorage (see relay/common/outbound_body.go). // a ReplayableBody. Callers must pass the original body because NewRequest
// // hides its dynamic type behind req.Body's io.ReadCloser wrapper.
// net/http.NewRequest only auto-detects ContentLength for *bytes.Reader, func ApplyUpstreamBodyMetadata(req *http.Request, body io.Reader) {
// *bytes.Buffer and *strings.Reader. When the body is a type-erased io.Reader replayable, ok := body.(common2.ReplayableBody)
// (which is the case for ReaderOnly(BodyStorage)), the Content-Length header if !ok {
// 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 {
return return
} }
if info.UpstreamRequestBodySize > 0 && req.ContentLength <= 0 {
req.ContentLength = info.UpstreamRequestBodySize
}
}
// applyUpstreamGetBody populates req.GetBody when the upstream body is wrapped // BodyStorage structurally satisfies ReplayableBody, but it also exposes
// in a BodyStorage (see relay/common/outbound_body.go). // io.Closer. If a caller passes the storage directly instead of using
// // NewReplayableBodyReader, hide Close before the transport takes ownership
// net/http.NewRequest only auto-populates GetBody for *bytes.Reader, // of req.Body so the shared replay source remains available to GetBody.
// *bytes.Buffer and *strings.Reader. When the body is a type-erased io.Reader if _, rawStorage := body.(common2.BodyStorage); rawStorage {
// (which is the case for ReaderOnly(BodyStorage)), GetBody would otherwise stay req.Body = io.NopCloser(body)
// 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
} }
req.ContentLength = replayable.Size()
if req.GetBody == nil { 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) { func SetupApiRequestHeader(info *common.RelayInfo, c *gin.Context, req *http.Header) {
if info.RelayMode == constant.RelayModeAudioTranscription || info.RelayMode == constant.RelayModeAudioTranslation { if info.RelayMode == constant.RelayModeAudioTranscription || info.RelayMode == constant.RelayModeAudioTranslation {
// multipart/form-data // multipart/form-data
@@ -341,7 +320,7 @@ func DoApiRequest(a Adaptor, c *gin.Context, info *common.RelayInfo, requestBody
if err != nil { if err != nil {
return nil, fmt.Errorf("new request failed: %w", err) return nil, fmt.Errorf("new request failed: %w", err)
} }
ApplyUpstreamBodyMetadata(req, info) ApplyUpstreamBodyMetadata(req, requestBody)
headers := req.Header headers := req.Header
err = a.SetupRequestHeader(c, &headers, info) err = a.SetupRequestHeader(c, &headers, info)
if err != nil { if err != nil {
@@ -371,7 +350,7 @@ func DoFormRequest(a Adaptor, c *gin.Context, info *common.RelayInfo, requestBod
if err != nil { if err != nil {
return nil, fmt.Errorf("new request failed: %w", err) return nil, fmt.Errorf("new request failed: %w", err)
} }
ApplyUpstreamBodyMetadata(req, info) ApplyUpstreamBodyMetadata(req, requestBody)
// set form data // set form data
req.Header.Set("Content-Type", c.Request.Header.Get("Content-Type")) req.Header.Set("Content-Type", c.Request.Header.Get("Content-Type"))
headers := req.Header headers := req.Header
@@ -588,13 +567,13 @@ func DoTaskApiRequest(a TaskAdaptor, c *gin.Context, info *common.RelayInfo, req
if err != nil { if err != nil {
return nil, fmt.Errorf("new request failed: %w", err) 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 // Do NOT wrap requestBody in a GetBody closure here: returning the same
// (already consumed) reader would make any transport-level retry silently // (already consumed) reader would make any transport-level retry silently
// replay an empty body. http.NewRequest already derives a correct, // replay an empty body. http.NewRequest already derives a correct,
// snapshot-based GetBody for *bytes.Reader/Buffer/strings.Reader bodies // snapshot-based GetBody for *bytes.Reader/Buffer/strings.Reader bodies
// (which most task adaptors pass in); for type-erased readers, // (which most task adaptors pass in); ApplyUpstreamBodyMetadata wires the
// ApplyUpstreamBodyMetadata wires a replayable body when one is available. // same contract for bodies that explicitly implement ReplayableBody.
// Otherwise GetBody stays nil so the transport fails the retry instead of // Otherwise GetBody stays nil so the transport fails the retry instead of
// sending a corrupted request. // sending a corrupted request.
+80 -89
View File
@@ -23,27 +23,29 @@ import (
"golang.org/x/net/http2/hpack" "golang.org/x/net/http2/hpack"
) )
func TestApplyUpstreamGetBody_SetsReplayableGetBody(t *testing.T) { func TestApplyUpstreamBodyMetadataSetsReplayableMetadata(t *testing.T) {
t.Parallel() t.Parallel()
payload := []byte(`{"model":"test-model","messages":[{"role":"user","content":"hi"}]}`) 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) require.NoError(t, err)
defer closer.Close() defer closer.Close()
// Mirror DoApiRequest: a type-erased io.Reader gives net/http neither // NewRequest hides the body's dynamic type behind req.Body, so metadata
// ContentLength nor GetBody. // extraction must use the original body passed to NewRequest.
req, err := http.NewRequest(http.MethodPost, "https://example.com/v1/chat/completions", body) req, err := http.NewRequest(http.MethodPost, "https://example.com/v1/chat/completions", body)
require.NoError(t, err) require.NoError(t, err)
assert.Nil(t, req.GetBody) assert.Nil(t, req.GetBody)
assert.Zero(t, req.ContentLength) assert.Zero(t, req.ContentLength)
_, requestBodyIsReplayable := req.Body.(common.ReplayableBody)
assert.False(t, requestBodyIsReplayable)
info := &relaycommon.RelayInfo{ ApplyUpstreamBodyMetadata(req, req.Body)
UpstreamRequestBodySize: size, assert.Nil(t, req.GetBody)
UpstreamRequestGetBody: getBody, assert.Zero(t, req.ContentLength)
}
ApplyUpstreamBodyMetadata(req, info) ApplyUpstreamBodyMetadata(req, body)
assert.EqualValues(t, len(payload), req.ContentLength) assert.EqualValues(t, len(payload), req.ContentLength)
require.NotNil(t, req.GetBody) 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 { tests := []struct {
name string name string
body func() io.Reader body func() io.Reader
@@ -79,17 +114,12 @@ func TestApplyUpstreamGetBody_KeepsExistingGetBody(t *testing.T) {
t.Run(test.name, func(t *testing.T) { t.Run(test.name, func(t *testing.T) {
t.Parallel() 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.NoError(t, err)
require.NotNil(t, req.GetBody, "net/http must derive GetBody for the concrete reader") require.NotNil(t, req.GetBody, "net/http must derive GetBody for the concrete reader")
info := &relaycommon.RelayInfo{ ApplyUpstreamBodyMetadata(req, body)
UpstreamRequestBodySize: 99,
UpstreamRequestGetBody: func() (io.ReadCloser, error) {
return io.NopCloser(bytes.NewReader([]byte("override"))), nil
},
}
ApplyUpstreamBodyMetadata(req, info)
rc, err := req.GetBody() rc, err := req.GetBody()
require.NoError(t, err) 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() t.Parallel()
storageBody, _, _, closer, err := relaycommon.NewOutboundJSONBody([]byte(`{}`)) payload := []byte(`{"model":"test-model"}`)
body, closer, err := relaycommon.NewOutboundJSONBody(payload)
require.NoError(t, err) require.NoError(t, err)
defer closer.Close() 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) require.NoError(t, err)
req.ContentLength = 99
req.GetBody = func() (io.ReadCloser, error) {
return io.NopCloser(bytes.NewReader([]byte("existing"))), nil
}
applyUpstreamGetBody(req, nil) ApplyUpstreamBodyMetadata(req, body)
assert.Nil(t, req.GetBody)
applyUpstreamGetBody(req, &relaycommon.RelayInfo{}) assert.EqualValues(t, len(payload), req.ContentLength)
assert.Nil(t, req.GetBody) 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() t.Parallel()
storage, err := common.CreateBodyStorage(nil) storage, err := common.CreateBodyStorage(nil)
require.NoError(t, err) require.NoError(t, err)
defer storage.Close() 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) require.NoError(t, err)
ApplyUpstreamBodyMetadata(req, &relaycommon.RelayInfo{ ApplyUpstreamBodyMetadata(req, body)
UpstreamRequestBodySize: storage.Size(),
UpstreamRequestGetBody: storage.NewReader,
})
assert.Zero(t, req.ContentLength) assert.Zero(t, req.ContentLength)
require.NotNil(t, req.GetBody) require.NotNil(t, req.GetBody)
@@ -143,45 +180,6 @@ func TestApplyUpstreamBodyMetadata_EmptyStorageRemainsReplayable(t *testing.T) {
assert.Empty(t, replay) 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. // stubTaskAdaptor implements just enough of TaskAdaptor for DoTaskApiRequest.
type stubTaskAdaptor struct { type stubTaskAdaptor struct {
TaskAdaptor TaskAdaptor
@@ -449,14 +447,11 @@ func newH2PriorKnowledgeClient(ln net.Listener) (*http.Client, *http2.Transport)
return &http.Client{Transport: transport, Timeout: 15 * time.Second}, 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() t.Helper()
storage, err := common.CreateBodyStorage(payload) storage, err := common.CreateBodyStorage(payload)
require.NoError(t, err) require.NoError(t, err)
return common.ReaderOnly(storage), &relaycommon.RelayInfo{ return common.NewReplayableBodyReader(storage), storage
UpstreamRequestBodySize: storage.Size(),
UpstreamRequestGetBody: storage.NewReader,
}, storage
} }
// TestUpstreamGetBody_HTTP2RetryAfterUpstreamStreamReset exercises the actual // TestUpstreamGetBody_HTTP2RetryAfterUpstreamStreamReset exercises the actual
@@ -475,19 +470,15 @@ func TestUpstreamGetBody_HTTP2RetryAfterUpstreamStreamReset(t *testing.T) {
client, transport := newH2PriorKnowledgeClient(ln) client, transport := newH2PriorKnowledgeClient(ln)
defer transport.CloseIdleConnections() defer transport.CloseIdleConnections()
// Build the upstream request exactly the way DoApiRequest does: a // Build the upstream request exactly the way DoApiRequest does: pass the
// type-erased BodyStorage reader plus the applyUpstream* helpers. // original replayable body to the metadata helper after NewRequest.
body, size, getBody, closer, err := relaycommon.NewOutboundJSONBody(payload) body, closer, err := relaycommon.NewOutboundJSONBody(payload)
require.NoError(t, err) require.NoError(t, err)
defer closer.Close() defer closer.Close()
req, err := http.NewRequest(http.MethodPost, "http://upstream.test/v1/chat/completions", body) req, err := http.NewRequest(http.MethodPost, "http://upstream.test/v1/chat/completions", body)
require.NoError(t, err) require.NoError(t, err)
info := &relaycommon.RelayInfo{ ApplyUpstreamBodyMetadata(req, body)
UpstreamRequestBodySize: size,
UpstreamRequestGetBody: getBody,
}
ApplyUpstreamBodyMetadata(req, info)
require.NotNil(t, req.GetBody) require.NotNil(t, req.GetBody)
resp, err := client.Do(req) resp, err := client.Do(req)
@@ -514,11 +505,11 @@ func TestUpstreamGetBody_HTTP2RetryAfterUpstreamStreamReset_PassThrough(t *testi
client, transport := newH2PriorKnowledgeClient(ln) client, transport := newH2PriorKnowledgeClient(ln)
defer transport.CloseIdleConnections() defer transport.CloseIdleConnections()
body, info, storage := newPassThroughBody(t, payload) body, storage := newPassThroughBody(t, payload)
defer storage.Close() defer storage.Close()
req, err := http.NewRequest(http.MethodPost, "http://upstream.test/v1/chat/completions", body) req, err := http.NewRequest(http.MethodPost, "http://upstream.test/v1/chat/completions", body)
require.NoError(t, err) require.NoError(t, err)
ApplyUpstreamBodyMetadata(req, info) ApplyUpstreamBodyMetadata(req, body)
require.NotNil(t, req.GetBody) require.NotNil(t, req.GetBody)
assert.EqualValues(t, len(payload), req.ContentLength) assert.EqualValues(t, len(payload), req.ContentLength)
@@ -546,11 +537,11 @@ func TestUpstreamGetBody_HTTP2RetryAfterGracefulGoAway_PassThrough(t *testing.T)
client, transport := newH2PriorKnowledgeClient(ln) client, transport := newH2PriorKnowledgeClient(ln)
defer transport.CloseIdleConnections() defer transport.CloseIdleConnections()
body, info, storage := newPassThroughBody(t, payload) body, storage := newPassThroughBody(t, payload)
defer storage.Close() defer storage.Close()
req, err := http.NewRequest(http.MethodPost, "http://upstream.test/v1/chat/completions", body) req, err := http.NewRequest(http.MethodPost, "http://upstream.test/v1/chat/completions", body)
require.NoError(t, err) require.NoError(t, err)
ApplyUpstreamBodyMetadata(req, info) ApplyUpstreamBodyMetadata(req, body)
require.NotNil(t, req.GetBody) require.NotNil(t, req.GetBody)
resp, err := client.Do(req) resp, err := client.Do(req)
@@ -580,13 +571,13 @@ func TestUpstreamGetBody_HTTP2CannotRetryWithoutGetBody(t *testing.T) {
client, transport := newH2PriorKnowledgeClient(ln) client, transport := newH2PriorKnowledgeClient(ln)
defer transport.CloseIdleConnections() defer transport.CloseIdleConnections()
body, size, _, closer, err := relaycommon.NewOutboundJSONBody(payload) body, closer, err := relaycommon.NewOutboundJSONBody(payload)
require.NoError(t, err) require.NoError(t, err)
defer closer.Close() defer closer.Close()
req, err := http.NewRequest(http.MethodPost, "http://upstream.test/v1/chat/completions", body) req, err := http.NewRequest(http.MethodPost, "http://upstream.test/v1/chat/completions", body)
require.NoError(t, err) require.NoError(t, err)
applyUpstreamContentLength(req, &relaycommon.RelayInfo{UpstreamRequestBodySize: size}) req.ContentLength = body.Size()
assert.Nil(t, req.GetBody) assert.Nil(t, req.GetBody)
resp, err := client.Do(req) //nolint:bodyclose // Do fails, no body to close 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 { if err != nil {
return nil, fmt.Errorf("new request failed: %w", err) return nil, fmt.Errorf("new request failed: %w", err)
} }
channel.ApplyUpstreamBodyMetadata(req, info) channel.ApplyUpstreamBodyMetadata(req, requestBody)
err = Sign(c, req, info.ApiKey) err = Sign(c, req, info.ApiKey)
if err != nil { if err != nil {
return nil, fmt.Errorf("setup request header failed: %w", err) 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 return &buf, nil
} }
info.UpstreamRequestBodySize = storage.Size() return common.NewReplayableBodyReader(storage), nil
info.UpstreamRequestGetBody = storage.NewReader
return common.ReaderOnly(storage), nil
} }
// DoRequest delegates to common helper. // DoRequest delegates to common helper.
+5 -4
View File
@@ -14,7 +14,7 @@ import (
"github.com/stretchr/testify/require" "github.com/stretchr/testify/require"
) )
func TestSoraBuildRequestBodyBindsReplayMetadataForPassThrough(t *testing.T) { func TestSoraBuildRequestBodyReturnsReplayablePassThroughBody(t *testing.T) {
payload := []byte("opaque-sora-request-body") payload := []byte("opaque-sora-request-body")
c, _ := gin.CreateTestContext(httptest.NewRecorder()) c, _ := gin.CreateTestContext(httptest.NewRecorder())
c.Request = httptest.NewRequest(http.MethodPost, "/v1/videos", bytes.NewReader(payload)) c.Request = httptest.NewRequest(http.MethodPost, "/v1/videos", bytes.NewReader(payload))
@@ -24,14 +24,15 @@ func TestSoraBuildRequestBodyBindsReplayMetadataForPassThrough(t *testing.T) {
info := &relaycommon.RelayInfo{} info := &relaycommon.RelayInfo{}
body, err := (&TaskAdaptor{}).BuildRequestBody(c, info) body, err := (&TaskAdaptor{}).BuildRequestBody(c, info)
require.NoError(t, err) require.NoError(t, err)
replayable, ok := body.(common.ReplayableBody)
require.True(t, ok)
sent, err := io.ReadAll(body) sent, err := io.ReadAll(body)
require.NoError(t, err) require.NoError(t, err)
assert.Equal(t, payload, sent) assert.Equal(t, payload, sent)
assert.EqualValues(t, len(payload), info.UpstreamRequestBodySize) assert.EqualValues(t, len(payload), replayable.Size())
require.NotNil(t, info.UpstreamRequestGetBody)
replayBody, err := info.UpstreamRequestGetBody() replayBody, err := replayable.NewReader()
require.NoError(t, err) require.NoError(t, err)
replay, err := io.ReadAll(replayBody) replay, err := io.ReadAll(replayBody)
require.NoError(t, err) 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()) 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 { if err != nil {
return nil, types.NewError(err, types.ErrorCodeConvertRequestFailed, types.ErrOptionWithSkipRetry()) return nil, types.NewError(err, types.ErrorCodeConvertRequestFailed, types.ErrOptionWithSkipRetry())
} }
defer closer.Close() defer closer.Close()
jsonData = nil jsonData = nil
info.UpstreamRequestBodySize = size
info.UpstreamRequestGetBody = getBody
var requestBody io.Reader = body var requestBody io.Reader = body
var httpResp *http.Response 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 { if err != nil {
return types.NewErrorWithStatusCode(err, types.ErrorCodeReadRequestBodyFailed, http.StatusBadRequest, types.ErrOptionWithSkipRetry()) return types.NewErrorWithStatusCode(err, types.ErrorCodeReadRequestBodyFailed, http.StatusBadRequest, types.ErrOptionWithSkipRetry())
} }
info.UpstreamRequestBodySize = storage.Size() requestBody = common.NewReplayableBodyReader(storage)
info.UpstreamRequestGetBody = storage.NewReader
requestBody = common.ReaderOnly(storage)
} else { } else {
convertedRequest, err := adaptor.ConvertClaudeRequest(c, info, request) convertedRequest, err := adaptor.ConvertClaudeRequest(c, info, request)
if err != nil { if err != nil {
@@ -188,14 +186,12 @@ func ClaudeHelper(c *gin.Context, info *relaycommon.RelayInfo) (newAPIError *typ
} }
logger.LogDebug(c, "requestBody: %s", jsonData) logger.LogDebug(c, "requestBody: %s", jsonData)
body, size, getBody, closer, err := relaycommon.NewOutboundJSONBody(jsonData) body, closer, err := relaycommon.NewOutboundJSONBody(jsonData)
if err != nil { if err != nil {
return types.NewError(err, types.ErrorCodeConvertRequestFailed, types.ErrOptionWithSkipRetry()) return types.NewError(err, types.ErrorCodeConvertRequestFailed, types.ErrOptionWithSkipRetry())
} }
defer closer.Close() defer closer.Close()
jsonData = nil jsonData = nil
info.UpstreamRequestBodySize = size
info.UpstreamRequestGetBody = getBody
requestBody = body requestBody = body
} }
+7 -19
View File
@@ -18,26 +18,14 @@ import (
// The caller MUST invoke closer.Close() once the upstream call has finished // The caller MUST invoke closer.Close() once the upstream call has finished
// (typically via defer) to release the disk file / memory accounting. // (typically via defer) to release the disk file / memory accounting.
// //
// The returned reader is wrapped with common.ReaderOnly to prevent the HTTP // The returned body exposes its size and replay capability without exposing
// transport from prematurely closing the underlying BodyStorage. The returned // io.Closer. Request construction uses that metadata to populate ContentLength
// size is meant to be propagated to http.Request.ContentLength because the // and GetBody, while the caller retains ownership of the underlying storage
// type-erased io.Reader prevents net/http from auto-detecting it. // through the separately returned closer.
// func NewOutboundJSONBody(data []byte) (body common.ReplayableBody, closer io.Closer, err error) {
// 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) {
storage, err := common.CreateBodyStorage(data) storage, err := common.CreateBodyStorage(data)
if err != nil { 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"}]}`) 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) require.NoError(t, err)
defer closer.Close() defer closer.Close()
assert.EqualValues(t, len(payload), size) assert.EqualValues(t, len(payload), body.Size())
require.NotNil(t, getBody)
// Consume the primary body, as the HTTP transport does on the first attempt. // Consume the primary body, as the HTTP transport does on the first attempt.
first, err := io.ReadAll(body) 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 // GetBody must hand out the complete body again — and repeatedly, since the
// transport may need more than one retry. // transport may need more than one retry.
for i := 0; i < 2; i++ { for i := 0; i < 2; i++ {
rc, err := getBody() rc, err := body.NewReader()
require.NoError(t, err) require.NoError(t, err)
replay, err := io.ReadAll(rc) replay, err := io.ReadAll(rc)
require.NoError(t, err) require.NoError(t, err)
@@ -43,7 +42,7 @@ func TestNewOutboundJSONBody_GetBodyAfterPartialRead(t *testing.T) {
payload := []byte(`{"model":"test-model","input":"0123456789"}`) payload := []byte(`{"model":"test-model","input":"0123456789"}`)
body, _, getBody, closer, err := NewOutboundJSONBody(payload) body, closer, err := NewOutboundJSONBody(payload)
require.NoError(t, err) require.NoError(t, err)
defer closer.Close() defer closer.Close()
@@ -52,7 +51,7 @@ func TestNewOutboundJSONBody_GetBodyAfterPartialRead(t *testing.T) {
_, err = io.ReadFull(body, partial) _, err = io.ReadFull(body, partial)
require.NoError(t, err) require.NoError(t, err)
rc, err := getBody() rc, err := body.NewReader()
require.NoError(t, err) require.NoError(t, err)
replay, err := io.ReadAll(rc) replay, err := io.ReadAll(rc)
require.NoError(t, err) 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 // Closing the replayed body must not close the underlying storage: the
// handler owns the storage lifetime via the returned closer. // handler owns the storage lifetime via the returned closer.
require.NoError(t, rc.Close()) require.NoError(t, rc.Close())
rc2, err := getBody() rc2, err := body.NewReader()
require.NoError(t, err) require.NoError(t, err)
replay2, err := io.ReadAll(rc2) replay2, err := io.ReadAll(rc2)
require.NoError(t, err) 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 // independent cursors, per the http.Request.GetBody contract of returning a
// new copy of the body: interleaved reads across two replay readers and the // new copy of the body: interleaved reads across two replay readers and the
// primary body each observe exactly their own byte stream. // 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() t.Helper()
half := len(payload) / 2 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 // Interleave two replay readers: A reads half, B reads everything, then A
// reads the rest. // reads the rest.
a, err := getBody() a, err := body.NewReader()
require.NoError(t, err) require.NoError(t, err)
b, err := getBody() b, err := body.NewReader()
require.NoError(t, err) require.NoError(t, err)
aHead := make([]byte, half) aHead := make([]byte, half)
@@ -118,16 +117,16 @@ func TestNewOutboundJSONBody_GetBodyReadersAreIndependent(t *testing.T) {
payload := []byte(`{"model":"test-model","input":"abcdefghijklmnopqrstuvwxyz"}`) payload := []byte(`{"model":"test-model","input":"abcdefghijklmnopqrstuvwxyz"}`)
body, _, getBody, closer, err := NewOutboundJSONBody(payload) body, closer, err := NewOutboundJSONBody(payload)
require.NoError(t, err) require.NoError(t, err)
defer closer.Close() defer closer.Close()
assertIndependentReplayReaders(t, payload, body, getBody) assertIndependentReplayReaders(t, payload, body)
// Once the handler releases the storage, GetBody must fail loudly instead // Once the handler releases the storage, GetBody must fail loudly instead
// of replaying stale data. // of replaying stale data.
require.NoError(t, closer.Close()) require.NoError(t, closer.Close())
_, err = getBody() _, err = body.NewReader()
require.ErrorIs(t, err, common.ErrStorageClosed) require.ErrorIs(t, err, common.ErrStorageClosed)
} }
@@ -147,7 +146,7 @@ func TestNewOutboundJSONBody_GetBodyReadersAreIndependent_DiskStorage(t *testing
payload := []byte(`{"model":"test-model","input":"abcdefghijklmnopqrstuvwxyz"}`) payload := []byte(`{"model":"test-model","input":"abcdefghijklmnopqrstuvwxyz"}`)
body, _, getBody, closer, err := NewOutboundJSONBody(payload) body, closer, err := NewOutboundJSONBody(payload)
require.NoError(t, err) require.NoError(t, err)
defer closer.Close() defer closer.Close()
@@ -155,9 +154,9 @@ func TestNewOutboundJSONBody_GetBodyReadersAreIndependent_DiskStorage(t *testing
require.True(t, ok) require.True(t, ok)
assert.True(t, storage.IsDisk(), "the payload must have taken the diskStorage path") 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()) require.NoError(t, closer.Close())
_, err = getBody() _, err = body.NewReader()
require.ErrorIs(t, err, common.ErrStorageClosed) require.ErrorIs(t, err, common.ErrStorageClosed)
} }
-24
View File
@@ -4,7 +4,6 @@ import (
"encoding/json" "encoding/json"
"errors" "errors"
"fmt" "fmt"
"io"
"strconv" "strconv"
"strings" "strings"
"time" "time"
@@ -150,23 +149,6 @@ type RelayInfo struct {
UseRuntimeHeadersOverride bool UseRuntimeHeadersOverride bool
ParamOverrideAudit []string 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 PriceData hosttypes.PriceData
// QuotaClamp is set (non-nil) when a quota conversion saturated at the // 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) { 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) channelType := common.GetContextKeyInt(c, constant.ContextKeyChannelType)
paramOverride := common.GetContextKeyStringMap(c, constant.ContextKeyChannelParamOverride) paramOverride := common.GetContextKeyStringMap(c, constant.ContextKeyChannelParamOverride)
headerOverride := common.GetContextKeyStringMap(c, constant.ContextKeyChannelHeaderOverride) headerOverride := common.GetContextKeyStringMap(c, constant.ContextKeyChannelHeaderOverride)
-18
View File
@@ -1,32 +1,14 @@
package common package common
import ( import (
"io"
"net/http/httptest"
"testing" "testing"
"github.com/QuantumNous/new-api/relaykit/relayconvert/convmeta" "github.com/QuantumNous/new-api/relaykit/relayconvert/convmeta"
"github.com/QuantumNous/new-api/relaykit/types" "github.com/QuantumNous/new-api/relaykit/types"
"github.com/gin-gonic/gin"
"github.com/stretchr/testify/assert" "github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require" "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) { func TestRelayInfoGetFinalRequestRelayFormatPrefersExplicitFinal(t *testing.T) {
info := &RelayInfo{ info := &RelayInfo{
RelayFormat: types.RelayFormatOpenAI, 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) logger.LogDebug(c, "requestBody: %s", debugBytes)
} }
} }
info.UpstreamRequestBodySize = storage.Size() requestBody = common.NewReplayableBodyReader(storage)
info.UpstreamRequestGetBody = storage.NewReader
requestBody = common.ReaderOnly(storage)
} else { } else {
convertedRequest, err := adaptor.ConvertOpenAIRequest(c, info, request) convertedRequest, err := adaptor.ConvertOpenAIRequest(c, info, request)
if err != nil { 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) 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 { if err != nil {
return types.NewError(err, types.ErrorCodeConvertRequestFailed, types.ErrOptionWithSkipRetry()) return types.NewError(err, types.ErrorCodeConvertRequestFailed, types.ErrOptionWithSkipRetry())
} }
defer closer.Close() defer closer.Close()
jsonData = nil jsonData = nil
info.UpstreamRequestBodySize = size
info.UpstreamRequestGetBody = getBody
requestBody = body 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) 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 { if err != nil {
return types.NewError(err, types.ErrorCodeConvertRequestFailed, types.ErrOptionWithSkipRetry()) return types.NewError(err, types.ErrorCodeConvertRequestFailed, types.ErrOptionWithSkipRetry())
} }
defer closer.Close() defer closer.Close()
jsonData = nil jsonData = nil
info.UpstreamRequestBodySize = size
info.UpstreamRequestGetBody = getBody
var requestBody io.Reader = body var requestBody io.Reader = body
statusCodeMappingStr := c.GetString("status_code_mapping") statusCodeMappingStr := c.GetString("status_code_mapping")
resp, err := adaptor.DoRequest(c, info, requestBody) 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 { if err != nil {
return types.NewErrorWithStatusCode(err, types.ErrorCodeReadRequestBodyFailed, http.StatusBadRequest, types.ErrOptionWithSkipRetry()) return types.NewErrorWithStatusCode(err, types.ErrorCodeReadRequestBodyFailed, http.StatusBadRequest, types.ErrOptionWithSkipRetry())
} }
info.UpstreamRequestBodySize = storage.Size() requestBody = common.NewReplayableBodyReader(storage)
info.UpstreamRequestGetBody = storage.NewReader
requestBody = common.ReaderOnly(storage)
} else { } else {
// 使用 ConvertGeminiRequest 转换请求格式 // 使用 ConvertGeminiRequest 转换请求格式
convertedRequest, err := adaptor.ConvertGeminiRequest(c, info, request) 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) 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 { if err != nil {
return types.NewError(err, types.ErrorCodeConvertRequestFailed, types.ErrOptionWithSkipRetry()) return types.NewError(err, types.ErrorCodeConvertRequestFailed, types.ErrOptionWithSkipRetry())
} }
defer closer.Close() defer closer.Close()
jsonData = nil jsonData = nil
info.UpstreamRequestBodySize = size
info.UpstreamRequestGetBody = getBody
requestBody = body requestBody = body
} }
@@ -272,14 +268,12 @@ func GeminiEmbeddingHandler(c *gin.Context, info *relaycommon.RelayInfo) (newAPI
} }
} }
logger.LogDebug(c, "Gemini embedding request body: %s", jsonData) 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 { if err != nil {
return types.NewError(err, types.ErrorCodeConvertRequestFailed, types.ErrOptionWithSkipRetry()) return types.NewError(err, types.ErrorCodeConvertRequestFailed, types.ErrOptionWithSkipRetry())
} }
defer closer.Close() defer closer.Close()
jsonData = nil jsonData = nil
info.UpstreamRequestBodySize = size
info.UpstreamRequestGetBody = getBody
requestBody = body requestBody = body
resp, err := adaptor.DoRequest(c, info, requestBody) 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 { if err != nil {
return types.NewErrorWithStatusCode(err, types.ErrorCodeReadRequestBodyFailed, http.StatusBadRequest, types.ErrOptionWithSkipRetry()) return types.NewErrorWithStatusCode(err, types.ErrorCodeReadRequestBodyFailed, http.StatusBadRequest, types.ErrOptionWithSkipRetry())
} }
info.UpstreamRequestBodySize = storage.Size() requestBody = common.NewReplayableBodyReader(storage)
info.UpstreamRequestGetBody = storage.NewReader
requestBody = common.ReaderOnly(storage)
} else { } else {
convertedRequest, err := adaptor.ConvertImageRequest(c, info, *request) convertedRequest, err := adaptor.ConvertImageRequest(c, info, *request)
if err != nil { 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) 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 { if err != nil {
return types.NewError(err, types.ErrorCodeConvertRequestFailed, types.ErrOptionWithSkipRetry()) return types.NewError(err, types.ErrorCodeConvertRequestFailed, types.ErrOptionWithSkipRetry())
} }
defer closer.Close() defer closer.Close()
jsonData = nil jsonData = nil
info.UpstreamRequestBodySize = size
info.UpstreamRequestGetBody = getBody
requestBody = body requestBody = body
} }
} }
+2 -6
View File
@@ -47,9 +47,7 @@ func RerankHelper(c *gin.Context, info *relaycommon.RelayInfo) (newAPIError *typ
if err != nil { if err != nil {
return types.NewErrorWithStatusCode(err, types.ErrorCodeReadRequestBodyFailed, http.StatusBadRequest, types.ErrOptionWithSkipRetry()) return types.NewErrorWithStatusCode(err, types.ErrorCodeReadRequestBodyFailed, http.StatusBadRequest, types.ErrOptionWithSkipRetry())
} }
info.UpstreamRequestBodySize = storage.Size() requestBody = common.NewReplayableBodyReader(storage)
info.UpstreamRequestGetBody = storage.NewReader
requestBody = common.ReaderOnly(storage)
} else { } else {
convertedRequest, err := adaptor.ConvertRerankRequest(c, info.RelayMode, *request) convertedRequest, err := adaptor.ConvertRerankRequest(c, info.RelayMode, *request)
if err != nil { 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) 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 { if err != nil {
return types.NewError(err, types.ErrorCodeConvertRequestFailed, types.ErrOptionWithSkipRetry()) return types.NewError(err, types.ErrorCodeConvertRequestFailed, types.ErrOptionWithSkipRetry())
} }
defer closer.Close() defer closer.Close()
jsonData = nil jsonData = nil
info.UpstreamRequestBodySize = size
info.UpstreamRequestGetBody = getBody
requestBody = body requestBody = body
} }
+2 -6
View File
@@ -82,9 +82,7 @@ func ResponsesHelper(c *gin.Context, info *relaycommon.RelayInfo) (newAPIError *
if err != nil { if err != nil {
return types.NewError(err, types.ErrorCodeReadRequestBodyFailed, types.ErrOptionWithSkipRetry()) return types.NewError(err, types.ErrorCodeReadRequestBodyFailed, types.ErrOptionWithSkipRetry())
} }
info.UpstreamRequestBodySize = storage.Size() requestBody = common.NewReplayableBodyReader(storage)
info.UpstreamRequestGetBody = storage.NewReader
requestBody = common.ReaderOnly(storage)
} else { } else {
convertedRequest, err := adaptor.ConvertOpenAIResponsesRequest(c, info, *request) convertedRequest, err := adaptor.ConvertOpenAIResponsesRequest(c, info, *request)
if err != nil { if err != nil {
@@ -111,14 +109,12 @@ func ResponsesHelper(c *gin.Context, info *relaycommon.RelayInfo) (newAPIError *
} }
logger.LogDebug(c, "requestBody: %s", jsonData) logger.LogDebug(c, "requestBody: %s", jsonData)
body, size, getBody, closer, err := relaycommon.NewOutboundJSONBody(jsonData) body, closer, err := relaycommon.NewOutboundJSONBody(jsonData)
if err != nil { if err != nil {
return types.NewError(err, types.ErrorCodeConvertRequestFailed, types.ErrOptionWithSkipRetry()) return types.NewError(err, types.ErrorCodeConvertRequestFailed, types.ErrOptionWithSkipRetry())
} }
defer closer.Close() defer closer.Close()
jsonData = nil jsonData = nil
info.UpstreamRequestBodySize = size
info.UpstreamRequestGetBody = getBody
requestBody = body requestBody = body
} }