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:
Lucas
2026-08-06 15:56:38 +08:00
committed by GitHub
co-authored by Claude Fable 5
parent 0ab0202060
commit d6b5ce99de
20 changed files with 1067 additions and 20 deletions
+15 -3
View File
@@ -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
}
+163
View File
@@ -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)
}
+17
View File
@@ -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)
+18
View File
@@ -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,