fix(relay): set Request.GetBody so the HTTP/2 transport can transparently retry after an upstream stream reset (#6249)
* fix(relay): set Request.GetBody so the HTTP/2 transport can transparently retry after an upstream stream reset
The outbound request body is a type-erased io.Reader over BodyStorage, so
net/http cannot derive Request.GetBody (it only does so for *bytes.Reader,
*bytes.Buffer and *strings.Reader). With GetBody nil, the HTTP/2 transport
cannot transparently retry a request once the body has been written and the
upstream resets the stream with a retryable error (REFUSED_STREAM, or a
connection-level GOAWAY); the relay request then fails with:
http2: Transport: cannot retry err [...] after Request.Body was written;
define Request.GetBody to avoid this error
This affects every relay path that goes through DoApiRequest (chat, claude,
gemini, responses, embedding, image, rerank).
BodyStorage (memory and disk) already implements io.Seeker, so replay support
only needed wiring:
- NewOutboundJSONBody additionally returns a getBody that rewinds the storage
and hands out a fresh non-closing reader. The transport only calls GetBody
after the previous attempt's body has been abandoned, so the rewind cannot
race an in-flight read.
- RelayInfo carries it in the new UpstreamRequestGetBody field, set alongside
UpstreamRequestBodySize by the handlers that build storage-backed bodies.
- applyUpstreamGetBody (symmetric with applyUpstreamContentLength) wires it
into DoApiRequest/DoFormRequest/DoTaskApiRequest, only when req.GetBody is
still nil.
Also remove the hand-rolled GetBody override in DoTaskApiRequest: it returned
the same already-consumed reader, so any transport-level replay would have
silently sent an empty body, and it clobbered the correct snapshot-based
GetBody that net/http derives from the *bytes.Reader bodies the task adaptors
pass in. For non-replayable bodies GetBody now stays nil, so a retry fails
loudly instead of corrupting the request.
Covered by unit tests plus an end-to-end raw-frame HTTP/2 test that resets
the first stream with REFUSED_STREAM after the body is written and asserts
the transport transparently retries with the complete body.
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
* fix(relay): hand out independent readers from GetBody (address review)
Per the http.Request.GetBody contract ("returns a new copy of Body"),
each call must yield a reader with its own cursor. The previous
implementation rewound and reused the shared BodyStorage, so two
consecutive GetBody readers would interfere with each other, and a
replay could disturb the primary body's offset under extreme transport
timing (e.g. attempt N's body write not yet fully abandoned when the
transport builds attempt N+1).
Instead of snapshotting the payload (an extra copy), add
BodyStorage.NewReader, which returns an independent zero-copy reader:
- memory mode: a fresh bytes.Reader over the same immutable backing
array;
- disk mode: a separate file descriptor over the cache file, so the
transport closing a replayed body only closes that descriptor.
NewOutboundJSONBody's getBody now simply hands out storage.NewReader,
and once the handler releases the storage, GetBody fails with
ErrStorageClosed instead of replaying stale data.
Tests: interleaved reads across two replay readers and the primary
body each observe exactly their own byte stream, for both the memory
and the disk-backed storage; the existing GetBody and HTTP/2 retry
suites still pass (h2 e2e tests flake-free with -count=20).
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
* fix(relay): bind replayable metadata on pass-through requests
* fix(relay): reset upstream body metadata between channels
* test(relay): cover replay across retries and channel attempts
* fix(relay): stop following upstream redirects
---------
Co-authored-by: Claude Fable 5 <noreply@anthropic.com>
This commit is contained in:
@@ -20,6 +20,13 @@ type BodyStorage interface {
|
|||||||
Size() int64
|
Size() int64
|
||||||
// IsDisk 是否是磁盘存储
|
// IsDisk 是否是磁盘存储
|
||||||
IsDisk() bool
|
IsDisk() bool
|
||||||
|
// NewReader returns an independent reader positioned at the start of the
|
||||||
|
// stored payload. Each call returns a reader with its own cursor, so
|
||||||
|
// callers (e.g. http.Request.GetBody) can replay the body concurrently
|
||||||
|
// with, or after, other readers without sharing seek state. Closing the
|
||||||
|
// returned reader releases only that reader, never the storage itself;
|
||||||
|
// after the storage has been closed, NewReader returns ErrStorageClosed.
|
||||||
|
NewReader() (io.ReadCloser, error)
|
||||||
}
|
}
|
||||||
|
|
||||||
// ErrStorageClosed 存储已关闭错误
|
// ErrStorageClosed 存储已关闭错误
|
||||||
@@ -80,6 +87,18 @@ func (m *memoryStorage) Bytes() ([]byte, error) {
|
|||||||
return m.data, nil
|
return m.data, nil
|
||||||
}
|
}
|
||||||
|
|
||||||
|
func (m *memoryStorage) NewReader() (io.ReadCloser, error) {
|
||||||
|
m.mu.Lock()
|
||||||
|
defer m.mu.Unlock()
|
||||||
|
if atomic.LoadInt32(&m.closed) == 1 {
|
||||||
|
return nil, ErrStorageClosed
|
||||||
|
}
|
||||||
|
// A fresh bytes.Reader over the shared immutable backing array: an
|
||||||
|
// independent cursor at zero copy cost. NopCloser keeps Close a no-op, so
|
||||||
|
// the storage lifecycle stays owned by whoever holds the storage itself.
|
||||||
|
return io.NopCloser(bytes.NewReader(m.data)), nil
|
||||||
|
}
|
||||||
|
|
||||||
func (m *memoryStorage) Size() int64 {
|
func (m *memoryStorage) Size() int64 {
|
||||||
return m.size
|
return m.size
|
||||||
}
|
}
|
||||||
@@ -229,6 +248,24 @@ func (d *diskStorage) Bytes() ([]byte, error) {
|
|||||||
return data, nil
|
return data, nil
|
||||||
}
|
}
|
||||||
|
|
||||||
|
func (d *diskStorage) NewReader() (io.ReadCloser, error) {
|
||||||
|
d.mu.Lock()
|
||||||
|
defer d.mu.Unlock()
|
||||||
|
if atomic.LoadInt32(&d.closed) == 1 {
|
||||||
|
return nil, ErrStorageClosed
|
||||||
|
}
|
||||||
|
// A separate file descriptor over the same cache file: an independent
|
||||||
|
// cursor at zero copy cost. Closing the returned reader closes only that
|
||||||
|
// descriptor; the storage keeps owning the primary descriptor and the
|
||||||
|
// file's lifetime. Readers opened before Close stay usable even after the
|
||||||
|
// file is unlinked, as the descriptor keeps the inode alive.
|
||||||
|
file, err := os.Open(d.filePath)
|
||||||
|
if err != nil {
|
||||||
|
return nil, fmt.Errorf("failed to open body cache file for replay: %w", err)
|
||||||
|
}
|
||||||
|
return file, nil
|
||||||
|
}
|
||||||
|
|
||||||
func (d *diskStorage) Size() int64 {
|
func (d *diskStorage) Size() int64 {
|
||||||
return d.size
|
return d.size
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -62,12 +62,13 @@ func AlphaSearchHelper(c *gin.Context, info *relaycommon.RelayInfo) (newAPIError
|
|||||||
}
|
}
|
||||||
|
|
||||||
logger.LogDebug(c, "requestBody: %s", jsonData)
|
logger.LogDebug(c, "requestBody: %s", jsonData)
|
||||||
body, size, closer, err := relaycommon.NewOutboundJSONBody(jsonData)
|
body, size, getBody, 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.UpstreamRequestBodySize = size
|
||||||
|
info.UpstreamRequestGetBody = getBody
|
||||||
|
|
||||||
adaptor := GetAdaptor(info.ApiType)
|
adaptor := GetAdaptor(info.ApiType)
|
||||||
if adaptor == nil {
|
if adaptor == nil {
|
||||||
|
|||||||
@@ -42,6 +42,33 @@ func applyUpstreamContentLength(req *http.Request, info *common.RelayInfo) {
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
// applyUpstreamGetBody populates req.GetBody when the upstream body is wrapped
|
||||||
|
// in a BodyStorage (see relay/common/outbound_body.go).
|
||||||
|
//
|
||||||
|
// net/http.NewRequest only auto-populates GetBody for *bytes.Reader,
|
||||||
|
// *bytes.Buffer and *strings.Reader. When the body is a type-erased io.Reader
|
||||||
|
// (which is the case for ReaderOnly(BodyStorage)), GetBody would otherwise stay
|
||||||
|
// nil, and the HTTP/2 transport cannot transparently retry the request once the
|
||||||
|
// upstream resets the stream after the body was already written; the request
|
||||||
|
// then fails with "http2: Transport: cannot retry err ... after Request.Body
|
||||||
|
// was written; define Request.GetBody to avoid this error".
|
||||||
|
func applyUpstreamGetBody(req *http.Request, info *common.RelayInfo) {
|
||||||
|
if info == nil || info.UpstreamRequestGetBody == nil {
|
||||||
|
return
|
||||||
|
}
|
||||||
|
if req.GetBody == nil {
|
||||||
|
req.GetBody = info.UpstreamRequestGetBody
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// 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
|
||||||
@@ -314,7 +341,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)
|
||||||
}
|
}
|
||||||
applyUpstreamContentLength(req, info)
|
ApplyUpstreamBodyMetadata(req, info)
|
||||||
headers := req.Header
|
headers := req.Header
|
||||||
err = a.SetupRequestHeader(c, &headers, info)
|
err = a.SetupRequestHeader(c, &headers, info)
|
||||||
if err != nil {
|
if err != nil {
|
||||||
@@ -344,7 +371,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)
|
||||||
}
|
}
|
||||||
applyUpstreamContentLength(req, info)
|
ApplyUpstreamBodyMetadata(req, info)
|
||||||
// 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
|
||||||
@@ -474,11 +501,24 @@ func sendPingData(c *gin.Context, mutex *sync.Mutex) error {
|
|||||||
func DoRequest(c *gin.Context, req *http.Request, info *common.RelayInfo) (*http.Response, error) {
|
func DoRequest(c *gin.Context, req *http.Request, info *common.RelayInfo) (*http.Response, error) {
|
||||||
return doRequest(c, req, info)
|
return doRequest(c, req, info)
|
||||||
}
|
}
|
||||||
|
|
||||||
|
// keepUpstreamRedirectResponse stops net/http from following redirects while
|
||||||
|
// returning the upstream 3xx response to the relay without an extra error.
|
||||||
|
func keepUpstreamRedirectResponse(_ *http.Request, _ []*http.Request) error {
|
||||||
|
return http.ErrUseLastResponse
|
||||||
|
}
|
||||||
|
|
||||||
func doRequest(c *gin.Context, req *http.Request, info *common.RelayInfo) (*http.Response, error) {
|
func doRequest(c *gin.Context, req *http.Request, info *common.RelayInfo) (*http.Response, error) {
|
||||||
client, err := service.GetHttpClientWithProxySettings(info.ChannelSetting.Proxy, info.ChannelSetting)
|
client, err := service.GetHttpClientWithProxySettings(info.ChannelSetting.Proxy, info.ChannelSetting)
|
||||||
if err != nil {
|
if err != nil {
|
||||||
return nil, fmt.Errorf("new proxy http client failed: %w", err)
|
return nil, fmt.Errorf("new proxy http client failed: %w", err)
|
||||||
}
|
}
|
||||||
|
// Clients are cached and shared across channels, so override redirect
|
||||||
|
// behavior on a shallow copy instead of mutating the cached client. This
|
||||||
|
// still reuses its transport and connection pools, including HTTP/2's
|
||||||
|
// transparent stream retries.
|
||||||
|
relayClient := *client
|
||||||
|
relayClient.CheckRedirect = keepUpstreamRedirectResponse
|
||||||
if common2.DebugEnabled && req != nil && req.URL != nil {
|
if common2.DebugEnabled && req != nil && req.URL != nil {
|
||||||
policy := service.NormalizeHTTPTransportPolicy(info.ChannelSetting)
|
policy := service.NormalizeHTTPTransportPolicy(info.ChannelSetting)
|
||||||
logger.LogDebug(c, fmt.Sprintf(
|
logger.LogDebug(c, fmt.Sprintf(
|
||||||
@@ -510,7 +550,7 @@ func doRequest(c *gin.Context, req *http.Request, info *common.RelayInfo) (*http
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
resp, err := client.Do(req)
|
resp, err := relayClient.Do(req)
|
||||||
if err != nil {
|
if err != nil {
|
||||||
logger.LogError(c, "do request failed: "+err.Error())
|
logger.LogError(c, "do request failed: "+err.Error())
|
||||||
return nil, types.NewError(err, types.ErrorCodeDoRequestFailed, types.ErrOptionWithHideErrMsg("upstream error: do request failed"))
|
return nil, types.NewError(err, types.ErrorCodeDoRequestFailed, types.ErrOptionWithHideErrMsg("upstream error: do request failed"))
|
||||||
@@ -548,10 +588,15 @@ 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)
|
||||||
}
|
}
|
||||||
applyUpstreamContentLength(req, info)
|
ApplyUpstreamBodyMetadata(req, info)
|
||||||
req.GetBody = func() (io.ReadCloser, error) {
|
// Do NOT wrap requestBody in a GetBody closure here: returning the same
|
||||||
return io.NopCloser(requestBody), nil
|
// (already consumed) reader would make any transport-level retry silently
|
||||||
}
|
// replay an empty body. http.NewRequest already derives a correct,
|
||||||
|
// snapshot-based GetBody for *bytes.Reader/Buffer/strings.Reader bodies
|
||||||
|
// (which most task adaptors pass in); for type-erased readers,
|
||||||
|
// ApplyUpstreamBodyMetadata wires a replayable body when one is available.
|
||||||
|
// Otherwise GetBody stays nil so the transport fails the retry instead of
|
||||||
|
// sending a corrupted request.
|
||||||
|
|
||||||
err = a.BuildRequestHeader(c, req, info)
|
err = a.BuildRequestHeader(c, req, info)
|
||||||
if err != nil {
|
if err != nil {
|
||||||
|
|||||||
@@ -0,0 +1,603 @@
|
|||||||
|
package channel
|
||||||
|
|
||||||
|
import (
|
||||||
|
"bytes"
|
||||||
|
"context"
|
||||||
|
"crypto/tls"
|
||||||
|
"fmt"
|
||||||
|
"io"
|
||||||
|
"net"
|
||||||
|
"net/http"
|
||||||
|
"net/http/httptest"
|
||||||
|
"strings"
|
||||||
|
"testing"
|
||||||
|
"time"
|
||||||
|
|
||||||
|
"github.com/QuantumNous/new-api/common"
|
||||||
|
relaycommon "github.com/QuantumNous/new-api/relay/common"
|
||||||
|
"github.com/QuantumNous/new-api/service"
|
||||||
|
"github.com/gin-gonic/gin"
|
||||||
|
"github.com/stretchr/testify/assert"
|
||||||
|
"github.com/stretchr/testify/require"
|
||||||
|
"golang.org/x/net/http2"
|
||||||
|
"golang.org/x/net/http2/hpack"
|
||||||
|
)
|
||||||
|
|
||||||
|
func TestApplyUpstreamGetBody_SetsReplayableGetBody(t *testing.T) {
|
||||||
|
t.Parallel()
|
||||||
|
|
||||||
|
payload := []byte(`{"model":"test-model","messages":[{"role":"user","content":"hi"}]}`)
|
||||||
|
|
||||||
|
body, size, getBody, closer, err := relaycommon.NewOutboundJSONBody(payload)
|
||||||
|
require.NoError(t, err)
|
||||||
|
defer closer.Close()
|
||||||
|
|
||||||
|
// Mirror DoApiRequest: a type-erased io.Reader gives net/http neither
|
||||||
|
// ContentLength nor GetBody.
|
||||||
|
req, err := http.NewRequest(http.MethodPost, "https://example.com/v1/chat/completions", body)
|
||||||
|
require.NoError(t, err)
|
||||||
|
assert.Nil(t, req.GetBody)
|
||||||
|
assert.Zero(t, req.ContentLength)
|
||||||
|
|
||||||
|
info := &relaycommon.RelayInfo{
|
||||||
|
UpstreamRequestBodySize: size,
|
||||||
|
UpstreamRequestGetBody: getBody,
|
||||||
|
}
|
||||||
|
ApplyUpstreamBodyMetadata(req, info)
|
||||||
|
|
||||||
|
assert.EqualValues(t, len(payload), req.ContentLength)
|
||||||
|
require.NotNil(t, req.GetBody)
|
||||||
|
|
||||||
|
// Drain the primary body as the transport does on the first attempt, then
|
||||||
|
// make sure GetBody can replay the complete payload repeatedly.
|
||||||
|
sent, err := io.ReadAll(req.Body)
|
||||||
|
require.NoError(t, err)
|
||||||
|
assert.Equal(t, payload, sent)
|
||||||
|
|
||||||
|
for i := 0; i < 2; i++ {
|
||||||
|
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, payload, replay, "replay %d must equal the original payload", i+1)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestApplyUpstreamGetBody_KeepsExistingGetBody(t *testing.T) {
|
||||||
|
tests := []struct {
|
||||||
|
name string
|
||||||
|
body func() io.Reader
|
||||||
|
}{
|
||||||
|
{name: "bytes reader", body: func() io.Reader { return bytes.NewReader([]byte("original")) }},
|
||||||
|
{name: "bytes buffer", body: func() io.Reader { return bytes.NewBufferString("original") }},
|
||||||
|
{name: "strings reader", body: func() io.Reader { return strings.NewReader("original") }},
|
||||||
|
}
|
||||||
|
|
||||||
|
for _, test := range tests {
|
||||||
|
test := test
|
||||||
|
t.Run(test.name, func(t *testing.T) {
|
||||||
|
t.Parallel()
|
||||||
|
|
||||||
|
req, err := http.NewRequest(http.MethodPost, "https://example.com/v1/chat/completions", test.body())
|
||||||
|
require.NoError(t, err)
|
||||||
|
require.NotNil(t, req.GetBody, "net/http must derive GetBody for the concrete reader")
|
||||||
|
|
||||||
|
info := &relaycommon.RelayInfo{
|
||||||
|
UpstreamRequestBodySize: 99,
|
||||||
|
UpstreamRequestGetBody: func() (io.ReadCloser, error) {
|
||||||
|
return io.NopCloser(bytes.NewReader([]byte("override"))), nil
|
||||||
|
},
|
||||||
|
}
|
||||||
|
ApplyUpstreamBodyMetadata(req, info)
|
||||||
|
|
||||||
|
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, "original", string(got), "an already correct GetBody must not be overwritten")
|
||||||
|
assert.EqualValues(t, len("original"), req.ContentLength, "native content length must not be overwritten")
|
||||||
|
})
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestApplyUpstreamGetBody_NoopWithoutReplaySource(t *testing.T) {
|
||||||
|
t.Parallel()
|
||||||
|
|
||||||
|
storageBody, _, _, closer, err := relaycommon.NewOutboundJSONBody([]byte(`{}`))
|
||||||
|
require.NoError(t, err)
|
||||||
|
defer closer.Close()
|
||||||
|
|
||||||
|
req, err := http.NewRequest(http.MethodPost, "https://example.com/v1/chat/completions", storageBody)
|
||||||
|
require.NoError(t, err)
|
||||||
|
|
||||||
|
applyUpstreamGetBody(req, nil)
|
||||||
|
assert.Nil(t, req.GetBody)
|
||||||
|
|
||||||
|
applyUpstreamGetBody(req, &relaycommon.RelayInfo{})
|
||||||
|
assert.Nil(t, req.GetBody)
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestApplyUpstreamBodyMetadata_EmptyStorageRemainsReplayable(t *testing.T) {
|
||||||
|
t.Parallel()
|
||||||
|
|
||||||
|
storage, err := common.CreateBodyStorage(nil)
|
||||||
|
require.NoError(t, err)
|
||||||
|
defer storage.Close()
|
||||||
|
|
||||||
|
req, err := http.NewRequest(http.MethodPost, "https://example.com/v1/chat/completions", common.ReaderOnly(storage))
|
||||||
|
require.NoError(t, err)
|
||||||
|
ApplyUpstreamBodyMetadata(req, &relaycommon.RelayInfo{
|
||||||
|
UpstreamRequestBodySize: storage.Size(),
|
||||||
|
UpstreamRequestGetBody: storage.NewReader,
|
||||||
|
})
|
||||||
|
|
||||||
|
assert.Zero(t, 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.Empty(t, replay)
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestUpstreamBodyMetadataIsReboundAcrossChannelAttempts(t *testing.T) {
|
||||||
|
firstPayload := []byte(`{"attempt":"first-with-longer-body"}`)
|
||||||
|
_, firstSize, firstGetBody, firstCloser, err := relaycommon.NewOutboundJSONBody(firstPayload)
|
||||||
|
require.NoError(t, err)
|
||||||
|
|
||||||
|
info := &relaycommon.RelayInfo{
|
||||||
|
UpstreamRequestBodySize: firstSize,
|
||||||
|
UpstreamRequestGetBody: firstGetBody,
|
||||||
|
}
|
||||||
|
require.NoError(t, firstCloser.Close())
|
||||||
|
_, err = firstGetBody()
|
||||||
|
require.ErrorIs(t, err, common.ErrStorageClosed)
|
||||||
|
|
||||||
|
c, _ := gin.CreateTestContext(httptest.NewRecorder())
|
||||||
|
info.InitChannelMeta(c)
|
||||||
|
assert.Zero(t, info.UpstreamRequestBodySize)
|
||||||
|
assert.Nil(t, info.UpstreamRequestGetBody)
|
||||||
|
|
||||||
|
secondPayload := []byte(`{"attempt":"second"}`)
|
||||||
|
secondStorage, err := common.CreateBodyStorage(secondPayload)
|
||||||
|
require.NoError(t, err)
|
||||||
|
defer secondStorage.Close()
|
||||||
|
info.UpstreamRequestBodySize = secondStorage.Size()
|
||||||
|
info.UpstreamRequestGetBody = secondStorage.NewReader
|
||||||
|
|
||||||
|
req, err := http.NewRequest(http.MethodPost, "https://example.com/v1/chat/completions", common.ReaderOnly(secondStorage))
|
||||||
|
require.NoError(t, err)
|
||||||
|
ApplyUpstreamBodyMetadata(req, info)
|
||||||
|
|
||||||
|
assert.EqualValues(t, len(secondPayload), req.ContentLength)
|
||||||
|
require.NotNil(t, req.GetBody)
|
||||||
|
rc, err := req.GetBody()
|
||||||
|
require.NoError(t, err)
|
||||||
|
replay, err := io.ReadAll(rc)
|
||||||
|
require.NoError(t, err)
|
||||||
|
require.NoError(t, rc.Close())
|
||||||
|
assert.Equal(t, secondPayload, replay)
|
||||||
|
}
|
||||||
|
|
||||||
|
// stubTaskAdaptor implements just enough of TaskAdaptor for DoTaskApiRequest.
|
||||||
|
type stubTaskAdaptor struct {
|
||||||
|
TaskAdaptor
|
||||||
|
baseURL string
|
||||||
|
capturedReq *http.Request
|
||||||
|
}
|
||||||
|
|
||||||
|
func (s *stubTaskAdaptor) BuildRequestURL(info *relaycommon.RelayInfo) (string, error) {
|
||||||
|
return s.baseURL + "/v1/video/generations", nil
|
||||||
|
}
|
||||||
|
|
||||||
|
func (s *stubTaskAdaptor) BuildRequestHeader(c *gin.Context, req *http.Request, info *relaycommon.RelayInfo) error {
|
||||||
|
s.capturedReq = req
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
|
||||||
|
// TestDoTaskApiRequest_KeepsReplayableGetBody guards against reintroducing the
|
||||||
|
// hand-rolled GetBody override that wrapped the already consumed request
|
||||||
|
// reader: any transport-level retry would then have silently replayed an empty
|
||||||
|
// body. net/http derives a correct snapshot-based GetBody from the
|
||||||
|
// *bytes.Reader bodies the task adaptors pass in, and it must be left intact.
|
||||||
|
func TestDoTaskApiRequest_KeepsReplayableGetBody(t *testing.T) {
|
||||||
|
service.InitHttpClient()
|
||||||
|
|
||||||
|
payload := []byte(`{"model":"test-model","prompt":"hello"}`)
|
||||||
|
|
||||||
|
type receivedBody struct {
|
||||||
|
body []byte
|
||||||
|
err error
|
||||||
|
}
|
||||||
|
receivedCh := make(chan receivedBody, 1)
|
||||||
|
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||||
|
body, err := io.ReadAll(r.Body)
|
||||||
|
receivedCh <- receivedBody{body: body, err: err}
|
||||||
|
w.WriteHeader(http.StatusOK)
|
||||||
|
}))
|
||||||
|
defer server.Close()
|
||||||
|
|
||||||
|
gin.SetMode(gin.TestMode)
|
||||||
|
recorder := httptest.NewRecorder()
|
||||||
|
ctx, _ := gin.CreateTestContext(recorder)
|
||||||
|
ctx.Request = httptest.NewRequest(http.MethodPost, "/v1/video/generations", bytes.NewReader(payload))
|
||||||
|
|
||||||
|
info := &relaycommon.RelayInfo{
|
||||||
|
ChannelMeta: &relaycommon.ChannelMeta{},
|
||||||
|
}
|
||||||
|
|
||||||
|
adaptor := &stubTaskAdaptor{baseURL: server.URL}
|
||||||
|
resp, err := DoTaskApiRequest(adaptor, ctx, info, bytes.NewReader(payload))
|
||||||
|
require.NoError(t, err)
|
||||||
|
defer resp.Body.Close()
|
||||||
|
assert.Equal(t, http.StatusOK, resp.StatusCode)
|
||||||
|
received := <-receivedCh
|
||||||
|
require.NoError(t, received.err)
|
||||||
|
assert.Equal(t, payload, received.body)
|
||||||
|
|
||||||
|
req := adaptor.capturedReq
|
||||||
|
require.NotNil(t, req)
|
||||||
|
require.NotNil(t, req.GetBody)
|
||||||
|
// Even after the request body has been fully written, GetBody must still
|
||||||
|
// return the complete payload, repeatedly.
|
||||||
|
for i := 0; i < 2; i++ {
|
||||||
|
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, payload, replay, "replay %d must equal the original payload", i+1)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
type h2ServerResult struct {
|
||||||
|
err error
|
||||||
|
streamCount int
|
||||||
|
attemptBodies [][]byte
|
||||||
|
}
|
||||||
|
|
||||||
|
func acceptH2TestConnection(ln net.Listener) (net.Conn, *http2.Framer, error) {
|
||||||
|
conn, err := ln.Accept()
|
||||||
|
if err != nil {
|
||||||
|
return nil, nil, err
|
||||||
|
}
|
||||||
|
_ = conn.SetDeadline(time.Now().Add(15 * time.Second))
|
||||||
|
|
||||||
|
preface := make([]byte, len(http2.ClientPreface))
|
||||||
|
if _, err := io.ReadFull(conn, preface); err != nil {
|
||||||
|
conn.Close()
|
||||||
|
return nil, nil, fmt.Errorf("read client preface: %w", err)
|
||||||
|
}
|
||||||
|
if !bytes.Equal(preface, []byte(http2.ClientPreface)) {
|
||||||
|
conn.Close()
|
||||||
|
return nil, nil, fmt.Errorf("unexpected client preface")
|
||||||
|
}
|
||||||
|
|
||||||
|
framer := http2.NewFramer(conn, conn)
|
||||||
|
framer.ReadMetaHeaders = hpack.NewDecoder(4096, nil)
|
||||||
|
if err := framer.WriteSettings(); err != nil {
|
||||||
|
conn.Close()
|
||||||
|
return nil, nil, err
|
||||||
|
}
|
||||||
|
return conn, framer, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
func readH2TestRequest(framer *http2.Framer) (uint32, []byte, error) {
|
||||||
|
var streamID uint32
|
||||||
|
var body []byte
|
||||||
|
for {
|
||||||
|
frame, err := framer.ReadFrame()
|
||||||
|
if err != nil {
|
||||||
|
return 0, nil, fmt.Errorf("read frame: %w", err)
|
||||||
|
}
|
||||||
|
switch f := frame.(type) {
|
||||||
|
case *http2.SettingsFrame:
|
||||||
|
if !f.IsAck() {
|
||||||
|
if err := framer.WriteSettingsAck(); err != nil {
|
||||||
|
return 0, nil, err
|
||||||
|
}
|
||||||
|
}
|
||||||
|
case *http2.MetaHeadersFrame:
|
||||||
|
streamID = f.Header().StreamID
|
||||||
|
if f.StreamEnded() {
|
||||||
|
return streamID, body, nil
|
||||||
|
}
|
||||||
|
case *http2.DataFrame:
|
||||||
|
if streamID == 0 {
|
||||||
|
streamID = f.Header().StreamID
|
||||||
|
}
|
||||||
|
if f.Header().StreamID != streamID {
|
||||||
|
continue
|
||||||
|
}
|
||||||
|
body = append(body, f.Data()...)
|
||||||
|
if f.StreamEnded() {
|
||||||
|
return streamID, body, nil
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func writeH2TestResponse(framer *http2.Framer, streamID uint32) error {
|
||||||
|
var hpackBuf bytes.Buffer
|
||||||
|
henc := hpack.NewEncoder(&hpackBuf)
|
||||||
|
if err := henc.WriteField(hpack.HeaderField{Name: ":status", Value: "200"}); err != nil {
|
||||||
|
return err
|
||||||
|
}
|
||||||
|
if err := framer.WriteHeaders(http2.HeadersFrameParam{
|
||||||
|
StreamID: streamID,
|
||||||
|
BlockFragment: hpackBuf.Bytes(),
|
||||||
|
EndHeaders: true,
|
||||||
|
}); err != nil {
|
||||||
|
return err
|
||||||
|
}
|
||||||
|
return framer.WriteData(streamID, true, []byte(`{}`))
|
||||||
|
}
|
||||||
|
|
||||||
|
func awaitH2ServerResult(t *testing.T, resultCh <-chan h2ServerResult) h2ServerResult {
|
||||||
|
t.Helper()
|
||||||
|
select {
|
||||||
|
case result := <-resultCh:
|
||||||
|
return result
|
||||||
|
case <-time.After(20 * time.Second):
|
||||||
|
t.Fatal("timed out waiting for HTTP/2 test server")
|
||||||
|
return h2ServerResult{}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// runResetOnFirstStreamServer speaks just enough raw HTTP/2 to emulate an
|
||||||
|
// upstream that accepts the first request, waits until the request body has
|
||||||
|
// been fully written, and then resets the stream with REFUSED_STREAM (the
|
||||||
|
// retry-safe reset some proxy/CDN-fronted upstreams send under load or during
|
||||||
|
// graceful shutdown, see RFC 9113 section 8.7). When expectRetry is true it
|
||||||
|
// serves the retried stream a 200 response; otherwise it stops after the reset.
|
||||||
|
func runResetOnFirstStreamServer(ln net.Listener, expectRetry bool) <-chan h2ServerResult {
|
||||||
|
resCh := make(chan h2ServerResult, 1)
|
||||||
|
go func() {
|
||||||
|
res := h2ServerResult{}
|
||||||
|
defer func() { resCh <- res }()
|
||||||
|
|
||||||
|
conn, framer, err := acceptH2TestConnection(ln)
|
||||||
|
if err != nil {
|
||||||
|
res.err = err
|
||||||
|
return
|
||||||
|
}
|
||||||
|
defer conn.Close()
|
||||||
|
|
||||||
|
attempts:
|
||||||
|
for attempt := 0; ; attempt++ {
|
||||||
|
streamID, body, err := readH2TestRequest(framer)
|
||||||
|
if err != nil {
|
||||||
|
res.err = err
|
||||||
|
return
|
||||||
|
}
|
||||||
|
res.streamCount++
|
||||||
|
res.attemptBodies = append(res.attemptBodies, body)
|
||||||
|
|
||||||
|
if attempt == 0 {
|
||||||
|
if err := framer.WriteRSTStream(streamID, http2.ErrCodeRefusedStream); err != nil {
|
||||||
|
res.err = err
|
||||||
|
return
|
||||||
|
}
|
||||||
|
if !expectRetry {
|
||||||
|
break attempts
|
||||||
|
}
|
||||||
|
continue
|
||||||
|
}
|
||||||
|
|
||||||
|
if err := writeH2TestResponse(framer, streamID); err != nil {
|
||||||
|
res.err = err
|
||||||
|
}
|
||||||
|
return
|
||||||
|
}
|
||||||
|
}()
|
||||||
|
return resCh
|
||||||
|
}
|
||||||
|
|
||||||
|
func runGoAwayAfterFirstRequestServer(ln net.Listener) <-chan h2ServerResult {
|
||||||
|
resCh := make(chan h2ServerResult, 1)
|
||||||
|
go func() {
|
||||||
|
res := h2ServerResult{}
|
||||||
|
defer func() { resCh <- res }()
|
||||||
|
|
||||||
|
for attempt := 0; attempt < 2; attempt++ {
|
||||||
|
conn, framer, err := acceptH2TestConnection(ln)
|
||||||
|
if err != nil {
|
||||||
|
res.err = err
|
||||||
|
return
|
||||||
|
}
|
||||||
|
streamID, body, err := readH2TestRequest(framer)
|
||||||
|
if err != nil {
|
||||||
|
conn.Close()
|
||||||
|
res.err = err
|
||||||
|
return
|
||||||
|
}
|
||||||
|
res.streamCount++
|
||||||
|
res.attemptBodies = append(res.attemptBodies, body)
|
||||||
|
|
||||||
|
if attempt == 0 {
|
||||||
|
err = framer.WriteGoAway(0, http2.ErrCodeNo, nil)
|
||||||
|
conn.Close()
|
||||||
|
if err != nil {
|
||||||
|
res.err = err
|
||||||
|
return
|
||||||
|
}
|
||||||
|
continue
|
||||||
|
}
|
||||||
|
|
||||||
|
err = writeH2TestResponse(framer, streamID)
|
||||||
|
conn.Close()
|
||||||
|
if err != nil {
|
||||||
|
res.err = err
|
||||||
|
}
|
||||||
|
return
|
||||||
|
}
|
||||||
|
}()
|
||||||
|
return resCh
|
||||||
|
}
|
||||||
|
|
||||||
|
func newH2PriorKnowledgeClient(ln net.Listener) (*http.Client, *http2.Transport) {
|
||||||
|
transport := &http2.Transport{
|
||||||
|
AllowHTTP: true,
|
||||||
|
DialTLSContext: func(ctx context.Context, network, _ string, _ *tls.Config) (net.Conn, error) {
|
||||||
|
var dialer net.Dialer
|
||||||
|
return dialer.DialContext(ctx, network, ln.Addr().String())
|
||||||
|
},
|
||||||
|
}
|
||||||
|
return &http.Client{Transport: transport, Timeout: 15 * time.Second}, transport
|
||||||
|
}
|
||||||
|
|
||||||
|
func newPassThroughBody(t *testing.T, payload []byte) (io.Reader, *relaycommon.RelayInfo, common.BodyStorage) {
|
||||||
|
t.Helper()
|
||||||
|
storage, err := common.CreateBodyStorage(payload)
|
||||||
|
require.NoError(t, err)
|
||||||
|
return common.ReaderOnly(storage), &relaycommon.RelayInfo{
|
||||||
|
UpstreamRequestBodySize: storage.Size(),
|
||||||
|
UpstreamRequestGetBody: storage.NewReader,
|
||||||
|
}, storage
|
||||||
|
}
|
||||||
|
|
||||||
|
// TestUpstreamGetBody_HTTP2RetryAfterUpstreamStreamReset exercises the actual
|
||||||
|
// failure this change fixes: an HTTP/2 upstream resets the stream with a
|
||||||
|
// retryable error after the request body has been written. With GetBody wired
|
||||||
|
// up the transport must transparently retry, and the retried request must
|
||||||
|
// carry the complete body.
|
||||||
|
func TestUpstreamGetBody_HTTP2RetryAfterUpstreamStreamReset(t *testing.T) {
|
||||||
|
payload := []byte(`{"model":"test-model","messages":[{"role":"user","content":"retry me"}]}`)
|
||||||
|
|
||||||
|
ln, err := net.Listen("tcp", "127.0.0.1:0")
|
||||||
|
require.NoError(t, err)
|
||||||
|
defer ln.Close()
|
||||||
|
resCh := runResetOnFirstStreamServer(ln, true)
|
||||||
|
|
||||||
|
client, transport := newH2PriorKnowledgeClient(ln)
|
||||||
|
defer transport.CloseIdleConnections()
|
||||||
|
|
||||||
|
// Build the upstream request exactly the way DoApiRequest does: a
|
||||||
|
// type-erased BodyStorage reader plus the applyUpstream* helpers.
|
||||||
|
body, size, getBody, closer, err := relaycommon.NewOutboundJSONBody(payload)
|
||||||
|
require.NoError(t, err)
|
||||||
|
defer closer.Close()
|
||||||
|
|
||||||
|
req, err := http.NewRequest(http.MethodPost, "http://upstream.test/v1/chat/completions", body)
|
||||||
|
require.NoError(t, err)
|
||||||
|
info := &relaycommon.RelayInfo{
|
||||||
|
UpstreamRequestBodySize: size,
|
||||||
|
UpstreamRequestGetBody: getBody,
|
||||||
|
}
|
||||||
|
ApplyUpstreamBodyMetadata(req, info)
|
||||||
|
require.NotNil(t, req.GetBody)
|
||||||
|
|
||||||
|
resp, err := client.Do(req)
|
||||||
|
require.NoError(t, err, "the transport must transparently retry after RST_STREAM")
|
||||||
|
defer resp.Body.Close()
|
||||||
|
assert.Equal(t, http.StatusOK, resp.StatusCode)
|
||||||
|
|
||||||
|
srv := awaitH2ServerResult(t, resCh)
|
||||||
|
require.NoError(t, srv.err)
|
||||||
|
assert.Equal(t, 2, srv.streamCount, "the request must have been attempted twice")
|
||||||
|
require.Len(t, srv.attemptBodies, 2)
|
||||||
|
assert.Equal(t, payload, srv.attemptBodies[0], "first attempt must carry the full body")
|
||||||
|
assert.Equal(t, payload, srv.attemptBodies[1], "the retried request must carry the complete body")
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestUpstreamGetBody_HTTP2RetryAfterUpstreamStreamReset_PassThrough(t *testing.T) {
|
||||||
|
payload := []byte(`{"model":"test-model","messages":[{"role":"user","content":"pass through"}]}`)
|
||||||
|
|
||||||
|
ln, err := net.Listen("tcp", "127.0.0.1:0")
|
||||||
|
require.NoError(t, err)
|
||||||
|
defer ln.Close()
|
||||||
|
resCh := runResetOnFirstStreamServer(ln, true)
|
||||||
|
|
||||||
|
client, transport := newH2PriorKnowledgeClient(ln)
|
||||||
|
defer transport.CloseIdleConnections()
|
||||||
|
|
||||||
|
body, info, storage := newPassThroughBody(t, payload)
|
||||||
|
defer storage.Close()
|
||||||
|
req, err := http.NewRequest(http.MethodPost, "http://upstream.test/v1/chat/completions", body)
|
||||||
|
require.NoError(t, err)
|
||||||
|
ApplyUpstreamBodyMetadata(req, info)
|
||||||
|
require.NotNil(t, req.GetBody)
|
||||||
|
assert.EqualValues(t, len(payload), req.ContentLength)
|
||||||
|
|
||||||
|
resp, err := client.Do(req)
|
||||||
|
require.NoError(t, err, "the transport must transparently retry a pass-through body after RST_STREAM")
|
||||||
|
defer resp.Body.Close()
|
||||||
|
assert.Equal(t, http.StatusOK, resp.StatusCode)
|
||||||
|
|
||||||
|
srv := awaitH2ServerResult(t, resCh)
|
||||||
|
require.NoError(t, srv.err)
|
||||||
|
assert.Equal(t, 2, srv.streamCount)
|
||||||
|
require.Len(t, srv.attemptBodies, 2)
|
||||||
|
assert.Equal(t, payload, srv.attemptBodies[0])
|
||||||
|
assert.Equal(t, payload, srv.attemptBodies[1])
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestUpstreamGetBody_HTTP2RetryAfterGracefulGoAway_PassThrough(t *testing.T) {
|
||||||
|
payload := []byte(`{"model":"test-model","messages":[{"role":"user","content":"go away"}]}`)
|
||||||
|
|
||||||
|
ln, err := net.Listen("tcp", "127.0.0.1:0")
|
||||||
|
require.NoError(t, err)
|
||||||
|
defer ln.Close()
|
||||||
|
resCh := runGoAwayAfterFirstRequestServer(ln)
|
||||||
|
|
||||||
|
client, transport := newH2PriorKnowledgeClient(ln)
|
||||||
|
defer transport.CloseIdleConnections()
|
||||||
|
|
||||||
|
body, info, storage := newPassThroughBody(t, payload)
|
||||||
|
defer storage.Close()
|
||||||
|
req, err := http.NewRequest(http.MethodPost, "http://upstream.test/v1/chat/completions", body)
|
||||||
|
require.NoError(t, err)
|
||||||
|
ApplyUpstreamBodyMetadata(req, info)
|
||||||
|
require.NotNil(t, req.GetBody)
|
||||||
|
|
||||||
|
resp, err := client.Do(req)
|
||||||
|
require.NoError(t, err, "the transport must retry on a new connection after graceful GOAWAY")
|
||||||
|
defer resp.Body.Close()
|
||||||
|
assert.Equal(t, http.StatusOK, resp.StatusCode)
|
||||||
|
|
||||||
|
srv := awaitH2ServerResult(t, resCh)
|
||||||
|
require.NoError(t, srv.err)
|
||||||
|
assert.Equal(t, 2, srv.streamCount)
|
||||||
|
require.Len(t, srv.attemptBodies, 2)
|
||||||
|
assert.Equal(t, payload, srv.attemptBodies[0])
|
||||||
|
assert.Equal(t, payload, srv.attemptBodies[1])
|
||||||
|
}
|
||||||
|
|
||||||
|
// TestUpstreamGetBody_HTTP2CannotRetryWithoutGetBody documents the pre-fix
|
||||||
|
// behavior: without GetBody the transport cannot safely retry once the body
|
||||||
|
// has been written, and the whole relay request fails.
|
||||||
|
func TestUpstreamGetBody_HTTP2CannotRetryWithoutGetBody(t *testing.T) {
|
||||||
|
payload := []byte(`{"model":"test-model","messages":[{"role":"user","content":"retry me"}]}`)
|
||||||
|
|
||||||
|
ln, err := net.Listen("tcp", "127.0.0.1:0")
|
||||||
|
require.NoError(t, err)
|
||||||
|
defer ln.Close()
|
||||||
|
resCh := runResetOnFirstStreamServer(ln, false)
|
||||||
|
|
||||||
|
client, transport := newH2PriorKnowledgeClient(ln)
|
||||||
|
defer transport.CloseIdleConnections()
|
||||||
|
|
||||||
|
body, size, _, closer, err := relaycommon.NewOutboundJSONBody(payload)
|
||||||
|
require.NoError(t, err)
|
||||||
|
defer closer.Close()
|
||||||
|
|
||||||
|
req, err := http.NewRequest(http.MethodPost, "http://upstream.test/v1/chat/completions", body)
|
||||||
|
require.NoError(t, err)
|
||||||
|
applyUpstreamContentLength(req, &relaycommon.RelayInfo{UpstreamRequestBodySize: size})
|
||||||
|
assert.Nil(t, req.GetBody)
|
||||||
|
|
||||||
|
resp, err := client.Do(req) //nolint:bodyclose // Do fails, no body to close
|
||||||
|
require.Error(t, err)
|
||||||
|
assert.Nil(t, resp)
|
||||||
|
require.ErrorContains(t, err, "cannot retry err")
|
||||||
|
require.ErrorContains(t, err, "Request.Body was written")
|
||||||
|
|
||||||
|
srv := awaitH2ServerResult(t, resCh)
|
||||||
|
require.NoError(t, srv.err)
|
||||||
|
assert.Equal(t, 1, srv.streamCount)
|
||||||
|
require.Len(t, srv.attemptBodies, 1)
|
||||||
|
assert.Equal(t, payload, srv.attemptBodies[0])
|
||||||
|
}
|
||||||
@@ -0,0 +1,88 @@
|
|||||||
|
package channel
|
||||||
|
|
||||||
|
import (
|
||||||
|
"bytes"
|
||||||
|
"io"
|
||||||
|
"net/http"
|
||||||
|
"net/http/httptest"
|
||||||
|
"reflect"
|
||||||
|
"sync/atomic"
|
||||||
|
"testing"
|
||||||
|
|
||||||
|
relaycommon "github.com/QuantumNous/new-api/relay/common"
|
||||||
|
"github.com/QuantumNous/new-api/service"
|
||||||
|
"github.com/gin-gonic/gin"
|
||||||
|
"github.com/stretchr/testify/assert"
|
||||||
|
"github.com/stretchr/testify/require"
|
||||||
|
)
|
||||||
|
|
||||||
|
func TestDoRequestReturnsUpstreamRedirectWithoutFollowing(t *testing.T) {
|
||||||
|
service.InitHttpClient()
|
||||||
|
gin.SetMode(gin.TestMode)
|
||||||
|
sharedClient := service.GetHttpClient()
|
||||||
|
require.NotNil(t, sharedClient)
|
||||||
|
require.NotNil(t, sharedClient.CheckRedirect)
|
||||||
|
originalRedirectPolicy := reflect.ValueOf(sharedClient.CheckRedirect).Pointer()
|
||||||
|
|
||||||
|
var targetRequests atomic.Int32
|
||||||
|
target := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
|
||||||
|
targetRequests.Add(1)
|
||||||
|
w.WriteHeader(http.StatusTeapot)
|
||||||
|
}))
|
||||||
|
defer target.Close()
|
||||||
|
|
||||||
|
const responseBody = "redirect response"
|
||||||
|
tests := []int{
|
||||||
|
http.StatusMovedPermanently,
|
||||||
|
http.StatusFound,
|
||||||
|
http.StatusSeeOther,
|
||||||
|
http.StatusTemporaryRedirect,
|
||||||
|
http.StatusPermanentRedirect,
|
||||||
|
}
|
||||||
|
|
||||||
|
for _, statusCode := range tests {
|
||||||
|
t.Run(http.StatusText(statusCode), func(t *testing.T) {
|
||||||
|
targetRequests.Store(0)
|
||||||
|
var sourceRequests atomic.Int32
|
||||||
|
type sourceResult struct {
|
||||||
|
body []byte
|
||||||
|
err error
|
||||||
|
}
|
||||||
|
sourceResultCh := make(chan sourceResult, 1)
|
||||||
|
source := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||||
|
sourceRequests.Add(1)
|
||||||
|
body, err := io.ReadAll(r.Body)
|
||||||
|
sourceResultCh <- sourceResult{body: body, err: err}
|
||||||
|
w.Header().Set("Location", target.URL+"/redirect-target")
|
||||||
|
w.WriteHeader(statusCode)
|
||||||
|
_, _ = io.WriteString(w, responseBody)
|
||||||
|
}))
|
||||||
|
defer source.Close()
|
||||||
|
|
||||||
|
recorder := httptest.NewRecorder()
|
||||||
|
ctx, _ := gin.CreateTestContext(recorder)
|
||||||
|
ctx.Request = httptest.NewRequest(http.MethodPost, "/relay", nil)
|
||||||
|
|
||||||
|
req, err := http.NewRequest(http.MethodPost, source.URL, bytes.NewReader([]byte("request body")))
|
||||||
|
require.NoError(t, err)
|
||||||
|
info := &relaycommon.RelayInfo{ChannelMeta: &relaycommon.ChannelMeta{}}
|
||||||
|
|
||||||
|
resp, err := doRequest(ctx, req, info)
|
||||||
|
require.NoError(t, err)
|
||||||
|
defer resp.Body.Close()
|
||||||
|
body, err := io.ReadAll(resp.Body)
|
||||||
|
require.NoError(t, err)
|
||||||
|
gotSource := <-sourceResultCh
|
||||||
|
require.NoError(t, gotSource.err)
|
||||||
|
|
||||||
|
assert.Equal(t, statusCode, resp.StatusCode)
|
||||||
|
assert.Equal(t, target.URL+"/redirect-target", resp.Header.Get("Location"))
|
||||||
|
assert.Equal(t, responseBody, string(body))
|
||||||
|
assert.Equal(t, []byte("request body"), gotSource.body)
|
||||||
|
assert.EqualValues(t, 1, sourceRequests.Load())
|
||||||
|
assert.Zero(t, targetRequests.Load())
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
|
assert.Equal(t, originalRedirectPolicy, reflect.ValueOf(sharedClient.CheckRedirect).Pointer(), "the cached client must not be mutated")
|
||||||
|
}
|
||||||
@@ -112,6 +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)
|
||||||
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,6 +216,8 @@ func (a *TaskAdaptor) BuildRequestBody(c *gin.Context, info *relaycommon.RelayIn
|
|||||||
return &buf, nil
|
return &buf, nil
|
||||||
}
|
}
|
||||||
|
|
||||||
|
info.UpstreamRequestBodySize = storage.Size()
|
||||||
|
info.UpstreamRequestGetBody = storage.NewReader
|
||||||
return common.ReaderOnly(storage), nil
|
return common.ReaderOnly(storage), nil
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|||||||
@@ -0,0 +1,40 @@
|
|||||||
|
package sora
|
||||||
|
|
||||||
|
import (
|
||||||
|
"bytes"
|
||||||
|
"io"
|
||||||
|
"net/http"
|
||||||
|
"net/http/httptest"
|
||||||
|
"testing"
|
||||||
|
|
||||||
|
"github.com/QuantumNous/new-api/common"
|
||||||
|
relaycommon "github.com/QuantumNous/new-api/relay/common"
|
||||||
|
"github.com/gin-gonic/gin"
|
||||||
|
"github.com/stretchr/testify/assert"
|
||||||
|
"github.com/stretchr/testify/require"
|
||||||
|
)
|
||||||
|
|
||||||
|
func TestSoraBuildRequestBodyBindsReplayMetadataForPassThrough(t *testing.T) {
|
||||||
|
payload := []byte("opaque-sora-request-body")
|
||||||
|
c, _ := gin.CreateTestContext(httptest.NewRecorder())
|
||||||
|
c.Request = httptest.NewRequest(http.MethodPost, "/v1/videos", bytes.NewReader(payload))
|
||||||
|
c.Request.Header.Set("Content-Type", "application/octet-stream")
|
||||||
|
defer common.CleanupBodyStorage(c)
|
||||||
|
|
||||||
|
info := &relaycommon.RelayInfo{}
|
||||||
|
body, err := (&TaskAdaptor{}).BuildRequestBody(c, info)
|
||||||
|
require.NoError(t, err)
|
||||||
|
|
||||||
|
sent, err := io.ReadAll(body)
|
||||||
|
require.NoError(t, err)
|
||||||
|
assert.Equal(t, payload, sent)
|
||||||
|
assert.EqualValues(t, len(payload), info.UpstreamRequestBodySize)
|
||||||
|
require.NotNil(t, info.UpstreamRequestGetBody)
|
||||||
|
|
||||||
|
replayBody, err := info.UpstreamRequestGetBody()
|
||||||
|
require.NoError(t, err)
|
||||||
|
replay, err := io.ReadAll(replayBody)
|
||||||
|
require.NoError(t, err)
|
||||||
|
require.NoError(t, replayBody.Close())
|
||||||
|
assert.Equal(t, payload, replay)
|
||||||
|
}
|
||||||
@@ -128,13 +128,14 @@ 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, closer, err := relaycommon.NewOutboundJSONBody(jsonData)
|
body, size, getBody, 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.UpstreamRequestBodySize = size
|
||||||
|
info.UpstreamRequestGetBody = getBody
|
||||||
var requestBody io.Reader = body
|
var requestBody io.Reader = body
|
||||||
|
|
||||||
var httpResp *http.Response
|
var httpResp *http.Response
|
||||||
|
|||||||
@@ -160,6 +160,7 @@ func ClaudeHelper(c *gin.Context, info *relaycommon.RelayInfo) (newAPIError *typ
|
|||||||
return types.NewErrorWithStatusCode(err, types.ErrorCodeReadRequestBodyFailed, http.StatusBadRequest, types.ErrOptionWithSkipRetry())
|
return types.NewErrorWithStatusCode(err, types.ErrorCodeReadRequestBodyFailed, http.StatusBadRequest, types.ErrOptionWithSkipRetry())
|
||||||
}
|
}
|
||||||
info.UpstreamRequestBodySize = storage.Size()
|
info.UpstreamRequestBodySize = storage.Size()
|
||||||
|
info.UpstreamRequestGetBody = storage.NewReader
|
||||||
requestBody = common.ReaderOnly(storage)
|
requestBody = common.ReaderOnly(storage)
|
||||||
} else {
|
} else {
|
||||||
convertedRequest, err := adaptor.ConvertClaudeRequest(c, info, request)
|
convertedRequest, err := adaptor.ConvertClaudeRequest(c, info, request)
|
||||||
@@ -187,13 +188,14 @@ func ClaudeHelper(c *gin.Context, info *relaycommon.RelayInfo) (newAPIError *typ
|
|||||||
}
|
}
|
||||||
|
|
||||||
logger.LogDebug(c, "requestBody: %s", jsonData)
|
logger.LogDebug(c, "requestBody: %s", jsonData)
|
||||||
body, size, closer, err := relaycommon.NewOutboundJSONBody(jsonData)
|
body, size, getBody, 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.UpstreamRequestBodySize = size
|
||||||
|
info.UpstreamRequestGetBody = getBody
|
||||||
requestBody = body
|
requestBody = body
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|||||||
@@ -22,10 +22,22 @@ import (
|
|||||||
// transport from prematurely closing the underlying BodyStorage. The returned
|
// transport from prematurely closing the underlying BodyStorage. The returned
|
||||||
// size is meant to be propagated to http.Request.ContentLength because the
|
// size is meant to be propagated to http.Request.ContentLength because the
|
||||||
// type-erased io.Reader prevents net/http from auto-detecting it.
|
// type-erased io.Reader prevents net/http from auto-detecting it.
|
||||||
func NewOutboundJSONBody(data []byte) (body io.Reader, size int64, 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, err
|
return nil, 0, nil, nil, err
|
||||||
}
|
}
|
||||||
return common.ReaderOnly(storage), storage.Size(), storage, nil
|
return common.ReaderOnly(storage), storage.Size(), storage.NewReader, storage, nil
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -0,0 +1,163 @@
|
|||||||
|
package common
|
||||||
|
|
||||||
|
import (
|
||||||
|
"io"
|
||||||
|
"testing"
|
||||||
|
|
||||||
|
"github.com/QuantumNous/new-api/common"
|
||||||
|
"github.com/stretchr/testify/assert"
|
||||||
|
"github.com/stretchr/testify/require"
|
||||||
|
)
|
||||||
|
|
||||||
|
func TestNewOutboundJSONBody_GetBodyReplaysFullBody(t *testing.T) {
|
||||||
|
t.Parallel()
|
||||||
|
|
||||||
|
payload := []byte(`{"model":"test-model","messages":[{"role":"user","content":"hello"}]}`)
|
||||||
|
|
||||||
|
body, size, getBody, closer, err := NewOutboundJSONBody(payload)
|
||||||
|
require.NoError(t, err)
|
||||||
|
defer closer.Close()
|
||||||
|
|
||||||
|
assert.EqualValues(t, len(payload), size)
|
||||||
|
require.NotNil(t, getBody)
|
||||||
|
|
||||||
|
// Consume the primary body, as the HTTP transport does on the first attempt.
|
||||||
|
first, err := io.ReadAll(body)
|
||||||
|
require.NoError(t, err)
|
||||||
|
assert.Equal(t, payload, first)
|
||||||
|
|
||||||
|
// 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()
|
||||||
|
require.NoError(t, err)
|
||||||
|
replay, err := io.ReadAll(rc)
|
||||||
|
require.NoError(t, err)
|
||||||
|
require.NoError(t, rc.Close())
|
||||||
|
assert.Equal(t, payload, replay, "replay %d must equal the original payload", i+1)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestNewOutboundJSONBody_GetBodyAfterPartialRead(t *testing.T) {
|
||||||
|
t.Parallel()
|
||||||
|
|
||||||
|
payload := []byte(`{"model":"test-model","input":"0123456789"}`)
|
||||||
|
|
||||||
|
body, _, getBody, closer, err := NewOutboundJSONBody(payload)
|
||||||
|
require.NoError(t, err)
|
||||||
|
defer closer.Close()
|
||||||
|
|
||||||
|
// Simulate an aborted first attempt that only wrote part of the body.
|
||||||
|
partial := make([]byte, 10)
|
||||||
|
_, err = io.ReadFull(body, partial)
|
||||||
|
require.NoError(t, err)
|
||||||
|
|
||||||
|
rc, err := getBody()
|
||||||
|
require.NoError(t, err)
|
||||||
|
replay, err := io.ReadAll(rc)
|
||||||
|
require.NoError(t, err)
|
||||||
|
assert.Equal(t, payload, replay)
|
||||||
|
|
||||||
|
// 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()
|
||||||
|
require.NoError(t, err)
|
||||||
|
replay2, err := io.ReadAll(rc2)
|
||||||
|
require.NoError(t, err)
|
||||||
|
require.NoError(t, rc2.Close())
|
||||||
|
assert.Equal(t, payload, replay2)
|
||||||
|
}
|
||||||
|
|
||||||
|
// assertIndependentReplayReaders proves that readers handed out by getBody own
|
||||||
|
// 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)) {
|
||||||
|
t.Helper()
|
||||||
|
|
||||||
|
half := len(payload) / 2
|
||||||
|
|
||||||
|
// Partially drain the primary body first, as if attempt N's body write
|
||||||
|
// were still in flight when the transport builds attempt N+1 via GetBody.
|
||||||
|
primaryHead := make([]byte, half)
|
||||||
|
_, err := io.ReadFull(body, primaryHead)
|
||||||
|
require.NoError(t, err)
|
||||||
|
assert.Equal(t, payload[:half], primaryHead)
|
||||||
|
|
||||||
|
// Interleave two replay readers: A reads half, B reads everything, then A
|
||||||
|
// reads the rest.
|
||||||
|
a, err := getBody()
|
||||||
|
require.NoError(t, err)
|
||||||
|
b, err := getBody()
|
||||||
|
require.NoError(t, err)
|
||||||
|
|
||||||
|
aHead := make([]byte, half)
|
||||||
|
_, err = io.ReadFull(a, aHead)
|
||||||
|
require.NoError(t, err)
|
||||||
|
assert.Equal(t, payload[:half], aHead)
|
||||||
|
|
||||||
|
bAll, err := io.ReadAll(b)
|
||||||
|
require.NoError(t, err)
|
||||||
|
require.NoError(t, b.Close())
|
||||||
|
assert.Equal(t, payload, bAll, "reader B must see the complete body even while A is mid-read")
|
||||||
|
|
||||||
|
aRest, err := io.ReadAll(a)
|
||||||
|
require.NoError(t, err)
|
||||||
|
require.NoError(t, a.Close())
|
||||||
|
assert.Equal(t, payload[half:], aRest, "reader A must resume from its own cursor, unaffected by B")
|
||||||
|
|
||||||
|
// The replays must not have disturbed the primary body's cursor either.
|
||||||
|
primaryRest, err := io.ReadAll(body)
|
||||||
|
require.NoError(t, err)
|
||||||
|
assert.Equal(t, payload[half:], primaryRest, "the primary body must be unaffected by replay readers")
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestNewOutboundJSONBody_GetBodyReadersAreIndependent(t *testing.T) {
|
||||||
|
t.Parallel()
|
||||||
|
|
||||||
|
payload := []byte(`{"model":"test-model","input":"abcdefghijklmnopqrstuvwxyz"}`)
|
||||||
|
|
||||||
|
body, _, getBody, closer, err := NewOutboundJSONBody(payload)
|
||||||
|
require.NoError(t, err)
|
||||||
|
defer closer.Close()
|
||||||
|
|
||||||
|
assertIndependentReplayReaders(t, payload, body, getBody)
|
||||||
|
|
||||||
|
// Once the handler releases the storage, GetBody must fail loudly instead
|
||||||
|
// of replaying stale data.
|
||||||
|
require.NoError(t, closer.Close())
|
||||||
|
_, err = getBody()
|
||||||
|
require.ErrorIs(t, err, common.ErrStorageClosed)
|
||||||
|
}
|
||||||
|
|
||||||
|
// TestNewOutboundJSONBody_GetBodyReadersAreIndependent_DiskStorage runs the
|
||||||
|
// same independence assertions against the disk-backed storage. Deliberately
|
||||||
|
// not parallel: it temporarily lowers the global disk-cache threshold so the
|
||||||
|
// payload takes the diskStorage path.
|
||||||
|
func TestNewOutboundJSONBody_GetBodyReadersAreIndependent_DiskStorage(t *testing.T) {
|
||||||
|
prev := common.GetDiskCacheConfig()
|
||||||
|
common.SetDiskCacheConfig(common.DiskCacheConfig{
|
||||||
|
Enabled: true,
|
||||||
|
ThresholdMB: 0,
|
||||||
|
MaxSizeMB: 64,
|
||||||
|
Path: t.TempDir(),
|
||||||
|
})
|
||||||
|
defer common.SetDiskCacheConfig(prev)
|
||||||
|
|
||||||
|
payload := []byte(`{"model":"test-model","input":"abcdefghijklmnopqrstuvwxyz"}`)
|
||||||
|
|
||||||
|
body, _, getBody, closer, err := NewOutboundJSONBody(payload)
|
||||||
|
require.NoError(t, err)
|
||||||
|
defer closer.Close()
|
||||||
|
|
||||||
|
storage, ok := closer.(common.BodyStorage)
|
||||||
|
require.True(t, ok)
|
||||||
|
assert.True(t, storage.IsDisk(), "the payload must have taken the diskStorage path")
|
||||||
|
|
||||||
|
assertIndependentReplayReaders(t, payload, body, getBody)
|
||||||
|
|
||||||
|
require.NoError(t, closer.Close())
|
||||||
|
_, err = getBody()
|
||||||
|
require.ErrorIs(t, err, common.ErrStorageClosed)
|
||||||
|
}
|
||||||
@@ -4,6 +4,7 @@ import (
|
|||||||
"encoding/json"
|
"encoding/json"
|
||||||
"errors"
|
"errors"
|
||||||
"fmt"
|
"fmt"
|
||||||
|
"io"
|
||||||
"strconv"
|
"strconv"
|
||||||
"strings"
|
"strings"
|
||||||
"time"
|
"time"
|
||||||
@@ -156,6 +157,16 @@ type RelayInfo struct {
|
|||||||
// *bytes.Reader/Buffer/strings.Reader). 0 means "let net/http decide".
|
// *bytes.Reader/Buffer/strings.Reader). 0 means "let net/http decide".
|
||||||
UpstreamRequestBodySize int64
|
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
|
||||||
@@ -193,6 +204,12 @@ 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)
|
||||||
|
|||||||
@@ -1,14 +1,32 @@
|
|||||||
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,6 +104,8 @@ 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()
|
||||||
|
info.UpstreamRequestGetBody = storage.NewReader
|
||||||
requestBody = common.ReaderOnly(storage)
|
requestBody = common.ReaderOnly(storage)
|
||||||
} else {
|
} else {
|
||||||
convertedRequest, err := adaptor.ConvertOpenAIRequest(c, info, request)
|
convertedRequest, err := adaptor.ConvertOpenAIRequest(c, info, request)
|
||||||
@@ -175,13 +177,14 @@ 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, closer, err := relaycommon.NewOutboundJSONBody(jsonData)
|
body, size, getBody, 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.UpstreamRequestBodySize = size
|
||||||
|
info.UpstreamRequestGetBody = getBody
|
||||||
requestBody = body
|
requestBody = body
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|||||||
@@ -58,13 +58,14 @@ 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, closer, err := relaycommon.NewOutboundJSONBody(jsonData)
|
body, size, getBody, 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.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,6 +141,8 @@ 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()
|
||||||
|
info.UpstreamRequestGetBody = storage.NewReader
|
||||||
requestBody = common.ReaderOnly(storage)
|
requestBody = common.ReaderOnly(storage)
|
||||||
} else {
|
} else {
|
||||||
// 使用 ConvertGeminiRequest 转换请求格式
|
// 使用 ConvertGeminiRequest 转换请求格式
|
||||||
@@ -164,13 +166,14 @@ 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, closer, err := relaycommon.NewOutboundJSONBody(jsonData)
|
body, size, getBody, 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.UpstreamRequestBodySize = size
|
||||||
|
info.UpstreamRequestGetBody = getBody
|
||||||
requestBody = body
|
requestBody = body
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -269,13 +272,14 @@ 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, closer, err := relaycommon.NewOutboundJSONBody(jsonData)
|
body, size, getBody, 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.UpstreamRequestBodySize = size
|
||||||
|
info.UpstreamRequestGetBody = getBody
|
||||||
requestBody = body
|
requestBody = body
|
||||||
|
|
||||||
resp, err := adaptor.DoRequest(c, info, requestBody)
|
resp, err := adaptor.DoRequest(c, info, requestBody)
|
||||||
|
|||||||
@@ -51,6 +51,8 @@ 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()
|
||||||
|
info.UpstreamRequestGetBody = storage.NewReader
|
||||||
requestBody = common.ReaderOnly(storage)
|
requestBody = common.ReaderOnly(storage)
|
||||||
} else {
|
} else {
|
||||||
convertedRequest, err := adaptor.ConvertImageRequest(c, info, *request)
|
convertedRequest, err := adaptor.ConvertImageRequest(c, info, *request)
|
||||||
@@ -77,13 +79,14 @@ 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, closer, err := relaycommon.NewOutboundJSONBody(jsonData)
|
body, size, getBody, 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.UpstreamRequestBodySize = size
|
||||||
|
info.UpstreamRequestGetBody = getBody
|
||||||
requestBody = body
|
requestBody = body
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -47,6 +47,8 @@ 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()
|
||||||
|
info.UpstreamRequestGetBody = storage.NewReader
|
||||||
requestBody = common.ReaderOnly(storage)
|
requestBody = common.ReaderOnly(storage)
|
||||||
} else {
|
} else {
|
||||||
convertedRequest, err := adaptor.ConvertRerankRequest(c, info.RelayMode, *request)
|
convertedRequest, err := adaptor.ConvertRerankRequest(c, info.RelayMode, *request)
|
||||||
@@ -68,13 +70,14 @@ 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, closer, err := relaycommon.NewOutboundJSONBody(jsonData)
|
body, size, getBody, 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.UpstreamRequestBodySize = size
|
||||||
|
info.UpstreamRequestGetBody = getBody
|
||||||
requestBody = body
|
requestBody = body
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|||||||
@@ -82,6 +82,8 @@ 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()
|
||||||
|
info.UpstreamRequestGetBody = storage.NewReader
|
||||||
requestBody = common.ReaderOnly(storage)
|
requestBody = common.ReaderOnly(storage)
|
||||||
} else {
|
} else {
|
||||||
convertedRequest, err := adaptor.ConvertOpenAIResponsesRequest(c, info, *request)
|
convertedRequest, err := adaptor.ConvertOpenAIResponsesRequest(c, info, *request)
|
||||||
@@ -109,13 +111,14 @@ func ResponsesHelper(c *gin.Context, info *relaycommon.RelayInfo) (newAPIError *
|
|||||||
}
|
}
|
||||||
|
|
||||||
logger.LogDebug(c, "requestBody: %s", jsonData)
|
logger.LogDebug(c, "requestBody: %s", jsonData)
|
||||||
body, size, closer, err := relaycommon.NewOutboundJSONBody(jsonData)
|
body, size, getBody, 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.UpstreamRequestBodySize = size
|
||||||
|
info.UpstreamRequestGetBody = getBody
|
||||||
requestBody = body
|
requestBody = body
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|||||||
Reference in New Issue
Block a user