mirror of
https://github.com/Wei-Shaw/sub2api.git
synced 2026-09-21 14:19:18 +08:00
fix: reuse OpenAI failover error body
This commit is contained in:
@@ -15,6 +15,14 @@ import (
|
||||
"github.com/tidwall/gjson"
|
||||
)
|
||||
|
||||
type panicOnReadCloser struct{}
|
||||
|
||||
func (panicOnReadCloser) Read(_ []byte) (int, error) {
|
||||
panic("response body should not be reread")
|
||||
}
|
||||
|
||||
func (panicOnReadCloser) Close() error { return nil }
|
||||
|
||||
func TestOpenAIGatewayService_Forward_FailoverReparsesCachedBodyForNextAccount(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
|
||||
@@ -102,6 +110,26 @@ func TestOpenAIGatewayService_Forward_FailoverReparsesCachedBodyForNextAccount(t
|
||||
}
|
||||
}
|
||||
|
||||
func TestOpenAIGatewayService_HandleFailoverSideEffects_DoesNotRereadResponseBody(t *testing.T) {
|
||||
svc := &OpenAIGatewayService{}
|
||||
account := &Account{
|
||||
ID: 88,
|
||||
Platform: PlatformOpenAI,
|
||||
Type: AccountTypeOAuth,
|
||||
}
|
||||
resp := &http.Response{
|
||||
StatusCode: http.StatusTooManyRequests,
|
||||
Header: http.Header{},
|
||||
Body: panicOnReadCloser{},
|
||||
}
|
||||
|
||||
require.NotPanics(t, func() {
|
||||
svc.handleFailoverSideEffects(context.Background(), resp, account, []byte(`{"error":{"type":"rate_limit_error","message":"rate limited"}}`))
|
||||
})
|
||||
|
||||
require.True(t, svc.isOpenAIAccountRuntimeBlocked(account))
|
||||
}
|
||||
|
||||
func TestGetOpenAIRequestBodyMap_IgnoresLegacyContextCache(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
|
||||
|
||||
@@ -2358,13 +2358,12 @@ func (s *OpenAIGatewayService) readUpstreamErrorBody(resp *http.Response) []byte
|
||||
return body
|
||||
}
|
||||
|
||||
func (s *OpenAIGatewayService) handleFailoverSideEffects(ctx context.Context, resp *http.Response, account *Account, requestedModel ...string) {
|
||||
body := s.readUpstreamErrorBody(resp)
|
||||
func (s *OpenAIGatewayService) handleFailoverSideEffects(ctx context.Context, resp *http.Response, account *Account, responseBody []byte, requestedModel ...string) {
|
||||
if len(requestedModel) > 0 {
|
||||
s.handleOpenAIAccountUpstreamError(ctx, account, resp.StatusCode, resp.Header, body, requestedModel[0])
|
||||
s.handleOpenAIAccountUpstreamError(ctx, account, resp.StatusCode, resp.Header, responseBody, requestedModel[0])
|
||||
return
|
||||
}
|
||||
s.handleOpenAIAccountUpstreamError(ctx, account, resp.StatusCode, resp.Header, body)
|
||||
s.handleOpenAIAccountUpstreamError(ctx, account, resp.StatusCode, resp.Header, responseBody)
|
||||
}
|
||||
|
||||
// Forward forwards request to OpenAI API
|
||||
@@ -3035,7 +3034,7 @@ func (s *OpenAIGatewayService) Forward(ctx context.Context, c *gin.Context, acco
|
||||
Detail: upstreamDetail,
|
||||
})
|
||||
|
||||
s.handleFailoverSideEffects(ctx, resp, account, upstreamModel)
|
||||
s.handleFailoverSideEffects(ctx, resp, account, respBody, upstreamModel)
|
||||
return nil, &UpstreamFailoverError{
|
||||
StatusCode: resp.StatusCode,
|
||||
ResponseBody: respBody,
|
||||
|
||||
@@ -638,7 +638,7 @@ func (s *OpenAIGatewayService) forwardOpenAIImagesAPIKey(
|
||||
Kind: "failover",
|
||||
Message: upstreamMsg,
|
||||
})
|
||||
s.handleFailoverSideEffects(upstreamCtx, resp, account, upstreamModel)
|
||||
s.handleFailoverSideEffects(upstreamCtx, resp, account, respBody, upstreamModel)
|
||||
return nil, &UpstreamFailoverError{
|
||||
StatusCode: resp.StatusCode,
|
||||
ResponseBody: respBody,
|
||||
|
||||
@@ -1358,7 +1358,7 @@ func (s *OpenAIGatewayService) forwardOpenAIImagesOAuth(
|
||||
Kind: "failover",
|
||||
Message: upstreamMsg,
|
||||
})
|
||||
s.handleFailoverSideEffects(upstreamCtx, resp, account, requestModel)
|
||||
s.handleFailoverSideEffects(upstreamCtx, resp, account, respBody, requestModel)
|
||||
return nil, &UpstreamFailoverError{
|
||||
StatusCode: resp.StatusCode,
|
||||
ResponseBody: respBody,
|
||||
|
||||
Reference in New Issue
Block a user