diff --git a/common/body_storage.go b/common/body_storage.go index 094dbda3..0c4a849d 100644 --- a/common/body_storage.go +++ b/common/body_storage.go @@ -20,6 +20,13 @@ type BodyStorage interface { Size() int64 // IsDisk 是否是磁盘存储 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 存储已关闭错误 @@ -80,6 +87,18 @@ func (m *memoryStorage) Bytes() ([]byte, error) { 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 { return m.size } @@ -229,6 +248,24 @@ func (d *diskStorage) Bytes() ([]byte, error) { 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 { return d.size } diff --git a/relay/alpha_search_handler.go b/relay/alpha_search_handler.go index 23ebafea..c66cb1de 100644 --- a/relay/alpha_search_handler.go +++ b/relay/alpha_search_handler.go @@ -62,12 +62,13 @@ func AlphaSearchHelper(c *gin.Context, info *relaycommon.RelayInfo) (newAPIError } logger.LogDebug(c, "requestBody: %s", jsonData) - body, size, closer, err := relaycommon.NewOutboundJSONBody(jsonData) + body, size, getBody, closer, err := relaycommon.NewOutboundJSONBody(jsonData) if err != nil { return types.NewError(err, types.ErrorCodeConvertRequestFailed, types.ErrOptionWithSkipRetry()) } defer closer.Close() info.UpstreamRequestBodySize = size + info.UpstreamRequestGetBody = getBody adaptor := GetAdaptor(info.ApiType) if adaptor == nil { diff --git a/relay/channel/api_request.go b/relay/channel/api_request.go index fc2855a8..2ed9f69f 100644 --- a/relay/channel/api_request.go +++ b/relay/channel/api_request.go @@ -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) { if info.RelayMode == constant.RelayModeAudioTranscription || info.RelayMode == constant.RelayModeAudioTranslation { // multipart/form-data @@ -314,7 +341,7 @@ func DoApiRequest(a Adaptor, c *gin.Context, info *common.RelayInfo, requestBody if err != nil { return nil, fmt.Errorf("new request failed: %w", err) } - applyUpstreamContentLength(req, info) + ApplyUpstreamBodyMetadata(req, info) headers := req.Header err = a.SetupRequestHeader(c, &headers, info) if err != nil { @@ -344,7 +371,7 @@ func DoFormRequest(a Adaptor, c *gin.Context, info *common.RelayInfo, requestBod if err != nil { return nil, fmt.Errorf("new request failed: %w", err) } - applyUpstreamContentLength(req, info) + ApplyUpstreamBodyMetadata(req, info) // set form data req.Header.Set("Content-Type", c.Request.Header.Get("Content-Type")) 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) { 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) { client, err := service.GetHttpClientWithProxySettings(info.ChannelSetting.Proxy, info.ChannelSetting) if err != nil { 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 { policy := service.NormalizeHTTPTransportPolicy(info.ChannelSetting) 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 { logger.LogError(c, "do request failed: "+err.Error()) 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 { return nil, fmt.Errorf("new request failed: %w", err) } - applyUpstreamContentLength(req, info) - req.GetBody = func() (io.ReadCloser, error) { - return io.NopCloser(requestBody), nil - } + ApplyUpstreamBodyMetadata(req, info) + // Do NOT wrap requestBody in a GetBody closure here: returning the same + // (already consumed) reader would make any transport-level retry silently + // replay an empty body. http.NewRequest already derives a correct, + // snapshot-based GetBody for *bytes.Reader/Buffer/strings.Reader bodies + // (which most task adaptors pass in); for type-erased readers, + // ApplyUpstreamBodyMetadata wires a replayable body when one is available. + // Otherwise GetBody stays nil so the transport fails the retry instead of + // sending a corrupted request. err = a.BuildRequestHeader(c, req, info) if err != nil { diff --git a/relay/channel/api_request_getbody_test.go b/relay/channel/api_request_getbody_test.go new file mode 100644 index 00000000..847468e1 --- /dev/null +++ b/relay/channel/api_request_getbody_test.go @@ -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]) +} diff --git a/relay/channel/api_request_redirect_test.go b/relay/channel/api_request_redirect_test.go new file mode 100644 index 00000000..e24685ad --- /dev/null +++ b/relay/channel/api_request_redirect_test.go @@ -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") +} diff --git a/relay/channel/jimeng/adaptor.go b/relay/channel/jimeng/adaptor.go index 658d9236..d2f9c389 100644 --- a/relay/channel/jimeng/adaptor.go +++ b/relay/channel/jimeng/adaptor.go @@ -112,6 +112,7 @@ func (a *Adaptor) DoRequest(c *gin.Context, info *relaycommon.RelayInfo, request if err != nil { return nil, fmt.Errorf("new request failed: %w", err) } + channel.ApplyUpstreamBodyMetadata(req, info) err = Sign(c, req, info.ApiKey) if err != nil { return nil, fmt.Errorf("setup request header failed: %w", err) diff --git a/relay/channel/task/sora/adaptor.go b/relay/channel/task/sora/adaptor.go index e9029aa2..59a7a258 100644 --- a/relay/channel/task/sora/adaptor.go +++ b/relay/channel/task/sora/adaptor.go @@ -216,6 +216,8 @@ func (a *TaskAdaptor) BuildRequestBody(c *gin.Context, info *relaycommon.RelayIn return &buf, nil } + info.UpstreamRequestBodySize = storage.Size() + info.UpstreamRequestGetBody = storage.NewReader return common.ReaderOnly(storage), nil } diff --git a/relay/channel/task/sora/adaptor_test.go b/relay/channel/task/sora/adaptor_test.go new file mode 100644 index 00000000..ef514952 --- /dev/null +++ b/relay/channel/task/sora/adaptor_test.go @@ -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) +} diff --git a/relay/chat_completions_via_responses.go b/relay/chat_completions_via_responses.go index ccc7fb28..660cc889 100644 --- a/relay/chat_completions_via_responses.go +++ b/relay/chat_completions_via_responses.go @@ -128,13 +128,14 @@ func chatCompletionsViaResponses(c *gin.Context, info *relaycommon.RelayInfo, ad 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 { return nil, types.NewError(err, types.ErrorCodeConvertRequestFailed, types.ErrOptionWithSkipRetry()) } defer closer.Close() jsonData = nil info.UpstreamRequestBodySize = size + info.UpstreamRequestGetBody = getBody var requestBody io.Reader = body var httpResp *http.Response diff --git a/relay/claude_handler.go b/relay/claude_handler.go index 63bc11f5..bcbe9908 100644 --- a/relay/claude_handler.go +++ b/relay/claude_handler.go @@ -160,6 +160,7 @@ func ClaudeHelper(c *gin.Context, info *relaycommon.RelayInfo) (newAPIError *typ return types.NewErrorWithStatusCode(err, types.ErrorCodeReadRequestBodyFailed, http.StatusBadRequest, types.ErrOptionWithSkipRetry()) } info.UpstreamRequestBodySize = storage.Size() + info.UpstreamRequestGetBody = storage.NewReader requestBody = common.ReaderOnly(storage) } else { 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) - body, size, closer, err := relaycommon.NewOutboundJSONBody(jsonData) + body, size, getBody, closer, err := relaycommon.NewOutboundJSONBody(jsonData) if err != nil { return types.NewError(err, types.ErrorCodeConvertRequestFailed, types.ErrOptionWithSkipRetry()) } defer closer.Close() jsonData = nil info.UpstreamRequestBodySize = size + info.UpstreamRequestGetBody = getBody requestBody = body } diff --git a/relay/common/outbound_body.go b/relay/common/outbound_body.go index 94ef8dde..0316cf74 100644 --- a/relay/common/outbound_body.go +++ b/relay/common/outbound_body.go @@ -22,10 +22,22 @@ import ( // 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. -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) 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 } diff --git a/relay/common/outbound_body_test.go b/relay/common/outbound_body_test.go new file mode 100644 index 00000000..23e2ff95 --- /dev/null +++ b/relay/common/outbound_body_test.go @@ -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) +} diff --git a/relay/common/relay_info.go b/relay/common/relay_info.go index 45ae30bf..52c74115 100644 --- a/relay/common/relay_info.go +++ b/relay/common/relay_info.go @@ -4,6 +4,7 @@ import ( "encoding/json" "errors" "fmt" + "io" "strconv" "strings" "time" @@ -156,6 +157,16 @@ type RelayInfo struct { // *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 @@ -193,6 +204,12 @@ type RelayInfo struct { } func (info *RelayInfo) InitChannelMeta(c *gin.Context) { + // RelayInfo is reused across channel attempts. Body metadata belongs to the + // current attempt and may reference storage that its handler has closed, so + // discard it before the next channel binds its outbound body. + info.UpstreamRequestBodySize = 0 + info.UpstreamRequestGetBody = nil + channelType := common.GetContextKeyInt(c, constant.ContextKeyChannelType) paramOverride := common.GetContextKeyStringMap(c, constant.ContextKeyChannelParamOverride) headerOverride := common.GetContextKeyStringMap(c, constant.ContextKeyChannelHeaderOverride) diff --git a/relay/common/relay_info_test.go b/relay/common/relay_info_test.go index 9e58f3f9..edf5fdbf 100644 --- a/relay/common/relay_info_test.go +++ b/relay/common/relay_info_test.go @@ -1,14 +1,32 @@ package common import ( + "io" + "net/http/httptest" "testing" "github.com/QuantumNous/new-api/relaykit/relayconvert/convmeta" "github.com/QuantumNous/new-api/relaykit/types" + "github.com/gin-gonic/gin" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" ) +func TestInitChannelMetaClearsUpstreamBodyMetadata(t *testing.T) { + c, _ := gin.CreateTestContext(httptest.NewRecorder()) + info := &RelayInfo{ + UpstreamRequestBodySize: 37, + UpstreamRequestGetBody: func() (io.ReadCloser, error) { + return nil, nil + }, + } + + info.InitChannelMeta(c) + + assert.Zero(t, info.UpstreamRequestBodySize) + assert.Nil(t, info.UpstreamRequestGetBody) +} + func TestRelayInfoGetFinalRequestRelayFormatPrefersExplicitFinal(t *testing.T) { info := &RelayInfo{ RelayFormat: types.RelayFormatOpenAI, diff --git a/relay/compatible_handler.go b/relay/compatible_handler.go index b1e25e03..24f355f9 100644 --- a/relay/compatible_handler.go +++ b/relay/compatible_handler.go @@ -104,6 +104,8 @@ func TextHelper(c *gin.Context, info *relaycommon.RelayInfo) (newAPIError *types logger.LogDebug(c, "requestBody: %s", debugBytes) } } + info.UpstreamRequestBodySize = storage.Size() + info.UpstreamRequestGetBody = storage.NewReader requestBody = common.ReaderOnly(storage) } else { 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) - body, size, closer, err := relaycommon.NewOutboundJSONBody(jsonData) + body, size, getBody, closer, err := relaycommon.NewOutboundJSONBody(jsonData) if err != nil { return types.NewError(err, types.ErrorCodeConvertRequestFailed, types.ErrOptionWithSkipRetry()) } defer closer.Close() jsonData = nil info.UpstreamRequestBodySize = size + info.UpstreamRequestGetBody = getBody requestBody = body } diff --git a/relay/embedding_handler.go b/relay/embedding_handler.go index 3be543bd..d3b617b6 100644 --- a/relay/embedding_handler.go +++ b/relay/embedding_handler.go @@ -58,13 +58,14 @@ func EmbeddingHelper(c *gin.Context, info *relaycommon.RelayInfo) (newAPIError * } 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 { return types.NewError(err, types.ErrorCodeConvertRequestFailed, types.ErrOptionWithSkipRetry()) } defer closer.Close() jsonData = nil info.UpstreamRequestBodySize = size + info.UpstreamRequestGetBody = getBody var requestBody io.Reader = body statusCodeMappingStr := c.GetString("status_code_mapping") resp, err := adaptor.DoRequest(c, info, requestBody) diff --git a/relay/gemini_handler.go b/relay/gemini_handler.go index 786feb8b..38ac0d87 100644 --- a/relay/gemini_handler.go +++ b/relay/gemini_handler.go @@ -141,6 +141,8 @@ func GeminiHelper(c *gin.Context, info *relaycommon.RelayInfo) (newAPIError *typ if err != nil { return types.NewErrorWithStatusCode(err, types.ErrorCodeReadRequestBodyFailed, http.StatusBadRequest, types.ErrOptionWithSkipRetry()) } + info.UpstreamRequestBodySize = storage.Size() + info.UpstreamRequestGetBody = storage.NewReader requestBody = common.ReaderOnly(storage) } else { // 使用 ConvertGeminiRequest 转换请求格式 @@ -164,13 +166,14 @@ func GeminiHelper(c *gin.Context, info *relaycommon.RelayInfo) (newAPIError *typ 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 { return types.NewError(err, types.ErrorCodeConvertRequestFailed, types.ErrOptionWithSkipRetry()) } defer closer.Close() jsonData = nil info.UpstreamRequestBodySize = size + info.UpstreamRequestGetBody = getBody requestBody = body } @@ -269,13 +272,14 @@ func GeminiEmbeddingHandler(c *gin.Context, info *relaycommon.RelayInfo) (newAPI } } 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 { return types.NewError(err, types.ErrorCodeConvertRequestFailed, types.ErrOptionWithSkipRetry()) } defer closer.Close() jsonData = nil info.UpstreamRequestBodySize = size + info.UpstreamRequestGetBody = getBody requestBody = body resp, err := adaptor.DoRequest(c, info, requestBody) diff --git a/relay/image_handler.go b/relay/image_handler.go index f970a43e..2d99a63c 100644 --- a/relay/image_handler.go +++ b/relay/image_handler.go @@ -51,6 +51,8 @@ func ImageHelper(c *gin.Context, info *relaycommon.RelayInfo) (newAPIError *type if err != nil { return types.NewErrorWithStatusCode(err, types.ErrorCodeReadRequestBodyFailed, http.StatusBadRequest, types.ErrOptionWithSkipRetry()) } + info.UpstreamRequestBodySize = storage.Size() + info.UpstreamRequestGetBody = storage.NewReader requestBody = common.ReaderOnly(storage) } else { 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) - body, size, closer, err := relaycommon.NewOutboundJSONBody(jsonData) + body, size, getBody, closer, err := relaycommon.NewOutboundJSONBody(jsonData) if err != nil { return types.NewError(err, types.ErrorCodeConvertRequestFailed, types.ErrOptionWithSkipRetry()) } defer closer.Close() jsonData = nil info.UpstreamRequestBodySize = size + info.UpstreamRequestGetBody = getBody requestBody = body } } diff --git a/relay/rerank_handler.go b/relay/rerank_handler.go index 769feb22..e2546dda 100644 --- a/relay/rerank_handler.go +++ b/relay/rerank_handler.go @@ -47,6 +47,8 @@ func RerankHelper(c *gin.Context, info *relaycommon.RelayInfo) (newAPIError *typ if err != nil { return types.NewErrorWithStatusCode(err, types.ErrorCodeReadRequestBodyFailed, http.StatusBadRequest, types.ErrOptionWithSkipRetry()) } + info.UpstreamRequestBodySize = storage.Size() + info.UpstreamRequestGetBody = storage.NewReader requestBody = common.ReaderOnly(storage) } else { 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) - body, size, closer, err := relaycommon.NewOutboundJSONBody(jsonData) + body, size, getBody, closer, err := relaycommon.NewOutboundJSONBody(jsonData) if err != nil { return types.NewError(err, types.ErrorCodeConvertRequestFailed, types.ErrOptionWithSkipRetry()) } defer closer.Close() jsonData = nil info.UpstreamRequestBodySize = size + info.UpstreamRequestGetBody = getBody requestBody = body } diff --git a/relay/responses_handler.go b/relay/responses_handler.go index c736c6a5..2531b3a6 100644 --- a/relay/responses_handler.go +++ b/relay/responses_handler.go @@ -82,6 +82,8 @@ func ResponsesHelper(c *gin.Context, info *relaycommon.RelayInfo) (newAPIError * if err != nil { return types.NewError(err, types.ErrorCodeReadRequestBodyFailed, types.ErrOptionWithSkipRetry()) } + info.UpstreamRequestBodySize = storage.Size() + info.UpstreamRequestGetBody = storage.NewReader requestBody = common.ReaderOnly(storage) } else { 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) - body, size, closer, err := relaycommon.NewOutboundJSONBody(jsonData) + body, size, getBody, closer, err := relaycommon.NewOutboundJSONBody(jsonData) if err != nil { return types.NewError(err, types.ErrorCodeConvertRequestFailed, types.ErrOptionWithSkipRetry()) } defer closer.Close() jsonData = nil info.UpstreamRequestBodySize = size + info.UpstreamRequestGetBody = getBody requestBody = body }