Commit ea4f0210 by CaIon

refactor(relay): move replay metadata onto request bodies

parent d6b5ce99
...@@ -29,6 +29,14 @@ type BodyStorage interface { ...@@ -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 ...@@ -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 清理旧的缓存文件(用于启动时清理残留)
......
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)
}
...@@ -62,13 +62,11 @@ func AlphaSearchHelper(c *gin.Context, info *relaycommon.RelayInfo) (newAPIError ...@@ -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 {
......
...@@ -25,50 +25,29 @@ import ( ...@@ -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 ...@@ -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 ...@@ -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 ...@@ -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.
......
...@@ -112,7 +112,7 @@ func (a *Adaptor) DoRequest(c *gin.Context, info *relaycommon.RelayInfo, request ...@@ -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)
......
...@@ -216,9 +216,7 @@ func (a *TaskAdaptor) BuildRequestBody(c *gin.Context, info *relaycommon.RelayIn ...@@ -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.
......
...@@ -14,7 +14,7 @@ import ( ...@@ -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) { ...@@ -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)
......
...@@ -128,14 +128,12 @@ func chatCompletionsViaResponses(c *gin.Context, info *relaycommon.RelayInfo, ad ...@@ -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
......
...@@ -159,9 +159,7 @@ func ClaudeHelper(c *gin.Context, info *relaycommon.RelayInfo) (newAPIError *typ ...@@ -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 ...@@ -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
} }
......
...@@ -18,26 +18,14 @@ import ( ...@@ -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
} }
...@@ -14,12 +14,11 @@ func TestNewOutboundJSONBody_GetBodyReplaysFullBody(t *testing.T) { ...@@ -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) { ...@@ -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) { ...@@ -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) { ...@@ -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) { ...@@ -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) { ...@@ -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 ...@@ -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) { ...@@ -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 ...@@ -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 ...@@ -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)
} }
...@@ -4,7 +4,6 @@ import ( ...@@ -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 { ...@@ -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 { ...@@ -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)
......
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,
......
...@@ -104,9 +104,7 @@ func TextHelper(c *gin.Context, info *relaycommon.RelayInfo) (newAPIError *types ...@@ -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 ...@@ -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
} }
......
...@@ -58,14 +58,12 @@ func EmbeddingHelper(c *gin.Context, info *relaycommon.RelayInfo) (newAPIError * ...@@ -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)
......
...@@ -141,9 +141,7 @@ func GeminiHelper(c *gin.Context, info *relaycommon.RelayInfo) (newAPIError *typ ...@@ -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 ...@@ -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 ...@@ -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)
......
...@@ -51,9 +51,7 @@ func ImageHelper(c *gin.Context, info *relaycommon.RelayInfo) (newAPIError *type ...@@ -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 ...@@ -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
} }
} }
......
...@@ -47,9 +47,7 @@ func RerankHelper(c *gin.Context, info *relaycommon.RelayInfo) (newAPIError *typ ...@@ -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 ...@@ -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
} }
......
...@@ -82,9 +82,7 @@ func ResponsesHelper(c *gin.Context, info *relaycommon.RelayInfo) (newAPIError * ...@@ -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 * ...@@ -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
} }
......
Markdown is supported
0% or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or sign in to comment