Jaymr/new-api
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

refactor(relay): move replay metadata onto request bodies

Browse Source
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
Download Patch File Download Diff File Expand all files Collapse all files
relay/common/outbound_body.go
+7 -19
View File
@@ -18,26 +18,14 @@ import (
// The caller MUST invoke closer.Close() once the upstream call has finished
// (typically via defer) to release the disk file / memory accounting.
//
// The returned reader is wrapped with common.ReaderOnly to prevent the HTTP
// transport from prematurely closing the underlying BodyStorage. The returned
// size is meant to be propagated to http.Request.ContentLength because the
// type-erased io.Reader prevents net/http from auto-detecting it.
//
// The returned getBody hands out a new, independent reader over the full body
// on every call, per the http.Request.GetBody contract of returning a fresh
// copy of the body. It is meant to be propagated to http.Request.GetBody
// (which net/http likewise cannot derive from a type-erased io.Reader) so the
// HTTP/2 transport can transparently retry the request when the upstream
// resets the stream after the body was already written ("http2: Transport:
// cannot retry err ... after Request.Body was written"). Each reader has its
// own cursor — in memory mode a fresh bytes.Reader over the shared immutable
// backing array, in disk mode a separate file descriptor — so replays never
// share seek state with the primary body or with each other, and closing a
// replayed reader never releases the underlying storage.
func NewOutboundJSONBody(data []byte) (body io.Reader, size int64, getBody func() (io.ReadCloser, error), closer io.Closer, err error) {
// The returned body exposes its size and replay capability without exposing
// io.Closer. Request construction uses that metadata to populate ContentLength
// and GetBody, while the caller retains ownership of the underlying storage
// through the separately returned closer.
func NewOutboundJSONBody(data []byte) (body common.ReplayableBody, closer io.Closer, err error) {
storage, err := common.CreateBodyStorage(data)
if err != nil {
return nil, 0, nil, nil, err
return nil, nil, err
}
return common.ReaderOnly(storage), storage.Size(), storage.NewReader, storage, nil
return common.NewReplayableBodyReader(storage), storage, nil
}
relay/common/outbound_body_test.go
+15 -16
View File
@@ -14,12 +14,11 @@ func TestNewOutboundJSONBody_GetBodyReplaysFullBody(t *testing.T) {
payload := []byte(`{"model":"test-model","messages":[{"role":"user","content":"hello"}]}`)
body, size, getBody, closer, err := NewOutboundJSONBody(payload)
body, closer, err := NewOutboundJSONBody(payload)
require.NoError(t, err)
defer closer.Close()
assert.EqualValues(t, len(payload), size)
require.NotNil(t, getBody)
assert.EqualValues(t, len(payload), body.Size())
// Consume the primary body, as the HTTP transport does on the first attempt.
first, err := io.ReadAll(body)
@@ -29,7 +28,7 @@ func TestNewOutboundJSONBody_GetBodyReplaysFullBody(t *testing.T) {
// GetBody must hand out the complete body again — and repeatedly, since the
// transport may need more than one retry.
for i := 0; i < 2; i++ {
rc, err := getBody()
rc, err := body.NewReader()
require.NoError(t, err)
replay, err := io.ReadAll(rc)
require.NoError(t, err)
@@ -43,7 +42,7 @@ func TestNewOutboundJSONBody_GetBodyAfterPartialRead(t *testing.T) {
payload := []byte(`{"model":"test-model","input":"0123456789"}`)
body, _, getBody, closer, err := NewOutboundJSONBody(payload)
body, closer, err := NewOutboundJSONBody(payload)
require.NoError(t, err)
defer closer.Close()
@@ -52,7 +51,7 @@ func TestNewOutboundJSONBody_GetBodyAfterPartialRead(t *testing.T) {
_, err = io.ReadFull(body, partial)
require.NoError(t, err)
rc, err := getBody()
rc, err := body.NewReader()
require.NoError(t, err)
replay, err := io.ReadAll(rc)
require.NoError(t, err)
@@ -61,7 +60,7 @@ func TestNewOutboundJSONBody_GetBodyAfterPartialRead(t *testing.T) {
// Closing the replayed body must not close the underlying storage: the
// handler owns the storage lifetime via the returned closer.
require.NoError(t, rc.Close())
rc2, err := getBody()
rc2, err := body.NewReader()
require.NoError(t, err)
replay2, err := io.ReadAll(rc2)
require.NoError(t, err)
@@ -73,7 +72,7 @@ func TestNewOutboundJSONBody_GetBodyAfterPartialRead(t *testing.T) {
// independent cursors, per the http.Request.GetBody contract of returning a
// new copy of the body: interleaved reads across two replay readers and the
// primary body each observe exactly their own byte stream.
func assertIndependentReplayReaders(t *testing.T, payload []byte, body io.Reader, getBody func() (io.ReadCloser, error)) {
func assertIndependentReplayReaders(t *testing.T, payload []byte, body common.ReplayableBody) {
t.Helper()
half := len(payload) / 2
@@ -87,9 +86,9 @@ func assertIndependentReplayReaders(t *testing.T, payload []byte, body io.Reader
// Interleave two replay readers: A reads half, B reads everything, then A
// reads the rest.
a, err := getBody()
a, err := body.NewReader()
require.NoError(t, err)
b, err := getBody()
b, err := body.NewReader()
require.NoError(t, err)
aHead := make([]byte, half)
@@ -118,16 +117,16 @@ func TestNewOutboundJSONBody_GetBodyReadersAreIndependent(t *testing.T) {
payload := []byte(`{"model":"test-model","input":"abcdefghijklmnopqrstuvwxyz"}`)
body, _, getBody, closer, err := NewOutboundJSONBody(payload)
body, closer, err := NewOutboundJSONBody(payload)
require.NoError(t, err)
defer closer.Close()
assertIndependentReplayReaders(t, payload, body, getBody)
assertIndependentReplayReaders(t, payload, body)
// Once the handler releases the storage, GetBody must fail loudly instead
// of replaying stale data.
require.NoError(t, closer.Close())
_, err = getBody()
_, err = body.NewReader()
require.ErrorIs(t, err, common.ErrStorageClosed)
}
@@ -147,7 +146,7 @@ func TestNewOutboundJSONBody_GetBodyReadersAreIndependent_DiskStorage(t *testing
payload := []byte(`{"model":"test-model","input":"abcdefghijklmnopqrstuvwxyz"}`)
body, _, getBody, closer, err := NewOutboundJSONBody(payload)
body, closer, err := NewOutboundJSONBody(payload)
require.NoError(t, err)
defer closer.Close()
@@ -155,9 +154,9 @@ func TestNewOutboundJSONBody_GetBodyReadersAreIndependent_DiskStorage(t *testing
require.True(t, ok)
assert.True(t, storage.IsDisk(), "the payload must have taken the diskStorage path")
assertIndependentReplayReaders(t, payload, body, getBody)
assertIndependentReplayReaders(t, payload, body)
require.NoError(t, closer.Close())
_, err = getBody()
_, err = body.NewReader()
require.ErrorIs(t, err, common.ErrStorageClosed)
}
relay/common/relay_info.go
-24
View File
@@ -4,7 +4,6 @@ import (
"encoding/json"
"errors"
"fmt"
"io"
"strconv"
"strings"
"time"
@@ -150,23 +149,6 @@ type RelayInfo struct {
UseRuntimeHeadersOverride bool
ParamOverrideAudit []string
// UpstreamRequestBodySize is the byte size of the marshaled upstream request
// body. It is set when the body is wrapped in a BodyStorage (see
// relay/common/outbound_body.go), so that DoApiRequest can populate
// http.Request.ContentLength manually (net/http only auto-detects it for
// *bytes.Reader/Buffer/strings.Reader). 0 means "let net/http decide".
UpstreamRequestBodySize int64
// UpstreamRequestGetBody returns a fresh reader over the full marshaled
// upstream request body. It is set alongside UpstreamRequestBodySize when
// the body is wrapped in a BodyStorage (see relay/common/outbound_body.go),
// so that DoApiRequest can populate http.Request.GetBody manually (net/http
// only auto-populates it for *bytes.Reader/Buffer/strings.Reader). Without
// GetBody the HTTP/2 transport cannot transparently retry a request whose
// stream was reset by the upstream after the body was already written.
// nil means "no safe replay available".
UpstreamRequestGetBody func() (io.ReadCloser, error)
PriceData hosttypes.PriceData
// QuotaClamp is set (non-nil) when a quota conversion saturated at the
@@ -204,12 +186,6 @@ type RelayInfo struct {
}
func (info *RelayInfo) InitChannelMeta(c *gin.Context) {
// RelayInfo is reused across channel attempts. Body metadata belongs to the
// current attempt and may reference storage that its handler has closed, so
// discard it before the next channel binds its outbound body.
info.UpstreamRequestBodySize = 0
info.UpstreamRequestGetBody = nil
channelType := common.GetContextKeyInt(c, constant.ContextKeyChannelType)
paramOverride := common.GetContextKeyStringMap(c, constant.ContextKeyChannelParamOverride)
headerOverride := common.GetContextKeyStringMap(c, constant.ContextKeyChannelHeaderOverride)
relay/common/relay_info_test.go
-18
View File
@@ -1,32 +1,14 @@
package common
import (
"io"
"net/http/httptest"
"testing"
"github.com/QuantumNous/new-api/relaykit/relayconvert/convmeta"
"github.com/QuantumNous/new-api/relaykit/types"
"github.com/gin-gonic/gin"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
)
func TestInitChannelMetaClearsUpstreamBodyMetadata(t *testing.T) {
c, _ := gin.CreateTestContext(httptest.NewRecorder())
info := &RelayInfo{
UpstreamRequestBodySize: 37,
UpstreamRequestGetBody: func() (io.ReadCloser, error) {
return nil, nil
},
}
info.InitChannelMeta(c)
assert.Zero(t, info.UpstreamRequestBodySize)
assert.Nil(t, info.UpstreamRequestGetBody)
}
func TestRelayInfoGetFinalRequestRelayFormatPrefersExplicitFinal(t *testing.T) {
info := &RelayInfo{
RelayFormat: types.RelayFormatOpenAI,