feat: sync deployed v1.0.0-rc.38-custom.42

Sync production source including routing, billing, Astra detection and public uptime. Update custom build instructions and license packaging.
This commit is contained in:
cneicy committed 2026-09-26 04:49:11 +08:00
1 parent aa6c3f50f4
commit ceddfd3d6b
828 files changed
+122253 -13702

No files matched your search

+25 -28
View File
@@ -77,15 +77,7 @@ func (a *Adaptor) ConvertClaudeRequest(c *gin.Context, info *relaycommon.RelayIn
case relayconvert.ConverterNone:
return a.claudeAdaptor.ConvertClaudeRequest(c, info, request)
case relayconvert.ConverterClaudeMessagesToOpenAIChat:
result, err := service.ConvertRequestByID(c, info, converter, request)
if err != nil {
return nil, err
}
chatRequest, ok := result.Value.(*dto.GeneralOpenAIRequest)
if !ok {
return nil, fmt.Errorf("expected OpenAI chat completions request, got %T", result.Value)
}
return a.convertOpenAICompatibleRequest(c, info, chatRequest)
return a.convertCrossProtocolChatRequest(c, info, converter, request)
default:
return nil, fmt.Errorf("converter %q does not support Anthropic Messages requests", converter)
}
@@ -101,15 +93,7 @@ func (a *Adaptor) ConvertGeminiRequest(c *gin.Context, info *relaycommon.RelayIn
case relayconvert.ConverterNone:
return a.geminiAdaptor.ConvertGeminiRequest(c, info, request)
case relayconvert.ConverterGeminiContentToOpenAIChat:
result, err := service.ConvertRequestByID(c, info, converter, request)
if err != nil {
return nil, err
}
chatRequest, ok := result.Value.(*dto.GeneralOpenAIRequest)
if !ok {
return nil, fmt.Errorf("expected OpenAI chat completions request, got %T", result.Value)
}
return a.convertOpenAICompatibleRequest(c, info, chatRequest)
return a.convertCrossProtocolChatRequest(c, info, converter, request)
default:
return nil, fmt.Errorf("converter %q does not support Gemini generateContent requests", converter)
}
@@ -124,15 +108,7 @@ func (a *Adaptor) ConvertOpenAIResponsesRequest(c *gin.Context, info *relaycommo
case relayconvert.ConverterNone:
return a.convertOpenAICompatibleResponsesRequest(c, info, request)
case relayconvert.ConverterOpenAIResponsesToOpenAIChat:
result, err := service.ConvertRequestByID(c, info, converter, request)
if err != nil {
return nil, err
}
chatRequest, ok := result.Value.(*dto.GeneralOpenAIRequest)
if !ok {
return nil, fmt.Errorf("expected OpenAI chat completions request, got %T", result.Value)
}
return a.convertOpenAICompatibleRequest(c, info, chatRequest)
return a.convertCrossProtocolChatRequest(c, info, converter, request)
case relayconvert.ConverterOpenAIResponsesToGemini:
result, err := service.ConvertRequestByID(c, info, converter, request)
if err != nil {
@@ -286,7 +262,7 @@ func (a *Adaptor) DoRequest(c *gin.Context, info *relaycommon.RelayInfo, request
if err := a.resolve(c, info); err != nil {
return nil, err
}
if !a.converted && a.converter != relayconvert.ConverterNone {
if !a.converted && !a.route.SupportsPassThroughBody() {
return nil, errors.New("advanced custom converter routes cannot be used with pass-through request body")
}
@@ -307,6 +283,8 @@ func (a *Adaptor) DoResponse(c *gin.Context, resp *http.Response, info *relaycom
}
switch a.converter {
case dto.AdvancedCustomConverterSGLangRerank:
return a.doSGLangRerankResponse(c, resp, info)
case relayconvert.ConverterNone:
return a.doNativeResponse(c, resp, info)
case relayconvert.ConverterClaudeMessagesToOpenAIChat,
@@ -534,6 +512,25 @@ func isJSONRequest(c *gin.Context) bool {
return strings.Contains(strings.ToLower(c.Request.Header.Get("Content-Type")), "application/json")
}
// convertCrossProtocolChatRequest converts a Claude, Gemini, or Responses request into an
// OpenAI chat completions request for an OpenAI-compatible upstream. Streaming requests
// always ask the upstream for usage, because the downstream protocol reports usage in its
// own format and never carries stream_options itself.
func (a *Adaptor) convertCrossProtocolChatRequest(c *gin.Context, info *relaycommon.RelayInfo, converter string, request any) (any, error) {
result, err := service.ConvertRequestByID(c, info, converter, request)
if err != nil {
return nil, err
}
chatRequest, ok := result.Value.(*dto.GeneralOpenAIRequest)
if !ok {
return nil, fmt.Errorf("expected OpenAI chat completions request, got %T", result.Value)
}
if info.SupportStreamOptions && info.IsStream {
chatRequest.StreamOptions = &dto.StreamOptions{IncludeUsage: true}
}
return a.convertOpenAICompatibleRequest(c, info, chatRequest)
}
func (a *Adaptor) convertOpenAICompatibleRequest(c *gin.Context, info *relaycommon.RelayInfo, request *dto.GeneralOpenAIRequest) (any, error) {
old := info.ChannelType
info.ChannelType = constant.ChannelTypeOpenAI
@@ -829,6 +829,143 @@ func TestAdaptorConvertsGeminiRequestToOpenAIChatUpstream(t *testing.T) {
assert.Equal(t, "user", chatReq.Messages[0].Role)
}
func TestAdaptorCrossProtocolChatUpstreamRequestsStreamUsage(t *testing.T) {
claudeReq := &dto.ClaudeRequest{
Model: "gpt-test",
Messages: []dto.ClaudeMessage{{Role: "user", Content: "hello"}},
}
geminiReq := &dto.GeminiChatRequest{
Contents: []dto.GeminiChatContent{{Role: "user", Parts: []dto.GeminiPart{{Text: "hello"}}}},
}
responsesReq := dto.OpenAIResponsesRequest{
Model: "gpt-test",
Input: mustAdvancedCustomRawMessage(t, "hello"),
}
tests := []struct {
name string
route dto.AdvancedCustomRoute
relayFormat types.RelayFormat
relayMode int
requestPath string
convert func(*Adaptor, *gin.Context, *relaycommon.RelayInfo) (any, error)
isStream bool
supportUsage bool
wantUsage bool
}{
{
name: "claude stream requests usage",
route: dto.AdvancedCustomRoute{
IncomingPath: "/v1/messages",
UpstreamPath: "/v1/chat/completions",
Converter: relayconvert.ConverterClaudeMessagesToOpenAIChat,
},
relayFormat: types.RelayFormatClaude,
relayMode: relayconstant.RelayModeChatCompletions,
requestPath: "/v1/messages",
convert: func(a *Adaptor, c *gin.Context, info *relaycommon.RelayInfo) (any, error) {
return a.ConvertClaudeRequest(c, info, claudeReq)
},
isStream: true,
supportUsage: true,
wantUsage: true,
},
{
name: "gemini stream requests usage",
route: dto.AdvancedCustomRoute{
IncomingPath: "/v1beta/models/{model}:generateContent",
UpstreamPath: "/v1/chat/completions",
Converter: relayconvert.ConverterGeminiContentToOpenAIChat,
},
relayFormat: types.RelayFormatGemini,
relayMode: relayconstant.RelayModeGemini,
requestPath: "/v1beta/models/gpt-test:generateContent",
convert: func(a *Adaptor, c *gin.Context, info *relaycommon.RelayInfo) (any, error) {
return a.ConvertGeminiRequest(c, info, geminiReq)
},
isStream: true,
supportUsage: true,
wantUsage: true,
},
{
name: "responses stream requests usage",
route: dto.AdvancedCustomRoute{
IncomingPath: "/v1/responses",
UpstreamPath: "/v1/chat/completions",
Converter: relayconvert.ConverterOpenAIResponsesToOpenAIChat,
},
relayFormat: types.RelayFormatOpenAIResponses,
relayMode: relayconstant.RelayModeResponses,
requestPath: "/v1/responses",
convert: func(a *Adaptor, c *gin.Context, info *relaycommon.RelayInfo) (any, error) {
return a.ConvertOpenAIResponsesRequest(c, info, responsesReq)
},
isStream: true,
supportUsage: true,
wantUsage: true,
},
{
name: "claude non-stream leaves stream options unset",
route: dto.AdvancedCustomRoute{
IncomingPath: "/v1/messages",
UpstreamPath: "/v1/chat/completions",
Converter: relayconvert.ConverterClaudeMessagesToOpenAIChat,
},
relayFormat: types.RelayFormatClaude,
relayMode: relayconstant.RelayModeChatCompletions,
requestPath: "/v1/messages",
convert: func(a *Adaptor, c *gin.Context, info *relaycommon.RelayInfo) (any, error) {
return a.ConvertClaudeRequest(c, info, claudeReq)
},
isStream: false,
supportUsage: true,
wantUsage: false,
},
{
name: "claude stream without stream options support",
route: dto.AdvancedCustomRoute{
IncomingPath: "/v1/messages",
UpstreamPath: "/v1/chat/completions",
Converter: relayconvert.ConverterClaudeMessagesToOpenAIChat,
},
relayFormat: types.RelayFormatClaude,
relayMode: relayconstant.RelayModeChatCompletions,
requestPath: "/v1/messages",
convert: func(a *Adaptor, c *gin.Context, info *relaycommon.RelayInfo) (any, error) {
return a.ConvertClaudeRequest(c, info, claudeReq)
},
isStream: true,
supportUsage: false,
wantUsage: false,
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
adaptor := &Adaptor{}
info := advancedCustomRelayInfo(&dto.AdvancedCustomConfig{Routes: []dto.AdvancedCustomRoute{tt.route}})
info.RelayFormat = tt.relayFormat
info.RelayMode = tt.relayMode
info.RequestURLPath = tt.requestPath
info.IsStream = tt.isStream
info.SupportStreamOptions = tt.supportUsage
c := advancedCustomGinContext(tt.requestPath)
converted, err := tt.convert(adaptor, c, info)
require.NoError(t, err)
chatReq, ok := converted.(*dto.GeneralOpenAIRequest)
require.True(t, ok)
if !tt.wantUsage {
assert.Nil(t, chatReq.StreamOptions)
return
}
require.NotNil(t, chatReq.StreamOptions)
assert.True(t, chatReq.StreamOptions.IncludeUsage)
})
}
}
func advancedCustomRelayInfo(config *dto.AdvancedCustomConfig) *relaycommon.RelayInfo {
return &relaycommon.RelayInfo{
RelayFormat: types.RelayFormatOpenAI,
@@ -0,0 +1,134 @@
package advancedcustom_test
import (
"bytes"
"context"
"io"
"net/http"
"net/http/httptest"
"strings"
"testing"
"github.com/QuantumNous/new-api/common"
"github.com/QuantumNous/new-api/constant"
"github.com/QuantumNous/new-api/relay"
"github.com/QuantumNous/new-api/relay/channel/advancedcustom"
relaycommon "github.com/QuantumNous/new-api/relay/common"
relayconstant "github.com/QuantumNous/new-api/relay/constant"
"github.com/QuantumNous/new-api/relaykit/dto"
"github.com/QuantumNous/new-api/relaykit/types"
"github.com/QuantumNous/new-api/service"
"github.com/gin-gonic/gin"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
)
func TestSGLangChannelProtocols(t *testing.T) {
apiType, ok := common.ChannelType2APIType(constant.ChannelTypeSGLang)
require.True(t, ok)
adaptor := relay.GetAdaptor(apiType)
require.NotNil(t, adaptor)
assert.Equal(t, "advanced_custom", adaptor.GetChannelName())
assert.Contains(t, common.GetEndpointTypesByChannelType(constant.ChannelTypeSGLang, "served-model"), constant.EndpointTypeAnthropic)
assert.Contains(t, common.GetEndpointTypesByChannelType(constant.ChannelTypeSGLang, "served-model"), constant.EndpointTypeOpenAIResponse)
info := &relaycommon.RelayInfo{ChannelMeta: &relaycommon.ChannelMeta{ChannelType: constant.ChannelTypeSGLang, ChannelBaseUrl: "https://inference.example/prefix", ApiKey: "test-key"}}
info.ChannelOtherSettings.AdvancedCustom = common.GetAdvancedCustomPreset(constant.ChannelTypeSGLang)
adaptor.Init(info)
for _, path := range []string{"/v1/chat/completions", "/v1/completions", "/v1/embeddings", "/v1/responses", "/v1/messages"} {
info.RequestURLPath = path
adaptor = &advancedcustom.Adaptor{}
adaptor.Init(info)
url, err := adaptor.GetRequestURL(info)
require.NoError(t, err)
assert.Equal(t, "https://inference.example/prefix"+path, url)
}
for _, format := range []types.RelayFormat{types.RelayFormatOpenAI, types.RelayFormatClaude} {
info.RelayFormat = format
info.RequestURLPath = "/v1/messages"
adaptor = &advancedcustom.Adaptor{}
adaptor.Init(info)
c, _ := gin.CreateTestContext(httptest.NewRecorder())
c.Request = httptest.NewRequest(http.MethodPost, "/v1/messages", nil)
headers := http.Header{}
require.NoError(t, adaptor.SetupRequestHeader(c, &headers, info))
assert.Equal(t, "Bearer test-key", headers.Get("Authorization"))
}
info.RequestURLPath = "/v1/chat/completions"
adaptor = &advancedcustom.Adaptor{}
adaptor.Init(info)
c, _ := gin.CreateTestContext(httptest.NewRecorder())
c.Request = httptest.NewRequest(http.MethodPost, "/v1/chat/completions", nil)
request := &dto.GeneralOpenAIRequest{Model: "served-model"}
converted, err := adaptor.ConvertOpenAIRequest(c, info, request)
require.NoError(t, err)
assert.Same(t, request, converted)
info.RequestURLPath = "/v1/messages"
adaptor = &advancedcustom.Adaptor{}
adaptor.Init(info)
c.Request = httptest.NewRequest(http.MethodPost, "/v1/messages", nil)
claudeRequest := &dto.ClaudeRequest{Model: "served-model", Thinking: &dto.Thinking{Type: "adaptive"}}
converted, err = adaptor.ConvertClaudeRequest(c, info, claudeRequest)
require.NoError(t, err)
assert.Same(t, claudeRequest, converted)
info.RequestURLPath = "/v1beta/models/test:generateContent"
adaptor = &advancedcustom.Adaptor{}
adaptor.Init(info)
_, err = adaptor.ConvertGeminiRequest(nil, info, &dto.GeminiChatRequest{})
require.Error(t, err)
}
func TestSGLangRerankRequestAndResponse(t *testing.T) {
service.InitHttpClient()
upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
assert.Equal(t, "/prefix/v1/rerank", r.URL.Path)
assert.Equal(t, "Bearer test-key", r.Header.Get("Authorization"))
var request dto.RerankRequest
require.NoError(t, common.DecodeJson(r.Body, &request))
assert.Equal(t, "served-reranker", request.Model)
assert.Equal(t, []any{"first", "second"}, request.Documents)
_, _ = io.WriteString(w, `[{"index":1,"score":0.9,"document":"second","meta_info":{"prompt_tokens":5}}]`)
}))
defer upstream.Close()
info := &relaycommon.RelayInfo{
ChannelMeta: &relaycommon.ChannelMeta{ChannelType: constant.ChannelTypeSGLang, ChannelBaseUrl: upstream.URL + "/prefix", ApiKey: "test-key"},
RequestURLPath: "/rerank", RelayMode: relayconstant.RelayModeRerank,
RelayFormat: types.RelayFormatOpenAI,
RerankerInfo: &relaycommon.RerankerInfo{Documents: []any{"first", "second"}, ReturnDocuments: true},
}
info.SetEstimatePromptTokens(12)
adaptor := &advancedcustom.Adaptor{}
info.ChannelOtherSettings.AdvancedCustom = common.GetAdvancedCustomPreset(constant.ChannelTypeSGLang)
adaptor.Init(info)
recorder := httptest.NewRecorder()
c, _ := gin.CreateTestContext(recorder)
c.Request = httptest.NewRequest(http.MethodPost, "/rerank", nil).WithContext(context.Background())
converted, err := adaptor.ConvertRerankRequest(c, info.RelayMode, dto.RerankRequest{Model: "served-reranker", Query: "query", Documents: info.Documents})
require.NoError(t, err)
body, err := common.Marshal(converted)
require.NoError(t, err)
response, err := adaptor.DoRequest(c, info, bytes.NewReader(body))
require.NoError(t, err)
usage, apiErr := adaptor.DoResponse(c, response.(*http.Response), info)
require.Nil(t, apiErr)
assert.Equal(t, 12, usage.(*dto.Usage).TotalTokens)
var decoded dto.RerankResponse
require.NoError(t, common.Unmarshal(recorder.Body.Bytes(), &decoded))
assert.Equal(t, []dto.RerankResponseResult{{Index: 1, RelevanceScore: 0.9, Document: map[string]any{"text": "second"}}}, decoded.Results)
assert.Equal(t, 12, decoded.Usage.PromptTokens)
assert.Zero(t, decoded.Usage.CompletionTokens)
}
func TestSGLangRerankRejectsMalformedResults(t *testing.T) {
for _, body := range []string{`null`, `{}`, `[{"index":99,"score":1}]`, `[{"index":-1,"score":0}]`, `[{"index":0}]`, `[{"score":1}]`} {
t.Run(body, func(t *testing.T) {
adaptor := &advancedcustom.Adaptor{}
info := &relaycommon.RelayInfo{ChannelMeta: &relaycommon.ChannelMeta{ChannelOtherSettings: dto.ChannelOtherSettings{AdvancedCustom: common.GetAdvancedCustomPreset(constant.ChannelTypeSGLang)}}, RequestURLPath: "/v1/rerank", RelayMode: relayconstant.RelayModeRerank, RerankerInfo: &relaycommon.RerankerInfo{Documents: []any{"only document"}}}
recorder := httptest.NewRecorder()
c, _ := gin.CreateTestContext(recorder)
usage, err := adaptor.DoResponse(c, &http.Response{Body: io.NopCloser(strings.NewReader(body))}, info)
require.NotNil(t, err)
assert.Nil(t, usage)
assert.Empty(t, recorder.Body.String())
})
}
}
+46
View File
@@ -0,0 +1,46 @@
package advancedcustom
import (
"errors"
"net/http"
"github.com/QuantumNous/new-api/common"
relaycommon "github.com/QuantumNous/new-api/relay/common"
"github.com/QuantumNous/new-api/relaykit/dto"
"github.com/QuantumNous/new-api/relaykit/types"
"github.com/QuantumNous/new-api/service"
"github.com/gin-gonic/gin"
)
func (a *Adaptor) doSGLangRerankResponse(c *gin.Context, resp *http.Response, info *relaycommon.RelayInfo) (any, *types.NewAPIError) {
defer service.CloseResponseBodyGracefully(resp)
// SGLang v0.5.19 returns a sorted array with score, not a Jina results envelope.
// https://github.com/sgl-project/sglang/blob/v0.5.19/python/sglang/srt/entrypoints/openai/serving_rerank.py
var results []struct {
Index *int `json:"index"`
Score *float64 `json:"score"`
}
if err := common.DecodeJson(resp.Body, &results); err != nil {
return nil, types.NewOpenAIError(err, types.ErrorCodeBadResponseBody, http.StatusBadGateway)
}
if results == nil {
return nil, types.NewOpenAIError(errors.New("invalid SGLang rerank response"), types.ErrorCodeBadResponseBody, http.StatusBadGateway)
}
response := dto.RerankResponse{Results: make([]dto.RerankResponseResult, 0, len(results))}
for _, result := range results {
if result.Index == nil || result.Score == nil || *result.Index < 0 || *result.Index >= len(info.Documents) {
return nil, types.NewOpenAIError(errors.New("invalid SGLang rerank result"), types.ErrorCodeBadResponseBody, http.StatusBadGateway)
}
item := dto.RerankResponseResult{Index: *result.Index, RelevanceScore: *result.Score}
if info.ReturnDocuments {
item.Document = dto.RerankDocument{Text: info.Documents[*result.Index]}
}
response.Results = append(response.Results, item)
}
// top_n can omit scored documents, and decoder rerankers can omit meta_info.
// Per-result counts therefore cannot account for the full request; use the
// existing prompt estimate, as with other rerank providers without usage.
response.Usage = dto.Usage{PromptTokens: info.GetEstimatePromptTokens(), TotalTokens: info.GetEstimatePromptTokens()}
c.JSON(http.StatusOK, response)
return &response.Usage, nil
}
+76
View File
@@ -1,10 +1,12 @@
package ali
import (
"bytes"
"encoding/base64"
"encoding/json"
"fmt"
"io"
"mime/multipart"
"net/http"
"net/http/httptest"
"strings"
@@ -19,6 +21,7 @@ import (
relayhelper "github.com/QuantumNous/new-api/relay/helper"
"github.com/QuantumNous/new-api/relaykit/dto"
"github.com/QuantumNous/new-api/service"
"github.com/QuantumNous/new-api/setting/ratio_setting"
"github.com/QuantumNous/new-api/setting/system_setting"
"github.com/gin-gonic/gin"
"github.com/stretchr/testify/assert"
@@ -26,6 +29,43 @@ import (
"github.com/tidwall/gjson"
)
func TestAliMultipartEditsUseValidatedProviderQuantity(t *testing.T) {
for _, tc := range []struct {
name string
convert func(*gin.Context, *relaycommon.RelayInfo, dto.ImageRequest) (*AliImageRequest, error)
}{
{"multimodal edit", oaiFormEdit2AliImageEdit},
{"legacy Wan edit", oaiFormEdit2WanxImageEdit},
} {
t.Run(tc.name, func(t *testing.T) {
var body bytes.Buffer
writer := multipart.NewWriter(&body)
require.NoError(t, writer.WriteField("model", "fixture-image"))
require.NoError(t, writer.WriteField("n", "2"))
require.NoError(t, writer.WriteField("parameters", `{"n":3,"prompt_extend":false}`))
part, err := writer.CreateFormFile("image", "input.png")
require.NoError(t, err)
_, err = part.Write([]byte("fixture image"))
require.NoError(t, err)
require.NoError(t, writer.Close())
c, _ := gin.CreateTestContext(httptest.NewRecorder())
c.Request = httptest.NewRequest(http.MethodPost, "/v1/images/edits", &body)
c.Request.Header.Set("Content-Type", writer.FormDataContentType())
common.SetContextKey(c, rootconstant.ContextKeyChannelType, rootconstant.ChannelTypeAli)
request, err := relayhelper.GetAndValidOpenAIImageRequest(c, constant.RelayModeImagesEdits)
require.NoError(t, err)
info := &relaycommon.RelayInfo{Request: request}
converted, err := tc.convert(c, info, *request)
require.NoError(t, err)
require.NotNil(t, converted.Parameters.N)
assert.Equal(t, uint(3), *converted.Parameters.N)
require.NotNil(t, converted.Parameters.PromptExtend)
assert.False(t, *converted.Parameters.PromptExtend)
assert.Equal(t, uint(2), *request.N, "conversion must preserve the incoming request")
})
}
}
func TestConvertOpenAIRequestFiltersThinkingBudgetByUpstreamModel(t *testing.T) {
tests := []struct {
name string
@@ -169,6 +209,42 @@ func TestMappedAliImageModelUsesUpstreamProtocol(t *testing.T) {
assert.IsType(t, &AliImageRequest{}, converted)
}
func TestAliImageCountMatchesLegacyReservation(t *testing.T) {
saved := ratio_setting.ModelPrice2JSONString()
require.NoError(t, ratio_setting.UpdateModelPriceByJSONString(`{"z-image":0.04}`))
t.Cleanup(func() { require.NoError(t, ratio_setting.UpdateModelPriceByJSONString(saved)) })
for _, tc := range []struct {
name, body string
count int
multiplier float64
}{
{"provider override", `{"model":"z-image","n":2,"parameters":{"n":4}}`, 4, 1},
{"empty parameters inherit top-level", `{"model":"z-image","n":2,"parameters":{}}`, 2, 1},
{"null count inherits top-level", `{"model":"z-image","n":2,"parameters":{"n":null}}`, 2, 1},
{"default count", `{"model":"z-image"}`, 1, 1},
{"prompt extension", `{"model":"z-image","parameters":{"n":3,"prompt_extend":true}}`, 3, 2},
} {
t.Run(tc.name, func(t *testing.T) {
c, _ := gin.CreateTestContext(httptest.NewRecorder())
c.Request = httptest.NewRequest(http.MethodPost, "/v1/images/generations", strings.NewReader(tc.body))
c.Request.Header.Set("Content-Type", "application/json")
common.SetContextKey(c, rootconstant.ContextKeyChannelType, rootconstant.ChannelTypeAli)
request, err := relayhelper.GetAndValidOpenAIImageRequest(c, constant.RelayModeImagesGenerations)
require.NoError(t, err)
info := &relaycommon.RelayInfo{Request: request, OriginModelName: request.Model, UserGroup: "default", UsingGroup: "default"}
price, err := relayhelper.ModelPriceHelper(c, info, 0, request.GetTokenCountMeta())
require.NoError(t, err)
assert.Equal(t, common.QuotaFromFloat(0.04*float64(tc.count)*tc.multiplier*common.QuotaPerUnit), price.QuotaToPreConsume)
info.ChannelMeta = &relaycommon.ChannelMeta{ChannelType: rootconstant.ChannelTypeAli, UpstreamModelName: request.Model}
converted, err := oaiImage2AliImageRequest(info, *request, true)
require.NoError(t, err)
encoded, err := common.Marshal(converted)
require.NoError(t, err)
assert.Equal(t, int64(tc.count), gjson.GetBytes(encoded, "parameters.n").Int())
})
}
}
func TestAliImageHandlerHonorsRequestResponseFormat(t *testing.T) {
imageBytes := []byte("ali-image")
var downloads atomic.Int32
+1 -1
View File
@@ -172,7 +172,7 @@ type AliImageRequest struct {
type AliImageParameters struct {
Size string `json:"size,omitempty"`
N int `json:"n,omitempty"`
N *uint `json:"n,omitempty"`
Steps string `json:"steps,omitempty"`
Scale string `json:"scale,omitempty"`
Watermark *bool `json:"watermark,omitempty"`
+43 -26
View File
@@ -18,25 +18,29 @@ import (
"github.com/QuantumNous/new-api/service"
"github.com/gin-gonic/gin"
"github.com/samber/lo"
)
func oaiImage2AliImageRequest(info *relaycommon.RelayInfo, request dto.ImageRequest, isSync bool) (*AliImageRequest, error) {
var imageRequest AliImageRequest
imageRequest.Model = request.Model
imageRequest.ResponseFormat = request.ResponseFormat
imageRequest.Parameters = AliImageParameters{
Size: strings.ReplaceAll(request.Size, "x", "*"),
Watermark: request.Watermark,
}
if request.Extra != nil {
if val, ok := request.Extra["parameters"]; ok {
imageRequest.Parameters = AliImageParameters{}
err := common.Unmarshal(val, &imageRequest.Parameters)
if err != nil {
return nil, fmt.Errorf("invalid parameters field: %w", err)
}
} else {
// 兼容没有parameters字段的情况,从openai标准字段中提取参数
imageRequest.Parameters = AliImageParameters{
Size: strings.Replace(request.Size, "x", "*", -1),
N: int(lo.FromPtrOr(request.N, uint(1))),
Watermark: request.Watermark,
// Direct adaptor callers can bypass ingress validation. Reuse the
// decoded provider scalars and the same quantity validation below.
if request.BillingParameters == nil {
request.BillingParameters = &dto.ImageBillingParameters{
N: imageRequest.Parameters.N, PromptExtend: imageRequest.Parameters.PromptExtend,
}
}
}
if val, ok := request.Extra["input"]; ok {
@@ -47,23 +51,21 @@ func oaiImage2AliImageRequest(info *relaycommon.RelayInfo, request dto.ImageRequ
}
}
if strings.Contains(request.Model, "z-image") {
// z-image 开启prompt_extend后,按2倍计费
if imageRequest.Parameters.PromptExtendValue() {
info.PriceData.AddOtherRatio("prompt_extend", 2)
count, err := request.ImageCount(true)
if err != nil {
return nil, err
}
imageRequest.Parameters.N = common.GetPointer(uint(count))
if request.BillingParameters != nil {
imageRequest.Parameters.PromptExtend = request.BillingParameters.PromptExtend
}
if info.TieredBillingSnapshot == nil {
info.PriceData.AddOtherRatio("n", float64(count))
if strings.Contains(request.Model, "z-image") && imageRequest.Parameters.PromptExtendValue() {
info.PriceData.AddOtherRatio("prompt_extend", common.ZImagePromptExtendMultiplier)
}
}
// Parameters may come from Extra["parameters"], bypassing the standard
// top-level n validation; enforce the same bound before it becomes a
// billing multiplier.
if imageRequest.Parameters.N < 0 || imageRequest.Parameters.N > dto.MaxImageN {
return nil, fmt.Errorf("parameters.n must be an integer between 1 and %d", dto.MaxImageN)
}
if imageRequest.Parameters.N != 0 {
info.PriceData.AddOtherRatio("n", float64(imageRequest.Parameters.N))
}
// 同步图片模型和异步图片模型请求格式不一样
if isSync {
if imageRequest.Input == nil {
@@ -159,6 +161,10 @@ func getImageBase64sFromForm(c *gin.Context, fieldName string) ([]string, error)
}
func oaiFormEdit2AliImageEdit(c *gin.Context, info *relaycommon.RelayInfo, request dto.ImageRequest) (*AliImageRequest, error) {
count, err := request.ImageCount(true)
if err != nil {
return nil, err
}
var imageRequest AliImageRequest
imageRequest.Model = request.Model
imageRequest.ResponseFormat = request.ResponseFormat
@@ -186,9 +192,12 @@ func oaiFormEdit2AliImageEdit(c *gin.Context, info *relaycommon.RelayInfo, reque
},
}
imageRequest.Parameters = AliImageParameters{
N: int(lo.FromPtrOr(request.N, uint(1))),
N: common.GetPointer(uint(count)),
Watermark: request.Watermark,
}
if request.BillingParameters != nil {
imageRequest.Parameters.PromptExtend = request.BillingParameters.PromptExtend
}
return &imageRequest, nil
}
@@ -336,10 +345,18 @@ func aliImageHandler(a *Adaptor, c *gin.Context, resp *http.Response, info *rela
}
imageResponses := responseAli2OpenAIImage(c, aliResponse, originRespBody, info, responseFormat)
if aliResponse.Usage.ImageCount != 0 {
info.PriceData.AddOtherRatio("n", float64(aliResponse.Usage.ImageCount))
} else if len(imageResponses.Data) != 0 {
info.PriceData.AddOtherRatio("n", float64(len(imageResponses.Data)))
count := int64(aliResponse.Usage.ImageCount)
if count < 0 || count > dto.MaxImageN {
logger.LogWarn(c, "invalid Ali image usage count %d; falling back to response images or requested quantity", count)
count = 0
}
if count == 0 {
count = int64(len(imageResponses.Data))
}
if count > dto.MaxImageN {
logger.LogWarn(c, "invalid Ali response image count %d; retaining requested quantity", count)
} else if count > 0 {
info.UpdateImageCount(count)
}
jsonResponse, err := common.Marshal(imageResponses)
if err != nil {
+8 -6
View File
@@ -13,7 +13,10 @@ import (
)
func oaiFormEdit2WanxImageEdit(c *gin.Context, info *relaycommon.RelayInfo, request dto.ImageRequest) (*AliImageRequest, error) {
var err error
count, err := request.ImageCount(true)
if err != nil {
return nil, err
}
var imageRequest AliImageRequest
imageRequest.Model = request.Model
imageRequest.ResponseFormat = request.ResponseFormat
@@ -27,14 +30,13 @@ func oaiFormEdit2WanxImageEdit(c *gin.Context, info *relaycommon.RelayInfo, requ
if wanInput.Images, err = getImageBase64sFromForm(c, "image"); err != nil {
return nil, fmt.Errorf("get image base64s from form failed: %w", err)
}
//wanParams := WanImageParameters{
// N: int(request.N),
//}
imageRequest.Input = wanInput
imageRequest.Parameters = AliImageParameters{
N: int(lo.FromPtrOr(request.N, uint(1))),
N: common.GetPointer(uint(count)),
}
if request.BillingParameters != nil {
imageRequest.Parameters.PromptExtend = request.BillingParameters.PromptExtend
}
info.PriceData.AddOtherRatio("n", float64(imageRequest.Parameters.N))
return &imageRequest, nil
}
+84 -6
View File
@@ -372,11 +372,16 @@ func DoFormRequest(a Adaptor, c *gin.Context, info *common.RelayInfo, requestBod
return resp, nil
}
// DoWssRequest dials the adaptor's upstream over WebSocket for the realtime
// and Responses WebSocket relays. It honors the channel proxy, is bound to the
// request context, and reports a rejected handshake as a *types.NewAPIError
// carrying the upstream status code (types.NewError preserves it).
func DoWssRequest(a Adaptor, c *gin.Context, info *common.RelayInfo, requestBody io.Reader) (*websocket.Conn, error) {
fullRequestURL, err := a.GetRequestURL(info)
if err != nil {
return nil, fmt.Errorf("get request url failed: %w", err)
}
fullRequestURL = toWebSocketURL(fullRequestURL)
targetHeader := http.Header{}
err = a.SetupRequestHeader(c, &targetHeader, info)
if err != nil {
@@ -392,16 +397,41 @@ func DoWssRequest(a Adaptor, c *gin.Context, info *common.RelayInfo, requestBody
targetHeader.Set(key, value)
}
targetHeader.Set("Content-Type", c.Request.Header.Get("Content-Type"))
targetConn, _, err := websocket.DefaultDialer.Dial(fullRequestURL, targetHeader)
dialer := *websocket.DefaultDialer
if info.ChannelSetting.Proxy != "" {
proxyURL, _, proxyErr := common2.ParseProxyURLRuntime(info.ChannelSetting.Proxy)
if proxyErr != nil {
return nil, proxyErr
}
dialer.Proxy = http.ProxyURL(proxyURL)
}
targetConn, resp, err := dialer.DialContext(c.Request.Context(), fullRequestURL, targetHeader)
if err != nil {
return nil, fmt.Errorf("dial failed to %s: %w", common.SanitizeURLForLog(fullRequestURL), err)
statusCode := http.StatusInternalServerError
if resp != nil {
statusCode = resp.StatusCode
if resp.Body != nil {
_ = resp.Body.Close()
}
}
return nil, types.NewErrorWithStatusCode(fmt.Errorf("dial failed to %s: %w", common.SanitizeURLForLog(fullRequestURL), err), types.ErrorCodeDoRequestFailed, statusCode)
}
// send request body
//all, err := io.ReadAll(requestBody)
//err = service.WssString(c, targetConn, string(all))
return targetConn, nil
}
// toWebSocketURL maps an http(s) endpoint to ws(s). Realtime adaptors already
// return ws(s) URLs, which pass through unchanged.
func toWebSocketURL(raw string) string {
switch {
case strings.HasPrefix(raw, "https://"):
return "wss://" + strings.TrimPrefix(raw, "https://")
case strings.HasPrefix(raw, "http://"):
return "ws://" + strings.TrimPrefix(raw, "http://")
default:
return raw
}
}
func startPingKeepAlive(c *gin.Context, pingInterval time.Duration) (context.CancelFunc, <-chan struct{}) {
pingerCtx, stopPinger := context.WithCancel(context.Background())
done := make(chan struct{})
@@ -488,6 +518,13 @@ func keepUpstreamRedirectResponse(_ *http.Request, _ []*http.Request) error {
}
func doRequest(c *gin.Context, req *http.Request, info *common.RelayInfo) (*http.Response, error) {
guardEnabled := info.ResponseGuard != nil && info.AstraAutoSwitch
if guardEnabled && info.UpstreamRequests >= common.ResponseGuardMaxAttempts {
return nil, types.NewErrorWithStatusCode(errors.New("response model attempt budget exhausted"), common.ResponseGuardErrorCode, http.StatusBadGateway, types.ErrOptionWithSkipRetry())
}
if guardEnabled && c.Request.Context().Err() != nil {
return nil, types.NewErrorWithStatusCode(c.Request.Context().Err(), common.ResponseGuardErrorCode, http.StatusBadGateway, types.ErrOptionWithSkipRetry())
}
client, err := service.GetHttpClientWithProxySettings(info.ChannelSetting.Proxy, info.ChannelSetting)
if err != nil {
return nil, fmt.Errorf("new proxy http client failed: %w", err)
@@ -511,7 +548,7 @@ func doRequest(c *gin.Context, req *http.Request, info *common.RelayInfo) (*http
var stopPinger context.CancelFunc
var pingerDone <-chan struct{}
if info.IsStream {
if info.IsStream && !guardEnabled {
helper.SetEventStreamHeaders(c)
// 处理流式请求的 ping 保活
generalSettings := operation_setting.GetGeneralSetting()
@@ -529,14 +566,55 @@ func doRequest(c *gin.Context, req *http.Request, info *common.RelayInfo) (*http
}
}
info.UpstreamRequests++
var guard *common.ResponseGuardAttempt
var guardUsage *service.GuardUsageCollector
if guardEnabled {
var upstream context.Context
upstream, guard = info.ResponseGuard.Begin(c.Request.Context(), info.UpstreamRequests, info.ChannelId)
if info.DurableAttemptID != 0 {
guardUsage = service.NewGuardUsageCollector(c, info)
guard.OnEvent = guardUsage.Observe
guard.OnClose = func() { guardUsage.CheckpointFinal(guard.Reason) }
}
req = req.WithContext(upstream)
}
resp, err := relayClient.Do(req)
if err != nil {
if guard != nil {
guard.Close()
guard.Finish(false)
if c.Request.Context().Err() != nil {
guard.Reject("client_cancelled")
return nil, guard.APIError()
}
if errors.Is(err, context.Canceled) || errors.Is(err, context.DeadlineExceeded) {
guard.Reject("confirmation_timeout")
return nil, guard.APIError()
}
guard.Reason = "transport_error"
}
logger.LogError(c, "do request failed: "+err.Error())
return nil, types.NewError(err, types.ErrorCodeDoRequestFailed, types.ErrOptionWithHideErrMsg("upstream error: do request failed"))
}
if resp == nil {
if guard != nil {
guard.Close()
}
return nil, errors.New("resp is nil")
}
if guard != nil {
if err := guard.Confirm(resp, c.Request.Context()); err != nil {
if guardUsage != nil && guardUsage.FinishRejected(info, guard, resp) {
guard.Finish(false)
return nil, guard.APIError()
}
guard.Close()
_ = resp.Body.Close()
guard.Finish(false)
return nil, guard.APIError()
}
}
if common2.DebugEnabled {
policy := service.NormalizeHTTPTransportPolicy(info.ChannelSetting)
logger.LogDebug(c, fmt.Sprintf(
+12
View File
@@ -8,6 +8,7 @@ import (
relaycommon "github.com/QuantumNous/new-api/relay/common"
"github.com/gin-gonic/gin"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
)
@@ -205,3 +206,14 @@ func TestProcessHeaderOverride_PassHeadersTemplateSetsRuntimeHeaders(t *testing.
require.Equal(t, "sess-123", upstreamReq.Header.Get("Session_id"))
require.Empty(t, upstreamReq.Header.Get("X-Codex-Beta-Features"))
}
func TestToWebSocketURL(t *testing.T) {
for input, want := range map[string]string{
"https://api.openai.com/v1/responses": "wss://api.openai.com/v1/responses",
"http://127.0.0.1:3000/v1/responses": "ws://127.0.0.1:3000/v1/responses",
"wss://chatgpt.com/backend-api/codex/responses": "wss://chatgpt.com/backend-api/codex/responses",
"ws://127.0.0.1:3000/backend-api/codex/responses": "ws://127.0.0.1:3000/backend-api/codex/responses",
} {
assert.Equal(t, want, toWebSocketURL(input), input)
}
}
+13 -4
View File
@@ -24,11 +24,12 @@ func stopReasonClaude2OpenAI(reason string) string {
return relayconvert.StopReasonClaudeToOpenAI(reason)
}
func maybeMarkClaudeRefusal(c *gin.Context, stopReason string) {
func maybeMarkClaudeRefusal(c *gin.Context, info *relaycommon.RelayInfo, stopReason string) {
if c == nil {
return
}
if strings.EqualFold(stopReason, "refusal") {
info.PerformanceBusinessRejection = true
common.SetContextKey(c, constant.ContextKeyAdminRejectReason, "claude_stop_reason=refusal")
}
}
@@ -94,11 +95,17 @@ func HandleStreamResponseData(c *gin.Context, info *relaycommon.RelayInfo, claud
if claudeError := claudeResponse.GetClaudeError(); claudeError != nil && claudeError.Type != "" {
return types.WithClaudeError(*claudeError, http.StatusInternalServerError)
}
if claudeResponse.Type == "message_start" && claudeResponse.Message != nil {
info.ObserveResponseModel(claudeResponse.Message.Model)
}
if claudeResponse.StopReason != "" {
maybeMarkClaudeRefusal(c, claudeResponse.StopReason)
maybeMarkClaudeRefusal(c, info, claudeResponse.StopReason)
}
if claudeResponse.Delta != nil && claudeResponse.Delta.StopReason != nil {
maybeMarkClaudeRefusal(c, *claudeResponse.Delta.StopReason)
maybeMarkClaudeRefusal(c, info, *claudeResponse.Delta.StopReason)
}
if claudeResponse.Type == "message_stop" {
info.StreamStatus.MarkCompleted()
}
if info.RelayFormat == types.RelayFormatClaude {
FormatClaudeResponseInfo(&claudeResponse, nil, claudeInfo)
@@ -299,6 +306,7 @@ func ClaudeStreamHandler(c *gin.Context, resp *http.Response, info *relaycommon.
sr.Stop(err)
}
})
info.StreamStatus.RequireTerminal()
if err != nil {
return nil, err
}
@@ -316,7 +324,8 @@ func HandleClaudeResponseData(c *gin.Context, info *relaycommon.RelayInfo, claud
if claudeError := claudeResponse.GetClaudeError(); claudeError != nil && claudeError.Type != "" {
return types.WithClaudeError(*claudeError, http.StatusInternalServerError)
}
maybeMarkClaudeRefusal(c, claudeResponse.StopReason)
info.ObserveResponseModel(claudeResponse.Model)
maybeMarkClaudeRefusal(c, info, claudeResponse.StopReason)
if claudeInfo.Usage == nil {
claudeInfo.Usage = &dto.Usage{}
}
+6 -2
View File
@@ -107,12 +107,16 @@ func ClaudeResponsesStreamHandler(c *gin.Context, resp *http.Response, info *rel
}
if claudeResponse.StopReason != "" {
maybeMarkClaudeRefusal(c, claudeResponse.StopReason)
maybeMarkClaudeRefusal(c, info, claudeResponse.StopReason)
}
if claudeResponse.Delta != nil && claudeResponse.Delta.StopReason != nil {
maybeMarkClaudeRefusal(c, *claudeResponse.Delta.StopReason)
maybeMarkClaudeRefusal(c, info, *claudeResponse.Delta.StopReason)
}
if claudeResponse.Type == "message_stop" {
info.StreamStatus.MarkCompleted()
}
if claudeResponse.Type == "message_start" && claudeResponse.Message != nil {
info.ObserveResponseModel(claudeResponse.Message.Model)
info.UpstreamModelName = claudeResponse.Message.Model
}
FormatClaudeResponseInfo(&claudeResponse, nil, claudeInfo)
@@ -15,6 +15,7 @@ import (
"github.com/QuantumNous/new-api/service"
"github.com/gin-gonic/gin"
"github.com/tidwall/gjson"
)
func GeminiTextGenerationHandler(c *gin.Context, info *relaycommon.RelayInfo, resp *http.Response) (*dto.Usage, *types.NewAPIError) {
@@ -34,9 +35,11 @@ func GeminiTextGenerationHandler(c *gin.Context, info *relaycommon.RelayInfo, re
if err != nil {
return nil, types.NewOpenAIError(err, types.ErrorCodeBadResponseBody, http.StatusInternalServerError)
}
info.ObserveResponseModel(gjson.GetBytes(responseBody, "modelVersion").Str)
countGeminiBillableFunctionCalls(info, &geminiResponse)
if len(geminiResponse.Candidates) == 0 && geminiResponse.PromptFeedback != nil && geminiResponse.PromptFeedback.BlockReason != nil {
info.PerformanceBusinessRejection = true
common.SetContextKey(c, constant.ContextKeyAdminRejectReason, fmt.Sprintf("gemini_block_reason=%s", *geminiResponse.PromptFeedback.BlockReason))
}
+16
View File
@@ -20,6 +20,7 @@ import (
"github.com/QuantumNous/new-api/relaykit/types"
"github.com/QuantumNous/new-api/service"
"github.com/gin-gonic/gin"
"github.com/tidwall/gjson"
)
func buildUsageFromGeminiMetadata(metadata *dto.GeminiUsageMetadata, fallbackPromptTokens int) dto.Usage {
@@ -181,8 +182,20 @@ func geminiStreamHandler(c *gin.Context, info *relaycommon.RelayInfo, resp *http
}
if len(geminiResponse.Candidates) == 0 && geminiResponse.PromptFeedback != nil && geminiResponse.PromptFeedback.BlockReason != nil {
info.PerformanceBusinessRejection = true
common.SetContextKey(c, constant.ContextKeyAdminRejectReason, fmt.Sprintf("gemini_block_reason=%s", *geminiResponse.PromptFeedback.BlockReason))
}
info.ObserveResponseModel(gjson.Get(data, "modelVersion").Str)
for _, candidate := range geminiResponse.Candidates {
if candidate.FinishReason == nil || *candidate.FinishReason == "" || *candidate.FinishReason == "FINISH_REASON_UNSPECIFIED" {
continue
}
switch *candidate.FinishReason {
case "SAFETY", "RECITATION", "BLOCKLIST", "PROHIBITED_CONTENT", "SPII", "IMAGE_SAFETY", "IMAGE_PROHIBITED_CONTENT", "IMAGE_RECITATION":
info.PerformanceBusinessRejection = true
}
info.StreamStatus.MarkCompleted()
}
markGeminiGoogleSearchCall(c, &geminiResponse)
countGeminiBillableFunctionCalls(info, &geminiResponse)
@@ -216,6 +229,7 @@ func geminiStreamHandler(c *gin.Context, info *relaycommon.RelayInfo, resp *http
sr.Stop(streamErr)
}
})
info.StreamStatus.RequireTerminal()
if !hasBillableUsageMetadata {
if info.ReceivedResponseCount > 0 {
@@ -370,6 +384,7 @@ func GeminiChatHandler(c *gin.Context, info *relaycommon.RelayInfo, resp *http.R
if err != nil {
return nil, types.NewOpenAIError(err, types.ErrorCodeBadResponseBody, http.StatusInternalServerError)
}
info.ObserveResponseModel(gjson.GetBytes(responseBody, "modelVersion").Str)
markGeminiGoogleSearchCall(c, &geminiResponse)
countGeminiBillableFunctionCalls(info, &geminiResponse)
if len(geminiResponse.Candidates) == 0 {
@@ -377,6 +392,7 @@ func GeminiChatHandler(c *gin.Context, info *relaycommon.RelayInfo, resp *http.R
var newAPIError *types.NewAPIError
if geminiResponse.PromptFeedback != nil && geminiResponse.PromptFeedback.BlockReason != nil {
info.PerformanceBusinessRejection = true
common.SetContextKey(c, constant.ContextKeyAdminRejectReason, fmt.Sprintf("gemini_block_reason=%s", *geminiResponse.PromptFeedback.BlockReason))
newAPIError = types.NewOpenAIError(
errors.New("request blocked by Gemini API: "+*geminiResponse.PromptFeedback.BlockReason),
+3
View File
@@ -16,6 +16,7 @@ import (
"github.com/QuantumNous/new-api/relaykit/types"
"github.com/QuantumNous/new-api/service"
"github.com/gin-gonic/gin"
"github.com/tidwall/gjson"
)
func GeminiResponsesHandler(c *gin.Context, info *relaycommon.RelayInfo, resp *http.Response) (*dto.Usage, *types.NewAPIError) {
@@ -31,11 +32,13 @@ func GeminiResponsesHandler(c *gin.Context, info *relaycommon.RelayInfo, resp *h
if err := common.Unmarshal(responseBody, &geminiResponse); err != nil {
return nil, types.NewOpenAIError(err, types.ErrorCodeBadResponseBody, http.StatusInternalServerError)
}
info.ObserveResponseModel(gjson.GetBytes(responseBody, "modelVersion").Str)
markGeminiGoogleSearchCall(c, &geminiResponse)
countGeminiBillableFunctionCalls(info, &geminiResponse)
if len(geminiResponse.Candidates) == 0 {
usage := buildUsageFromGeminiResponse(c, info, &geminiResponse)
if geminiResponse.PromptFeedback != nil && geminiResponse.PromptFeedback.BlockReason != nil {
info.PerformanceBusinessRejection = true
common.SetContextKey(c, constant.ContextKeyAdminRejectReason, fmt.Sprintf("gemini_block_reason=%s", *geminiResponse.PromptFeedback.BlockReason))
return &usage, types.NewOpenAIError(
errors.New("request blocked by Gemini API: "+*geminiResponse.PromptFeedback.BlockReason),
+14
View File
@@ -38,6 +38,9 @@ func (a *Adaptor) ConvertImageRequest(c *gin.Context, info *relaycommon.RelayInf
}
func (a *Adaptor) Init(info *relaycommon.RelayInfo) {
// OpenAI-compatible chat responses are handled by the OpenAI adaptor, which
// relies on the thinking-to-content state initialized here.
(&openai.Adaptor{}).Init(info)
}
func (a *Adaptor) GetRequestURL(info *relaycommon.RelayInfo) (string, error) {
@@ -55,6 +58,9 @@ func (a *Adaptor) GetRequestURL(info *relaycommon.RelayInfo) (string, error) {
case relayconstant.RelayModeCompletions:
return fmt.Sprintf("%s/api/generate", info.ChannelBaseUrl), nil
default:
if info.ChannelOtherSettings.OllamaOpenAIChat {
return fmt.Sprintf("%s/v1/chat/completions", info.ChannelBaseUrl), nil
}
return fmt.Sprintf("%s/api/chat", info.ChannelBaseUrl), nil
}
}
@@ -83,6 +89,9 @@ func (a *Adaptor) ConvertOpenAIRequest(c *gin.Context, info *relaycommon.RelayIn
case relayconstant.RelayModeCompletions:
return openAIToGenerate(c, request)
default:
if info.ChannelOtherSettings.OllamaOpenAIChat {
return (&openai.Adaptor{}).ConvertOpenAIRequest(c, info, request)
}
return openAIChatToOllamaChat(c, request)
}
}
@@ -117,6 +126,11 @@ func (a *Adaptor) DoResponse(c *gin.Context, resp *http.Response, info *relaycom
adaptor := openai.Adaptor{}
return adaptor.DoResponse(c, resp, info)
default:
// /api/generate always returns native NDJSON; chat follows the channel setting.
if info.RelayMode != relayconstant.RelayModeCompletions && info.ChannelOtherSettings.OllamaOpenAIChat {
adaptor := openai.Adaptor{}
return adaptor.DoResponse(c, resp, info)
}
if info.IsStream {
return ollamaStreamHandler(c, info, resp)
}
+51 -37
View File
@@ -79,6 +79,51 @@ func ollamaToolCallsToOpenAI(toolCalls []OllamaToolCall, startIndex int, include
return result, startIndex
}
func buildOllamaStreamDelta(chunk *ollamaChatStreamChunk, responseID string, created int64, model string, toolCallIndex *int) (dto.ChatCompletionsStreamResponse, bool) {
delta := dto.ChatCompletionsStreamResponse{
Id: responseID,
Object: "chat.completion.chunk",
Created: created,
Model: model,
Choices: []dto.ChatCompletionsStreamResponseChoice{{
Index: 0,
Delta: dto.ChatCompletionsStreamResponseChoiceDelta{Role: "assistant"},
}},
}
var content string
if chunk.Message != nil {
content = chunk.Message.Content
} else {
content = chunk.Response
}
if content != "" {
delta.Choices[0].Delta.SetContentString(content)
}
hasPayload := content != ""
if chunk.Message != nil && len(chunk.Message.Thinking) > 0 {
raw := strings.TrimSpace(string(chunk.Message.Thinking))
if raw != "" && raw != "null" {
// Unmarshal the JSON string to get the actual content without quotes
var thinkingContent string
if err := common.Unmarshal(chunk.Message.Thinking, &thinkingContent); err == nil {
delta.Choices[0].Delta.SetReasoningContent(thinkingContent)
} else {
// Fallback to raw string if it's not a JSON string
delta.Choices[0].Delta.SetReasoningContent(raw)
}
hasPayload = true
}
}
if chunk.Message != nil && len(chunk.Message.ToolCalls) > 0 {
delta.Choices[0].Delta.ToolCalls, *toolCallIndex = ollamaToolCallsToOpenAI(chunk.Message.ToolCalls, *toolCallIndex, true)
hasPayload = true
}
return delta, hasPayload
}
func toUnix(ts string) int64 {
if ts == "" {
return time.Now().Unix()
@@ -130,49 +175,18 @@ func ollamaStreamHandler(c *gin.Context, info *relaycommon.RelayInfo, resp *http
created = toUnix(chunk.CreatedAt)
if !chunk.Done {
// delta content
var content string
if chunk.Message != nil {
content = chunk.Message.Content
} else {
content = chunk.Response
}
delta := dto.ChatCompletionsStreamResponse{
Id: responseId,
Object: "chat.completion.chunk",
Created: created,
Model: model,
Choices: []dto.ChatCompletionsStreamResponseChoice{{
Index: 0,
Delta: dto.ChatCompletionsStreamResponseChoiceDelta{Role: "assistant"},
}},
}
if content != "" {
delta.Choices[0].Delta.SetContentString(content)
}
if chunk.Message != nil && len(chunk.Message.Thinking) > 0 {
raw := strings.TrimSpace(string(chunk.Message.Thinking))
if raw != "" && raw != "null" {
// Unmarshal the JSON string to get the actual content without quotes
var thinkingContent string
if err := common.Unmarshal(chunk.Message.Thinking, &thinkingContent); err == nil {
delta.Choices[0].Delta.SetReasoningContent(thinkingContent)
} else {
// Fallback to raw string if it's not a JSON string
delta.Choices[0].Delta.SetReasoningContent(raw)
}
}
}
// tool calls
if chunk.Message != nil && len(chunk.Message.ToolCalls) > 0 {
delta.Choices[0].Delta.ToolCalls, toolCallIndex = ollamaToolCallsToOpenAI(chunk.Message.ToolCalls, toolCallIndex, true)
}
delta, _ := buildOllamaStreamDelta(&chunk, responseId, created, model, &toolCallIndex)
if data, err := common.Marshal(delta); err == nil {
_ = helper.StringData(c, string(data))
}
continue
}
// done frame
if delta, hasPayload := buildOllamaStreamDelta(&chunk, responseId, created, model, &toolCallIndex); hasPayload {
if data, err := common.Marshal(delta); err == nil {
_ = helper.StringData(c, string(data))
}
}
// finalize once and break loop
usage.PromptTokens = chunk.PromptEvalCount
usage.CompletionTokens = chunk.EvalCount
+145
View File
@@ -97,3 +97,148 @@ func TestOllamaChatHandlerNonStreamToolCalls(t *testing.T) {
})
}
}
func TestBuildOllamaStreamDelta(t *testing.T) {
tests := []struct {
name string
raw string
startIndex int
wantContent string
wantReasoning string
wantToolCalls []struct {
id string
name string
index int
arguments string
}
wantNextIndex int
wantPayload bool
}{
{
name: "chat content",
raw: `{"message":{"content":"hello"}}`,
wantContent: "hello",
wantPayload: true,
},
{
name: "generate content",
raw: `{"response":"hello"}`,
wantContent: "hello",
wantPayload: true,
},
{
name: "thinking json string",
raw: `{"message":{"thinking":"consider this"}}`,
wantReasoning: "consider this",
wantPayload: true,
},
{
name: "thinking raw fallback",
raw: `{"message":{"thinking":{"step":"consider this"}}}`,
wantReasoning: `{"step":"consider this"}`,
wantPayload: true,
},
{
name: "multiple tool calls",
raw: `{"message":{"tool_calls":[{"id":"call_upstream","function":{"name":"get_weather","arguments":{"city":"Paris"}}},{"function":{"name":"get_time","arguments":{"timezone":"UTC"}}}]}}`,
startIndex: 2,
wantToolCalls: []struct {
id string
name string
index int
arguments string
}{
{id: "call_upstream", name: "get_weather", index: 2, arguments: `{"city":"Paris"}`},
{id: "call_3", name: "get_time", index: 3, arguments: `{"timezone":"UTC"}`},
},
wantNextIndex: 4,
wantPayload: true,
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
var chunk ollamaChatStreamChunk
require.NoError(t, common.Unmarshal([]byte(tt.raw), &chunk))
toolCallIndex := tt.startIndex
delta, hasPayload := buildOllamaStreamDelta(&chunk, "response-id", 123, "model", &toolCallIndex)
assert.Equal(t, tt.wantPayload, hasPayload)
assert.Equal(t, tt.wantContent, delta.Choices[0].Delta.GetContentString())
assert.Equal(t, tt.wantReasoning, delta.Choices[0].Delta.GetReasoningContent())
assert.Equal(t, tt.wantNextIndex, toolCallIndex)
if tt.wantToolCalls == nil {
assert.Empty(t, delta.Choices[0].Delta.ToolCalls)
return
}
require.Len(t, delta.Choices[0].Delta.ToolCalls, len(tt.wantToolCalls))
for i, want := range tt.wantToolCalls {
got := delta.Choices[0].Delta.ToolCalls[i]
assert.Equal(t, want.id, got.ID)
assert.Equal(t, "function", got.Type)
assert.Equal(t, want.name, got.Function.Name)
assert.Equal(t, want.arguments, got.Function.Arguments)
require.NotNil(t, got.Index)
assert.Equal(t, want.index, *got.Index)
}
})
}
}
func TestOllamaStreamHandlerDoneToolCalls(t *testing.T) {
gin.SetMode(gin.TestMode)
w := httptest.NewRecorder()
c, _ := gin.CreateTestContext(w)
resp := &http.Response{
StatusCode: http.StatusOK,
Header: make(http.Header),
Body: io.NopCloser(strings.NewReader(`{"model":"qwen3-coder","created_at":"2026-05-27T12:00:00Z","message":{"role":"assistant","content":"","tool_calls":[{"function":{"name":"get_weather","arguments":{"city":"北京"}}},{"function":{"name":"get_time","arguments":{"timezone":"Asia/Shanghai"}}}]},"done":true,"done_reason":"stop","prompt_eval_count":5,"eval_count":7}`)),
}
usage, apiErr := ollamaStreamHandler(c, &relaycommon.RelayInfo{
ChannelMeta: &relaycommon.ChannelMeta{UpstreamModelName: "fallback-model"},
}, resp)
require.Nil(t, apiErr)
require.NotNil(t, usage)
assert.Equal(t, 12, usage.TotalTokens)
var chunks []dto.ChatCompletionsStreamResponse
for _, line := range strings.Split(w.Body.String(), "\n") {
data, ok := strings.CutPrefix(line, "data: ")
if !ok || data == "[DONE]" {
continue
}
var chunk dto.ChatCompletionsStreamResponse
require.NoError(t, common.Unmarshal([]byte(data), &chunk))
chunks = append(chunks, chunk)
}
var toolChunkIndex, stopChunkIndex int = -1, -1
for i := range chunks {
if len(chunks[i].Choices) == 0 {
continue
}
if len(chunks[i].Choices[0].Delta.ToolCalls) > 0 {
toolChunkIndex = i
}
if chunks[i].Choices[0].FinishReason != nil {
stopChunkIndex = i
}
}
require.NotEqual(t, -1, toolChunkIndex)
require.NotEqual(t, -1, stopChunkIndex)
assert.Less(t, toolChunkIndex, stopChunkIndex)
toolCalls := chunks[toolChunkIndex].Choices[0].Delta.ToolCalls
require.Len(t, toolCalls, 2)
assert.Equal(t, "call_0", toolCalls[0].ID)
assert.Equal(t, "call_1", toolCalls[1].ID)
require.NotNil(t, toolCalls[0].Index)
require.NotNil(t, toolCalls[1].Index)
assert.Equal(t, 0, *toolCalls[0].Index)
assert.Equal(t, 1, *toolCalls[1].Index)
assert.Equal(t, constant.FinishReasonToolCalls, *chunks[stopChunkIndex].Choices[0].FinishReason)
assert.Contains(t, w.Body.String(), "data: [DONE]")
}
+36 -8
View File
@@ -51,6 +51,11 @@ func (a *Adaptor) ConvertGeminiRequest(c *gin.Context, info *relaycommon.RelayIn
if !ok {
return nil, fmt.Errorf("expected OpenAI chat completions request, got %T", result.Value)
}
if info.SupportStreamOptions && info.IsStream {
openaiRequest.StreamOptions = &dto.StreamOptions{
IncludeUsage: true,
}
}
return a.ConvertOpenAIRequest(c, info, openaiRequest)
}
@@ -262,23 +267,26 @@ func (a *Adaptor) ConvertOpenAIRequest(c *gin.Context, info *relaycommon.RelayIn
info.SetReasoningEffort(request.ReasoningEffort)
}
if info.ChannelType == constant.ChannelTypeOpenRouter {
initialIntent, err := kitreasoning.FromOpenAIChat(request)
initialIntent, diagnostics, err := kitreasoning.FromOpenAIChat(request)
if err != nil {
return nil, kitreasoning.AsClientError(err)
}
recordReasoningDiagnostics(c, info, diagnostics)
if request.THINKING != nil && strings.HasPrefix(info.UpstreamModelName, "anthropic") {
var thinking dto.Thinking
if err := common.Unmarshal(request.THINKING, &thinking); err != nil {
return nil, fmt.Errorf("error Unmarshal thinking: %w", err)
}
legacyIntent, err := kitreasoning.FromClaude(&dto.ClaudeRequest{Thinking: &thinking})
legacyIntent, diagnostics, err := kitreasoning.FromClaude(&dto.ClaudeRequest{Thinking: &thinking})
if err != nil {
return nil, kitreasoning.AsClientError(err)
}
initialIntent, err = kitreasoning.MergeExplicit(initialIntent, legacyIntent, request.Model)
recordReasoningDiagnostics(c, info, diagnostics)
initialIntent, diagnostics, err = kitreasoning.MergeExplicit(initialIntent, legacyIntent, request.Model)
if err != nil {
return nil, kitreasoning.AsClientError(err)
}
recordReasoningDiagnostics(c, info, diagnostics)
request.THINKING = nil
}
if len(request.Usage) == 0 {
@@ -298,7 +306,9 @@ func (a *Adaptor) ConvertOpenAIRequest(c *gin.Context, info *relaycommon.RelayIn
if effort == kitreasoning.EffortNone {
mode = kitreasoning.ModeDisabled
}
initialIntent, err = kitreasoning.MergeExplicitAndSuffix(initialIntent, kitreasoning.Intent{Mode: mode, Effort: effort, Source: kitreasoning.SourceSuffix}, modelName)
var diagnostics []types.ConversionDiagnostic
initialIntent, diagnostics, err = kitreasoning.MergeExplicitAndSuffix(initialIntent, kitreasoning.Intent{Mode: mode, Effort: effort, Source: kitreasoning.SourceSuffix}, modelName)
recordReasoningDiagnostics(c, info, diagnostics)
return err
}
if !preserveSuffix {
@@ -362,10 +372,11 @@ func (a *Adaptor) ConvertOpenAIRequest(c *gin.Context, info *relaycommon.RelayIn
if preserveSuffix {
effort = ""
}
currentIntent, err := kitreasoning.FromOpenAIChat(request)
currentIntent, diagnostics, err := kitreasoning.FromOpenAIChat(request)
if err != nil {
return nil, kitreasoning.AsClientError(err)
}
recordReasoningDiagnostics(c, info, diagnostics)
mergeSuffix := func(modelName, rawEffort string) error {
if rawEffort == "" {
return nil
@@ -378,7 +389,9 @@ func (a *Adaptor) ConvertOpenAIRequest(c *gin.Context, info *relaycommon.RelayIn
if suffixEffort == kitreasoning.EffortNone {
mode = kitreasoning.ModeDisabled
}
currentIntent, err = kitreasoning.MergeExplicitAndSuffix(currentIntent, kitreasoning.Intent{Mode: mode, Effort: suffixEffort, Source: kitreasoning.SourceSuffix}, modelName)
var diagnostics []types.ConversionDiagnostic
currentIntent, diagnostics, err = kitreasoning.MergeExplicitAndSuffix(currentIntent, kitreasoning.Intent{Mode: mode, Effort: suffixEffort, Source: kitreasoning.SourceSuffix}, modelName)
recordReasoningDiagnostics(c, info, diagnostics)
return err
}
if err := mergeSuffix(info.UpstreamModelName, effort); err != nil {
@@ -684,10 +697,11 @@ func (a *Adaptor) ConvertOpenAIResponsesRequest(c *gin.Context, info *relaycommo
}
return request, nil
}
currentIntent, err := kitreasoning.FromOpenAIResponses(&request)
currentIntent, diagnostics, err := kitreasoning.FromOpenAIResponses(&request)
if err != nil {
return nil, kitreasoning.AsClientError(err)
}
recordReasoningDiagnostics(c, info, diagnostics)
mergeSuffix := func(modelName, rawEffort string) error {
if rawEffort == "" {
return nil
@@ -700,7 +714,9 @@ func (a *Adaptor) ConvertOpenAIResponsesRequest(c *gin.Context, info *relaycommo
if suffixEffort == kitreasoning.EffortNone {
mode = kitreasoning.ModeDisabled
}
currentIntent, err = kitreasoning.MergeExplicitAndSuffix(currentIntent, kitreasoning.Intent{Mode: mode, Effort: suffixEffort, Source: kitreasoning.SourceSuffix}, modelName)
var diagnostics []types.ConversionDiagnostic
currentIntent, diagnostics, err = kitreasoning.MergeExplicitAndSuffix(currentIntent, kitreasoning.Intent{Mode: mode, Effort: suffixEffort, Source: kitreasoning.SourceSuffix}, modelName)
recordReasoningDiagnostics(c, info, diagnostics)
return err
}
if err := mergeSuffix(request.Model, effort); err != nil {
@@ -811,3 +827,15 @@ func (a *Adaptor) GetChannelName() string {
return ChannelName
}
}
// recordReasoningDiagnostics attaches best-effort reasoning resolutions made
// while merging request fields with model-name effort tails to the request log.
func recordReasoningDiagnostics(c *gin.Context, info *relaycommon.RelayInfo, diagnostics []types.ConversionDiagnostic) {
if info == nil || len(diagnostics) == 0 {
return
}
for i := range diagnostics {
diagnostics[i].From = info.RelayFormat
}
info.RecordConversionDiagnostics(c, diagnostics)
}
+11 -1
View File
@@ -41,6 +41,7 @@ func OaiResponsesToChatHandler(c *gin.Context, info *relaycommon.RelayInfo, resp
return nil, types.WithOpenAIError(*oaiError, resp.StatusCode)
}
info.ObserveResponseModel(responsesResp.Model)
responseValue, usage, err := convertResponsesResponseForClient(c, info, &responsesResp)
if err != nil {
return nil, types.NewOpenAIError(err, types.ErrorCodeBadResponseBody, http.StatusInternalServerError)
@@ -60,6 +61,8 @@ func OaiResponsesToChatBufferedStreamHandler(c *gin.Context, info *relaycommon.R
}
defer service.CloseResponseBodyGracefully(resp)
info.StreamStatus = relaycommon.NewStreamStatus()
info.StreamStatus.RequireTerminal()
accumulator := relayconvert.NewResponsesBufferedAccumulator()
var finalResponse *dto.OpenAIResponsesResponse
var streamErr *types.NewAPIError
@@ -86,6 +89,10 @@ func OaiResponsesToChatBufferedStreamHandler(c *gin.Context, info *relaycommon.R
streamErr = types.NewOpenAIError(err, types.ErrorCodeBadResponseBody, http.StatusInternalServerError)
break
}
if streamResp.Response != nil {
info.ObserveResponseModel(streamResp.Response.Model)
}
service.ObserveResponsesOutcome(info, &streamResp)
accumulator.ProcessEvent(&streamResp)
switch streamResp.Type {
case "response.completed", "response.done", "response.incomplete":
@@ -120,7 +127,7 @@ func OaiResponsesToChatBufferedStreamHandler(c *gin.Context, info *relaycommon.R
if finalResponse == nil {
finalResponse = &dto.OpenAIResponsesResponse{
ID: helper.GetResponseID(c),
CreatedAt: int(time.Now().Unix()),
CreatedAt: dto.IntValue(time.Now().Unix()),
Model: info.UpstreamModelName,
Status: []byte(`"completed"`),
}
@@ -257,6 +264,9 @@ func OaiResponsesToChatStreamHandler(c *gin.Context, info *relaycommon.RelayInfo
return
}
if streamResp.Response != nil {
info.ObserveResponseModel(streamResp.Response.Model)
}
if streamResp.Type == "response.error" || streamResp.Type == "response.failed" {
if streamResp.Response != nil {
if oaiErr := streamResp.Response.GetOpenAIError(); oaiErr != nil && oaiErr.Type != "" {
+3 -2
View File
@@ -137,13 +137,14 @@ func ProcessStreamResponse(streamResponse dto.ChatCompletionsStreamResponse, res
return nil
}
func processTokenData(relayMode int, data string, responseTextBuilder *strings.Builder, toolCount *int) error {
switch relayMode {
func processTokenData(info *relaycommon.RelayInfo, data string, responseTextBuilder *strings.Builder, toolCount *int) error {
switch info.RelayMode {
case relayconstant.RelayModeChatCompletions:
var streamResponse dto.ChatCompletionsStreamResponse
if err := common.UnmarshalJsonStr(data, &streamResponse); err != nil {
return err
}
info.ObserveResponseModel(streamResponse.Model)
return ProcessStreamResponse(streamResponse, responseTextBuilder, toolCount)
case relayconstant.RelayModeCompletions:
var streamResponse dto.CompletionsStreamResponse
+184
View File
@@ -2,6 +2,7 @@ package openai
import (
"context"
"fmt"
"io"
"net/http"
"net/http/httptest"
@@ -10,9 +11,12 @@ import (
"testing"
"github.com/QuantumNous/new-api/constant"
"github.com/QuantumNous/new-api/pkg/billingexpr"
relaycommon "github.com/QuantumNous/new-api/relay/common"
relayconstant "github.com/QuantumNous/new-api/relay/constant"
"github.com/QuantumNous/new-api/relaykit/dto"
"github.com/gin-gonic/gin"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
)
@@ -35,6 +39,112 @@ func newImageTestContext(t *testing.T, body, contentType string, isStream bool)
return c, recorder, resp, info
}
func TestImageExpressionUsesCompletedCountAndProtectsAbortedStreams(t *testing.T) {
oldTimeout := constant.StreamingTimeout
constant.StreamingTimeout = 30
t.Cleanup(func() { constant.StreamingTimeout = oldTimeout })
for _, tc := range []struct {
name, body string
stream, abort bool
requested, want int
}{
{"JSON uses actual count", `{"data":[{"b64_json":"first"},{"b64_json":"second"}]}`, false, false, 3, 2},
{"JSON wrapped as SSE uses actual count", `{"data":[{"b64_json":"first"}]}`, true, false, 3, 1},
{"empty response retains request", `{"data":[]}`, false, false, 3, 3},
{"JSON object data counts one image", `{"data":{"url":"https://example.com/a.png","b64_json":"first"}}`, false, false, 3, 1},
{"JSON wrapped as SSE counts object data once", `{"data":{"b64_json":"first"}}`, true, false, 3, 1},
{"completed stream refunds missing images", "data: {\"type\":\"image_generation.completed\",\"b64_json\":\"first\"}\n\ndata: [DONE]\n\n", true, false, 3, 1},
{"client abort cannot reduce count", "data: {\"type\":\"image_generation.completed\",\"b64_json\":\"first\"}\n\n", true, true, 3, 3},
} {
t.Run(tc.name, func(t *testing.T) {
contentType := "application/json"
if strings.HasPrefix(tc.body, "data:") {
contentType = "text/event-stream"
}
c, _, resp, info := newImageTestContext(t, tc.body, contentType, tc.stream)
if tc.abort {
c, _, resp, info = newDisconnectingImageStream(t, tc.body, "first")
}
info.TieredBillingSnapshot = &billingexpr.BillingSnapshot{EstimatedImageCount: &tc.requested}
if tc.stream {
_, err := OpenaiImageStreamHandler(c, info, resp)
require.Nil(t, err)
} else {
_, err := OpenaiImageHandler(c, info, resp)
require.Nil(t, err)
}
count := info.RequestedImageCount()
if info.BillingImageCount != nil {
count = *info.BillingImageCount
}
assert.Equal(t, tc.want, count)
assert.Empty(t, info.PriceData.OtherRatios(), "expression quantities must not add a legacy multiplier")
})
}
}
func TestImageCacheUsageAcrossResponseFormats(t *testing.T) {
previousTimeout := constant.StreamingTimeout
constant.StreamingTimeout = 30
t.Cleanup(func() { constant.StreamingTimeout = previousTimeout })
// Contract fixture for compatible upstreams. Live OpenAI cache response
// verification is a separate release check; this is not a captured response.
const usageJSON = `{"input_tokens":1000,"output_tokens":100,"total_tokens":1100,"input_tokens_details":{"text_tokens":400,"image_tokens":600,"cached_tokens":300,"cached_tokens_details":{"text_tokens":100,"image_tokens":200}}}`
for _, mode := range []string{"image_generation", "image_edit"} {
for _, stream := range []bool{false, true} {
t.Run(fmt.Sprintf("%s/stream=%t", mode, stream), func(t *testing.T) {
body := `{"data":[{"b64_json":"image"}],"usage":` + usageJSON + `}`
contentType := "application/json"
if stream {
body = "data: {\"type\":\"" + mode + ".completed\",\"usage\":" + usageJSON + "}\n\ndata: [DONE]\n\n"
contentType = "text/event-stream"
}
ctx, recorder, response, info := newImageTestContext(t, body, contentType, stream)
info.RelayMode = relayconstant.RelayModeImagesGenerations
if mode == "image_edit" {
info.RelayMode = relayconstant.RelayModeImagesEdits
ctx.Request.URL.Path = "/v1/images/edits"
}
result, apiErr := (&Adaptor{}).DoResponse(ctx, response, info)
require.Nil(t, apiErr)
usage := result.(*dto.Usage)
assert.Equal(t, 1000, usage.PromptTokens)
assert.Equal(t, 300, usage.PromptTokensDetails.CachedTokens)
require.NotNil(t, usage.PromptTokensDetails.CachedTokensDetails)
assert.Equal(t, 200, *usage.PromptTokensDetails.CachedTokensDetails.ImageTokens)
assert.Contains(t, recorder.Body.String(), `"cached_tokens_details":{"text_tokens":100,"image_tokens":200}`)
})
}
}
}
func TestNormalizeOpenAIUsageMapsOutputImageTokens(t *testing.T) {
previousTimeout := constant.StreamingTimeout
constant.StreamingTimeout = 30
t.Cleanup(func() { constant.StreamingTimeout = previousTimeout })
const usageJSON = `{"input_tokens":15,"output_tokens":1352,"total_tokens":1367,"input_tokens_details":{"text_tokens":15,"image_tokens":0},"output_tokens_details":{"image_tokens":1120,"text_tokens":232}}`
for _, stream := range []bool{false, true} {
t.Run(fmt.Sprintf("stream=%t", stream), func(t *testing.T) {
body := `{"data":[{"b64_json":"image"}],"usage":` + usageJSON + `}`
contentType := "application/json"
if stream {
body = "data: {\"type\":\"image_generation.completed\",\"usage\":" + usageJSON + "}\n\ndata: [DONE]\n\n"
contentType = "text/event-stream"
}
ctx, _, response, info := newImageTestContext(t, body, contentType, stream)
info.RelayMode = relayconstant.RelayModeImagesGenerations
result, apiErr := (&Adaptor{}).DoResponse(ctx, response, info)
require.Nil(t, apiErr)
usage := result.(*dto.Usage)
assert.Equal(t, 15, usage.PromptTokens)
assert.Equal(t, 1352, usage.CompletionTokens)
assert.Equal(t, 1120, usage.CompletionTokenDetails.ImageTokens)
assert.Equal(t, 232, usage.CompletionTokenDetails.TextTokens)
})
}
}
func TestOpenaiImageDoResponseUsesInfoIsStream(t *testing.T) {
oldMode := gin.Mode()
gin.SetMode(gin.TestMode)
@@ -344,6 +454,24 @@ func TestOpenaiImageHandlerUsesPositiveActualCountForFixedPrice(t *testing.T) {
usePrice: false,
wantCount: 3,
},
{
name: "object data with url and b64_json counts one image",
body: `{"data":{"url":"https://example.com/a.png","b64_json":"` + longImage + `"}}`,
usePrice: true,
wantCount: 1,
},
{
name: "url and b64_json split across entries count one image",
body: `{"data":[{"url":"https://example.com/a.png"},{"b64_json":"` + longImage + `"}]}`,
usePrice: true,
wantCount: 1,
},
{
name: "entries without image payload keep requested count",
body: `{"data":[{"revised_prompt":"draw a cat"}]}`,
usePrice: true,
wantCount: 3,
},
}
for _, tt := range tests {
@@ -361,6 +489,62 @@ func TestOpenaiImageHandlerUsesPositiveActualCountForFixedPrice(t *testing.T) {
}
}
// TestOpenaiImageStreamHandlerWrapsNonStandardDataShapes covers the JSON-to-SSE
// fallback when upstream returns data as an object, splits one image across a
// url entry and a b64_json entry, or returns entries without any image payload.
// Forwarded events and the billed quantity must both follow the real images.
func TestOpenaiImageStreamHandlerWrapsNonStandardDataShapes(t *testing.T) {
oldMode := gin.Mode()
gin.SetMode(gin.TestMode)
t.Cleanup(func() { gin.SetMode(oldMode) })
for _, tc := range []struct {
name string
body string
wantEvents int
wantCount float64
wantBody []string
}{
{
name: "object data is forwarded as one completed event",
body: `{"data":{"url":"https://example.com/a.png","b64_json":"first"}}`,
wantEvents: 1,
wantCount: 1,
wantBody: []string{`"url":"https://example.com/a.png"`, `"b64_json":"first"`},
},
{
name: "split url and b64_json entries bill one image",
body: `{"data":[{"url":"https://example.com/a.png"},{"b64_json":"first"}]}`,
wantEvents: 2,
wantCount: 1,
wantBody: []string{`"url":"https://example.com/a.png"`, `"b64_json":"first"`},
},
{
name: "entries without image payload are dropped and keep requested count",
body: `{"data":[{"revised_prompt":"draw a cat"}]}`,
wantEvents: 0,
wantCount: 3,
},
} {
t.Run(tc.name, func(t *testing.T) {
c, recorder, resp, info := newImageTestContext(t, tc.body, "application/json", true)
info.PriceData.UsePrice = true
info.PriceData.AddOtherRatio("n", 3)
_, err := OpenaiImageStreamHandler(c, info, resp)
require.Nil(t, err)
out := recorder.Body.String()
assert.Equal(t, tc.wantEvents, strings.Count(out, "event: image_generation.completed"))
assert.Equal(t, tc.wantEvents, info.ReceivedResponseCount)
for _, want := range tc.wantBody {
assert.Contains(t, out, want)
}
assert.True(t, strings.HasSuffix(out, "data: [DONE]\n\n"))
assert.Equal(t, tc.wantCount, info.PriceData.OtherRatios()["n"])
})
}
}
// TestOpenaiImageHandlersReturnJSONError covers JSON error responses for both
// entry points: the non-streaming handler and the stream handler's non-SSE
// fallback. Neither must leak the error body to the client.
+15 -3
View File
@@ -134,14 +134,16 @@ func OaiStreamHandler(c *gin.Context, info *relaycommon.RelayInfo, resp *http.Re
}
lastStreamData = data
collectStreamFunctionCallNames(data, seenStreamToolCalls, &streamFunctionCallNames)
if err := processTokenData(info.RelayMode, data, &responseTextBuilder, &toolCount); err != nil {
observeStreamChoices(info, data, seenStreamToolCalls, &streamFunctionCallNames)
if err := processTokenData(info, data, &responseTextBuilder, &toolCount); err != nil {
logger.LogError(c, "error processing stream token data: "+err.Error())
sr.Error(err)
}
}
})
info.StreamStatus.RequireTerminal()
// 处理最后的响应
shouldSendLastResp := true
if err := handleLastResponse(lastStreamData, &responseId, &createAt, &systemFingerprint, &model, &usage,
@@ -194,12 +196,20 @@ func OaiStreamHandler(c *gin.Context, info *relaycommon.RelayInfo, resp *http.Re
return usage, nil
}
func collectStreamFunctionCallNames(data string, seen map[string]struct{}, names *[]string) {
// observeStreamChoices collects billable function call names and records the
// finish reason facts used by health sampling from one parsed chunk.
func observeStreamChoices(info *relaycommon.RelayInfo, data string, seen map[string]struct{}, names *[]string) {
var streamResponse dto.ChatCompletionsStreamResponse
if err := common.UnmarshalJsonStr(data, &streamResponse); err != nil {
return
}
for _, choice := range streamResponse.Choices {
if choice.FinishReason != nil && *choice.FinishReason != "" {
if *choice.FinishReason == constant.FinishReasonContentFilter {
info.PerformanceBusinessRejection = true
}
info.StreamStatus.MarkCompleted()
}
for i, tc := range choice.Delta.ToolCalls {
name := strings.TrimSpace(tc.Function.Name)
if name == "" {
@@ -272,8 +282,10 @@ func OpenaiHandler(c *gin.Context, info *relaycommon.RelayInfo, resp *http.Respo
return nil, types.WithOpenAIError(*oaiError, resp.StatusCode)
}
info.ObserveResponseModel(simpleResponse.Model)
for _, choice := range simpleResponse.Choices {
if choice.FinishReason == constant.FinishReasonContentFilter {
info.PerformanceBusinessRejection = true
common.SetContextKey(c, constant.ContextKeyAdminRejectReason, "openai_finish_reason=content_filter")
break
}
+66 -33
View File
@@ -5,7 +5,6 @@ import (
"fmt"
"io"
"net/http"
"strconv"
"strings"
"time"
@@ -22,13 +21,6 @@ import (
"github.com/tidwall/sjson"
)
func updateOpenAIImageCount(info *relaycommon.RelayInfo, count int64) {
if info == nil || !info.PriceData.UsePrice || count <= 0 || count > int64(dto.MaxImageN) {
return
}
info.PriceData.AddOtherRatio("n", float64(count))
}
// OpenaiImageHandler handles non-streaming OpenAI image responses
// (generations/edits), returning the parsed usage for billing.
func OpenaiImageHandler(c *gin.Context, info *relaycommon.RelayInfo, resp *http.Response) (*dto.Usage, *types.NewAPIError) {
@@ -49,7 +41,7 @@ func OpenaiImageHandler(c *gin.Context, info *relaycommon.RelayInfo, resp *http.
return nil, types.WithOpenAIError(*oaiError, resp.StatusCode)
}
updateOpenAIImageCount(info, gjson.GetBytes(responseBody, "data.#").Int())
info.UpdateImageCount(openaiImageResponseCount(responseBody))
// 写入新的 response body
service.IOCopyBytesGracefully(c, resp, responseBody)
@@ -59,14 +51,54 @@ func OpenaiImageHandler(c *gin.Context, info *relaycommon.RelayInfo, resp *http.
return &usageResp.Usage, nil
}
// openaiImageResponseCount counts billable images in an OpenAI-format image
// response body. An object-shaped data is one image when it carries url or
// b64_json. For arrays the count is the larger of the url-bearing and the
// b64_json-bearing entry counts: a standard response uses one response_format
// so this equals the entry count, an upstream that splits one image into a url
// entry and a b64_json entry bills once, and entries without any image payload
// bill nothing. A zero result leaves the requested quantity in place because
// UpdateImageCount ignores non-positive counts.
func openaiImageResponseCount(responseBody []byte) int64 {
data := gjson.GetBytes(responseBody, "data")
if data.IsObject() {
if openaiImageDataHasField(data, "url") || openaiImageDataHasField(data, "b64_json") {
return 1
}
return 0
}
if !data.IsArray() {
return 0
}
var urls, b64s int64
data.ForEach(func(_, item gjson.Result) bool {
if openaiImageDataHasField(item, "url") {
urls++
}
if openaiImageDataHasField(item, "b64_json") {
b64s++
}
return true
})
return max(urls, b64s)
}
// openaiImageDataHasField reports whether an image data entry carries a
// non-empty string value for field.
func openaiImageDataHasField(item gjson.Result, field string) bool {
value := item.Get(field)
return value.Type == gjson.String && value.Raw != `""`
}
// normalizeOpenAIUsage maps the OpenAI Images usage shape (input_tokens /
// output_tokens / input_tokens_details) onto the canonical prompt/completion
// fields. It is used only on the OpenAI image relay paths (generations/edits,
// streaming and non-streaming): the image API never returns prompt_tokens /
// completion_tokens, so the overwrite (=) semantics here are equivalent to the
// previous additive (+=) behavior while avoiding any future double-counting if
// both field sets are ever populated. Do not reuse this on chat/embedding paths
// without revisiting the overwrite semantics.
// output_tokens / input_tokens_details / output_tokens_details) onto the
// canonical prompt/completion fields. It is used only on the OpenAI image
// relay paths (generations/edits, streaming and non-streaming): the image
// API never returns prompt_tokens / completion_tokens, so the overwrite (=)
// semantics here are equivalent to the previous additive (+=) behavior while
// avoiding any future double-counting if both field sets are ever populated.
// Do not reuse this on chat/embedding paths without revisiting the overwrite
// semantics.
func normalizeOpenAIUsage(usage *dto.Usage) {
if usage == nil {
return
@@ -78,12 +110,10 @@ func normalizeOpenAIUsage(usage *dto.Usage) {
usage.CompletionTokens = usage.OutputTokens
}
if usage.InputTokensDetails != nil {
usage.PromptTokensDetails.CachedTokens = usage.InputTokensDetails.CachedTokens
usage.PromptTokensDetails.CachedCreationTokens = usage.InputTokensDetails.CachedCreationTokens
usage.PromptTokensDetails.CacheWriteTokens = usage.InputTokensDetails.CacheWriteTokens
usage.PromptTokensDetails.ImageTokens = usage.InputTokensDetails.ImageTokens
usage.PromptTokensDetails.TextTokens = usage.InputTokensDetails.TextTokens
usage.PromptTokensDetails.AudioTokens = usage.InputTokensDetails.AudioTokens
usage.PromptTokensDetails = usage.InputTokensDetails.Clone()
}
if usage.OutputTokensDetails != nil {
usage.CompletionTokenDetails = *usage.OutputTokensDetails
}
if usage.TotalTokens == 0 {
usage.TotalTokens = usage.PromptTokens + usage.CompletionTokens
@@ -156,12 +186,8 @@ func OpenaiImageStreamHandler(c *gin.Context, info *relaycommon.RelayInfo, resp
if info.StreamStatus != nil {
upstreamFinished := info.StreamStatus.EndReason == relaycommon.StreamEndReasonDone ||
info.StreamStatus.EndReason == relaycommon.StreamEndReasonEOF
requestedN := 1.0
if n, ok := info.PriceData.OtherRatios()["n"]; ok {
requestedN = n
}
if upstreamFinished || float64(completedImages) > requestedN {
updateOpenAIImageCount(info, completedImages)
if upstreamFinished || completedImages > int64(info.RequestedImageCount()) {
info.UpdateImageCount(completedImages)
}
}
return usage, nil
@@ -252,8 +278,7 @@ func openaiImageJSONAsStreamHandler(c *gin.Context, info *relaycommon.RelayInfo,
normalizeOpenAIUsage(&usageResp.Usage)
applyUsagePostProcessing(info, &usageResp.Usage, responseBody)
imageCount := gjson.GetBytes(responseBody, "data.#").Int()
updateOpenAIImageCount(info, imageCount)
info.UpdateImageCount(openaiImageResponseCount(responseBody))
helper.SetEventStreamHeaders(c)
c.Status(http.StatusOK)
@@ -275,8 +300,15 @@ func openaiImageJSONAsStreamHandler(c *gin.Context, info *relaycommon.RelayInfo,
}
}
for i := range imageCount {
image := gjson.GetBytes(responseBody, "data."+strconv.FormatInt(i, 10))
// gjson.Result.Array returns the element list for an array and a single
// element for an object-shaped data, keeping document-relative indexes so
// the zero-copy field forwarding below stays valid for both shapes.
// Entries without url or b64_json carry no image and are not forwarded.
emitted := 0
for _, image := range gjson.GetBytes(responseBody, "data").Array() {
if !openaiImageDataHasField(image, "url") && !openaiImageDataHasField(image, "b64_json") {
continue
}
payload := []byte(`{"type":"image_generation.completed"}`)
payload, err = sjson.SetBytes(payload, "created_at", created)
if err != nil {
@@ -311,6 +343,7 @@ func openaiImageJSONAsStreamHandler(c *gin.Context, info *relaycommon.RelayInfo,
}
return &usageResp.Usage, nil
}
emitted++
}
if err := writeOpenaiImageStreamDone(c); err != nil {
if info != nil && info.StreamStatus != nil {
@@ -319,7 +352,7 @@ func openaiImageJSONAsStreamHandler(c *gin.Context, info *relaycommon.RelayInfo,
return &usageResp.Usage, nil
}
if info != nil {
info.ReceivedResponseCount += int(imageCount)
info.ReceivedResponseCount += emitted
if info.StreamStatus == nil {
info.StreamStatus = relaycommon.NewStreamStatus()
}
+94 -72
View File
@@ -1,12 +1,15 @@
package openai
import (
"bufio"
"bytes"
"fmt"
"io"
"net/http"
"strings"
"github.com/QuantumNous/new-api/common"
"github.com/QuantumNous/new-api/constant"
"github.com/QuantumNous/new-api/logger"
relaycommon "github.com/QuantumNous/new-api/relay/common"
"github.com/QuantumNous/new-api/relay/helper"
@@ -16,10 +19,25 @@ import (
"github.com/QuantumNous/new-api/service"
"github.com/gin-gonic/gin"
"github.com/tidwall/gjson"
"github.com/tidwall/sjson"
)
func OaiResponsesHandler(c *gin.Context, info *relaycommon.RelayInfo, resp *http.Response) (*dto.Usage, *types.NewAPIError) {
defer service.CloseResponseBodyGracefully(resp)
if strings.HasPrefix(strings.ToLower(resp.Header.Get("Content-Type")), "text/event-stream") {
body, err := readBufferedResponsesSSE(resp.Body)
if err != nil {
return nil, types.NewOpenAIError(err, types.ErrorCodeBadResponseBody, http.StatusBadGateway)
}
converted := *resp
converted.Header = resp.Header.Clone()
converted.Header.Set("Content-Type", "application/json")
converted.Header.Del("Content-Length")
converted.Body = io.NopCloser(bytes.NewReader(body))
converted.ContentLength = int64(len(body))
return OaiResponsesHandler(c, info, &converted)
}
// read response body
var responsesResponse dto.OpenAIResponsesResponse
@@ -35,11 +53,15 @@ func OaiResponsesHandler(c *gin.Context, info *relaycommon.RelayInfo, resp *http
return nil, types.WithOpenAIError(*oaiError, resp.StatusCode)
}
info.ObserveResponseModel(responsesResponse.Model)
responseBody = rewriteSGLangResponsesCreatedAt(info, responseBody, "created_at", responsesResponse.CreatedAt)
// 写入新的 response body
service.IOCopyBytesGracefully(c, resp, responseBody)
// compute usage
usage := relayconvert.NormalizeResponsesUsage(responsesResponse.Usage)
usage := &dto.Usage{}
service.ApplyResponsesUsage(usage, responsesResponse.Usage)
// Count actual tool invocations from Output (not tool declarations).
for _, output := range responsesResponse.Output {
switch output.Type {
@@ -64,6 +86,58 @@ func OaiResponsesHandler(c *gin.Context, info *relaycommon.RelayInfo, resp *http
return usage, nil
}
// Some upstreams return SSE even for stream=false. Require a terminal response
// and reuse JSON delivery and settlement so upstream usage is charged once.
func readBufferedResponsesSSE(body io.Reader) ([]byte, error) {
scanner := helper.NewStreamScanner(body)
scanner.Split(bufio.ScanLines)
accumulator := relayconvert.NewResponsesBufferedAccumulator()
var data []string
for scanner.Scan() {
line := strings.TrimSuffix(scanner.Text(), "\r")
if payload, ok := strings.CutPrefix(line, "data:"); ok {
data = append(data, strings.TrimPrefix(payload, " "))
}
if line != "" || len(data) == 0 {
continue
}
payload := strings.Join(data, "\n")
data = nil
if payload == "[DONE]" {
break
}
var event dto.ResponsesStreamResponse
if err := common.UnmarshalJsonStr(payload, &event); err != nil {
return nil, err
}
accumulator.ProcessEvent(&event)
switch event.Type {
case "error", "response.error", "response.failed", "response.cancelled", "response.canceled":
return nil, fmt.Errorf("responses stream ended with %s", event.Type)
case "response.completed", "response.done", "response.incomplete":
if event.Response == nil {
return nil, fmt.Errorf("responses terminal event has no response")
}
result := []byte(gjson.Get(payload, "response").Raw)
if len(event.Response.Output) == 0 {
accumulator.SupplementResponseOutput(event.Response)
if len(event.Response.Output) != 0 {
output, err := common.Marshal(event.Response.Output)
if err != nil {
return nil, err
}
return sjson.SetRawBytes(result, "output", output)
}
}
return result, nil
}
}
if err := scanner.Err(); err != nil {
return nil, err
}
return nil, fmt.Errorf("responses stream ended without a terminal response")
}
func OaiResponsesStreamHandler(c *gin.Context, info *relaycommon.RelayInfo, resp *http.Response) (*dto.Usage, *types.NewAPIError) {
if resp == nil || resp.Body == nil {
logger.LogError(c, "invalid response or response body")
@@ -72,10 +146,7 @@ func OaiResponsesStreamHandler(c *gin.Context, info *relaycommon.RelayInfo, resp
defer service.CloseResponseBodyGracefully(resp)
var usage = &dto.Usage{}
var responseTextBuilder strings.Builder
imageCounter := &relaycommon.ImageGenerationCallCounter{}
imageCommitted := false
accumulator := service.NewResponsesUsageAccumulator(info)
helper.StreamScannerHandler(c, resp, info, func(data string, sr *helper.StreamResult) {
@@ -86,77 +157,28 @@ func OaiResponsesStreamHandler(c *gin.Context, info *relaycommon.RelayInfo, resp
sr.Error(err)
return
}
sendResponsesStreamData(c, streamResponse, data)
switch streamResponse.Type {
case "response.completed", "response.done":
if streamResponse.Response != nil {
if streamResponse.Response.Usage != nil {
incomingUsage := relayconvert.NormalizeResponsesUsage(streamResponse.Response.Usage)
usage = dto.MergeUsageNonZero(usage, incomingUsage)
}
if !imageCommitted {
if relaycommon.IsNonBillableResponsesStatus(streamResponse.Response.Status) {
imageCounter.Reset()
imageCounter.Commit(info)
imageCommitted = true
} else {
for i := range streamResponse.Response.Output {
idx := i
imageCounter.Observe(&streamResponse.Response.Output[i], &idx)
}
imageCounter.Commit(info)
imageCommitted = true
}
}
} else if !imageCommitted {
imageCounter.Commit(info)
imageCommitted = true
}
case "response.failed", "response.incomplete", "response.cancelled", "response.canceled":
if !imageCommitted {
imageCounter.Reset()
imageCounter.Commit(info)
imageCommitted = true
}
case "response.output_text.delta":
// 处理输出文本
responseTextBuilder.WriteString(streamResponse.Delta)
case dto.ResponsesOutputTypeItemDone:
if streamResponse.Item != nil {
switch streamResponse.Item.Type {
case dto.BuildInCallWebSearchCall:
info.CountBillableToolCall(dto.BuildInCallWebSearchCall, "")
case dto.BuildInCallFileSearchCall:
info.CountBillableToolCall(dto.BuildInCallFileSearchCall, "")
case dto.BuildInCallFunctionCall:
info.CountBillableToolCall(dto.BuildInCallFunctionCall, streamResponse.Item.Name)
case dto.ResponsesOutputTypeImageGenerationCall:
if !imageCommitted {
imageCounter.Observe(streamResponse.Item, streamResponse.OutputIndex)
}
}
}
if streamResponse.Response != nil {
data = string(rewriteSGLangResponsesCreatedAt(info, []byte(data), "response.created_at", streamResponse.Response.CreatedAt))
}
sendResponsesStreamData(c, streamResponse, data)
accumulator.Observe(&streamResponse)
})
if usage.CompletionTokens == 0 {
// 计算输出文本的 token 数量
tempStr := responseTextBuilder.String()
if len(tempStr) > 0 {
// 非正常结束,使用输出文本的 token 数量
completionTokens := service.CountTextToken(tempStr, info.UpstreamModelName)
usage.CompletionTokens = completionTokens
}
common.SetContextKey(c, constant.ContextKeyResponseStreamStatus, info.StreamStatus)
info.StreamStatus.RequireTerminal()
return accumulator.Finish(), nil
}
func rewriteSGLangResponsesCreatedAt(info *relaycommon.RelayInfo, payload []byte, path string, createdAt dto.IntValue) []byte {
if info.GetChannelType() != constant.ChannelTypeSGLang {
return payload
}
if usage.PromptTokens == 0 && usage.CompletionTokens != 0 {
usage.PromptTokens = info.GetEstimatePromptTokens()
if !gjson.GetBytes(payload, path).Exists() {
return payload
}
usage.TotalTokens = usage.PromptTokens + usage.CompletionTokens
if usage.BillingUsage != nil {
usage.BillingUsage = dto.CloneBillingUsageWithEstimatedCompletion(usage.BillingUsage, usage.CompletionTokens)
patched, err := sjson.SetBytes(payload, path, int(createdAt))
if err != nil {
return payload
}
return usage, nil
return patched
}
@@ -18,6 +18,44 @@ import (
"github.com/stretchr/testify/require"
)
func TestOaiResponsesHandlerBuffersUnexpectedSSEAndPreservesUsage(t *testing.T) {
terminal := `{"id":"resp_test","object":"response","model":"deepseek-v4.1-flash","status":"completed","extra_provider_field":true,"output":[{"type":"message","role":"assistant","content":[{"type":"output_text","text":"OK"}]}],"usage":{"input_tokens":86926,"output_tokens":163,"total_tokens":87089,"input_tokens_details":{"cached_tokens":86784}}}`
for _, tc := range []struct {
name string
body string
wantError bool
}{
{"completed", ": ping\n\nevent: response.completed\ndata: {\"type\":\"response.completed\",\"response\":" + terminal + "}\n\ndata: [DONE]\n\n", false},
{"incomplete", "data: {\"type\":\"response.incomplete\",\"response\":" + strings.Replace(terminal, `"completed"`, `"incomplete"`, 1) + "}\n\n", false},
{"truncated", "data: {\"type\":\"response.created\",\"response\":{\"model\":\"deepseek-v4.1-flash\"}}\n\ndata: [DONE]\n\n", true},
{"failed", "data: {\"type\":\"response.failed\",\"response\":{\"status\":\"failed\"}}\n\n", true},
{"invalid", "data: not-json\n\n", true},
} {
t.Run(tc.name, func(t *testing.T) {
w := httptest.NewRecorder()
c, _ := gin.CreateTestContext(w)
c.Request = httptest.NewRequest(http.MethodPost, "/v1/responses", nil)
info := &relaycommon.RelayInfo{OriginModelName: "deepseek-v4.1-flash", ChannelMeta: &relaycommon.ChannelMeta{}}
resp := &http.Response{StatusCode: 200, Header: http.Header{"Content-Type": {"text/event-stream; charset=utf-8"}}, Body: io.NopCloser(strings.NewReader(tc.body))}
usage, apiErr := OaiResponsesHandler(c, info, resp)
if tc.wantError {
require.NotNil(t, apiErr)
assert.Nil(t, usage)
assert.False(t, c.Writer.Written())
return
}
require.Nil(t, apiErr)
require.NotNil(t, usage)
assert.Equal(t, 86926, usage.PromptTokens)
assert.Equal(t, 163, usage.CompletionTokens)
assert.Equal(t, 86784, usage.PromptTokensDetails.CachedTokens)
assert.Equal(t, "application/json", w.Header().Get("Content-Type"))
assert.Contains(t, w.Body.String(), `"extra_provider_field":true`)
assert.Contains(t, w.Body.String(), `"text":"OK"`)
})
}
}
func TestOaiResponsesHandlerCountsOutputCallsNotDeclarations(t *testing.T) {
gin.SetMode(gin.TestMode)
operation_setting.SetToolPriceForTest("priced_fn", 5.0)
@@ -265,3 +303,92 @@ func TestOaiResponsesStreamHandlerDoesNotCountPartialImageEvent(t *testing.T) {
assert.Equal(t, 0, info.ResponsesUsageInfo.BuiltInTools[dto.BuildInToolImageGeneration].CallCount)
}
func TestOaiResponsesHandlerRewritesSGLangCreatedAtToInt(t *testing.T) {
gin.SetMode(gin.TestMode)
w := httptest.NewRecorder()
c, _ := gin.CreateTestContext(w)
c.Request = httptest.NewRequest(http.MethodPost, "/v1/responses", nil)
info := &relaycommon.RelayInfo{
ChannelMeta: &relaycommon.ChannelMeta{ChannelType: constant.ChannelTypeSGLang},
}
resp := &http.Response{
StatusCode: http.StatusOK,
Body: io.NopCloser(strings.NewReader(`{"id":"resp_1","object":"response","created_at":1786588600.0,"status":"completed","output":[]}`)),
Header: http.Header{"Content-Type": []string{"application/json"}},
}
_, apiErr := OaiResponsesHandler(c, info, resp)
require.Nil(t, apiErr)
assert.Contains(t, w.Body.String(), `"created_at":1786588600`)
assert.NotContains(t, w.Body.String(), `1786588600.0`)
}
func TestOaiResponsesStreamHandlerRewritesSGLangCreatedAtToInt(t *testing.T) {
gin.SetMode(gin.TestMode)
oldTimeout := constant.StreamingTimeout
constant.StreamingTimeout = 30
t.Cleanup(func() {
constant.StreamingTimeout = oldTimeout
})
w := httptest.NewRecorder()
c, _ := gin.CreateTestContext(w)
c.Request = httptest.NewRequest(http.MethodPost, "/v1/responses", nil)
c.Set(common.RequestIdKey, "sglang-created-at-test")
info := &relaycommon.RelayInfo{
DisablePing: true,
ChannelMeta: &relaycommon.ChannelMeta{
ChannelType: constant.ChannelTypeSGLang,
UpstreamModelName: "served-model",
},
}
resp := &http.Response{
StatusCode: http.StatusOK,
Body: io.NopCloser(strings.NewReader(
"data: {\"type\":\"response.completed\",\"response\":{\"id\":\"resp_1\",\"created_at\":1.7865886E9,\"status\":\"completed\",\"output\":[]}}\n\n" +
"data: [DONE]\n\n",
)),
Header: http.Header{"Content-Type": []string{"text/event-stream"}},
}
_, apiErr := OaiResponsesStreamHandler(c, info, resp)
require.Nil(t, apiErr)
assert.Contains(t, w.Body.String(), `"created_at":1786588600`)
assert.NotContains(t, w.Body.String(), `1.7865886E9`)
assert.NotContains(t, w.Body.String(), `1786588600.0`)
}
func TestOaiResponsesStreamHandlerKeepsNonSGLangCreatedAt(t *testing.T) {
gin.SetMode(gin.TestMode)
oldTimeout := constant.StreamingTimeout
constant.StreamingTimeout = 30
t.Cleanup(func() {
constant.StreamingTimeout = oldTimeout
})
w := httptest.NewRecorder()
c, _ := gin.CreateTestContext(w)
c.Request = httptest.NewRequest(http.MethodPost, "/v1/responses", nil)
c.Set(common.RequestIdKey, "openai-created-at-test")
info := &relaycommon.RelayInfo{
DisablePing: true,
ChannelMeta: &relaycommon.ChannelMeta{ChannelType: constant.ChannelTypeOpenAI, UpstreamModelName: "gpt-test"},
}
resp := &http.Response{
StatusCode: http.StatusOK,
Body: io.NopCloser(strings.NewReader(
"data: {\"type\":\"response.completed\",\"response\":{\"id\":\"resp_1\",\"created_at\":1786588600.0,\"status\":\"completed\",\"output\":[]}}\n\n" +
"data: [DONE]\n\n",
)),
Header: http.Header{"Content-Type": []string{"text/event-stream"}},
}
_, apiErr := OaiResponsesStreamHandler(c, info, resp)
require.Nil(t, apiErr)
assert.Contains(t, w.Body.String(), `1786588600.0`)
}
@@ -35,6 +35,7 @@ func OaiChatToResponsesHandler(c *gin.Context, info *relaycommon.RelayInfo, resp
return nil, types.WithOpenAIError(*oaiError, resp.StatusCode)
}
info.ObserveResponseModel(chatResp.Model)
if responseID := helper.GetResponseID(c); responseID != "" {
chatResp.Id = responseID
}
@@ -140,6 +141,7 @@ func OaiChatToResponsesStreamHandler(c *gin.Context, info *relaycommon.RelayInfo
return
}
info.ObserveResponseModel(chunk.Model)
results, err := service.ConvertStreamResponseChunk(c, info, state, &chunk)
if err != nil {
if failResponsesStream(err) {
@@ -3,11 +3,13 @@ package openai
import (
"testing"
relaycommon "github.com/QuantumNous/new-api/relay/common"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
)
func TestCollectStreamFunctionCallNamesDedupesSameIndex(t *testing.T) {
func TestObserveStreamChoicesDedupesFunctionCallNames(t *testing.T) {
info := &relaycommon.RelayInfo{StreamStatus: relaycommon.NewStreamStatus()}
seen := make(map[string]struct{})
var names []string
@@ -19,9 +21,18 @@ func TestCollectStreamFunctionCallNamesDedupesSameIndex(t *testing.T) {
`{"choices":[{"index":0,"delta":{"tool_calls":[{"index":1,"function":{"arguments":"{}"}}]}}]}`,
}
for _, chunk := range chunks {
collectStreamFunctionCallNames(chunk, seen, &names)
observeStreamChoices(info, chunk, seen, &names)
}
require.Len(t, names, 2)
assert.Equal(t, []string{"get_weather", "get_time"}, names)
assert.Empty(t, info.StreamStatus.ResponseOutcome(), "tool call deltas carry no finish reason")
observeStreamChoices(info, `{"choices":[{"index":0,"delta":{},"finish_reason":"tool_calls"}]}`, seen, &names)
assert.Equal(t, "completed", info.StreamStatus.ResponseOutcome())
assert.False(t, info.PerformanceBusinessRejection)
filtered := &relaycommon.RelayInfo{StreamStatus: relaycommon.NewStreamStatus()}
observeStreamChoices(filtered, `{"choices":[{"index":0,"delta":{},"finish_reason":"content_filter"}]}`, map[string]struct{}{}, &names)
assert.True(t, filtered.PerformanceBusinessRejection)
}
+212 -75
View File
@@ -2,6 +2,7 @@ package jsplugin
import (
"bytes"
"cmp"
"context"
"encoding/base64"
"fmt"
@@ -18,6 +19,7 @@ import (
"strconv"
"strings"
"time"
"unicode/utf8"
"github.com/QuantumNous/new-api/common"
"github.com/QuantumNous/new-api/constant"
@@ -33,6 +35,7 @@ import (
)
type requestDescriptor struct {
ResponseType string `json:"responseType"`
URL string `json:"url"`
Method string `json:"method"`
Headers map[string]string `json:"headers"`
@@ -120,14 +123,24 @@ func (a *TaskAdaptor) ValidateRequestAndSetAction(c *gin.Context, info *relaycom
return taskErr
}
}
if request, exists := c.Get("task_request"); exists {
if err := a.validateResolvedUsageRequest(request); err != nil {
request, hasRequest := c.Get("task_request")
hasUsageProfiles := len(a.plugin.Meta.UsageProfiles) > 0
if hasRequest && !hasUsageProfiles {
if err := a.validateResolvedUsageRequest(request, ""); err != nil {
return service.TaskErrorWrapperLocal(err, "plugin_usage_invalid", http.StatusBadRequest)
}
}
if _, err := a.buildSubmit(c, info); err != nil {
return service.TaskErrorWrapperLocal(err, "plugin_request_invalid", http.StatusBadRequest)
}
// The descriptor may rewrite the model. Validate profiled requests against
// that final model before any quota calculation or upstream submission.
if hasRequest && hasUsageProfiles {
usageModel := cmp.Or(info.UpstreamModelName, info.OriginModelName)
if err := a.validateResolvedUsageRequest(request, usageModel); err != nil {
return service.TaskErrorWrapperLocal(err, "plugin_usage_invalid", http.StatusBadRequest)
}
}
return nil
}
@@ -143,7 +156,7 @@ func (a *TaskAdaptor) EstimateBilling(c *gin.Context, info *relaycommon.RelayInf
func (a *TaskAdaptor) EstimateBillingValidated(c *gin.Context, info *relaycommon.RelayInfo) (map[string]float64, error) {
usageContext := a.submitContext(c, info)
usageContext["usagePurpose"] = "billing_ratios"
return a.usageRatios(c.Request.Context(), "extractUsage", usageContext)
return a.usageRatios(c.Request.Context(), cmp.Or(info.UpstreamModelName, info.OriginModelName), "extractUsage", usageContext)
}
func (a *TaskAdaptor) ExtractUsageFacts(c *gin.Context, info *relaycommon.RelayInfo) map[string]any {
@@ -172,7 +185,7 @@ func (a *TaskAdaptor) ExtractUsageFactsValidated(c *gin.Context, info *relaycomm
if !ok {
return nil, fmt.Errorf("plugin usage hook must return an object")
}
if _, err = a.validatedUsageRatios(facts); err != nil {
if _, err = a.validatedUsageRatios(facts, cmp.Or(info.UpstreamModelName, info.OriginModelName)); err != nil {
return nil, err
}
return facts, nil
@@ -183,7 +196,7 @@ func (a *TaskAdaptor) AdjustBillingOnSubmit(info *relaycommon.RelayInfo, taskDat
if err := common.Unmarshal(taskData, &data); err != nil {
data = string(taskData)
}
ratios, err := a.usageRatios(context.Background(), "extractUsageOnSubmit", a.submitContext(nil, info), data)
ratios, err := a.usageRatios(context.Background(), cmp.Or(info.UpstreamModelName, info.OriginModelName), "extractUsageOnSubmit", a.submitContext(nil, info), data)
if err != nil {
a.logRejectedUsage("extractUsageOnSubmit", err)
return nil
@@ -199,7 +212,11 @@ func (a *TaskAdaptor) AdjustBillingOnComplete(task *model.Task, result *relaycom
if err != nil {
return 0
}
a.applyCompletionUsageFacts(result, value)
usageModel := ""
if task != nil {
usageModel = cmp.Or(task.Properties.UpstreamModelName, task.Properties.OriginModelName)
}
a.applyCompletionUsageFacts(result, value, usageModel)
return 0
}
@@ -423,6 +440,9 @@ func encodeFilePlaceholder(placeholder map[string]any, form *multipart.Form, lim
}
func (a *TaskAdaptor) DoRequest(c *gin.Context, info *relaycommon.RelayInfo, body io.Reader) (*http.Response, error) {
wasStream := info.IsStream
info.IsStream = false
defer func() { info.IsStream = wasStream }()
if a.submit != nil && strings.TrimSpace(a.submit.Method) != "" {
originalMethod := c.Request.Method
c.Request.Method = strings.ToUpper(a.submit.Method)
@@ -431,24 +451,42 @@ func (a *TaskAdaptor) DoRequest(c *gin.Context, info *relaycommon.RelayInfo, bod
return channel.DoTaskApiRequest(a, c, info, body)
}
func (a *TaskAdaptor) ParseResponse(c *gin.Context, resp *http.Response, info *relaycommon.RelayInfo) (*channel.TaskSubmitResponse, *dto.TaskError) {
func (a *TaskAdaptor) ParseResponse(c *gin.Context, resp *http.Response, info *relaycommon.RelayInfo) (_ *channel.TaskSubmitResponse, taskErr *dto.TaskError) {
started := time.Now()
body, err := io.ReadAll(resp.Body)
if err != nil {
logger.LogDebug(c, "task_plugin subsystem=adaptor event=parse_submit_failed plugin=%q stage=read_response reason=read_failed status=%d", a.plugin.Meta.Key, resp.StatusCode)
return nil, service.TaskErrorWrapper(err, "read_response_body_failed", http.StatusInternalServerError)
var responseBody any
var err error
streaming := a.submit != nil && a.submit.ResponseType == "sse"
mediaType, _, _ := mime.ParseMediaType(resp.Header.Get("Content-Type"))
acceptedStream := streaming || mediaType == "text/event-stream"
defer func() {
if acceptedStream && taskErr != nil {
taskErr.NoRetry = true
}
}()
if !streaming && acceptedStream {
return nil, service.TaskErrorWrapperLocal(fmt.Errorf("unexpected SSE response for a JSON submission"), "plugin_submit_response_invalid", http.StatusBadGateway)
}
logger.LogDebug(
c,
"task_plugin subsystem=adaptor event=submit_response_received plugin=%q status=%d body_bytes=%d",
a.plugin.Meta.Key,
resp.StatusCode,
len(body),
)
responseBody := any(string(body))
var decoded any
if common.Unmarshal(body, &decoded) == nil {
responseBody = decoded
if streaming {
// Every failure after accepting this stream is non-retryable: the
// upstream may already have performed billable work.
defer resp.Body.Close()
responseBody, err = a.readSubmitEvents(c.Request.Context(), resp, a.submitContext(c, info))
} else {
var body []byte
body, err = io.ReadAll(io.LimitReader(resp.Body, maxTaskPluginPersistedJSONBytes+1))
if err == nil && len(body) > maxTaskPluginPersistedJSONBytes {
err = fmt.Errorf("task submit response exceeds size limit")
}
responseBody = string(body)
var decoded any
if err == nil && common.Unmarshal(body, &decoded) == nil {
responseBody = decoded
}
}
if err != nil {
failure := service.TaskErrorWrapper(err, "read_response_body_failed", http.StatusBadGateway)
failure.NoRetry = streaming
return nil, failure
}
headers := make(map[string][]string, len(resp.Header))
maps.Copy(headers, resp.Header)
@@ -485,34 +523,57 @@ func (a *TaskAdaptor) ParseResponse(c *gin.Context, resp *http.Response, info *r
var taskData []byte
if parsed.TaskData != nil {
taskData, err = common.Marshal(parsed.TaskData)
if err != nil {
if err != nil || len(taskData) > maxTaskPluginPersistedJSONBytes {
if err == nil {
err = fmt.Errorf("task data exceeds size limit")
}
return nil, service.TaskErrorWrapper(err, "plugin_submit_response_invalid", http.StatusBadGateway)
}
}
var immediate *relaycommon.TaskInfo
if parsed.Immediate != nil {
immediate = &relaycommon.TaskInfo{Code: parsed.Immediate.Code, TaskID: parsed.Immediate.TaskID, Status: parsed.Immediate.Status, Progress: parsed.Immediate.Progress, Reason: parsed.Immediate.Reason, Url: parsed.Immediate.URL, RemoteUrl: parsed.Immediate.RemoteURL}
}
logger.LogDebug(
c,
"task_plugin subsystem=adaptor event=parse_submit_complete plugin=%q status=%d task_data_bytes=%d immediate=%t elapsed_ms=%d",
a.plugin.Meta.Key,
resp.StatusCode,
len(taskData),
immediate != nil,
time.Since(started).Milliseconds(),
)
pluginState, _ := encodeReturnedPluginState(value)
if len(pluginState) > maxTaskPluginPersistedJSONBytes {
logger.LogWarn(c, fmt.Sprintf("task plugin %s rejected oversized submit state (%d bytes)", a.plugin.Meta.Key, len(pluginState)))
pluginState = nil
}
return &channel.TaskSubmitResponse{
response := &channel.TaskSubmitResponse{
UpstreamTaskID: parsed.TaskID,
TaskData: taskData,
Immediate: immediate,
PluginState: pluginState,
}, nil
}
logger.LogDebug(c, "task_plugin subsystem=adaptor event=parse_submit_complete plugin=%q status=%d task_data_bytes=%d immediate=%t elapsed_ms=%d",
a.plugin.Meta.Key, resp.StatusCode, len(taskData), parsed.Immediate != nil, time.Since(started).Milliseconds())
if parsed.Immediate == nil {
return response, nil
}
immediate := &relaycommon.TaskInfo{
Code: parsed.Immediate.Code, TaskID: parsed.Immediate.TaskID,
Status: parsed.Immediate.Status, Progress: parsed.Immediate.Progress,
Reason: parsed.Immediate.Reason, Url: parsed.Immediate.URL, RemoteUrl: parsed.Immediate.RemoteURL,
CompletionTokens: positiveInt(parsed.Immediate.CompletionTokens), TotalTokens: positiveInt(parsed.Immediate.TotalTokens),
}
response.Immediate = immediate
if immediate.Status != model.TaskStatusSuccess || !a.hasHook(c.Request.Context(), "extractUsageOnComplete") {
return response, nil
}
// Match the polling hook context without persisting a temporary task.
task := &model.Task{TaskID: info.PublicTaskID, Action: info.Action, Data: taskData,
Properties: model.Properties{OriginModelName: info.OriginModelName, UpstreamModelName: info.UpstreamModelName}}
task.PrivateData.UpstreamTaskID = parsed.TaskID
task.PrivateData.PluginState = pluginState
ctx, err := a.queryContext(task, info.ApiKey, info.ChannelBaseUrl, info.ChannelSetting.Proxy)
if err != nil {
logger.LogWarn(c, fmt.Sprintf("task plugin %s completion context failed; retaining reserved quota: %v", a.plugin.Meta.Key, err))
return response, nil
}
facts, err := a.plugin.Engine.Call(c.Request.Context(), "extractUsageOnComplete", ctx, jsonValue(immediate), parsed.TaskData)
if err != nil {
logger.LogWarn(c, fmt.Sprintf("task plugin %s completion usage failed; retaining reserved quota: %v", a.plugin.Meta.Key, err))
return response, nil
}
if err := a.applyCompletionUsageFacts(immediate, facts, cmp.Or(info.UpstreamModelName, info.OriginModelName)); err != nil {
logger.LogWarn(c, fmt.Sprintf("task plugin %s completion usage rejected; retaining reserved quota: %v", a.plugin.Meta.Key, err))
}
return response, nil
}
func (a *TaskAdaptor) GetModelList() []string { return append([]string(nil), a.plugin.Meta.Models...) }
@@ -680,7 +741,8 @@ func (a *TaskAdaptor) ParseBatchResult(tasks []*model.Task, resp *http.Response,
}
facts, hookErr := a.plugin.Engine.Call(context.Background(), "extractUsageOnComplete", itemCtx, jsonValue(&info), usageBody)
if hookErr == nil {
a.applyCompletionUsageFacts(&info, facts)
usageModel, _ := itemCtx["upstreamModel"].(string)
a.applyCompletionUsageFacts(&info, facts, usageModel)
}
}
results[item.TaskID] = &service.BatchTaskResult{TaskInfo: info, Action: item.Action, SubmitTime: item.SubmitTime, StartTime: item.StartTime, FinishTime: item.FinishTime, Data: item.Data}
@@ -745,7 +807,8 @@ func (a *TaskAdaptor) ParseTaskResult(task *model.Task, resp *http.Response, bod
if a.hasHook(context.Background(), "extractUsageOnComplete") {
facts, hookErr := a.plugin.Engine.Call(context.Background(), "extractUsageOnComplete", ctx, jsonValue(result), input)
if hookErr == nil {
a.applyCompletionUsageFacts(result, facts)
usageModel, _ := ctx["upstreamModel"].(string)
a.applyCompletionUsageFacts(result, facts, usageModel)
}
}
taskStatus := model.TaskStatus(result.Status)
@@ -760,20 +823,20 @@ func (a *TaskAdaptor) ParseTaskResult(task *model.Task, resp *http.Response, bod
return result, nil
}
func (a *TaskAdaptor) applyCompletionUsageFacts(result *relaycommon.TaskInfo, facts any) {
values, err := a.validatedCompletionUsageFacts(facts)
func (a *TaskAdaptor) applyCompletionUsageFacts(result *relaycommon.TaskInfo, facts any, modelName string) error {
values, err := a.validatedCompletionUsageFacts(facts, modelName)
if err != nil {
a.logRejectedUsage("extractUsageOnComplete", err)
return
return err
}
if len(values) == 0 {
return
return nil
}
result.UsageFacts = values
if units := positiveInt(values["upstreamUnits"]); units > 0 {
result.CompletionTokens = units
result.TotalTokens = units
return
return nil
}
if completionTokens, exists := values["completionTokens"]; exists {
result.CompletionTokens = positiveInt(completionTokens)
@@ -781,6 +844,7 @@ func (a *TaskAdaptor) applyCompletionUsageFacts(result *relaycommon.TaskInfo, fa
if totalTokens, exists := values["totalTokens"]; exists {
result.TotalTokens = positiveInt(totalTokens)
}
return nil
}
func (a *TaskAdaptor) ConvertToOpenAIVideo(task *model.Task) ([]byte, error) {
@@ -805,26 +869,24 @@ func (a *TaskAdaptor) ConvertToOpenAIVideo(task *model.Task) ([]byte, error) {
if err != nil {
return nil, err
}
rendered := kitdto.NewOpenAIVideo()
if err = common.Unmarshal(encoded, rendered); err != nil {
var rendered map[string]any
if err = common.Unmarshal(encoded, &rendered); err != nil || rendered == nil {
return nil, fmt.Errorf("plugin returned an invalid OpenAI video object")
}
// Keep provider extensions intact while the host owns the public task's
// identity and lifecycle, including completion timestamps after settlement.
host := task.ToOpenAIVideo()
rendered.ID = host.ID
rendered.Object = host.Object
rendered.TaskID = ""
rendered.Status = host.Status
rendered.Progress = host.Progress
rendered.CreatedAt = host.CreatedAt
rendered.Model = host.Model
rendered.CompletedAt = host.CompletedAt
for key := range rendered.Metadata {
if strings.EqualFold(key, "url") {
delete(rendered.Metadata, key)
}
}
if len(rendered.Metadata) == 0 {
rendered.Metadata = nil
rendered["id"] = host.ID
rendered["object"] = host.Object
delete(rendered, "task_id")
rendered["status"] = host.Status
rendered["progress"] = host.Progress
rendered["created_at"] = host.CreatedAt
rendered["model"] = host.Model
if host.CompletedAt != 0 {
rendered["completed_at"] = host.CompletedAt
} else {
delete(rendered, "completed_at")
}
return common.Marshal(rendered)
}
@@ -1139,6 +1201,16 @@ func (a *TaskAdaptor) buildSubmit(c *gin.Context, info *relaycommon.RelayInfo) (
)
return nil, err
}
if descriptor.ResponseType == "" {
descriptor.ResponseType = "json"
}
allowed := a.plugin.Meta.SubmitResponseTypes
if len(allowed) == 0 {
allowed = []string{"json"}
}
if !slices.Contains(allowed, descriptor.ResponseType) {
return nil, fmt.Errorf("plugin does not support submit response type %q", descriptor.ResponseType)
}
if strings.TrimSpace(descriptor.URL) == "" {
logger.LogDebug(c, "task_plugin subsystem=adaptor event=build_submit_failed plugin=%q stage=validate_url reason=empty_url", a.plugin.Meta.Key)
return nil, fmt.Errorf("plugin returned an empty submit URL")
@@ -1294,7 +1366,7 @@ func (a *TaskAdaptor) submitContext(c *gin.Context, info *relaycommon.RelayInfo)
return ctx
}
func (a *TaskAdaptor) usageRatios(ctx context.Context, hook string, args ...any) (map[string]float64, error) {
func (a *TaskAdaptor) usageRatios(ctx context.Context, modelName, hook string, args ...any) (map[string]float64, error) {
if !a.hasHook(ctx, hook) {
return nil, nil
}
@@ -1312,7 +1384,7 @@ func (a *TaskAdaptor) usageRatios(ctx context.Context, hook string, args ...any)
logger.LogDebug(ctx, "task_plugin subsystem=adaptor event=usage_hook_failed plugin=%q hook=%q reason=result_not_object elapsed_ms=%d", a.plugin.Meta.Key, hook, time.Since(started).Milliseconds())
return nil, fmt.Errorf("plugin usage hook must return an object")
}
ratios, err := a.validatedUsageRatios(facts)
ratios, err := a.validatedUsageRatios(facts, modelName)
if err != nil {
logger.LogDebug(ctx, "task_plugin subsystem=adaptor event=usage_hook_failed plugin=%q hook=%q reason=invalid_usage elapsed_ms=%d", a.plugin.Meta.Key, hook, time.Since(started).Milliseconds())
return nil, err
@@ -1329,15 +1401,16 @@ func (a *TaskAdaptor) usageRatios(ctx context.Context, hook string, args ...any)
return ratios, nil
}
func (a *TaskAdaptor) validateResolvedUsageRequest(request any) error {
return a.validateResolvedUsageValue(jsonValue(request))
func (a *TaskAdaptor) validateResolvedUsageRequest(request any, modelName string) error {
schema, _ := a.plugin.Meta.UsageForModel(modelName)
return a.validateResolvedUsageValue(jsonValue(request), schema)
}
func (a *TaskAdaptor) validateResolvedUsageValue(value any) error {
func (a *TaskAdaptor) validateResolvedUsageValue(value any, usageSchema map[string]pluginruntime.UsageFieldSchema) error {
switch typed := value.(type) {
case map[string]any:
for key, item := range typed {
if schema, declared := a.plugin.Meta.UsageSchema[key]; declared {
if schema, declared := usageSchema[key]; declared {
if _, err := validateUsageValue(item, schema, true); err != nil {
return err
}
@@ -1346,13 +1419,13 @@ func (a *TaskAdaptor) validateResolvedUsageValue(value any) error {
return err
}
}
if err := a.validateResolvedUsageValue(item); err != nil {
if err := a.validateResolvedUsageValue(item, usageSchema); err != nil {
return err
}
}
case []any:
for _, item := range typed {
if err := a.validateResolvedUsageValue(item); err != nil {
if err := a.validateResolvedUsageValue(item, usageSchema); err != nil {
return err
}
}
@@ -1360,10 +1433,11 @@ func (a *TaskAdaptor) validateResolvedUsageValue(value any) error {
return nil
}
func (a *TaskAdaptor) validatedUsageRatios(facts map[string]any) (map[string]float64, error) {
func (a *TaskAdaptor) validatedUsageRatios(facts map[string]any, modelName string) (map[string]float64, error) {
usageSchema, _ := a.plugin.Meta.UsageForModel(modelName)
ratios := make(map[string]float64)
for key, value := range facts {
if schema, declared := a.plugin.Meta.UsageSchema[key]; declared {
if schema, declared := usageSchema[key]; declared {
number, err := validateUsageValue(value, schema, false)
if err != nil {
return nil, err
@@ -1390,6 +1464,9 @@ func (a *TaskAdaptor) validatedUsageRatios(facts map[string]any) (map[string]flo
if err := validateUsageNumberLimit(number, limit); err != nil {
return nil, err
}
// Keep numeric facts usable by existing expressions even when the
// selected profile no longer declares that legacy extension field.
facts[key] = number
if number > 0 {
ratios[key] = number
}
@@ -1397,7 +1474,8 @@ func (a *TaskAdaptor) validatedUsageRatios(facts map[string]any) (map[string]flo
return ratios, nil
}
func (a *TaskAdaptor) validatedCompletionUsageFacts(facts any) (map[string]any, error) {
func (a *TaskAdaptor) validatedCompletionUsageFacts(facts any, modelName string) (map[string]any, error) {
usageSchema, _ := a.plugin.Meta.UsageForModel(modelName)
if facts == nil {
return nil, nil
}
@@ -1408,7 +1486,7 @@ func (a *TaskAdaptor) validatedCompletionUsageFacts(facts any) (map[string]any,
validated := make(map[string]any, len(values))
for key, value := range values {
validated[key] = value
if schema, declared := a.plugin.Meta.UsageSchema[key]; declared {
if schema, declared := usageSchema[key]; declared {
number, err := validateUsageValue(value, schema, false)
if err != nil {
return nil, err
@@ -1436,6 +1514,10 @@ func (a *TaskAdaptor) validatedCompletionUsageFacts(facts any) (map[string]any,
return nil, fmt.Errorf("plugin usage value must be a finite non-negative number")
}
validated[key] = float64(common.QuotaFromFloat(number))
default:
if number, numeric := usageNumber(value, false); numeric {
validated[key] = number
}
}
}
return validated, nil
@@ -1549,6 +1631,9 @@ func convert(value any, target any) error {
}
func jsonValue(value any) any {
if normalized, ok := cloneJSONValue(value, 0); ok {
return normalized
}
data, err := common.Marshal(value)
if err != nil {
return value
@@ -1559,6 +1644,58 @@ func jsonValue(value any) any {
}
return normalized
}
// Already-parsed JSON and Sobek's plain exported values do not need another
// encode/decode pass. Copy containers to isolate hooks, preserving the codec's
// float64 numbers and nulls. Other Go types still use the configured codec.
func cloneJSONValue(value any, depth int) (any, bool) {
if depth > 64 {
return nil, false
}
switch typed := value.(type) {
case nil, bool:
return typed, true
case string:
return typed, utf8.ValidString(typed)
case float64:
return typed, !math.IsNaN(typed) && !math.IsInf(typed, 0)
case int64:
return float64(typed), true
case int:
return float64(typed), true
case map[string]any:
if typed == nil {
return nil, true
}
cloned := make(map[string]any, len(typed))
for key, item := range typed {
if !utf8.ValidString(key) {
return nil, false
}
child, ok := cloneJSONValue(item, depth+1)
if !ok {
return nil, false
}
cloned[key] = child
}
return cloned, true
case []any:
if typed == nil {
return nil, true
}
cloned := make([]any, len(typed))
for index, item := range typed {
child, ok := cloneJSONValue(item, depth+1)
if !ok {
return nil, false
}
cloned[index] = child
}
return cloned, true
default:
return nil, false
}
}
func positiveInt(value any) int {
switch number := value.(type) {
case int64:
+485 -57
View File
@@ -4,6 +4,7 @@ import (
"bytes"
"context"
"encoding/base64"
"encoding/json"
"io"
"math"
"mime/multipart"
@@ -16,7 +17,9 @@ import (
"github.com/QuantumNous/new-api/common"
"github.com/QuantumNous/new-api/constant"
"github.com/QuantumNous/new-api/model"
"github.com/QuantumNous/new-api/pkg/billingexpr"
pluginruntime "github.com/QuantumNous/new-api/pkg/jsplugin"
"github.com/QuantumNous/new-api/plugins"
"github.com/QuantumNous/new-api/relay/channel"
relaycommon "github.com/QuantumNous/new-api/relay/common"
"github.com/QuantumNous/new-api/relay/helper"
@@ -532,69 +535,90 @@ func TestTaskAdaptorMapsJSContract(t *testing.T) {
require.NoError(t, err)
}
func TestTaskAdaptorSanitizesOpenAIVideoRendererOutput(t *testing.T) {
source := `
export const meta = {
apiVersion: 1, key: "safe-video", name: "Safe Video", version: "1.0.0",
author: {name: "Test"}, models: ["model"], fetchMode: "per_task", protocols: ["openai_video"],
};
export function buildSubmitRequest(ctx) { return {url: ctx.baseUrl + "/submit"}; }
export function parseSubmitResponse() { return {taskId: "upstream"}; }
export function buildQueryRequest(ctx) { return {url: ctx.baseUrl + "/query"}; }
export function parseTaskResult() { return {status: "SUCCESS"}; }
export function listArtifacts() { return []; }
export function buildContentRequest() { throw new Error("artifact_not_found"); }
export const protocols = {openai_video: {
decodeRequest: function(ctx) { return {kind:"submit", model:ctx.model, requestBody:ctx.body.value}; },
render: function() {
return {
id: "upstream-id",
task_id: "upstream-task-id",
object: "provider-video",
model: "model",
status: "completed",
progress: 100,
created_at: 10,
completed_at: 20,
metadata: {
url: "https://upstream.example/video.mp4",
URL: "https://upstream.example/uppercase.mp4",
label: "safe",
},
url: "https://upstream.example/top-level.mp4",
upstream_url: "https://upstream.example/unknown.mp4",
provider_payload: {task_id: "upstream-task-id"},
};
}}};
`
func TestTaskAdaptorPreservesSoraVideoResponseFields(t *testing.T) {
source, err := plugins.Source("sora")
require.NoError(t, err)
plugin, err := pluginruntime.NewRegistry().Register(source, pluginruntime.Options{})
require.NoError(t, err)
adaptor := New(plugin)
rendered, err := adaptor.ConvertToOpenAIVideo(&model.Task{
TaskID: "task_public",
Status: model.TaskStatusInProgress,
Properties: model.Properties{OriginModelName: "origin-model"},
})
require.NoError(t, err)
for _, tc := range []struct {
status model.TaskStatus
want string
}{
{model.TaskStatusInProgress, "in_progress"},
{model.TaskStatusSuccess, "completed"},
{model.TaskStatusFailure, "failed"},
} {
t.Run(string(tc.status), func(t *testing.T) {
task := &model.Task{
TaskID: "task_public",
Status: tc.status,
Progress: "42%",
CreatedAt: 100,
FinishTime: 200,
Properties: model.Properties{OriginModelName: "origin-model"},
PrivateData: model.TaskPrivateData{UpstreamTaskID: "upstream-task-id"},
Data: []byte(`{
"id":"upstream-task-id","task_id":"upstream-task-id",
"object":"provider-video","model":"provider-model",
"status":"completed","progress":100,"created_at":10,"completed_at":20,
"url":"https://cdn.example/video.mp4",
"metadata":{"url":"https://cdn.example/video.mp4","URL":"https://cdn.example/uppercase.mp4"},
"provider_payload":{"task_id":"upstream-task-id","items":[{"enabled":false,"count":0,"value":null}]},
"seconds":8,"resolution":"720p","aspect_ratio":"16:9",
"reference_images":["https://cdn.example/reference.png"],
"error":{"code":"provider_error","message":"provider rejected request","detail":{"retryable":false}}
}`),
}
rendered, err := adaptor.ConvertToOpenAIVideo(task)
require.NoError(t, err)
var video dto.OpenAIVideo
require.NoError(t, common.Unmarshal(rendered, &video))
assert.Equal(t, "task_public", video.ID)
assert.Equal(t, "video", video.Object)
assert.Empty(t, video.TaskID)
assert.Equal(t, "origin-model", video.Model)
assert.Zero(t, video.CompletedAt)
assert.Equal(t, map[string]any{"label": "safe"}, video.Metadata)
var fields map[string]any
require.NoError(t, common.Unmarshal(rendered, &fields))
assert.Equal(t, "task_public", fields["id"])
assert.NotContains(t, fields, "task_id")
assert.Equal(t, "video", fields["object"])
assert.Equal(t, "origin-model", fields["model"])
assert.Equal(t, tc.want, fields["status"])
assert.Equal(t, float64(42), fields["progress"])
assert.Equal(t, float64(100), fields["created_at"])
if tc.status == model.TaskStatusSuccess {
assert.Equal(t, float64(200), fields["completed_at"])
} else {
assert.NotContains(t, fields, "completed_at")
}
assert.Equal(t, "https://cdn.example/video.mp4", fields["url"])
assert.Equal(t, map[string]any{
"url": "https://cdn.example/video.mp4",
"URL": "https://cdn.example/uppercase.mp4",
}, fields["metadata"])
assert.Equal(t, map[string]any{
"task_id": "task_public",
"items": []any{map[string]any{"enabled": false, "count": float64(0), "value": nil}},
}, fields["provider_payload"])
assert.Equal(t, float64(8), fields["seconds"])
assert.Equal(t, "720p", fields["resolution"])
assert.Equal(t, "16:9", fields["aspect_ratio"])
assert.Equal(t, []any{"https://cdn.example/reference.png"}, fields["reference_images"])
assert.Equal(t, map[string]any{
"code": "provider_error", "message": "provider rejected request",
"detail": map[string]any{"retryable": false},
}, fields["error"])
})
}
}
var fields map[string]any
require.NoError(t, common.Unmarshal(rendered, &fields))
assert.NotContains(t, fields, "url")
assert.NotContains(t, fields, "upstream_url")
assert.NotContains(t, fields, "provider_payload")
assert.NotContains(t, fields, "completed_at")
assert.NotContains(t, string(rendered), "upstream.example")
assert.NotContains(t, string(rendered), "upstream-task-id")
func TestTaskAdaptorRejectsNonObjectOpenAIVideoRendererOutput(t *testing.T) {
for _, value := range []string{"null", "[]", `"video"`, "42", "false"} {
t.Run(value, func(t *testing.T) {
source := strings.Replace(mockPlugin, `return {id: task.task_id, status: "completed"};`, "return "+value+";", 1)
plugin, err := pluginruntime.NewRegistry().Register(source, pluginruntime.Options{})
require.NoError(t, err)
_, err = New(plugin).ConvertToOpenAIVideo(&model.Task{TaskID: "task_public"})
require.ErrorContains(t, err, "invalid OpenAI video object")
})
}
}
func TestTaskAdaptorPreservesOpenAIVideoFailureSlotsAndOwnsLifecycle(t *testing.T) {
@@ -1413,3 +1437,407 @@ export function parseBatchResult(){return [];}
assert.Equal(t, []any{"model-a", "model-b"}, captured["models"])
assert.Equal(t, false, captured["hasRequestBody"])
}
func TestTaskAdaptorUsageProfilesFollowExecutionModel(t *testing.T) {
const source = `
export const meta = {
apiVersion:1, key:"profile-usage", name:"Profile Usage", version:"1.0.0", author:{name:"Test"},
models:["image", "video"], fetchMode:"batch",
usageSchema:{units:{type:"number",unit:"second"}},
usageProfiles:[
{models:["image"],schema:{units:{type:"number",unit:"count"},mode:{enum:["image"]}}},
{models:["video"],schema:{units:{type:"number",unit:"token"},mode:{enum:["video"]}},
examples:[{label:"video",facts:{units:1,mode:"video"}}]}
]
};
export function buildSubmitRequest(ctx) {
return {url:ctx.baseUrl+"/submit",rewriteModel:ctx.requestBody.rewriteTo || ""};
}
export function parseSubmitResponse(){return {taskId:"task"};}
export function buildQueryRequest(){return {url:"https://provider.example/task"};}
export function buildBatchQueryRequest(){return {url:"https://provider.example/tasks"};}
export function parseTaskResult(){return {status:"SUCCESS"};}
export function parseBatchResult(ctx,body){return body.items;}
export function extractUsage(ctx){return {units:ctx.requestBody.hookUnits,mode:ctx.requestBody.hookMode,legacyRatio:2};}
export function extractUsageOnSubmit(ctx,body){return body.usage;}
export function extractUsageOnComplete(ctx,result,body){return body.usage;}
`
plugin, err := pluginruntime.NewRegistry().Register(source, pluginruntime.Options{})
require.NoError(t, err)
newRequest := func(t *testing.T, upstream string, body map[string]any) (*TaskAdaptor, *gin.Context, *relaycommon.RelayInfo) {
t.Helper()
info := &relaycommon.RelayInfo{
OriginModelName: "public-alias",
ChannelMeta: &relaycommon.ChannelMeta{
UpstreamModelName: upstream, ChannelBaseUrl: "https://provider.example",
},
TaskRelayInfo: &relaycommon.TaskRelayInfo{},
}
adaptor := New(plugin)
adaptor.Init(info)
c, _ := gin.CreateTestContext(httptest.NewRecorder())
c.Request = httptest.NewRequest(http.MethodPost, "/native/submit", nil)
c.Set("task_request", body)
return adaptor, c, info
}
for _, tc := range []struct {
name, upstream, rewrite, mode string
units float64
wantError bool
}{
{"mapped image", "image", "", "image", 2, false},
{"mapped video token units", "video", "", "video", 500000, false},
{"rewritten model", "image", "video", "video", 500000, false},
{"image count ceiling", "image", "", "image", float64(dto.MaxImageN + 1), true},
{"profile enum", "image", "", "video", 2, true},
} {
t.Run(tc.name, func(t *testing.T) {
adaptor, c, info := newRequest(t, tc.upstream, map[string]any{
"hookUnits": tc.units, "hookMode": tc.mode, "rewriteTo": tc.rewrite,
})
require.Nil(t, adaptor.ValidateRequestAndSetAction(c, info))
facts, err := adaptor.ExtractUsageFactsValidated(c, info)
if tc.wantError {
require.Error(t, err)
return
}
require.NoError(t, err)
assert.Equal(t, map[string]any{"units": tc.units, "mode": tc.mode, "legacyRatio": 2.0}, facts)
ratios, err := adaptor.EstimateBillingValidated(c, info)
require.NoError(t, err)
assert.Equal(t, map[string]float64{"units": tc.units, "legacyRatio": 2}, ratios)
body, err := common.Marshal(map[string]any{"usage": facts})
require.NoError(t, err)
assert.Equal(t, ratios, adaptor.AdjustBillingOnSubmit(info, body))
})
}
t.Run("request revalidated after model rewrite", func(t *testing.T) {
adaptor, c, info := newRequest(t, "video", map[string]any{
"rewriteTo": "image", "metadata": map[string]any{"units": dto.MaxImageN + 1},
})
taskErr := adaptor.ValidateRequestAndSetAction(c, info)
require.NotNil(t, taskErr)
assert.Equal(t, "plugin_usage_invalid", taskErr.Code)
})
t.Run("request enum belongs to the rewritten model", func(t *testing.T) {
adaptor, c, info := newRequest(t, "video", map[string]any{
"rewriteTo": "image", "mode": "image", "units": 2,
})
require.Nil(t, adaptor.ValidateRequestAndSetAction(c, info))
assert.Equal(t, "image", info.UpstreamModelName)
})
t.Run("completion selects each persisted model including legacy origin fallback", func(t *testing.T) {
adaptor, _, _ := newRequest(t, "image", map[string]any{})
for _, tc := range []struct {
upstream, origin, mode string
units float64
accepted bool
}{
{"video", "public-alias", "video", 500000, true},
{"image", "public-alias", "image", 2, true},
{"", "video", "video", 500000, true},
{"image", "public-alias", "image", float64(dto.MaxImageN + 1), false},
{"video", "public-alias", "image", 1, false},
} {
task := &model.Task{Properties: model.Properties{OriginModelName: tc.origin, UpstreamModelName: tc.upstream}}
body, err := common.Marshal(map[string]any{"usage": map[string]any{"units": tc.units, "mode": tc.mode}})
require.NoError(t, err)
result, err := adaptor.ParseTaskResult(task, &http.Response{StatusCode: http.StatusOK}, body)
require.NoError(t, err)
if tc.accepted {
assert.Equal(t, map[string]any{"units": tc.units, "mode": tc.mode}, result.UsageFacts)
} else {
assert.Nil(t, result.UsageFacts)
}
}
})
t.Run("legacy completion facts remain usable by stored expressions", func(t *testing.T) {
updated, err := pluginruntime.NewRegistry().Register(strings.Replace(source,
`export function extractUsageOnComplete(ctx,result,body){return body.usage;}`,
`export function extractUsageOnComplete(){return {legacyRatio:3};}`, 1), pluginruntime.Options{})
require.NoError(t, err)
_, _, info := newRequest(t, "image", map[string]any{})
adaptor := New(updated)
adaptor.Init(info)
result, err := adaptor.ParseTaskResult(&model.Task{Properties: model.Properties{UpstreamModelName: "image"}},
&http.Response{StatusCode: http.StatusOK}, []byte(`{}`))
require.NoError(t, err)
cost, _, err := billingexpr.RunExprWithRequest(`tier("legacy", u("legacyRatio") * 2)`,
billingexpr.TokenParams{}, billingexpr.RequestInput{Usage: result.UsageFacts})
require.NoError(t, err)
assert.Equal(t, 6.0, cost)
})
t.Run("mixed batch selects schema by task rather than adaptor", func(t *testing.T) {
adaptor, _, _ := newRequest(t, "image", map[string]any{})
tasks := []*model.Task{
{TaskID: "video-task", Properties: model.Properties{UpstreamModelName: "video"}},
{TaskID: "image-task", Properties: model.Properties{UpstreamModelName: "image"}},
{TaskID: "invalid-image", Properties: model.Properties{UpstreamModelName: "image"}},
}
body := []byte(`{"items":[
{"taskId":"image-task","status":"SUCCESS","data":{"usage":{"units":2,"mode":"image"}}},
{"taskId":"video-task","status":"SUCCESS","data":{"usage":{"units":500000,"mode":"video"}}},
{"taskId":"invalid-image","status":"SUCCESS","data":{"usage":{"units":129,"mode":"image"}}}
]}`)
results, err := adaptor.ParseBatchResult(tasks, &http.Response{StatusCode: http.StatusOK}, body)
require.NoError(t, err)
require.Len(t, results, 3)
assert.Equal(t, map[string]any{"units": 500000.0, "mode": "video"}, results["video-task"].TaskInfo.UsageFacts)
assert.Equal(t, map[string]any{"units": 2.0, "mode": "image"}, results["image-task"].TaskInfo.UsageFacts)
assert.Nil(t, results["invalid-image"].TaskInfo.UsageFacts)
})
}
const documentStreamPlugin = `
export const meta = {apiVersion:1,key:"document-stream",name:"Document stream",version:"1.0.0",author:{name:"Test"},models:["document"],fetchMode:"per_task",submitResponseTypes:["json","sse"],usageSchema:{units:{type:"number",unit:"count"}}};
export function buildSubmitRequest(ctx){return {url:ctx.baseUrl+"/compile",responseType:ctx.requestBody.responseType || "sse"};}
export function parseSubmitEvent(ctx,event,previous) {
const chunk = JSON.parse(event.data);
if (chunk.error) throw new Error("provider stream failure");
if (chunk.badState) return {state:null};
if (chunk.largeState) return {state:{document:"x".repeat(1048577)},done:chunk.complete === true};
if (chunk.escapedState) return {state:{document:"<".repeat(200000)},done:true};
if (chunk.resetState) return {state:{document:"",units:0},done:true};
const state = Object.assign({}, previous || {document:"",units:0});
state.document += chunk.part || "";
if (chunk.units !== undefined) state.units = chunk.units;
state.event = event.event; state.id = event.id;
return {state:state,done:chunk.complete === true};
}
export function parseSubmitResponse(ctx,response){return {taskId:"vendor-document",taskData:response.body,immediate:{status:"SUCCESS"},state:{revision:1}};}
export function parseTaskResult(){return {status:"SUCCESS"};}
export function buildQueryRequest(ctx){return {url:ctx.baseUrl+"/query"};}
export function extractUsageOnComplete(ctx,result,body){
if(ctx.model!=="alias" || ctx.upstreamModel!=="document" || ctx.taskId!=="vendor-document" || ctx.publicTaskId!=="public-document" || ctx.state.revision!==1 || ctx.data.document!==body.document) throw new Error("invalid completion context");
return {units:body.units};
}
`
func TestTaskSubmitStreamContract(t *testing.T) {
plugin, err := pluginruntime.CompilePlugin(documentStreamPlugin, pluginruntime.Options{})
require.NoError(t, err)
for _, tc := range []struct {
name, body, contentType string
valid bool
}{
{"multiline and CRLF", ": heartbeat\r\nid: doc-1\r\nevent: update\r\ndata: {\"part\":\r\ndata: \"hello\",\"units\":2}\r\n\r\ndata: {\"part\":\"world\",\"units\":3,\"complete\":true}\n\n", "text/event-stream; charset=utf-8", true},
{"zero actual units", "data: {\"part\":\"free\",\"units\":0,\"complete\":true}\n\n", "text/event-stream", true},
{"premature EOF", "data: {\"part\":\"partial\"}\n\n", "text/event-stream", false},
{"unterminated event", "data: {\"complete\":true}", "text/event-stream", false},
{"provider error", "data: {\"error\":true}\n\n", "text/event-stream", false},
{"invalid hook result", "data: {\"badState\":true}\n\n", "text/event-stream", false},
{"state limit", "data: {\"largeState\":true}\n\n", "text/event-stream", false},
{"intermediate state limit", "data: {\"largeState\":true}\n\ndata: {\"resetState\":true}\n\n", "text/event-stream", false},
{"encoded state limit", "data: {\"escapedState\":true}\n\n", "text/event-stream", false},
{"event limit", "data: " + strings.Repeat("x", maxTaskPluginPersistedJSONBytes) + "\n\n", "text/event-stream", false},
{"wrong response type", "{}", "application/json", false},
{"undeclared stream", "data: {}\n\n", "text/event-stream", false},
} {
t.Run(tc.name, func(t *testing.T) {
info := &relaycommon.RelayInfo{OriginModelName: "alias", ChannelMeta: &relaycommon.ChannelMeta{UpstreamModelName: "document", ChannelBaseUrl: "https://provider.example"}, TaskRelayInfo: &relaycommon.TaskRelayInfo{PublicTaskID: "public-document"}}
adaptor := New(plugin)
adaptor.Init(info)
c, recorder := gin.CreateTestContext(httptest.NewRecorder())
_ = recorder
c.Request = httptest.NewRequest(http.MethodPost, "/compile", nil)
requestBody := map[string]any{}
if tc.name == "undeclared stream" {
requestBody["responseType"] = "json"
}
c.Set("task_request", requestBody)
require.Nil(t, adaptor.ValidateRequestAndSetAction(c, info))
response := &http.Response{StatusCode: 200, Header: http.Header{"Content-Type": {tc.contentType}}, Body: io.NopCloser(strings.NewReader(tc.body))}
parsed, taskErr := adaptor.ParseResponse(c, response, info)
assert.False(t, c.Writer.Written())
if !tc.valid {
require.NotNil(t, taskErr)
assert.True(t, taskErr.NoRetry)
return
}
require.Nil(t, taskErr)
require.NotNil(t, parsed.Immediate)
if tc.name == "zero actual units" {
assert.Equal(t, map[string]any{"units": float64(0)}, parsed.Immediate.UsageFacts)
return
}
assert.Equal(t, map[string]any{"units": float64(3)}, parsed.Immediate.UsageFacts)
assert.JSONEq(t, `{"document":"helloworld","units":3,"event":"message","id":"doc-1"}`, string(parsed.TaskData))
})
}
}
func TestTaskSubmitStreamCancellationClosesReader(t *testing.T) {
plugin, err := pluginruntime.CompilePlugin(documentStreamPlugin, pluginruntime.Options{})
require.NoError(t, err)
reader, writer := io.Pipe()
defer writer.Close()
ctx, cancel := context.WithCancel(t.Context())
cancel()
_, err = New(plugin).readSubmitEvents(ctx, &http.Response{Header: http.Header{"Content-Type": {"text/event-stream"}}, Body: reader}, nil)
require.ErrorIs(t, err, context.Canceled)
_, err = writer.Write([]byte("data: {}\n\n"))
require.Error(t, err)
}
func TestTaskSubmitStreamIdleTimeout(t *testing.T) {
previous := constant.StreamingTimeout
constant.StreamingTimeout = 1
t.Cleanup(func() { constant.StreamingTimeout = previous })
plugin, err := pluginruntime.CompilePlugin(documentStreamPlugin, pluginruntime.Options{})
require.NoError(t, err)
reader, writer := io.Pipe()
defer writer.Close()
_, err = New(plugin).readSubmitEvents(t.Context(), &http.Response{Header: http.Header{"Content-Type": {"text/event-stream"}}, Body: reader}, nil)
require.ErrorContains(t, err, "idle timeout")
}
func TestPluginJSONValuesPreserveCodecNormalizationAndIsolation(t *testing.T) {
for _, value := range []any{
map[string]any{"units": int64(3), "nested": []any{true, "<image> / 图像", math.Copysign(0, -1), nil}},
map[string]any{"empty": []any{}, "null": []any(nil), "object": map[string]any(nil)},
map[string]any{"large": int64(math.MaxInt64), "invalid UTF-8": string([]byte{0xff, 0xfe})},
map[string]any{string([]byte{0xff}): "invalid key"},
map[string]any{"bytes": []byte{0, 1, 255}, "number": json.Number("9007199254740993")},
json.RawMessage(`{"units":2,"enabled":false}`),
struct {
Units int `json:"units"`
}{Units: 0},
} {
encoded, err := common.Marshal(value)
require.NoError(t, err)
var expected any
require.NoError(t, common.Unmarshal(encoded, &expected))
assert.Equal(t, expected, jsonValue(value))
}
source := map[string]any{"items": []any{map[string]any{"label": "original"}}}
copy := jsonValue(source).(map[string]any)
copy["items"].([]any)[0].(map[string]any)["label"] = "changed"
assert.Equal(t, "original", source["items"].([]any)[0].(map[string]any)["label"])
}
func TestTaskSubmitHooksReceiveIndependentRequestSnapshots(t *testing.T) {
source := `
export const meta={apiVersion:1,key:"request-copy",name:"Request copy",version:"1.0.0",author:{name:"Test"},models:["copy"],fetchMode:"per_task",usageSchema:{units:{type:"number",unit:"count"}}};
export function buildSubmitRequest(ctx){ctx.requestBody.units=3;ctx.requestBody.nested.label="built";return {url:ctx.baseUrl+"/submit",body:ctx.requestBody};}
export function extractUsage(ctx){if(ctx.requestBody.nested.label!=="original")throw new Error("request changed");return {units:ctx.requestBody.units};}
export function parseSubmitResponse(){return {taskId:"copy"};}
export function buildQueryRequest(ctx){return {url:ctx.baseUrl+"/query"};}
export function parseTaskResult(){return {status:"SUCCESS"};}
`
plugin, err := pluginruntime.CompilePlugin(source, pluginruntime.Options{})
require.NoError(t, err)
request := map[string]any{"units": int64(2), "nested": map[string]any{"label": "original"}}
c, _ := gin.CreateTestContext(httptest.NewRecorder())
c.Request = httptest.NewRequest(http.MethodPost, "/submit", nil)
c.Set("task_request", request)
info := &relaycommon.RelayInfo{OriginModelName: "copy", ChannelMeta: &relaycommon.ChannelMeta{UpstreamModelName: "copy", ChannelBaseUrl: "https://provider.example"}, TaskRelayInfo: &relaycommon.TaskRelayInfo{}}
adaptor := New(plugin)
adaptor.Init(info)
require.Nil(t, adaptor.ValidateRequestAndSetAction(c, info))
facts, err := adaptor.ExtractUsageFactsValidated(c, info)
require.NoError(t, err)
assert.Equal(t, map[string]any{"units": float64(2)}, facts)
body, err := adaptor.BuildRequestBody(c, info)
require.NoError(t, err)
encoded, err := io.ReadAll(body)
require.NoError(t, err)
assert.JSONEq(t, `{"units":3,"nested":{"label":"built"}}`, string(encoded))
assert.Equal(t, int64(2), request["units"])
assert.Equal(t, "original", request["nested"].(map[string]any)["label"])
}
func TestTaskSubmitDeltaStreamContract(t *testing.T) {
source := strings.Replace(documentStreamPlugin, `submitResponseTypes:["json","sse"],`, `submitResponseTypes:["json","sse"],requiredCapabilities:["submit-sse-delta@1"],`, 1) + `
export function parseSubmitEventDelta(ctx,event,previous) {
if(previous && previous.document !== undefined) throw new Error("full result leaked into control state");
const chunk=JSON.parse(event.data);
let changes=chunk.changes;
if(chunk.largeResult) changes=[{op:"set",path:[],value:{document:"<".repeat(200000)}}];
if(chunk.resetAfterLarge) changes.push({op:"set",path:[],value:{document:"",units:0}});
const result={changes:changes,state:{events:(previous ? previous.events : 0)+1},done:chunk.complete===true};
if(chunk.largeControl) result.state={text:"x".repeat(65537)};
if(chunk.extra) result.extra=true;
return result;
}`
plugin, err := pluginruntime.CompilePlugin(source, pluginruntime.Options{})
require.NoError(t, err)
const first = `{"changes":[{"op":"set","path":[],"value":{"document":"hello","units":2}}]}`
const last = `{"changes":[{"op":"appendText","path":["document"],"value":"world"},{"op":"set","path":["units"],"value":0}],"complete":true}`
for _, tc := range []struct {
name string
frames []string
valid bool
}{
{"control state and zero usage", []string{first, last}, true},
{"oversized intermediate result", []string{`{"largeResult":true}`, last}, false},
{"oversized operation before reset", []string{`{"largeResult":true,"resetAfterLarge":true,"complete":true}`}, false},
{"control state limit", []string{first, `{"changes":[],"largeControl":true,"complete":true}`}, false},
{"missing changes", []string{`{"complete":true}`}, false},
{"extra result fields", []string{first, `{"changes":[],"extra":true,"complete":true}`}, false},
{"unfinished delta stream", []string{first}, false},
} {
t.Run(tc.name, func(t *testing.T) {
var stream strings.Builder
for _, frame := range tc.frames {
stream.WriteString("data: " + frame + "\n\n")
}
info := &relaycommon.RelayInfo{OriginModelName: "alias", ChannelMeta: &relaycommon.ChannelMeta{UpstreamModelName: "document", ChannelBaseUrl: "https://provider.example"}, TaskRelayInfo: &relaycommon.TaskRelayInfo{PublicTaskID: "public-document"}}
adaptor := New(plugin)
adaptor.Init(info)
c, _ := gin.CreateTestContext(httptest.NewRecorder())
c.Request = httptest.NewRequest(http.MethodPost, "/compile", nil)
c.Set("task_request", map[string]any{})
require.Nil(t, adaptor.ValidateRequestAndSetAction(c, info))
response := &http.Response{StatusCode: http.StatusOK, Header: http.Header{"Content-Type": {"text/event-stream"}}, Body: io.NopCloser(strings.NewReader(stream.String()))}
parsed, taskErr := adaptor.ParseResponse(c, response, info)
assert.False(t, c.Writer.Written())
if !tc.valid {
require.NotNil(t, taskErr)
assert.True(t, taskErr.NoRetry)
return
}
require.Nil(t, taskErr)
require.NotNil(t, parsed.Immediate)
assert.JSONEq(t, `{"document":"helloworld","units":0}`, string(parsed.TaskData))
assert.Equal(t, map[string]any{"units": float64(0)}, parsed.Immediate.UsageFacts)
})
}
}
func TestAlibabaSubmitDeltaDoesNotMutateControlState(t *testing.T) {
source, err := plugins.Source("alibaba")
require.NoError(t, err)
plugin, err := pluginruntime.CompilePlugin(source, pluginruntime.Options{})
require.NoError(t, err)
const before = `{"request_id":"old","usage":{"image_count":1},"output":{"finished":false,"choices":[{"message":{"role":"assistant","content":[{"text":"hello"}]}}]}}`
const control = `{"choices":[{"count":1,"lastText":true,"finishReason":""}],"hasUsage":true}`
var previous, initial any
require.NoError(t, common.UnmarshalJsonStr(control, &previous))
require.NoError(t, common.UnmarshalJsonStr(before, &initial))
accumulated := pluginruntime.NewJSONState(maxTaskPluginPersistedJSONBytes)
require.NoError(t, accumulated.Apply(t.Context(), []any{map[string]any{"op": "set", "path": []any{}, "value": initial}}))
value, err := plugin.Engine.Call(t.Context(), "parseSubmitEventDelta", map[string]any{}, map[string]any{
"event": "message", "data": `{"request_id":"new","usage":{"image_count":2},"output":{"choices":[{"finish_reason":"stop","message":{"content":[{"text":" world"},{"image":"https://cdn.example/image.png"}]}}]}}`,
}, previous)
require.NoError(t, err)
unchanged, err := common.Marshal(previous)
require.NoError(t, err)
assert.JSONEq(t, control, string(unchanged))
result := value.(map[string]any)
assert.Equal(t, true, result["done"])
require.NoError(t, accumulated.Apply(t.Context(), result["changes"]))
final, err := accumulated.Value()
require.NoError(t, err)
encoded, err := common.Marshal(final)
require.NoError(t, err)
assert.JSONEq(t, `{"request_id":"new","usage":{"image_count":2},"output":{"finished":true,"choices":[{"finish_reason":"stop","message":{"role":"assistant","content":[{"text":"hello world"},{"image":"https://cdn.example/image.png"}]}}]}}`, string(encoded))
encoded, err = common.Marshal(result["state"])
require.NoError(t, err)
assert.JSONEq(t, `{"choices":[{"count":2,"lastText":false,"finishReason":"stop"}],"hasUsage":true}`, string(encoded))
}
@@ -0,0 +1,190 @@
package jsplugin
import (
"context"
"fmt"
"io"
"net/http"
"net/http/httptest"
"runtime"
"strings"
"sync/atomic"
"testing"
"github.com/QuantumNous/new-api/common"
pluginruntime "github.com/QuantumNous/new-api/pkg/jsplugin"
"github.com/QuantumNous/new-api/plugins"
relaycommon "github.com/QuantumNous/new-api/relay/common"
"github.com/gin-gonic/gin"
"github.com/stretchr/testify/require"
)
func benchmarkAlibaba(b *testing.B) (*pluginruntime.LoadedPlugin, *atomic.Int64) {
b.Helper()
source, err := plugins.Source("alibaba")
require.NoError(b, err)
initializations := &atomic.Int64{}
plugin, err := pluginruntime.CompilePlugin(source+"\n"+`console.log("benchmark_runtime_initialized");`, pluginruntime.Options{
Log: func(message string) {
if strings.HasSuffix(message, "benchmark_runtime_initialized") {
initializations.Add(1)
}
},
})
require.NoError(b, err)
return plugin, initializations
}
func BenchmarkTaskPluginRuntime(b *testing.B) {
for _, name := range []string{"HasExport", "MissingExport", "HookCall", "ParallelHookCall"} {
b.Run(name, func(b *testing.B) {
plugin, initializations := benchmarkAlibaba(b)
ctx := context.Background()
input := map[string]any{"model": "wan2.5-i2v-preview", "upstreamModel": "wan2.5-i2v-preview", "requestBody": map[string]any{
"prompt": "animate the landscape", "image": "https://cdn.example/image.png", "duration": 5, "size": "720P",
}}
before := initializations.Load()
b.ReportAllocs()
if name == "ParallelHookCall" {
b.ResetTimer()
b.RunParallel(func(pb *testing.PB) {
for pb.Next() {
if _, err := plugin.Engine.Call(ctx, "extractUsage", input); err != nil {
b.Error(err)
return
}
}
})
} else {
for b.Loop() {
var err error
switch name {
case "HasExport":
_, err = plugin.Engine.HasExport(ctx, "extractUsage")
case "MissingExport":
_, err = plugin.Engine.HasExport(ctx, "absentHook")
default:
_, err = plugin.Engine.Call(ctx, "extractUsage", input)
}
if err != nil {
b.Fatal(err)
}
}
}
b.ReportMetric(float64(initializations.Load()-before)/float64(b.N), "runtime_inits/op")
})
}
}
// Run with -benchtime=50x: GC is excluded so this measures the first hook after
// reclamation, not the runtime.GC operation or an artificial throughput workload.
func BenchmarkTaskPluginAfterGC(b *testing.B) {
plugin, initializations := benchmarkAlibaba(b)
input := map[string]any{"model": "wan2.5-i2v-preview", "upstreamModel": "wan2.5-i2v-preview", "requestBody": map[string]any{
"prompt": "animate the landscape", "image": "https://cdn.example/image.png", "duration": 5, "size": "720P",
}}
before := initializations.Load()
b.ReportAllocs()
b.ResetTimer()
for range b.N {
b.StopTimer()
runtime.GC()
runtime.GC()
b.StartTimer()
if _, err := plugin.Engine.Call(b.Context(), "extractUsage", input); err != nil {
b.Fatal(err)
}
}
b.ReportMetric(float64(initializations.Load()-before)/float64(b.N), "runtime_inits/op")
}
func BenchmarkTaskPluginSubmit(b *testing.B) {
for _, size := range []int{4 << 10, 1 << 20, 8 << 20} {
b.Run(fmt.Sprintf("%dKiB", size>>10), func(b *testing.B) {
plugin, initializations := benchmarkAlibaba(b)
const modelName = "wan2.5-i2v-preview"
wire, err := common.Marshal(map[string]any{
"model": modelName, "duration": 5, "size": "720P",
"input": []any{map[string]any{"type": "input_text", "text": "animate the landscape"}, map[string]any{"type": "input_image", "image_url": "data:image/png;base64," + strings.Repeat("A", size)}},
})
require.NoError(b, err)
before := initializations.Load()
b.SetBytes(int64(len(wire)))
b.ReportAllocs()
for b.Loop() {
var body any
if err := common.Unmarshal(wire, &body); err != nil {
b.Fatal(err)
}
protocol := pluginruntime.ProtocolRequestContext{
RouteRequestContext: pluginruntime.RouteRequestContext{Path: "/v1/responses", Method: http.MethodPost, Body: map[string]any{"kind": "json", "value": body}},
Protocol: "openai_responses", Operation: "create", Model: modelName,
}
// The middleware decoder runs before channel selection; the adaptor
// decodes again after the executing provider has been pinned.
decoded, err := plugin.Engine.CallPath(b.Context(), "protocols", []string{"openai_responses", "decodeRequest"}, protocol.JSValue())
if err != nil {
b.Fatal(err)
}
c, _ := gin.CreateTestContext(httptest.NewRecorder())
c.Request = httptest.NewRequest(http.MethodPost, "/v1/responses", nil)
c.Set("task_request", decoded.(map[string]any)["requestBody"])
c.Set(pluginruntime.ContextKeyRouteRequest, protocol.RouteRequestContext)
c.Set(pluginruntime.ContextKeyProtocolRequest, protocol)
c.Set(pluginruntime.ContextKeyPinnedEndpoint, pluginruntime.PinnedEndpoint{Plugin: plugin, Protocol: protocol.Protocol, Model: modelName})
info := &relaycommon.RelayInfo{OriginModelName: modelName, ChannelMeta: &relaycommon.ChannelMeta{UpstreamModelName: modelName, ChannelBaseUrl: "https://provider.example"}, TaskRelayInfo: &relaycommon.TaskRelayInfo{PublicTaskID: "benchmark-task"}}
adaptor := New(plugin)
adaptor.Init(info)
if taskErr := adaptor.ValidateRequestAndSetAction(c, info); taskErr != nil {
b.Fatal(taskErr)
}
if _, err := adaptor.ExtractUsageFactsValidated(c, info); err != nil {
b.Fatal(err)
}
upstream, err := adaptor.BuildRequestBody(c, info)
if err != nil {
b.Fatal(err)
}
if _, err := io.Copy(io.Discard, upstream); err != nil {
b.Fatal(err)
}
}
b.ReportMetric(float64(initializations.Load()-before)/float64(b.N), "runtime_inits/op")
})
}
}
func BenchmarkTaskPluginSSE(b *testing.B) {
for _, tc := range []struct {
name string
events, textBytes int
}{{"small_state", 32, 48}, {"growing_state", 64, 8 << 10}} {
b.Run(tc.name, func(b *testing.B) {
plugin, initializations := benchmarkAlibaba(b)
var stream strings.Builder
for index := range tc.events {
choice := map[string]any{"message": map[string]any{"role": "assistant", "content": []any{map[string]any{"text": strings.Repeat("a", tc.textBytes)}}}}
if index == tc.events-1 {
choice["finish_reason"] = "stop"
}
frame, err := common.Marshal(map[string]any{"request_id": "benchmark-sse", "output": map[string]any{"choices": []any{choice}}, "usage": map[string]any{"image_count": 1}})
require.NoError(b, err)
stream.WriteString("data: ")
stream.Write(frame)
stream.WriteString("\n\n")
}
wire := stream.String()
input := map[string]any{"model": "wan2.6-image"}
before := initializations.Load()
b.SetBytes(int64(len(wire)))
b.ReportAllocs()
for b.Loop() {
response := &http.Response{Header: http.Header{"Content-Type": {"text/event-stream"}}, Body: io.NopCloser(strings.NewReader(wire))}
if _, err := New(plugin).readSubmitEvents(b.Context(), response, input); err != nil {
b.Fatal(err)
}
}
b.ReportMetric(float64(initializations.Load()-before)/float64(b.N), "runtime_inits/op")
})
}
}
@@ -0,0 +1,146 @@
package jsplugin
import (
"context"
"fmt"
"mime"
"net/http"
"slices"
"strings"
"time"
"github.com/QuantumNous/new-api/common"
"github.com/QuantumNous/new-api/constant"
pluginruntime "github.com/QuantumNous/new-api/pkg/jsplugin"
"github.com/QuantumNous/new-api/relay/helper"
)
// readSubmitEvents owns only SSE framing and bounded JSON state. The plugin
// interprets data, accumulates provider output and decides when it is complete.
// No response bytes are written to the client at this boundary.
func (a *TaskAdaptor) readSubmitEvents(parent context.Context, resp *http.Response, driverContext map[string]any) (any, error) {
mediaType, _, err := mime.ParseMediaType(resp.Header.Get("Content-Type"))
if err != nil || mediaType != "text/event-stream" {
return nil, fmt.Errorf("expected a text/event-stream submit response")
}
ctx, cancel := context.WithCancelCause(parent)
defer cancel(nil)
stopClose := context.AfterFunc(ctx, func() { _ = resp.Body.Close() })
defer stopClose()
idleTimeout := time.Duration(constant.StreamingTimeout) * time.Second
if idleTimeout <= 0 {
idleTimeout = 30 * time.Second
}
timer := time.AfterFunc(idleTimeout, func() { cancel(fmt.Errorf("task submit stream idle timeout")) })
defer timer.Stop()
scanner := helper.NewStreamScanner(resp.Body, maxTaskPluginPersistedJSONBytes+1)
hook := "parseSubmitEvent"
var accumulated *pluginruntime.JSONState
if slices.Contains(a.plugin.Meta.RequiredCapabilities, pluginruntime.CapabilitySubmitSSEDelta) {
hook = "parseSubmitEventDelta"
accumulated = pluginruntime.NewJSONState(maxTaskPluginPersistedJSONBytes)
}
var state any
var data []string
var eventName, lastID string
frameBytes := 0
firstLine := true
for scanner.Scan() {
if err := context.Cause(ctx); err != nil {
return nil, err
}
timer.Reset(idleTimeout)
line := scanner.Text()
if firstLine {
line = strings.TrimPrefix(line, "\ufeff")
firstLine = false
}
frameBytes += len(line) + 1
if frameBytes > maxTaskPluginPersistedJSONBytes {
return nil, fmt.Errorf("task submit SSE event exceeds size limit")
}
if line != "" {
field, value, _ := strings.Cut(line, ":")
value, _ = strings.CutPrefix(value, " ")
switch field {
case "data":
data = append(data, value)
case "event":
eventName = value
case "id":
if !strings.ContainsRune(value, '\x00') {
lastID = value
}
}
continue
}
frameBytes = 0
if len(data) == 0 {
eventName = ""
continue
}
if eventName == "" {
eventName = "message"
}
event := map[string]any{"event": eventName, "id": lastID, "data": strings.Join(data, "\n")}
value, err := a.plugin.Engine.Call(ctx, hook, driverContext, event, state)
if err != nil {
return nil, err
}
result, ok := value.(map[string]any)
if !ok {
return nil, fmt.Errorf("%s must return an object", hook)
}
nextState, hasState := result["state"]
done, hasDone := result["done"].(bool)
if !hasState || !hasDone {
return nil, fmt.Errorf("%s must return state and a boolean done", hook)
}
encoded, err := common.Marshal(nextState)
if err != nil {
return nil, fmt.Errorf("invalid submit stream state: %w", err)
}
if len(encoded) > maxTaskPluginPersistedJSONBytes {
return nil, fmt.Errorf("task submit stream state exceeds size limit")
}
if accumulated != nil {
if len(result) != 3 {
return nil, fmt.Errorf("parseSubmitEventDelta must return only changes, state and done")
}
// Control state is separate from the full result and is the only
// snapshot passed back to JS on the next event.
if len(encoded) > 64<<10 {
return nil, fmt.Errorf("task submit stream control state exceeds size limit")
}
if err = accumulated.Apply(ctx, result["changes"]); err != nil {
return nil, fmt.Errorf("invalid submit stream changes: %w", err)
}
}
// Keep the exact encoded-byte limit above. Plain JSON state can be
// isolated without parsing large strings again; codec-specific values
// (for example exported typed arrays) retain the old normalization.
var plainJSON bool
state, plainJSON = cloneJSONValue(nextState, 0)
if !plainJSON {
if err = common.Unmarshal(encoded, &state); err != nil {
return nil, err
}
}
if done {
if accumulated != nil {
return accumulated.Value()
}
return state, nil
}
data = nil
eventName = ""
}
if err := context.Cause(ctx); err != nil {
return nil, err
}
if err := scanner.Err(); err != nil {
return nil, fmt.Errorf("read task submit stream: %w", err)
}
return nil, fmt.Errorf("task submit stream ended before the plugin reported completion")
}
@@ -14,6 +14,7 @@ import (
relaycommon "github.com/QuantumNous/new-api/relay/common"
relayconstant "github.com/QuantumNous/new-api/relay/constant"
"github.com/QuantumNous/new-api/relaykit/dto"
"github.com/QuantumNous/new-api/relaykit/relayconvert"
relaytypes "github.com/QuantumNous/new-api/relaykit/types"
hosttypes "github.com/QuantumNous/new-api/types"
"github.com/gin-gonic/gin"
@@ -21,6 +22,59 @@ import (
"github.com/stretchr/testify/require"
)
func TestPrepareResponsesRequestRetainsConvertedAdaptor(t *testing.T) {
type capturedRequest struct {
path string
body []byte
}
captured := make(chan capturedRequest, 1)
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
body, err := io.ReadAll(r.Body)
if err != nil {
http.Error(w, err.Error(), http.StatusBadRequest)
return
}
captured <- capturedRequest{path: r.URL.Path, body: body}
w.Header().Set("Content-Type", "application/json")
_, _ = w.Write([]byte(`{"id":"chatcmpl_1","object":"chat.completion","model":"gpt-4o","choices":[{"index":0,"message":{"role":"assistant","content":"ok"},"finish_reason":"stop"}],"usage":{"prompt_tokens":3,"completion_tokens":2,"total_tokens":5}}`))
}))
defer server.Close()
c, _ := gin.CreateTestContext(httptest.NewRecorder())
c.Request = httptest.NewRequest(http.MethodPost, "/v1/responses", nil)
c.Request.Header.Set("Content-Type", "application/json")
common.SetContextKey(c, constant.ContextKeyOriginalModel, "gpt-4o")
common.SetContextKey(c, constant.ContextKeyChannelType, constant.ChannelTypeAdvancedCustom)
common.SetContextKey(c, constant.ContextKeyChannelBaseUrl, server.URL)
common.SetContextKey(c, constant.ContextKeyChannelKey, "test-key")
common.SetContextKey(c, constant.ContextKeyChannelOtherSetting, dto.ChannelOtherSettings{
AdvancedCustom: &dto.AdvancedCustomConfig{Routes: []dto.AdvancedCustomRoute{{
IncomingPath: "/v1/responses", UpstreamPath: "/v1/chat/completions", Converter: relayconvert.ConverterOpenAIResponsesToOpenAIChat,
}}},
})
request := &dto.OpenAIResponsesRequest{Model: "gpt-4o", Input: []byte(`"hello"`)}
info := relaycommon.GenRelayInfoResponses(c, request)
adaptor, body, closer, apiErr := PrepareResponsesRequest(c, info, request)
require.Nil(t, apiErr)
defer closer.Close()
response, err := adaptor.DoRequest(c, info, body)
require.NoError(t, err)
httpResponse, ok := response.(*http.Response)
require.True(t, ok)
usage, apiErr := adaptor.DoResponse(c, httpResponse, info)
require.Nil(t, apiErr)
require.IsType(t, &dto.Usage{}, usage)
assert.Equal(t, 5, usage.(*dto.Usage).TotalTokens)
upstream := <-captured
assert.Equal(t, "/v1/chat/completions", upstream.path)
var upstreamRequest dto.GeneralOpenAIRequest
require.NoError(t, common.Unmarshal(upstream.body, &upstreamRequest))
require.Len(t, upstreamRequest.Messages, 1)
assert.Equal(t, "hello", upstreamRequest.Messages[0].StringContent())
assert.Equal(t, []relaytypes.RelayFormat{relaytypes.RelayFormatOpenAIResponses, relaytypes.RelayFormatOpenAI}, info.RequestConversionChain)
}
func TestIsResponsesEventStreamContentType(t *testing.T) {
tests := []struct {
name string
+68 -5
View File
@@ -99,6 +99,11 @@ type RelayInfo struct {
UsePrice bool
RelayMode int
OriginModelName string
ResponseModel *ResponseModel
ResponseGuard *ResponseModelGuard
AstraDetectionDisabled bool
AstraAutoSwitch bool
UpstreamRequests int
// BillingModelName is the pricing identity for this request. It is kept
// separate from OriginModelName and UpstreamModelName so virtual pricing
@@ -138,7 +143,11 @@ type RelayInfo struct {
ForcePreConsume bool
// Billing 是计费会话,封装了预扣费/结算/退款的统一生命周期。
// 初始免费组可为 nil;若 auto 重试切换到付费组,会在发送前创建。
Billing BillingSettler
Billing BillingSettler
DurableAttemptBilling bool
DurableAttemptID int64
AttemptDetached bool
GuardDeadline time.Time
// BillingSource indicates whether this request is billed from wallet quota or subscription.
// "" or "wallet" => wallet; "subscription" => subscription
BillingSource string
@@ -176,6 +185,11 @@ type RelayInfo struct {
// and again before settlement. Non-nil only when billing mode is "tiered_expr".
TieredBillingSnapshot *billingexpr.BillingSnapshot
BillingRequestInput *billingexpr.RequestInput
BillingImageCount *int
// ImageRequestCount is the effective quantity sent on the current attempt;
// ImageQuotaBeforeGroup is the frozen legacy estimate before request ratios.
ImageRequestCount int
ImageQuotaBeforeGroup float64
Request dto.Request
@@ -187,6 +201,10 @@ type RelayInfo struct {
FinalRequestRelayFormat types.RelayFormat
StreamStatus *StreamStatus
// PerformanceOutputTokens is captured by settlement and sampled once at
// the request boundary, independently of billing success or failure.
PerformanceOutputTokens int64
PerformanceBusinessRejection bool
// convOptions caches the converter settings snapshot (see ConvOptions).
convOptions *convmeta.Options
@@ -204,7 +222,36 @@ type RelayInfo struct {
*TaskRelayInfo
}
// UpdateImageCount replaces the billable quantity without changing the frozen
// request parameters or multiplying the legacy and expression prices together.
func (info *RelayInfo) UpdateImageCount(count int64) {
if info == nil || count <= 0 || count > int64(dto.MaxImageN) {
return
}
if info.PriceData.UsePrice {
info.PriceData.AddOtherRatio("n", float64(count))
}
if info.TieredBillingSnapshot != nil && info.TieredBillingSnapshot.EstimatedImageCount != nil {
n := int(count)
info.BillingImageCount = &n
}
}
func (info *RelayInfo) RequestedImageCount() int {
if info.ImageRequestCount > 0 {
return info.ImageRequestCount
}
if info.TieredBillingSnapshot != nil && info.TieredBillingSnapshot.EstimatedImageCount != nil {
return *info.TieredBillingSnapshot.EstimatedImageCount
}
if count, ok := info.PriceData.OtherRatios()["n"]; ok && count >= 1 && count <= dto.MaxImageN {
return int(count)
}
return 1
}
func (info *RelayInfo) InitChannelMeta(c *gin.Context) {
info.ResponseModel = nil
info.FinalRequestRelayFormat = ""
info.RequestConversionChain = nil
info.InitRequestConversionChain()
@@ -252,6 +299,15 @@ func (info *RelayInfo) InitChannelMeta(c *gin.Context) {
channelMeta.ChannelOtherSettings = channelOtherSettings
}
if channelType == constant.ChannelTypeAdvancedCustom &&
!channelMeta.ChannelSetting.PassThroughBodyEnabled &&
c.Request != nil && c.Request.URL != nil {
route, matched := channelMeta.ChannelOtherSettings.AdvancedCustom.MatchPathForModel(c.Request.URL.Path, info.OriginModelName)
if matched && route.PassThroughBodyEnabled {
channelMeta.ChannelSetting.PassThroughBodyEnabled = true
}
}
if streamSupportedChannels[channelMeta.ChannelType] {
channelMeta.SupportStreamOptions = true
}
@@ -377,6 +433,8 @@ var streamSupportedChannels = map[int]bool{
constant.ChannelTypeAdvancedCustom: true,
constant.ChannelTypeSub2API: true,
constant.ChannelTypeNewAPI: true,
constant.ChannelTypeVLLM: true,
constant.ChannelTypeSGLang: true,
constant.ChannelTypeTencent: true,
}
@@ -498,6 +556,9 @@ func reasoningEffortFromRequest(request dto.Request) string {
if req != nil && req.GenerationConfig.ThinkingConfig != nil {
config := req.GenerationConfig.ThinkingConfig
effort = config.ThinkingLevel
if canonical, err := kitreasoning.ParseEffort(effort); err == nil {
effort = string(canonical)
}
if effort == "" && config.ThinkingBudget != nil {
effort = string(kitreasoning.EffortFromBudget(*config.ThinkingBudget))
}
@@ -551,10 +612,12 @@ func genBaseRelayInfo(c *gin.Context, request dto.Request) *RelayInfo {
OriginModelName: originModelName,
TokenId: common.GetContextKeyInt(c, constant.ContextKeyTokenId),
TokenKey: common.GetContextKeyString(c, constant.ContextKeyTokenKey),
TokenUnlimited: common.GetContextKeyBool(c, constant.ContextKeyTokenUnlimited),
TokenGroup: tokenGroup,
TokenId: common.GetContextKeyInt(c, constant.ContextKeyTokenId),
TokenKey: common.GetContextKeyString(c, constant.ContextKeyTokenKey),
TokenUnlimited: common.GetContextKeyBool(c, constant.ContextKeyTokenUnlimited),
TokenGroup: tokenGroup,
AstraAutoSwitch: common.GetContextKeyBool(c, constant.ContextKeyTokenAstraAutoSwitch),
AstraDetectionDisabled: common.GetContextKeyBool(c, constant.ContextKeyTokenAstraDetectionDisabled),
isFirstResponse: true,
RelayMode: relayconstant.Path2RelayMode(c.Request.URL.Path),
+98
View File
@@ -6,6 +6,8 @@ import (
"net/http/httptest"
"testing"
"github.com/QuantumNous/new-api/common"
"github.com/QuantumNous/new-api/constant"
"github.com/QuantumNous/new-api/relaykit/dto"
"github.com/QuantumNous/new-api/relaykit/relayconvert"
"github.com/QuantumNous/new-api/relaykit/relayconvert/convmeta"
@@ -147,6 +149,24 @@ func TestGenRelayInfoCapturesRequestReasoningEffort(t *testing.T) {
}},
expected: "low",
},
{
name: "Gemini uppercase enum thinking level is canonicalized",
path: "/v1beta/models/gemini-3.7-flash:generateContent",
relayFormat: types.RelayFormatGemini,
request: &dto.GeminiChatRequest{GenerationConfig: dto.GeminiChatGenerationConfig{
ThinkingConfig: &dto.GeminiThinkingConfig{ThinkingLevel: " MEDIUM "},
}},
expected: "medium",
},
{
name: "Gemini unknown thinking level is recorded as sent",
path: "/v1beta/models/gemini-3.7-flash:generateContent",
relayFormat: types.RelayFormatGemini,
request: &dto.GeminiChatRequest{GenerationConfig: dto.GeminiChatGenerationConfig{
ThinkingConfig: &dto.GeminiThinkingConfig{ThinkingLevel: "ULTRA"},
}},
expected: "ULTRA",
},
}
for _, tt := range tests {
@@ -195,6 +215,84 @@ func TestInitChannelMetaRestoresRequestReasoningEffortForRetry(t *testing.T) {
assert.Equal(t, "max", info.ReasoningEffort)
}
func TestInitChannelMetaAppliesAdvancedCustomRoutePassThrough(t *testing.T) {
gin.SetMode(gin.TestMode)
advancedCustom := &dto.AdvancedCustomConfig{
Routes: []dto.AdvancedCustomRoute{
{
IncomingPath: "/v1/chat/completions",
UpstreamPath: "/v1/chat/completions",
Models: []string{"native-model"},
PassThroughBodyEnabled: true,
},
{
IncomingPath: "/v1/chat/completions",
UpstreamPath: "/v1/messages",
Converter: relayconvert.ConverterOpenAIChatToClaudeMessages,
},
},
}
require.NoError(t, advancedCustom.Validate())
tests := []struct {
name string
channelType int
channelSetting dto.ChannelSettings
model string
wantPassThrough bool
wantEffort string
}{
{
name: "matched route enables pass-through and drops reasoning effort",
channelType: constant.ChannelTypeAdvancedCustom,
model: "native-model",
wantPassThrough: true,
wantEffort: "",
},
{
name: "fallback converter route keeps conversion",
channelType: constant.ChannelTypeAdvancedCustom,
model: "other-model",
wantPassThrough: false,
wantEffort: "high",
},
{
name: "channel-level pass-through still applies to converter route",
channelType: constant.ChannelTypeAdvancedCustom,
channelSetting: dto.ChannelSettings{PassThroughBodyEnabled: true},
model: "other-model",
wantPassThrough: true,
wantEffort: "",
},
{
name: "route flag is ignored for other channel types",
channelType: constant.ChannelTypeOpenAI,
model: "native-model",
wantPassThrough: false,
wantEffort: "high",
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
ctx, _ := gin.CreateTestContext(httptest.NewRecorder())
ctx.Request = httptest.NewRequest("POST", "/v1/chat/completions", nil)
ctx.Set("original_model", tt.model)
common.SetContextKey(ctx, constant.ContextKeyChannelType, tt.channelType)
common.SetContextKey(ctx, constant.ContextKeyChannelSetting, tt.channelSetting)
common.SetContextKey(ctx, constant.ContextKeyChannelOtherSetting, dto.ChannelOtherSettings{AdvancedCustom: advancedCustom})
info, err := GenRelayInfo(ctx, types.RelayFormatOpenAI, &dto.GeneralOpenAIRequest{Model: tt.model, ReasoningEffort: "high"}, nil)
require.NoError(t, err)
info.InitChannelMeta(ctx)
assert.Equal(t, tt.wantPassThrough, info.ChannelSetting.PassThroughBodyEnabled)
assert.Equal(t, tt.wantEffort, info.ReasoningEffort)
})
}
}
func TestInitChannelMetaResetsPerAttemptStreamStateAndPreservesRequestState(t *testing.T) {
gin.SetMode(gin.TestMode)
ctx, _ := gin.CreateTestContext(httptest.NewRecorder())
+48
View File
@@ -0,0 +1,48 @@
package common
import "strings"
// ResponseModel records upstream declarations before response conversion. It is
// diagnostic only: it must never change routing, pricing, or downstream output.
type ResponseModel struct {
RequestedModel string `json:"requested_model"`
UpstreamModel string `json:"upstream_model"`
ReturnedModel string `json:"returned_model"`
Mismatch bool `json:"mismatch"`
}
// ObserveResponseModel retains the first differing model for inspection, with
// mismatches taking priority over prefix or case-only differences. A later
// matching or empty event cannot erase it. Only observe upstream declarations,
// never models synthesized by a response converter.
func (info *RelayInfo) ObserveResponseModel(model string) {
if info == nil || strings.TrimSpace(model) == "" || (info.OriginModelName == "gpt-6-astra" && info.AstraDetectionDisabled) {
return
}
if info.ResponseModel == nil {
info.ResponseModel = &ResponseModel{
RequestedModel: info.OriginModelName,
UpstreamModel: info.GetUpstreamModelName(),
}
}
observation := info.ResponseModel
if observation.Mismatch {
return
}
mismatch := true
for _, expected := range []string{observation.RequestedModel, observation.UpstreamModel} {
if expected != "" && (strings.HasPrefix(model, expected) || strings.EqualFold(model, expected)) {
mismatch = false
break
}
}
if observation.RequestedModel == "gpt-6-astra" {
mismatch = strings.TrimSpace(model) != "gpt-6-astra"
}
if !mismatch && observation.ReturnedModel != "" &&
observation.ReturnedModel != observation.RequestedModel && observation.ReturnedModel != observation.UpstreamModel {
return
}
observation.ReturnedModel = model
observation.Mismatch = mismatch
}
+358
View File
@@ -0,0 +1,358 @@
package common
import (
"bufio"
"bytes"
"context"
"errors"
"io"
"net/http"
"strings"
"sync"
"time"
appcommon "github.com/QuantumNous/new-api/common"
"github.com/QuantumNous/new-api/constant"
"github.com/QuantumNous/new-api/relaykit/types"
)
const ResponseGuardMaxAttempts = 3
const ResponseGuardErrorCode types.ErrorCode = "response_model_rejected"
// The guard is request-scoped, unlike channel metadata which is reset on retry.
type ResponseModelGuard struct {
Attempts []*ResponseGuardAttempt `json:"attempts"`
Deadline time.Time `json:"-"`
}
type ResponseGuardUsage struct {
Input int `json:"input_tokens"`
Output int `json:"output_tokens"`
Total int `json:"total_tokens"`
}
type ResponseGuardAttempt struct {
OnEvent func([]byte, bool) `json:"-"`
OnClose func() `json:"-"`
closeOnce sync.Once
Drain func() `json:"-"`
Number int `json:"number"`
ChannelID int `json:"channel_id"`
Expected string `json:"expected_model"`
Returned string `json:"returned_model,omitempty"`
Reason string `json:"reason"`
ElapsedMS int64 `json:"elapsed_ms"`
Usage *ResponseGuardUsage `json:"usage,omitempty"`
Delivered bool `json:"delivered"`
Status int `json:"status,omitempty"`
ResponseID string `json:"-"`
Sequence int `json:"-"`
started time.Time
timeout time.Duration
confirmed bool
err error
cancel context.CancelFunc
timer *time.Timer
stopClient func() bool
}
func (g *ResponseModelGuard) Last() *ResponseGuardAttempt {
if g == nil || len(g.Attempts) == 0 {
return nil
}
return g.Attempts[len(g.Attempts)-1]
}
func (g *ResponseModelGuard) Begin(ctx context.Context, number, channel int) (context.Context, *ResponseGuardAttempt) {
upstream, cancel := context.WithCancel(context.WithoutCancel(ctx))
a := &ResponseGuardAttempt{Number: number, ChannelID: channel, Expected: "gpt-6-astra", Reason: "pending", started: time.Now(), cancel: cancel}
// Model confirmation must allow the same initial wait as an ordinary stream.
// A separate 10-second cutoff rejects valid slow upstream responses.
a.timeout = time.Duration(constant.StreamingTimeout) * time.Second
if a.timeout <= 0 {
a.timeout = 300 * time.Second
}
if !g.Deadline.IsZero() {
a.timeout = max(time.Until(g.Deadline), time.Nanosecond)
}
a.timer = time.AfterFunc(a.timeout, cancel)
a.stopClient = context.AfterFunc(ctx, cancel)
g.Attempts = append(g.Attempts, a)
return upstream, a
}
func (a *ResponseGuardAttempt) Finish(delivered bool) {
a.ElapsedMS = time.Since(a.started).Milliseconds()
a.Delivered = delivered
}
func (a *ResponseGuardAttempt) Close() {
a.closeOnce.Do(func() {
if a.OnClose != nil {
a.OnClose()
}
a.timer.Stop()
a.stopClient()
a.cancel()
})
}
func (a *ResponseGuardAttempt) Reject(reason string) error {
if a.err == nil {
a.Reason = reason
a.err = errors.New("response model guard: " + reason)
}
return a.err
}
func (a *ResponseGuardAttempt) APIError() *types.NewAPIError {
if a == nil || a.err == nil {
return nil
}
return types.NewErrorWithStatusCode(errors.New(a.PublicErrorMessage()), ResponseGuardErrorCode, http.StatusConflict, types.ErrOptionWithSkipRetry())
}
func (a *ResponseGuardAttempt) PublicErrorMessage() string {
if a != nil && a.Reason == "model_mismatch" && a.Returned != "" {
return "Requested model " + a.Expected + ", but the upstream response declared " + a.Returned + ". No verified result could be delivered. Usage from each upstream attempt is charged independently, including rejected results."
}
if a != nil && a.Expected != "" {
return "Could not verify that the upstream response used " + a.Expected + ". Usage from each upstream attempt is charged independently, including rejected results."
}
return "Could not verify the upstream response model."
}
func (g *ResponseModelGuard) APIError() *types.NewAPIError {
if g == nil {
return nil
}
for i := len(g.Attempts) - 1; i >= 0; i-- {
attempt := g.Attempts[i]
if attempt != nil && attempt.Reason == "model_mismatch" && attempt.Returned != "" {
return attempt.APIError()
}
}
last := g.Last()
if last == nil {
return nil
}
return last.APIError()
}
// Only read the provider's declaration and numeric usage, never generated text.
func (a *ResponseGuardAttempt) observe(payload []byte, stream bool) error {
if a.OnEvent != nil {
a.OnEvent(payload, stream)
}
var event struct {
Model string `json:"model"`
ID string `json:"id"`
Sequence int `json:"sequence_number"`
Usage *ResponseGuardUsage `json:"usage"`
Response *struct {
Model string `json:"model"`
ID string `json:"id"`
Usage *ResponseGuardUsage `json:"usage"`
} `json:"response"`
}
if err := appcommon.Unmarshal(payload, &event); err != nil {
return a.Reject("invalid_response")
}
model, id, usage := event.Model, event.ID, event.Usage
if stream {
model, id, usage = "", "", nil
if event.Response != nil {
model, id, usage = event.Response.Model, event.Response.ID, event.Response.Usage
}
}
if usage != nil {
a.Usage = usage
}
a.Sequence = max(a.Sequence, event.Sequence)
model = strings.TrimSpace(model)
if model == "" {
return nil
}
// Bound diagnostics from an untrusted upstream.
a.Returned = string([]rune(model)[:min(len([]rune(model)), 128)])
if model != a.Expected {
return a.Reject("model_mismatch")
}
if id != "" {
a.ResponseID = string([]rune(id)[:min(len([]rune(id)), 256)])
}
if !a.confirmed {
if !a.timer.Stop() {
return a.Reject("confirmation_timeout")
}
if !a.stopClient() {
return a.Reject("client_cancelled")
}
a.confirmed = true
a.Reason = "accepted"
}
return nil
}
// Confirm buffers only the prefix until a provider model is known. No caller
// may commit headers or start a downstream pinger before it returns success.
func (a *ResponseGuardAttempt) Confirm(resp *http.Response, ctx context.Context) error {
a.Status = resp.StatusCode
if resp.StatusCode != http.StatusOK {
a.timer.Stop()
a.stopClient()
a.Reason = "http_error"
resp.Body = &guardCloseBody{Reader: resp.Body, body: resp.Body, close: a.Close}
return nil
}
if !strings.HasPrefix(strings.ToLower(resp.Header.Get("Content-Type")), "text/event-stream") {
body, err := io.ReadAll(io.LimitReader(resp.Body, (1<<20)+1))
if err != nil {
return a.readError(ctx, err)
}
if len(body) > 1<<20 {
return a.Reject("confirmation_buffer_limit")
}
if err := a.observe(body, false); err != nil {
return err
}
if !a.confirmed {
return a.Reject("missing_model")
}
_ = resp.Body.Close()
resp.Body = &guardCloseBody{Reader: bytes.NewReader(body), body: io.NopCloser(bytes.NewReader(nil)), close: a.Close}
return nil
}
b := &guardStreamBody{reader: bufio.NewReader(resp.Body), body: resp.Body, attempt: a}
a.Drain = func() {
b.draining = true
for {
if _, err := b.readFrame(8 << 20); err != nil {
return
}
}
}
var prefix bytes.Buffer
for !a.confirmed {
frame, err := b.readFrame((1 << 20) - prefix.Len())
if err != nil {
return a.readError(ctx, err)
}
prefix.Write(frame)
}
b.pending = prefix.Bytes()
resp.Body = b
return nil
}
func (a *ResponseGuardAttempt) readError(ctx context.Context, err error) error {
if a.err != nil {
return a.err
}
if ctx.Err() != nil {
return a.Reject("client_cancelled")
}
if time.Since(a.started) >= a.timeout {
return a.Reject("confirmation_timeout")
}
if errors.Is(err, io.EOF) {
return a.Reject("missing_model")
}
return a.Reject("response_read_error")
}
type guardCloseBody struct {
io.Reader
body io.Closer
close func()
once sync.Once
}
func (b *guardCloseBody) Close() error { b.once.Do(b.close); return b.body.Close() }
// Validate complete SSE records before exposing their bytes to the existing
// scanner, so a late wrong-model event cannot race past the check.
type guardStreamBody struct {
draining bool
reader *bufio.Reader
body io.ReadCloser
attempt *ResponseGuardAttempt
pending []byte
once sync.Once
}
func (b *guardStreamBody) Close() error { b.once.Do(b.attempt.Close); return b.body.Close() }
func (b *guardStreamBody) Read(p []byte) (int, error) {
if len(p) == 0 {
return 0, nil
}
if len(b.pending) == 0 {
frame, err := b.readFrame(128 << 20)
if err != nil {
return 0, err
}
b.pending = frame
}
n := copy(p, b.pending)
b.pending = b.pending[n:]
return n, nil
}
func (b *guardStreamBody) readFrame(limit int) ([]byte, error) {
var frame, line, data bytes.Buffer
for {
part, err := b.reader.ReadSlice('\n')
if frame.Len()+len(part) > limit {
if b.draining {
return nil, io.ErrShortBuffer
}
return nil, b.attempt.Reject("confirmation_buffer_limit")
}
frame.Write(part)
line.Write(part)
if errors.Is(err, bufio.ErrBufferFull) {
continue
}
if err != nil && !errors.Is(err, io.EOF) {
return nil, err
}
text := strings.TrimSuffix(strings.TrimSuffix(line.String(), "\n"), "\r")
line.Reset()
if value, ok := strings.CutPrefix(text, "data:"); ok {
if data.Len() > 0 {
data.WriteByte('\n')
}
data.WriteString(strings.TrimPrefix(value, " "))
}
if text == "" || err != nil {
if frame.Len() == 0 {
return nil, io.EOF
}
payload := bytes.TrimSpace(data.Bytes())
if len(payload) > 0 && !bytes.Equal(payload, []byte("[DONE]")) {
if b.draining {
if b.attempt.OnEvent != nil {
b.attempt.OnEvent(payload, true)
}
} else if e := b.attempt.observe(payload, true); e != nil {
return nil, e
}
}
if b.draining && bytes.Equal(payload, []byte("[DONE]")) {
return nil, io.EOF
}
if !b.attempt.confirmed && bytes.Equal(payload, []byte("[DONE]")) {
return nil, b.attempt.Reject("missing_model")
}
return frame.Bytes(), nil
}
}
}
// Detach keeps only this upstream attempt alive while a different channel is
// tried; its numeric usage is collected for at most one minute.
func (a *ResponseGuardAttempt) Detach() {
a.timer.Stop()
a.stopClient()
a.timer = time.AfterFunc(60*time.Second, a.cancel)
}
+159 -10
View File
@@ -21,6 +21,18 @@ const (
StreamEndReasonPingFail StreamEndReason = "ping_fail"
)
// ResponseOutcome is the protocol-level result of one response, independent of
// how the transport ended. Adaptors mark it from the events they already parse.
type ResponseOutcome string
const (
ResponseOutcomeUnknown ResponseOutcome = ""
ResponseOutcomeCompleted ResponseOutcome = "completed"
ResponseOutcomeFailed ResponseOutcome = "failed"
ResponseOutcomeIncomplete ResponseOutcome = "incomplete"
ResponseOutcomeCancelled ResponseOutcome = "cancelled"
)
const maxStreamErrorEntries = 20
type StreamErrorEntry struct {
@@ -31,11 +43,32 @@ type StreamErrorEntry struct {
type StreamStatus struct {
EndReason StreamEndReason
EndError error
endOnce sync.Once
endMu sync.Mutex
endSet bool
mu sync.Mutex
Errors []StreamErrorEntry
ErrorCount int
response ResponseOutcome
errorCode string
errorType string
errorStatus int
incompleteReason string
expectsTerminal bool
}
// StreamOutcome holds classification facts only; upstream messages never
// enter it because they may contain credentials or request content.
type StreamOutcome struct {
EndReason StreamEndReason
HasErrors bool
ExpectsTerminal bool
Response ResponseOutcome
ErrorCode string
ErrorType string
ErrorStatus int
IncompleteReason string
}
func NewStreamStatus() *StreamStatus {
@@ -46,10 +79,30 @@ func (s *StreamStatus) SetEndReason(reason StreamEndReason, err error) {
if s == nil {
return
}
s.endOnce.Do(func() {
s.EndReason = reason
s.endMu.Lock()
defer s.endMu.Unlock()
if s.endSet {
return
}
s.EndReason = reason
s.EndError = err
s.endSet = true
}
// SetClientGone preserves the client disconnect as the diagnostic reason when
// it races with a normal upstream terminal event. Fatal and handler-stop reasons
// remain authoritative.
func (s *StreamStatus) SetClientGone(err error) {
if s == nil {
return
}
s.endMu.Lock()
defer s.endMu.Unlock()
if !s.endSet || s.EndReason == StreamEndReasonDone || s.EndReason == StreamEndReasonEOF {
s.EndReason = StreamEndReasonClientGone
s.EndError = err
})
s.endSet = true
}
}
func (s *StreamStatus) RecordError(msg string) {
@@ -67,6 +120,95 @@ func (s *StreamStatus) RecordError(msg string) {
}
}
// RequireTerminal declares that the protocol always ends with an explicit
// terminal event, so a stream that ends without one was cut short.
func (s *StreamStatus) RequireTerminal() {
if s == nil {
return
}
s.mu.Lock()
defer s.mu.Unlock()
s.expectsTerminal = true
}
// MarkCompleted, MarkIncomplete and MarkCancelled keep the first terminal seen;
// MarkFailed always wins because an error after completion is still a failure.
func (s *StreamStatus) MarkCompleted() {
s.markTerminal(ResponseOutcomeCompleted, "")
}
func (s *StreamStatus) MarkIncomplete(reason string) {
s.markTerminal(ResponseOutcomeIncomplete, reason)
}
func (s *StreamStatus) MarkCancelled() {
s.markTerminal(ResponseOutcomeCancelled, "")
}
func (s *StreamStatus) markTerminal(outcome ResponseOutcome, incompleteReason string) {
if s == nil {
return
}
s.mu.Lock()
defer s.mu.Unlock()
if s.response != ResponseOutcomeUnknown {
return
}
s.response = outcome
s.incompleteReason = incompleteReason
}
// MarkFailed records a protocol failure. Empty details never erase details
// recorded earlier, so a bare error envelope keeps the structured error.
func (s *StreamStatus) MarkFailed(code, errorType string, status int) {
if s == nil {
return
}
s.mu.Lock()
defer s.mu.Unlock()
s.response = ResponseOutcomeFailed
if code != "" {
s.errorCode = code
}
if errorType != "" {
s.errorType = errorType
}
if status != 0 {
s.errorStatus = status
}
}
func (s *StreamStatus) ResponseOutcome() string {
if s == nil {
return ""
}
s.mu.Lock()
defer s.mu.Unlock()
return string(s.response)
}
func (s *StreamStatus) ResponseFailed() bool {
return s.ResponseOutcome() == string(ResponseOutcomeFailed)
}
func (s *StreamStatus) OutcomeSnapshot() StreamOutcome {
if s == nil {
return StreamOutcome{}
}
s.mu.Lock()
defer s.mu.Unlock()
return StreamOutcome{
EndReason: s.EndReason,
HasErrors: s.ErrorCount > 0,
ExpectsTerminal: s.expectsTerminal,
Response: s.response,
ErrorCode: s.errorCode,
ErrorType: s.errorType,
ErrorStatus: s.errorStatus,
IncompleteReason: s.incompleteReason,
}
}
func (s *StreamStatus) HasErrors() bool {
if s == nil {
return false
@@ -89,9 +231,12 @@ func (s *StreamStatus) IsNormalEnd() bool {
if s == nil {
return true
}
return s.EndReason == StreamEndReasonDone ||
s.EndReason == StreamEndReasonEOF ||
s.EndReason == StreamEndReasonHandlerStop
s.endMu.Lock()
reason := s.EndReason
s.endMu.Unlock()
return reason == StreamEndReasonDone ||
reason == StreamEndReasonEOF ||
reason == StreamEndReasonHandlerStop
}
func (s *StreamStatus) Summary() string {
@@ -99,9 +244,13 @@ func (s *StreamStatus) Summary() string {
return "StreamStatus<nil>"
}
b := &strings.Builder{}
fmt.Fprintf(b, "reason=%s", s.EndReason)
if s.EndError != nil {
fmt.Fprintf(b, " end_error=%q", s.EndError.Error())
s.endMu.Lock()
reason := s.EndReason
endError := s.EndError
s.endMu.Unlock()
fmt.Fprintf(b, "reason=%s", reason)
if endError != nil {
fmt.Fprintf(b, " end_error=%q", endError.Error())
}
s.mu.Lock()
if s.ErrorCount > 0 {
+11
View File
@@ -20,6 +20,17 @@ func TestStreamStatus_SetEndReason_FirstWins(t *testing.T) {
assert.Nil(t, s.EndError)
}
func TestStreamStatus_SetClientGoneOverridesNormalTerminal(t *testing.T) {
s := NewStreamStatus()
s.SetEndReason(StreamEndReasonDone, nil)
s.SetClientGone(fmt.Errorf("context canceled"))
assert.Equal(t, StreamEndReasonClientGone, s.EndReason)
assert.EqualError(t, s.EndError, "context canceled")
s.SetEndReason(StreamEndReasonHandlerStop, fmt.Errorf("fatal"))
assert.Equal(t, StreamEndReasonClientGone, s.EndReason)
}
func TestStreamStatus_SetEndReason_WithError(t *testing.T) {
t.Parallel()
s := NewStreamStatus()
+203
View File
@@ -1,22 +1,138 @@
package relay
import (
"errors"
"io"
"net/http"
"net/http/httptest"
"strings"
"testing"
"github.com/QuantumNous/new-api/common"
"github.com/QuantumNous/new-api/constant"
"github.com/QuantumNous/new-api/pkg/billingexpr"
relaycommon "github.com/QuantumNous/new-api/relay/common"
relayconstant "github.com/QuantumNous/new-api/relay/constant"
"github.com/QuantumNous/new-api/relay/helper"
"github.com/QuantumNous/new-api/relaykit/dto"
kitreasoning "github.com/QuantumNous/new-api/relaykit/relayconvert/reasoning"
"github.com/QuantumNous/new-api/relaykit/types"
"github.com/QuantumNous/new-api/service"
"github.com/QuantumNous/new-api/setting/model_setting"
hosttypes "github.com/QuantumNous/new-api/types"
"github.com/gin-gonic/gin"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
"github.com/tidwall/gjson"
)
type imageReservation struct {
held, limit int
}
func (s *imageReservation) Reserve(quota int) error {
if quota > s.limit {
return errors.New("insufficient image quota")
}
s.held = max(s.held, quota)
return nil
}
func (s *imageReservation) GetPreConsumedQuota() int { return s.held }
func (*imageReservation) Settle(int) error { return nil }
func (*imageReservation) Refund(*gin.Context) {}
func (*imageReservation) NeedsRefund() bool { return false }
func TestImageRequestReservesFinalQuantityBeforeUpstream(t *testing.T) {
service.InitHttpClient()
for _, tc := range []struct {
name, body string
tiered, passThrough, insufficient, retryToOpenAI bool
override any
count, status int
}{
{name: "Ali nested count", body: `{"model":"z-image","n":1,"parameters":{"n":4}}`, count: 4},
{name: "Ali empty parameters", body: `{"model":"z-image","n":2,"parameters":{}}`, count: 2},
{name: "legacy channel quantity override", body: `{"model":"z-image","n":1}`, override: 4, count: 4},
{name: "expression channel quantity override", body: `{"model":"z-image","n":1}`, override: 4, count: 4, tiered: true},
{name: "pass-through quantity", body: `{"model":"z-image","n":2,"parameters":{}}`, count: 2, passThrough: true},
{name: "zero override rejected", body: `{"model":"z-image","n":1}`, override: 0, status: http.StatusBadRequest},
{name: "oversized override rejected", body: `{"model":"z-image","n":1}`, override: 129, status: http.StatusBadRequest},
{name: "insufficient reservation blocks upstream", body: `{"model":"z-image","n":1}`, override: 4, count: 4, insufficient: true, status: http.StatusForbidden},
{name: "retry drops Ali quantity and surcharge", body: `{"model":"z-image","n":1,"parameters":{"n":4,"prompt_extend":true}}`, count: 1, retryToOpenAI: true},
} {
t.Run(tc.name, func(t *testing.T) {
received := make(chan []byte, 1)
upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
body, err := io.ReadAll(r.Body)
if err == nil {
received <- body
}
w.Header().Set("Content-Type", "application/json")
w.WriteHeader(http.StatusBadGateway)
_, _ = io.WriteString(w, `{"error":{"message":"fixture upstream failure","type":"upstream_error"}}`)
}))
t.Cleanup(upstream.Close)
c, _ := gin.CreateTestContext(httptest.NewRecorder())
c.Request = httptest.NewRequest(http.MethodPost, "/v1/images/generations", strings.NewReader(tc.body))
c.Request.Header.Set("Content-Type", "application/json")
channel := constant.ChannelTypeAli
if tc.retryToOpenAI {
channel = constant.ChannelTypeOpenAI
}
common.SetContextKey(c, constant.ContextKeyChannelType, channel)
common.SetContextKey(c, constant.ContextKeyChannelBaseUrl, upstream.URL)
common.SetContextKey(c, constant.ContextKeyChannelSetting, dto.ChannelSettings{PassThroughBodyEnabled: tc.passThrough})
if tc.override != nil {
common.SetContextKey(c, constant.ContextKeyChannelParamOverride, map[string]any{"operations": []any{map[string]any{"path": "parameters.n", "mode": "set", "value": tc.override}}})
}
request, err := helper.GetAndValidOpenAIImageRequest(c, relayconstant.RelayModeImagesGenerations)
require.NoError(t, err)
reservation := &imageReservation{held: 20000, limit: 500000}
if tc.insufficient {
reservation.limit = reservation.held
}
info := &relaycommon.RelayInfo{Request: request, OriginModelName: "z-image", RelayMode: relayconstant.RelayModeImagesGenerations,
RequestURLPath: c.Request.URL.Path, Billing: reservation,
PriceData: hosttypes.PriceData{UsePrice: true, ModelPrice: 0.04, GroupRatioInfo: hosttypes.GroupRatioInfo{GroupRatio: 1}},
}
if tc.tiered {
expr := `tier("image", fixed(0.04)) * image_count`
info.TieredBillingSnapshot = &billingexpr.BillingSnapshot{BillingMode: "tiered_expr", ExprString: expr, ExprHash: billingexpr.ExprHashString(expr), GroupRatio: 1, QuotaPerUnit: common.QuotaPerUnit, EstimatedImageCount: common.GetPointer(1)}
info.BillingRequestInput = &billingexpr.RequestInput{Body: []byte(tc.body), ImageCount: common.GetPointer(1)}
info.PriceData.UsePrice = false
}
if tc.retryToOpenAI {
reservation.held = 160000
info.PriceData.AddOtherRatio("n", 4)
info.PriceData.AddOtherRatio("prompt_extend", 2)
info.BillingImageCount = common.GetPointer(3)
}
apiErr := ImageHelper(c, info)
require.NotNil(t, apiErr)
if tc.status != 0 {
assert.Equal(t, tc.status, apiErr.StatusCode)
assert.Empty(t, received, "no upstream submission before quantity validation and reservation")
return
}
assert.Equal(t, http.StatusBadGateway, apiErr.StatusCode)
require.Len(t, received, 1)
body := <-received
path := "parameters.n"
if tc.retryToOpenAI {
path = "n"
}
assert.Equal(t, int64(tc.count), gjson.GetBytes(body, path).Int())
assert.Equal(t, tc.count*20000, info.PriceData.QuotaToPreConsume)
assert.GreaterOrEqual(t, reservation.held, info.PriceData.QuotaToPreConsume)
assert.Nil(t, info.BillingImageCount)
if tc.tiered {
assert.Equal(t, tc.body, string(info.BillingRequestInput.Body))
assert.Equal(t, 1, *info.BillingRequestInput.ImageCount)
}
})
}
}
func TestOptInSafeToolLossRejectedAsBadRequestWithAdminDiagnostics(t *testing.T) {
gin.SetMode(gin.TestMode)
c, _ := gin.CreateTestContext(httptest.NewRecorder())
@@ -96,3 +212,90 @@ func hasHostDiagnosticCode(diagnostics []types.ConversionDiagnostic, code string
}
return false
}
func TestGeminiThinkingControlsConvertBestEffort(t *testing.T) {
gin.SetMode(gin.TestMode)
settings := model_setting.GetGeminiSettings()
originalAdapter := settings.ThinkingAdapterEnabled
t.Cleanup(func() { settings.ThinkingAdapterEnabled = originalAdapter })
t.Run("openai thinking_budget on gemini 3 becomes thinkingLevel", func(t *testing.T) {
settings.ThinkingAdapterEnabled = false
c, _ := gin.CreateTestContext(httptest.NewRecorder())
c.Request = httptest.NewRequest(http.MethodPost, "/v1/chat/completions", nil)
req := &dto.GeneralOpenAIRequest{
Model: "gemini-3.1-pro-preview-thinking",
Messages: []dto.Message{{Role: "user", Content: "hello"}},
ExtraBody: []byte(`{"google":{"thinking_config":{"thinking_budget":8192}}}`),
}
info := &relaycommon.RelayInfo{
RelayFormat: types.RelayFormatOpenAI,
OriginModelName: "gemini-3.1-pro-preview-thinking",
Request: req,
ChannelMeta: &relaycommon.ChannelMeta{UpstreamModelName: "gemini-3.1-pro-preview"},
}
require.NoError(t, helper.ApplyReasoningModelSuffix(c, info, req))
result, err := service.ConvertRequest(c, info, types.RelayFormatGemini, req)
require.NoError(t, err)
converted, ok := result.Value.(*dto.GeminiChatRequest)
require.True(t, ok)
require.NotNil(t, converted.GenerationConfig.ThinkingConfig)
assert.Equal(t, "medium", converted.GenerationConfig.ThinkingConfig.ThinkingLevel)
assert.Nil(t, converted.GenerationConfig.ThinkingConfig.ThinkingBudget)
assert.True(t, hasHostDiagnosticCode(info.ConversionDiagnostics(), "gemini_budget_to_level"))
})
t.Run("thinking alias with native thinking_level keeps the level", func(t *testing.T) {
settings.ThinkingAdapterEnabled = true
c, _ := gin.CreateTestContext(httptest.NewRecorder())
c.Request = httptest.NewRequest(http.MethodPost, "/v1/chat/completions", nil)
req := &dto.GeneralOpenAIRequest{
Model: "gemini-3.1-pro-preview-thinking",
Messages: []dto.Message{{Role: "user", Content: "hello"}},
ExtraBody: []byte(`{"google":{"thinking_config":{"thinking_level":"low"}}}`),
}
info := &relaycommon.RelayInfo{
RelayFormat: types.RelayFormatOpenAI,
OriginModelName: "gemini-3.1-pro-preview-thinking",
Request: req,
ChannelMeta: &relaycommon.ChannelMeta{UpstreamModelName: "gemini-3.1-pro-preview-thinking"},
}
require.NoError(t, helper.ApplyReasoningModelSuffix(c, info, req))
assert.Equal(t, "gemini-3.1-pro-preview", info.UpstreamModelName)
result, err := service.ConvertRequest(c, info, types.RelayFormatGemini, req)
require.NoError(t, err)
converted, ok := result.Value.(*dto.GeminiChatRequest)
require.True(t, ok)
require.NotNil(t, converted.GenerationConfig.ThinkingConfig)
assert.Equal(t, "low", converted.GenerationConfig.ThinkingConfig.ThinkingLevel)
assert.Nil(t, converted.GenerationConfig.ThinkingConfig.ThinkingBudget)
})
t.Run("gemini thinkingBudget converts to openai reasoning_effort", func(t *testing.T) {
settings.ThinkingAdapterEnabled = false
c, _ := gin.CreateTestContext(httptest.NewRecorder())
c.Request = httptest.NewRequest(http.MethodPost, "/v1beta/models/gemini-3.1-pro-preview:generateContent", nil)
budget := 8192
req := &dto.GeminiChatRequest{
Contents: []dto.GeminiChatContent{{Role: "user", Parts: []dto.GeminiPart{{Text: "hello"}}}},
GenerationConfig: dto.GeminiChatGenerationConfig{
ThinkingConfig: &dto.GeminiThinkingConfig{ThinkingBudget: &budget},
},
}
info := &relaycommon.RelayInfo{
RelayFormat: types.RelayFormatGemini,
OriginModelName: "gemini-3.1-pro-preview",
Request: req,
ChannelMeta: &relaycommon.ChannelMeta{UpstreamModelName: "gpt-5"},
}
result, err := service.ConvertRequest(c, info, types.RelayFormatOpenAI, req)
require.NoError(t, err)
converted, ok := result.Value.(*dto.GeneralOpenAIRequest)
require.True(t, ok)
assert.Equal(t, "medium", converted.ReasoningEffort)
assert.True(t, hasHostDiagnosticCode(info.ConversionDiagnostics(), "gemini_budget_to_level"))
})
}
+38
View File
@@ -5,6 +5,7 @@ import (
"strings"
"github.com/QuantumNous/new-api/common"
"github.com/QuantumNous/new-api/constant"
"github.com/QuantumNous/new-api/pkg/billingexpr"
relaycommon "github.com/QuantumNous/new-api/relay/common"
"github.com/QuantumNous/new-api/relaykit/dto"
@@ -33,6 +34,39 @@ func ResolveIncomingBillingExprRequestInput(c *gin.Context, info *relaycommon.Re
return input, nil
}
// ResolveImageBillingRequestInput freezes only the validated scalar image
// parameters needed by pricing. Image files, prompts and base64 payloads are
// deliberately excluded, including for multipart edits.
func ResolveImageBillingRequestInput(c *gin.Context, info *relaycommon.RelayInfo, input billingexpr.RequestInput) (billingexpr.RequestInput, error) {
request, ok := info.Request.(*dto.ImageRequest)
if !ok {
return input, nil
}
channelType := common.GetContextKeyInt(c, constant.ContextKeyChannelType)
if info.ChannelMeta != nil {
channelType = info.ChannelType
}
count, err := request.ImageCount(channelType == constant.ChannelTypeAli)
if err != nil {
return input, err
}
topLevelCount, err := request.ImageCount(false)
if err != nil {
return input, err
}
body := map[string]any{"model": request.Model, "n": topLevelCount, "size": request.Size, "quality": request.Quality}
if request.BillingParameters != nil {
body["parameters"] = request.BillingParameters
}
encoded, err := common.Marshal(body)
if err != nil {
return input, err
}
input.Body = encoded
input.ImageCount = &count
return input, nil
}
func BuildBillingExprRequestInputFromRequest(request dto.Request, headers map[string]string) (billingexpr.RequestInput, error) {
input := billingexpr.RequestInput{
Headers: cloneStringMap(headers),
@@ -64,6 +98,10 @@ func cloneRequestInput(src billingexpr.RequestInput) billingexpr.RequestInput {
input := billingexpr.RequestInput{
Headers: cloneStringMap(src.Headers),
}
if src.ImageCount != nil {
count := *src.ImageCount
input.ImageCount = &count
}
if len(src.Body) > 0 {
input.Body = append([]byte(nil), src.Body...)
}
+22
View File
@@ -70,3 +70,25 @@ func TestMaxTokensBounds(t *testing.T) {
require.Contains(t, err.Error(), "max_output_tokens is invalid")
})
}
func TestSGLangMinTokensBounds(t *testing.T) {
for _, tc := range []struct {
value string
valid bool
}{
{"0", true}, {"8192", true}, {"-1", false}, {"18446744073686646784", false},
} {
t.Run(tc.value, func(t *testing.T) {
c, _ := gin.CreateTestContext(httptest.NewRecorder())
c.Request = httptest.NewRequest(http.MethodPost, "/v1/chat/completions", bytes.NewBufferString(`{"model":"served-model","messages":[{"role":"user","content":"hi"}],"min_tokens":`+tc.value+`}`))
c.Request.Header.Set("Content-Type", "application/json")
request, err := GetAndValidateTextRequest(c, relayconstant.RelayModeChatCompletions)
if !tc.valid {
require.Error(t, err)
return
}
require.NoError(t, err)
require.NotNil(t, request.MinTokens)
})
}
}
+81
View File
@@ -11,12 +11,50 @@ import (
"testing"
"github.com/QuantumNous/new-api/common"
"github.com/QuantumNous/new-api/constant"
"github.com/QuantumNous/new-api/pkg/billingexpr"
relaycommon "github.com/QuantumNous/new-api/relay/common"
relayconstant "github.com/QuantumNous/new-api/relay/constant"
"github.com/QuantumNous/new-api/relaykit/dto"
"github.com/gin-gonic/gin"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
)
func TestLegacyDalleValidationAndPricesRemainCompatible(t *testing.T) {
for _, tc := range []struct {
body, size, quality string
ratio float64
invalid bool
}{
{body: `{"model":"dall-e-2"}`, size: "1024x1024", ratio: 1},
{body: `{"model":"dall-e"}`, size: "1024x1024", ratio: 1},
{body: `{"model":"dall-e-3"}`, size: "1024x1024", quality: "standard", ratio: 1},
{body: `{"model":"dall-e-2","size":"256x256"}`, size: "256x256", ratio: 0.4},
{body: `{"model":"dall-e-2","size":"512x512","quality":"hd"}`, size: "512x512", quality: "hd", ratio: 0.45},
{body: `{"model":"dall-e-3","quality":"hd"}`, size: "1024x1024", quality: "hd", ratio: 2},
{body: `{"model":"dall-e-3","size":"1024x1792","quality":"hd"}`, size: "1024x1792", quality: "hd", ratio: 3},
{body: `{"model":"dall-e-3","size":"1792x1024"}`, size: "1792x1024", quality: "standard", ratio: 2},
{body: `{"model":"dall-e-3","size":"256x256"}`, invalid: true},
{body: `{"model":"dall-e-2","size":"1024x1792"}`, invalid: true},
} {
t.Run(tc.body, func(t *testing.T) {
c, _ := gin.CreateTestContext(httptest.NewRecorder())
c.Request = httptest.NewRequest(http.MethodPost, "/v1/images/generations", bytes.NewBufferString(tc.body))
c.Request.Header.Set("Content-Type", "application/json")
request, err := GetAndValidOpenAIImageRequest(c, relayconstant.RelayModeImagesGenerations)
if tc.invalid {
require.Error(t, err)
return
}
require.NoError(t, err)
assert.Equal(t, tc.size, request.Size)
assert.Equal(t, tc.quality, request.Quality)
assert.Equal(t, tc.ratio, request.GetTokenCountMeta().ImagePriceRatio)
})
}
}
// TestGetAndValidOpenAIImageRequestMultipartStream verifies multipart image
// edit parsing: the stream field is parsed and validated, and the request body
// stays replayable for the upstream request.
@@ -61,6 +99,11 @@ func TestGetAndValidOpenAIImageRequestMultipartStream(t *testing.T) {
require.NoError(t, err)
require.Equal(t, "true", url.Values(form.Value).Get("stream"))
require.Len(t, form.File["image"], 1)
billing, err := ResolveImageBillingRequestInput(c, &relaycommon.RelayInfo{Request: req}, billingexpr.RequestInput{})
require.NoError(t, err)
require.Equal(t, 1, *billing.ImageCount)
require.NotContains(t, string(billing.Body), "fake image")
require.NotContains(t, string(billing.Body), "edit this image")
})
t.Run("invalid stream value is rejected", func(t *testing.T) {
@@ -72,6 +115,44 @@ func TestGetAndValidOpenAIImageRequestMultipartStream(t *testing.T) {
})
}
func TestImageBillingRequestUsesValidatedProviderCount(t *testing.T) {
for _, tc := range []struct {
body string
channel, count int
invalid bool
}{
{`{"model":"z-image","n":2,"parameters":{"n":3,"prompt_extend":true}}`, constant.ChannelTypeAli, 3, false},
{`{"model":"gpt-image-2","n":2,"parameters":{"n":3}}`, constant.ChannelTypeOpenAI, 2, false},
{`{"model":"z-image","n":2,"parameters":{"n":129}}`, constant.ChannelTypeAli, 0, true},
{`{"model":"z-image","parameters":{"n":-1}}`, constant.ChannelTypeAli, 0, true},
{`{"model":"z-image","n":2,"parameters":{}}`, constant.ChannelTypeAli, 2, false},
{`{"model":"z-image","n":2,"parameters":{"n":null}}`, constant.ChannelTypeAli, 2, false},
{`{"model":"z-image","n":2,"parameters":{"n":0}}`, constant.ChannelTypeAli, 0, true},
{`{"model":"z-image","parameters":{"n":1.5}}`, constant.ChannelTypeAli, 0, true},
{`{"model":"z-image","parameters":{"n":18446744073686646784}}`, constant.ChannelTypeAli, 0, true},
{`{"model":"z-image","parameters":{"n":128}}`, constant.ChannelTypeAli, 128, false},
{`{"model":"z-image","n":0}`, constant.ChannelTypeAli, 1, false},
{`{"model":"gpt-image-2","n":2,"parameters":{"n":0}}`, constant.ChannelTypeOpenAI, 2, false},
} {
c, _ := gin.CreateTestContext(httptest.NewRecorder())
c.Request = httptest.NewRequest(http.MethodPost, "/v1/images/generations", bytes.NewBufferString(tc.body))
c.Request.Header.Set("Content-Type", "application/json")
common.SetContextKey(c, constant.ContextKeyChannelType, tc.channel)
request, err := GetAndValidOpenAIImageRequest(c, relayconstant.RelayModeImagesGenerations)
if tc.invalid {
require.Error(t, err)
continue
}
require.NoError(t, err)
input, err := ResolveImageBillingRequestInput(c, &relaycommon.RelayInfo{Request: request}, billingexpr.RequestInput{})
require.NoError(t, err)
require.Equal(t, tc.count, *input.ImageCount)
cost, _, err := billingexpr.RunExprWithRequest(`tier("image", fixed(0.04)) * image_count`, billingexpr.TokenParams{}, input)
require.NoError(t, err)
require.Equal(t, float64(tc.count)*40000, cost)
}
}
// TestGetAndValidOpenAIImageRequestNBounds guards the billing invariant that
// the image generation count can never reach quota calculation with a value
// large enough to overflow int64 into a negative charge.
+52 -2
View File
@@ -5,10 +5,12 @@ import (
"strings"
"github.com/QuantumNous/new-api/common"
"github.com/QuantumNous/new-api/constant"
"github.com/QuantumNous/new-api/logger"
"github.com/QuantumNous/new-api/model"
"github.com/QuantumNous/new-api/pkg/billingexpr"
relaycommon "github.com/QuantumNous/new-api/relay/common"
"github.com/QuantumNous/new-api/relaykit/dto"
"github.com/QuantumNous/new-api/relaykit/relayconvert/reasoning"
"github.com/QuantumNous/new-api/relaykit/types"
"github.com/QuantumNous/new-api/setting/billing_setting"
@@ -132,10 +134,16 @@ func ModelPriceHelper(c *gin.Context, info *relaycommon.RelayInfo, promptTokens
return hosttypes.PriceData{}, err
}
preConsumedQuota = quota
if _, image := info.Request.(*dto.ImageRequest); image {
info.ImageQuotaBeforeGroup = float64(preConsumedTokens) * modelRatio
}
} else {
if meta.ImagePriceRatio != 0 {
modelPrice = modelPrice * meta.ImagePriceRatio
}
if _, image := info.Request.(*dto.ImageRequest); image {
info.ImageQuotaBeforeGroup = modelPrice * common.QuotaPerUnit
}
}
// check if free model pre-consume is disabled
@@ -177,6 +185,36 @@ func ModelPriceHelper(c *gin.Context, info *relaycommon.RelayInfo, promptTokens
for name, ratio := range meta.BillingRatios {
priceData.AddOtherRatio(name, ratio)
}
}
if request, image := info.Request.(*dto.ImageRequest); image {
channelType := common.GetContextKeyInt(c, constant.ContextKeyChannelType)
count, err := request.ImageCount(channelType == constant.ChannelTypeAli)
if err != nil {
return hosttypes.PriceData{}, err
}
if usePrice || channelType == constant.ChannelTypeAli {
priceData.AddOtherRatio("n", float64(count))
}
if channelType == constant.ChannelTypeAli && request.BillingParameters != nil && request.BillingParameters.PromptExtend != nil && *request.BillingParameters.PromptExtend {
// Resolve only routing identity; do not initialize ChannelMeta on the
// real request, which also distinguishes the first channel attempt.
mapped := &relaycommon.RelayInfo{OriginModelName: info.OriginModelName, ChannelMeta: &relaycommon.ChannelMeta{UpstreamModelName: info.OriginModelName}}
if err := ModelMappedHelper(c, mapped, nil); err != nil {
return hosttypes.PriceData{}, err
}
if strings.Contains(mapped.UpstreamModelName, "z-image") {
priceData.AddOtherRatio("prompt_extend", common.ZImagePromptExtendMultiplier)
}
}
if !usePrice {
quota, err := common.QuotaFromFloatStrict(priceData.ApplyOtherRatiosToFloat(info.ImageQuotaBeforeGroup * groupRatioInfo.GroupRatio))
if err != nil {
return hosttypes.PriceData{}, err
}
priceData.QuotaToPreConsume = quota
}
}
if usePrice {
quotaToPreConsume := priceData.ApplyOtherRatiosToFloat(modelPrice * common.QuotaPerUnit * groupRatioInfo.GroupRatio)
quota, err := common.QuotaFromFloatStrict(quotaToPreConsume)
if err != nil {
@@ -324,6 +362,10 @@ func modelPriceHelperTiered(c *gin.Context, info *relaycommon.RelayInfo, billing
if !ok {
return hosttypes.PriceData{}, fmt.Errorf("model %s is configured as tiered_expr but has no billing expression", billingModelName)
}
exprHash := billingexpr.ExprHashString(exprStr)
if info.RelayFormat == types.RelayFormatOpenAIRealtime && billingexpr.UsesFixedPricingByHash(exprStr, exprHash) {
return hosttypes.PriceData{}, fmt.Errorf("fixed pricing is not supported for Realtime requests")
}
estimatedCompletionTokens := meta.MaxTokens
if estimatedCompletionTokens == 0 && groupRatioInfo.GroupRatio != 0 {
@@ -334,8 +376,14 @@ func modelPriceHelperTiered(c *gin.Context, info *relaycommon.RelayInfo, billing
if err != nil {
return hosttypes.PriceData{}, err
}
if billingexpr.UsedVarsByHash(exprStr, exprHash)["image_count"] {
requestInput, err = ResolveImageBillingRequestInput(c, info, requestInput)
if err != nil {
return hosttypes.PriceData{}, err
}
}
rawCost, trace, err := billingexpr.RunExprWithRequest(exprStr, billingexpr.TokenParams{
rawCost, trace, err := billingexpr.RunExprByHashWithRequest(exprStr, exprHash, billingexpr.TokenParams{
P: float64(promptTokens),
C: float64(estimatedCompletionTokens),
Len: float64(promptTokens),
@@ -359,8 +407,8 @@ func modelPriceHelperTiered(c *gin.Context, info *relaycommon.RelayInfo, billing
}
}
exprHash := billingexpr.ExprHashString(exprStr)
snapshot := &billingexpr.BillingSnapshot{
EstimatedImageCount: trace.ImageCount,
BillingMode: billing_setting.BillingModeTieredExpr,
ModelName: billingModelName,
ExprString: exprStr,
@@ -371,6 +419,8 @@ func modelPriceHelperTiered(c *gin.Context, info *relaycommon.RelayInfo, billing
EstimatedQuotaBeforeGroup: quotaBeforeGroup,
EstimatedQuotaAfterGroup: preConsumedQuota,
EstimatedTier: trace.MatchedTier,
EstimatedBillingUnit: trace.BillingUnit,
EstimatedFixedPrice: trace.FixedPrice,
QuotaPerUnit: common.QuotaPerUnit,
ExprVersion: billingexpr.ExprVersion(exprStr),
}
+41
View File
@@ -68,6 +68,47 @@ func TestModelPriceHelperTieredUsesPreloadedRequestInput(t *testing.T) {
require.Equal(t, common.QuotaPerUnit, info.TieredBillingSnapshot.QuotaPerUnit)
}
func TestFixedPricePreConsumeAndRealtimeRejection(t *testing.T) {
saved := map[string]string{}
require.NoError(t, config.GlobalConfig.SaveToDB(func(key, value string) error {
saved[key] = value
return nil
}))
t.Cleanup(func() { require.NoError(t, config.GlobalConfig.LoadFromDB(saved)) })
require.NoError(t, config.GlobalConfig.LoadFromDB(map[string]string{
"billing_setting.billing_mode": `{"fixed-test":"tiered_expr"}`,
"billing_setting.billing_expr": `{"fixed-test":"len <= 32000 ? tier(\"short\", fixed(0.01)) : tier(\"long\", p * 2)"}`,
"group_ratio_setting.group_ratio": `{"default":1.5}`,
}))
for _, tc := range []struct {
name string
format types.RelayFormat
prompt int
wantError bool
}{
{"HTTP charges once", types.RelayFormatOpenAI, 0, false},
{"Realtime rejects even unselected fixed branch", types.RelayFormatOpenAIRealtime, 50000, true},
} {
t.Run(tc.name, func(t *testing.T) {
ctx, _ := gin.CreateTestContext(httptest.NewRecorder())
ctx.Request = httptest.NewRequest(http.MethodPost, "/v1/chat/completions", nil)
info := &relaycommon.RelayInfo{OriginModelName: "fixed-test", UserGroup: "default", UsingGroup: "default", RelayFormat: tc.format, BillingRequestInput: &billingexpr.RequestInput{}}
price, err := ModelPriceHelper(ctx, info, tc.prompt, &types.TokenCountMeta{})
if tc.wantError {
require.ErrorContains(t, err, "Realtime")
assert.Nil(t, info.TieredBillingSnapshot)
return
}
require.NoError(t, err)
assert.Equal(t, 7500, price.QuotaToPreConsume)
require.NotNil(t, info.TieredBillingSnapshot)
assert.Equal(t, billingexpr.BillingUnitRequest, info.TieredBillingSnapshot.EstimatedBillingUnit)
require.NotNil(t, info.TieredBillingSnapshot.EstimatedFixedPrice)
assert.Equal(t, 0.01, *info.TieredBillingSnapshot.EstimatedFixedPrice)
})
}
}
func TestModelPriceHelperTieredPreConsumeMaxTokensFallback(t *testing.T) {
gin.SetMode(gin.TestMode)
+4 -3
View File
@@ -59,10 +59,11 @@ func ApplyReasoningModelSuffix(c *gin.Context, info *relaycommon.RelayInfo, outb
}
if selected.hasThinking {
explicit, err := explicitIntentFromRequest(info.Request)
explicit, explicitDiagnostics, err := explicitIntentFromRequest(info.Request)
if err != nil {
return reasoning.AsClientError(err)
}
diagnostics = append(diagnostics, explicitDiagnostics...)
diagnostics = append(diagnostics, modifierRequestOverrideDiagnostics(explicit, selected.intent)...)
if selected.intent.IncludeThoughts == nil {
selected.intent.IncludeThoughts = explicit.IncludeThoughts
@@ -215,7 +216,7 @@ func parseHostModelSuffix(name string, opts *convmeta.Options) (string, reasonin
)
}
func explicitIntentFromRequest(req dto.Request) (reasoning.Intent, error) {
func explicitIntentFromRequest(req dto.Request) (reasoning.Intent, []types.ConversionDiagnostic, error) {
switch r := req.(type) {
case *dto.ClaudeRequest:
return reasoning.FromClaude(r)
@@ -226,6 +227,6 @@ func explicitIntentFromRequest(req dto.Request) (reasoning.Intent, error) {
case *dto.OpenAIResponsesRequest:
return reasoning.FromOpenAIResponses(r)
default:
return reasoning.Intent{}, nil
return reasoning.Intent{}, nil, nil
}
}
+8 -3
View File
@@ -1,7 +1,10 @@
package helper
import (
"errors"
relaycommon "github.com/QuantumNous/new-api/relay/common"
"github.com/QuantumNous/new-api/relaykit/types"
)
// StreamResult is passed to each dataHandler invocation, providing methods
@@ -23,13 +26,15 @@ func (r *StreamResult) Error(err error) {
return
}
r.status.RecordError(err.Error())
var apiErr *types.NewAPIError
if errors.As(err, &apiErr) {
r.status.MarkFailed(string(apiErr.GetErrorCode()), apiErr.ToOpenAIError().Type, apiErr.StatusCode)
}
}
// Stop records a fatal error and marks the stream to stop after this chunk.
func (r *StreamResult) Stop(err error) {
if err != nil {
r.status.RecordError(err.Error())
}
r.Error(err)
r.status.SetEndReason(relaycommon.StreamEndReasonHandlerStop, err)
r.stopped = true
}
+46 -12
View File
@@ -30,7 +30,8 @@ const (
// unconditional wg.Wait() in cleanup can always finish. Without it, a slow
// but connected client (full TCP buffer, no server WriteTimeout) could hang
// the handler forever.
streamWriteTimeout = 30 * time.Second
streamWriteTimeout = 30 * time.Second
clientGoneDrainTimeout = 2 * time.Minute
)
func getScannerBufferSize() int {
@@ -40,9 +41,15 @@ func getScannerBufferSize() int {
return DefaultMaxScannerBufferSize
}
func NewStreamScanner(reader io.Reader) *bufio.Scanner {
// NewStreamScanner shares relay scanner configuration. Callers buffering bounded
// task state may additionally cap a line without increasing the configured limit.
func NewStreamScanner(reader io.Reader, maxBytes ...int) *bufio.Scanner {
limit := getScannerBufferSize()
if len(maxBytes) > 0 && maxBytes[0] > 0 {
limit = min(limit, maxBytes[0])
}
scanner := bufio.NewScanner(reader)
scanner.Buffer(make([]byte, InitialScannerBufferSize), getScannerBufferSize())
scanner.Buffer(make([]byte, min(InitialScannerBufferSize, limit)), limit)
return scanner
}
@@ -121,11 +128,11 @@ func StreamScannerHandler(c *gin.Context, resp *http.Response, info *relaycommon
logger.LogDebug(c, "streaming timeout seconds: %d", int64(streamingTimeout.Seconds()))
logger.LogDebug(c, "ping interval seconds: %d", int64(pingInterval.Seconds()))
cleanup := func() {
cleanup := func(closeUpstream bool) {
cleanupOnce.Do(func() {
cancel()
stop()
if resp.Body != nil {
if closeUpstream && resp.Body != nil {
_ = resp.Body.Close()
}
@@ -138,7 +145,9 @@ func StreamScannerHandler(c *gin.Context, resp *http.Response, info *relaycommon
})
}
// Ensure gin.Context is not returned to Gin's pool while any stream goroutine can still use it.
defer cleanup()
defer func() {
cleanup(true)
}()
scanner.Split(bufio.ScanLines)
copyCodexSSEHeaders(c, resp)
@@ -289,19 +298,44 @@ func StreamScannerHandler(c *gin.Context, resp *http.Response, info *relaycommon
info.StreamStatus.SetEndReason(relaycommon.StreamEndReasonEOF, nil)
})
// 客户端断开后仍读取上游终止事件中的 usage,避免上游已计费而
// NewAPI 无法结算。各写入函数会自行跳过已断开的客户端。
drainClientGone := func(err error) {
info.StreamStatus.SetClientGone(err)
logger.LogInfo(c, "client disconnected; draining upstream for billing usage")
drainTimer := time.NewTimer(clientGoneDrainTimeout)
select {
case <-stopChan:
case <-drainTimer.C:
logger.LogError(c, "client_gone upstream drain timeout; usage may be unavailable")
}
if !drainTimer.Stop() {
select {
case <-drainTimer.C:
default:
}
}
}
// 主循环等待完成或超时
select {
case <-ticker.C:
info.StreamStatus.SetEndReason(relaycommon.StreamEndReasonTimeout, nil)
if err := c.Request.Context().Err(); err != nil {
drainClientGone(err)
} else {
info.StreamStatus.SetEndReason(relaycommon.StreamEndReasonTimeout, nil)
}
case <-stopChan:
// EndReason already set by the goroutine that triggered stopChan
// A client disconnect can race with upstream [DONE]. Preserve it as the
// diagnostic reason while keeping fatal handler/scanner reasons intact.
if err := c.Request.Context().Err(); err != nil {
info.StreamStatus.SetClientGone(err)
}
case <-c.Request.Context().Done():
// 客户端断开:立即 cleanup 关闭上游 resp.Body,解除 scanner 阻塞并让上游停止生成,
// 避免为已放弃的请求继续消费上游 token。
info.StreamStatus.SetEndReason(relaycommon.StreamEndReasonClientGone, c.Request.Context().Err())
drainClientGone(c.Request.Context().Err())
}
cleanup()
cleanup(true)
if info.StreamStatus.IsNormalEnd() && !info.StreamStatus.HasErrors() {
logger.LogInfo(c, fmt.Sprintf("stream ended: %s", info.StreamStatus.Summary()))
} else {
+20 -22
View File
@@ -211,12 +211,9 @@ func TestStreamScannerHandler_DataWithExtraSpaces(t *testing.T) {
assert.Equal(t, "{\"trimmed\":true}", got)
}
// TestStreamScannerHandler_ClientCancelAbortsUpstreamAndReturns pins the
// disconnect contract: when the client goes away, the handler must return
// promptly (all goroutines joined, so the gin.Context can never leak into a
// pooled reuse), the upstream body must be closed to stop token generation,
// and no data received after the disconnect may be processed or written.
func TestStreamScannerHandler_ClientCancelAbortsUpstreamAndReturns(t *testing.T) {
// TestStreamScannerHandler_ClientCancelDrainsUpstreamUsage verifies that a client
// disconnect stops writes but leaves the upstream readable through its terminal event.
func TestStreamScannerHandler_ClientCancelDrainsUpstreamUsage(t *testing.T) {
ctx, cancel := context.WithCancel(context.Background())
defer cancel()
@@ -252,7 +249,6 @@ func TestStreamScannerHandler_ClientCancelAbortsUpstreamAndReturns(t *testing.T)
_, err := fmt.Fprint(pw, "data: first\n")
require.NoError(t, err)
select {
case <-firstHandled:
case <-time.After(2 * time.Second):
@@ -260,31 +256,21 @@ func TestStreamScannerHandler_ClientCancelAbortsUpstreamAndReturns(t *testing.T)
}
cancel()
_, err = fmt.Fprint(pw, "data: usage\ndata: [DONE]\n")
require.NoError(t, err, "client disconnect must not close the upstream body")
require.NoError(t, pw.Close())
// The handler must return without any further upstream input: cleanup
// closes resp.Body, which unblocks the scanner goroutine.
select {
case <-done:
case <-time.After(2 * time.Second):
t.Fatal("handler did not return after client disconnect")
t.Fatal("handler did not finish draining the upstream stream")
}
// Upstream read side must be closed so the provider stops generating
// (and billing) for a request nobody is listening to.
_, err = fmt.Fprint(pw, "data: second\n")
require.ErrorIs(t, err, io.ErrClosedPipe, "upstream body should be closed after client disconnect")
assert.Equal(t, int64(1), count.Load(), "no chunk after disconnect should be processed")
assert.Equal(t, int64(2), count.Load(), "terminal upstream data should be processed after client disconnect")
require.NotNil(t, info.StreamStatus)
assert.Equal(t, relaycommon.StreamEndReasonClientGone, info.StreamStatus.EndReason)
body := recorder.Body.String()
assert.Contains(t, body, "first")
assert.NotContains(t, body, "second")
}
// ---------- Ping tests ----------
func TestStreamScannerHandler_PingSentDuringSlowUpstream(t *testing.T) {
setting := operation_setting.GetGeneralSetting()
oldEnabled := setting.PingIntervalEnabled
@@ -571,3 +557,15 @@ func TestStreamScannerHandler_StreamStatus_ReplacesPreInitialized(t *testing.T)
assert.Equal(t, relaycommon.StreamEndReasonDone, info.StreamStatus.EndReason)
assert.Equal(t, 0, info.StreamStatus.TotalErrorCount())
}
func TestNewStreamScannerCallerLimit(t *testing.T) {
// The smaller buffer must actually constrain a line; a preallocated 64 KiB
// buffer would otherwise bypass this caller's 1 KiB limit in bufio.Scanner.
scanner := NewStreamScanner(strings.NewReader(strings.Repeat("x", 2048)+"\n"), 1024)
assert.False(t, scanner.Scan())
require.Error(t, scanner.Err())
scanner = NewStreamScanner(strings.NewReader("data: ok\n"), 1024)
require.True(t, scanner.Scan())
assert.Equal(t, "data: ok", scanner.Text())
require.NoError(t, scanner.Err())
}
+31 -25
View File
@@ -1,6 +1,7 @@
package helper
import (
"encoding/json"
"errors"
"fmt"
"math"
@@ -9,6 +10,7 @@ import (
"strings"
"github.com/QuantumNous/new-api/common"
"github.com/QuantumNous/new-api/constant"
"github.com/QuantumNous/new-api/logger"
relayconstant "github.com/QuantumNous/new-api/relay/constant"
"github.com/QuantumNous/new-api/relaykit/dto"
@@ -121,7 +123,8 @@ func GetAndValidateEmbeddingRequest(c *gin.Context, relayMode int) (*dto.Embeddi
// overflow the conversion and corrupt billing.
const maxTokensLimit = math.MaxInt32 / 2
func exceedsMaxTokensLimit(values ...*uint) bool {
// ExceedsMaxTokensLimit checks token limits before they reach billing arithmetic.
func ExceedsMaxTokensLimit(values ...*uint) bool {
for _, v := range values {
if lo.FromPtrOr(v, uint(0)) > maxTokensLimit {
return true
@@ -142,7 +145,7 @@ func GetAndValidateResponsesRequest(c *gin.Context) (*dto.OpenAIResponsesRequest
if request.Input == nil {
return nil, errors.New("input is required")
}
if exceedsMaxTokensLimit(request.MaxOutputTokens) {
if ExceedsMaxTokensLimit(request.MaxOutputTokens) {
return nil, errors.New("max_output_tokens is invalid")
}
return request, nil
@@ -203,6 +206,9 @@ func GetAndValidOpenAIImageRequest(c *gin.Context, relayMode int) (*dto.ImageReq
}
imageRequest.Quality = formData.Get("quality")
imageRequest.Size = formData.Get("size")
if parameters := formData.Get("parameters"); parameters != "" {
imageRequest.Extra = map[string]json.RawMessage{"parameters": json.RawMessage(parameters)}
}
if streamValue := strings.TrimSpace(formData.Get("stream")); streamValue != "" {
stream, err := strconv.ParseBool(streamValue)
if err != nil {
@@ -238,7 +244,6 @@ func GetAndValidOpenAIImageRequest(c *gin.Context, relayMode int) (*dto.ImageReq
}
if imageRequest.Model == "" {
//imageRequest.Model = "dall-e-3"
return nil, errors.New("model is required")
}
@@ -250,25 +255,10 @@ func GetAndValidOpenAIImageRequest(c *gin.Context, relayMode int) (*dto.ImageReq
return nil, fmt.Errorf("n must be an integer between 1 and %d", dto.MaxImageN)
}
// Not "256x256", "512x512", or "1024x1024"
if imageRequest.Model == "dall-e-2" || imageRequest.Model == "dall-e" {
if imageRequest.Size != "" && imageRequest.Size != "256x256" && imageRequest.Size != "512x512" && imageRequest.Size != "1024x1024" {
return nil, errors.New("size must be one of 256x256, 512x512, or 1024x1024 for dall-e-2 or dall-e")
}
if imageRequest.Size == "" {
imageRequest.Size = "1024x1024"
}
} else if imageRequest.Model == "dall-e-3" {
if imageRequest.Size != "" && imageRequest.Size != "1024x1024" && imageRequest.Size != "1024x1792" && imageRequest.Size != "1792x1024" {
return nil, errors.New("size must be one of 1024x1024, 1024x1792 or 1792x1024 for dall-e-3")
}
if imageRequest.Quality == "" {
imageRequest.Quality = "standard"
}
if imageRequest.Size == "" {
imageRequest.Size = "1024x1024"
}
} else if imageRequest.Model == "gpt-image-1" {
if err := imageRequest.NormalizeLegacyDalleImageRequest(); err != nil {
return nil, err
}
if imageRequest.Model == "gpt-image-1" {
if imageRequest.Quality == "" {
imageRequest.Quality = "auto"
}
@@ -283,6 +273,19 @@ func GetAndValidOpenAIImageRequest(c *gin.Context, relayMode int) (*dto.ImageReq
}
}
// Provider parameters can override the top-level count. Validate before
// pricing so malformed multipliers return a client error, not a pricing
// failure after reservation has started.
if raw, exists := imageRequest.Extra["parameters"]; exists {
parameters := &dto.ImageBillingParameters{}
if err := common.Unmarshal(raw, parameters); err != nil {
return nil, fmt.Errorf("invalid image parameters: %w", err)
}
imageRequest.BillingParameters = parameters
}
if _, err := imageRequest.ImageCount(common.GetContextKeyInt(c, constant.ContextKeyChannelType) == constant.ChannelTypeAli); err != nil {
return nil, err
}
return imageRequest, nil
}
@@ -298,7 +301,7 @@ func GetAndValidateClaudeRequest(c *gin.Context) (textRequest *dto.ClaudeRequest
if textRequest.Model == "" {
return nil, errors.New("field model is required")
}
if exceedsMaxTokensLimit(textRequest.MaxTokens, textRequest.MaxTokensToSample) {
if ExceedsMaxTokensLimit(textRequest.MaxTokens, textRequest.MaxTokensToSample) {
return nil, errors.New("max_tokens is invalid")
}
@@ -323,9 +326,12 @@ func GetAndValidateTextRequest(c *gin.Context, relayMode int) (*dto.GeneralOpenA
textRequest.Model = c.Param("model")
}
if exceedsMaxTokensLimit(textRequest.MaxTokens, textRequest.MaxCompletionTokens) {
if ExceedsMaxTokensLimit(textRequest.MaxTokens, textRequest.MaxCompletionTokens) {
return nil, errors.New("max_tokens is invalid")
}
if ExceedsMaxTokensLimit(textRequest.MinTokens) {
return nil, errors.New("min_tokens is invalid")
}
if textRequest.Model == "" {
return nil, errors.New("model is required")
}
@@ -376,7 +382,7 @@ func GetAndValidateGeminiRequest(c *gin.Context) (*dto.GeminiChatRequest, error)
if len(request.Contents) == 0 && len(request.Requests) == 0 {
return nil, errors.New("contents is required")
}
if exceedsMaxTokensLimit(request.GenerationConfig.MaxOutputTokens) {
if ExceedsMaxTokensLimit(request.GenerationConfig.MaxOutputTokens) {
return nil, errors.New("maxOutputTokens is invalid")
}
+53 -7
View File
@@ -18,10 +18,13 @@ import (
"github.com/QuantumNous/new-api/setting/model_setting"
"github.com/gin-gonic/gin"
"github.com/tidwall/sjson"
)
func ImageHelper(c *gin.Context, info *relaycommon.RelayInfo) (newAPIError *types.NewAPIError) {
info.InitChannelMeta(c)
info.BillingImageCount = nil
info.ImageRequestCount = 0
imageReq, ok := info.Request.(*dto.ImageRequest)
if !ok {
@@ -43,15 +46,28 @@ func ImageHelper(c *gin.Context, info *relaycommon.RelayInfo) (newAPIError *type
return types.NewError(fmt.Errorf("invalid api type: %d", info.ApiType), types.ErrorCodeInvalidApiType, types.ErrOptionWithSkipRetry())
}
adaptor.Init(info)
imageCount, err := request.ImageCount(info.ChannelType == constant.ChannelTypeAli)
if err != nil {
return types.NewErrorWithStatusCode(err, types.ErrorCodeInvalidRequest, http.StatusBadRequest, types.ErrOptionWithSkipRetry())
}
promptExtend := request.BillingParameters != nil && request.BillingParameters.PromptExtend != nil && *request.BillingParameters.PromptExtend
var requestBody io.Reader
var jsonData []byte
if model_setting.GetGlobalSettings().PassThroughRequestEnabled || info.ChannelSetting.PassThroughBodyEnabled {
storage, err := common.GetBodyStorage(c)
if err != nil {
return types.NewErrorWithStatusCode(err, types.ErrorCodeReadRequestBodyFailed, http.StatusBadRequest, types.ErrOptionWithSkipRetry())
}
requestBody = common.NewReplayableBodyReader(storage)
if strings.Contains(c.Request.Header.Get("Content-Type"), "multipart/form-data") {
requestBody = common.NewReplayableBodyReader(storage)
} else {
jsonData, err = storage.Bytes()
if err != nil {
return types.NewErrorWithStatusCode(err, types.ErrorCodeReadRequestBodyFailed, http.StatusBadRequest, types.ErrOptionWithSkipRetry())
}
}
} else {
convertedRequest, err := adaptor.ConvertImageRequest(c, info, *request)
if err != nil {
@@ -63,7 +79,7 @@ func ImageHelper(c *gin.Context, info *relaycommon.RelayInfo) (newAPIError *type
case *bytes.Buffer:
requestBody = convertedRequest.(io.Reader)
default:
jsonData, err := common.Marshal(convertedRequest)
jsonData, err = common.Marshal(convertedRequest)
if err != nil {
return types.NewError(err, types.ErrorCodeConvertRequestFailed, types.ErrOptionWithSkipRetry())
}
@@ -76,15 +92,45 @@ func ImageHelper(c *gin.Context, info *relaycommon.RelayInfo) (newAPIError *type
}
}
logger.LogDebug(c, "image request body: %s", jsonData)
body, closer, err := relaycommon.NewOutboundJSONBody(jsonData)
}
}
if jsonData != nil {
// This is a different trust boundary from ingress: channel overrides
// and pass-through bodies can change the quantity actually submitted.
var outbound struct {
N *uint `json:"n"`
Parameters *dto.ImageBillingParameters `json:"parameters"`
}
if err := common.Unmarshal(jsonData, &outbound); err != nil {
return types.NewErrorWithStatusCode(fmt.Errorf("invalid image billing parameters: %w", err), types.ErrorCodeInvalidRequest, http.StatusBadRequest, types.ErrOptionWithSkipRetry())
}
quantityRequest := dto.ImageRequest{N: outbound.N, BillingParameters: outbound.Parameters}
if quantityRequest.N == nil {
quantityRequest.N = common.GetPointer(uint(imageCount))
}
imageCount, err = quantityRequest.ImageCount(info.ChannelType == constant.ChannelTypeAli)
if err != nil {
return types.NewErrorWithStatusCode(err, types.ErrorCodeInvalidRequest, http.StatusBadRequest, types.ErrOptionWithSkipRetry())
}
promptExtend = outbound.Parameters != nil && outbound.Parameters.PromptExtend != nil && *outbound.Parameters.PromptExtend
if info.ChannelType == constant.ChannelTypeAli {
// Always send the same explicit quantity that is reserved, including
// when an empty parameters object accompanies a top-level n.
jsonData, err = sjson.SetBytes(jsonData, "parameters.n", imageCount)
if err != nil {
return types.NewError(err, types.ErrorCodeConvertRequestFailed, types.ErrOptionWithSkipRetry())
}
defer closer.Close()
jsonData = nil
requestBody = body
}
logger.LogDebug(c, "image request body: %s", jsonData)
body, closer, err := relaycommon.NewOutboundJSONBody(jsonData)
if err != nil {
return types.NewError(err, types.ErrorCodeConvertRequestFailed, types.ErrOptionWithSkipRetry())
}
defer closer.Close()
requestBody = body
}
if billingErr := service.PrepareImageBillingForRequest(c, info, imageCount, promptExtend); billingErr != nil {
return billingErr
}
statusCodeMappingStr := c.GetString("status_code_mapping")
+25 -11
View File
@@ -2,7 +2,6 @@ package relay
import (
"bytes"
"encoding/json"
"fmt"
"io"
"log"
@@ -125,7 +124,7 @@ func RelayMidjourneyNotify(c *gin.Context) *dto.MidjourneyResponse {
midjourneyTask.FinishTime = midjRequest.FinishTime
midjourneyTask.ImageUrl = midjRequest.ImageUrl
midjourneyTask.VideoUrl = midjRequest.VideoUrl
videoUrlsStr, _ := json.Marshal(midjRequest.VideoUrls)
videoUrlsStr, _ := common.Marshal(midjRequest.VideoUrls)
midjourneyTask.VideoUrls = string(videoUrlsStr)
midjourneyTask.Status = midjRequest.Status
midjourneyTask.FailReason = midjRequest.FailReason
@@ -167,21 +166,21 @@ func coverMidjourneyTaskDto(c *gin.Context, originTask *model.Midjourney) (midjo
midjourneyTask.Prompt = originTask.Prompt
if originTask.Buttons != "" {
var buttons []dto.ActionButton
err := json.Unmarshal([]byte(originTask.Buttons), &buttons)
err := common.Unmarshal([]byte(originTask.Buttons), &buttons)
if err == nil {
midjourneyTask.Buttons = buttons
}
}
if originTask.VideoUrls != "" {
var videoUrls []dto.ImgUrls
err := json.Unmarshal([]byte(originTask.VideoUrls), &videoUrls)
err := common.Unmarshal([]byte(originTask.VideoUrls), &videoUrls)
if err == nil {
midjourneyTask.VideoUrls = videoUrls
}
}
if originTask.Properties != "" {
var properties dto.Properties
err := json.Unmarshal([]byte(originTask.Properties), &properties)
err := common.Unmarshal([]byte(originTask.Properties), &properties)
if err == nil {
midjourneyTask.Properties = &properties
}
@@ -228,7 +227,9 @@ func RelaySwapFace(c *gin.Context, info *relaycommon.RelayInfo) *dto.MidjourneyR
requestURL := getMjRequestPath(c.Request.URL.String())
baseURL := c.GetString("base_url")
fullRequestURL := fmt.Sprintf("%s%s", baseURL, requestURL)
service.RequestPolicy(c).BeginAttempt(&model.Channel{Id: c.GetInt("channel_id")}, info.UsingGroup)
mjResp, _, err := service.DoMidjourneyHttpRequest(c, time.Second*60, fullRequestURL)
accepted := service.RecordMidjourneyPolicyResponse(c, mjResp, err)
if err != nil {
return &mjResp.Response
}
@@ -268,11 +269,15 @@ func RelaySwapFace(c *gin.Context, info *relaycommon.RelayInfo) *dto.MidjourneyR
if billingErr != nil {
common.SysLog("error settling Midjourney quota: " + billingErr.Error())
}
if accepted {
service.MarkRequestPolicySuccess(c, nil)
}
if billingApplied {
billingChannelId := midjourneyTask.GetBillingChannelId()
tokenName := c.GetString("token_name")
logContent := fmt.Sprintf("模型固定价格 %.2f,分组倍率 %.2f,操作 %s", priceData.ModelPrice, priceData.GroupRatioInfo.GroupRatio, constant.MjActionSwapFace)
other := service.GenerateMjOtherInfo(info, priceData)
service.AppendRelayLogAdminInfo(c, info, other)
model.RecordConsumeLog(c, info.UserId, model.RecordConsumeLogParams{
ChannelId: billingChannelId,
ModelName: modelName,
@@ -287,7 +292,7 @@ func RelaySwapFace(c *gin.Context, info *relaycommon.RelayInfo) *dto.MidjourneyR
model.UpdateChannelUsedQuota(billingChannelId, midjourneyTask.Quota)
}
c.Writer.WriteHeader(mjResp.StatusCode)
respBody, err := json.Marshal(midjResponse)
respBody, err := common.Marshal(midjResponse)
if err != nil {
return service.MidjourneyErrorWrapper(constant.MjRequestError, "unmarshal_response_body_failed")
}
@@ -323,7 +328,7 @@ func RelayMidjourneyTaskImageSeed(c *gin.Context) *dto.MidjourneyResponse {
}
midjResponse := &midjResponseWithStatus.Response
c.Writer.WriteHeader(midjResponseWithStatus.StatusCode)
respBody, err := json.Marshal(midjResponse)
respBody, err := common.Marshal(midjResponse)
if err != nil {
return service.MidjourneyErrorWrapper(constant.MjRequestError, "unmarshal_response_body_failed")
}
@@ -346,7 +351,7 @@ func RelayMidjourneyTask(c *gin.Context, relayMode int) *dto.MidjourneyResponse
}
}
midjourneyTask := coverMidjourneyTaskDto(c, originTask)
respBody, err = json.Marshal(midjourneyTask)
respBody, err = common.Marshal(midjourneyTask)
if err != nil {
return &dto.MidjourneyResponse{
Code: 4,
@@ -375,7 +380,7 @@ func RelayMidjourneyTask(c *gin.Context, relayMode int) *dto.MidjourneyResponse
if tasks == nil {
tasks = make([]dto.MidjourneyDto, 0)
}
respBody, err = json.Marshal(tasks)
respBody, err = common.Marshal(tasks)
if err != nil {
return &dto.MidjourneyResponse{
Code: 4,
@@ -539,7 +544,9 @@ func RelayMidjourneySubmit(c *gin.Context, relayInfo *relaycommon.RelayInfo) *dt
}
}
service.RequestPolicy(c).BeginAttempt(&model.Channel{Id: c.GetInt("channel_id")}, relayInfo.UsingGroup)
midjResponseWithStatus, responseBody, err := service.DoMidjourneyHttpRequest(c, time.Second*60, fullRequestURL)
accepted := service.RecordMidjourneyPolicyResponse(c, midjResponseWithStatus, err)
if err != nil {
return &midjResponseWithStatus.Response
}
@@ -576,8 +583,11 @@ func RelayMidjourneySubmit(c *gin.Context, relayInfo *relaycommon.RelayInfo) *dt
if err != nil {
common.SysLog("get_channel_null: " + err.Error())
}
if channel.GetAutoBan() && common.AutomaticDisableChannelEnabled {
model.UpdateChannelStatus(midjourneyTask.ChannelId, "", 2, "No available account instance")
if channel != nil && channel.GetAutoBan() && common.AutomaticDisableChannelEnabled {
if model.UpdateChannelStatus(midjourneyTask.ChannelId, "", common.ChannelStatusManuallyDisabled, "No available account instance") {
service.CloseActiveWebSocketsForChannel(midjourneyTask.ChannelId, service.ChannelDisabledCloseReason)
service.RequestPolicy(c).AddEvent(service.PolicyEvent{ChannelID: midjourneyTask.ChannelId, Decision: service.PolicyDecision{Action: "stop", Reason: "task_failed", Source: "upstream"}, Health: "channel_disable_requested"})
}
}
}
if midjResponse.Code != 1 && midjResponse.Code != 21 && midjResponse.Code != 22 {
@@ -633,11 +643,15 @@ func RelayMidjourneySubmit(c *gin.Context, relayInfo *relaycommon.RelayInfo) *dt
if billingErr != nil {
common.SysLog("error settling Midjourney quota: " + billingErr.Error())
}
if accepted {
service.MarkRequestPolicySuccess(c, nil)
}
if billingApplied {
billingChannelId := midjourneyTask.GetBillingChannelId()
tokenName := c.GetString("token_name")
logContent := fmt.Sprintf("模型固定价格 %.2f,分组倍率 %.2f,操作 %s,ID %s", priceData.ModelPrice, priceData.GroupRatioInfo.GroupRatio, midjRequest.Action, midjResponse.Result)
other := service.GenerateMjOtherInfo(relayInfo, priceData)
service.AppendRelayLogAdminInfo(c, relayInfo, other)
model.RecordConsumeLog(c, relayInfo.UserId, model.RecordConsumeLogParams{
ChannelId: billingChannelId,
ModelName: modelName,
+121
View File
@@ -0,0 +1,121 @@
package relay
import (
"context"
"fmt"
"github.com/QuantumNous/new-api/common"
"github.com/QuantumNous/new-api/constant"
"github.com/QuantumNous/new-api/model"
"github.com/QuantumNous/new-api/pkg/billingexpr"
"github.com/QuantumNous/new-api/relay/channel"
"github.com/QuantumNous/new-api/relay/channel/openai"
relaycommon "github.com/QuantumNous/new-api/relay/common"
"github.com/QuantumNous/new-api/relaykit/dto"
"github.com/QuantumNous/new-api/service"
hosttypes "github.com/QuantumNous/new-api/types"
"github.com/gin-gonic/gin"
"github.com/glebarez/sqlite"
"github.com/stretchr/testify/require"
"gorm.io/gorm"
"io"
"net/http"
"net/http/httptest"
"strings"
"testing"
"time"
)
func TestRelayAttemptMismatchCollectsUsageAndChargesRequestedPrice(t *testing.T) {
service.InitTokenEncoders()
db, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{})
require.NoError(t, err)
sqlDB, err := db.DB()
require.NoError(t, err)
sqlDB.SetMaxOpenConns(1)
previousDB, previousLog, previousRedis := model.DB, model.LOG_DB, common.RedisEnabled
model.DB = db
model.LOG_DB = db
common.RedisEnabled = false
t.Cleanup(func() {
model.DB = previousDB
model.LOG_DB = previousLog
common.RedisEnabled = previousRedis
require.NoError(t, sqlDB.Close())
})
require.NoError(t, db.AutoMigrate(&model.User{}, &model.Token{}, &model.Channel{}, &model.RelayAttempt{}, &model.Log{}, &model.QuotaData{}))
user := model.User{Username: "attempt-relay", Quota: 10000}
require.NoError(t, db.Create(&user).Error)
token := model.Token{UserId: user.Id, Key: "attempt-relay", Status: common.TokenStatusEnabled, RemainQuota: 10000}
require.NoError(t, db.Create(&token).Error)
expression := `p * 2 + c * 4 + cr * 0.5`
for i, ratio := range []float64{1, 3, 1} {
selected := model.Channel{Name: fmt.Sprint(i), Status: common.ChannelStatusEnabled}
require.NoError(t, db.Create(&selected).Error)
// The final usage is buffered behind the rejected prefix: Drain must reuse
// the same reader rather than lose bytes already read from the socket.
body := "data: {\"type\":\"response.created\",\"response\":{\"model\":\"gpt-5.6-luna\"}}\n\n" +
"data: {\"type\":\"response.output_text.delta\",\"delta\":\"hidden generated text\"}\n\n" +
"data: {\"type\":\"response.completed\",\"response\":{\"model\":\"gpt-5.6-luna\",\"usage\":{\"input_tokens\":10,\"output_tokens\":2,\"total_tokens\":12,\"input_tokens_details\":{\"cached_tokens\":4}}}}\n\ndata: [DONE]\n\n"
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
w.Header().Set("Content-Type", "text/event-stream")
_, _ = io.WriteString(w, body)
}))
ctx, cancel := context.WithCancel(context.Background())
c, writer := gin.CreateTestContext(httptest.NewRecorder())
_ = writer
c.Request = httptest.NewRequest("POST", "/v1/responses", strings.NewReader(`{"model":"gpt-6-astra","input":"secret prompt"}`)).WithContext(ctx)
c.Set("channel_id", selected.Id)
c.Set("id", user.Id)
c.Set("token_id", token.Id)
common.SetContextKey(c, constant.ContextKeyUserName, user.Username)
info := &relaycommon.RelayInfo{RequestId: fmt.Sprintf("attempt-request-%d", i), UserId: user.Id, TokenId: token.Id, TokenKey: token.Key, OriginModelName: "gpt-6-astra", UsingGroup: fmt.Sprint(ratio), UserSetting: dto.UserSetting{BillingPreference: "wallet_only"}, StartTime: time.Now(), DurableAttemptBilling: true, ResponseGuard: &relaycommon.ResponseModelGuard{}, ChannelMeta: &relaycommon.ChannelMeta{ChannelId: selected.Id, UpstreamModelName: "gpt-6-astra"}, PriceData: hosttypes.PriceData{GroupRatioInfo: hosttypes.GroupRatioInfo{GroupRatio: ratio}}, TieredBillingSnapshot: &billingexpr.BillingSnapshot{BillingMode: "tiered_expr", ExprString: expression, ExprHash: billingexpr.ExprHashString(expression), QuotaPerUnit: common.QuotaPerUnit, GroupRatio: ratio, EstimatedQuotaAfterGroup: 100}}
info.SetEstimatePromptTokens(30)
info.AstraAutoSwitch = i < 2
if !info.AstraAutoSwitch {
info.ResponseGuard = nil
}
require.Nil(t, service.BeginRelayAttempt(c, info))
req, err := http.NewRequest("POST", server.URL, strings.NewReader("{}"))
require.NoError(t, err)
resp, err := channel.DoRequest(c, req, info)
if info.AstraAutoSwitch {
require.Error(t, err)
require.Nil(t, resp)
require.False(t, c.Writer.Written())
require.True(t, info.AttemptDetached)
require.False(t, service.RequestPolicy(c).Successful, "collector cannot mark the client request successful")
} else {
require.NoError(t, err)
usage, apiErr := openai.OaiResponsesHandler(c, info, resp)
require.Nil(t, apiErr)
require.True(t, c.Writer.Written())
require.True(t, info.ResponseModel.Mismatch)
ConsumeResponsesQuota(c, info, usage)
}
var attempt model.RelayAttempt
require.Eventually(t, func() bool {
return db.First(&attempt, info.DurableAttemptID).Error == nil && attempt.State == "settled"
}, time.Second*5, time.Millisecond*10)
require.Equal(t, int(11*ratio), attempt.Quota)
if info.AstraAutoSwitch {
require.Equal(t, "mismatch_rejected", attempt.GuardResult)
} else {
require.Equal(t, "mismatch_allowed", attempt.GuardResult)
require.Contains(t, attempt.LogJSON, "mismatch_allowed")
}
require.Equal(t, "upstream", attempt.UsageSource)
require.NotContains(t, attempt.LogJSON, "secret prompt")
require.NotContains(t, attempt.LogJSON, "hidden generated text")
require.NoError(t, model.DeliverRelayAttemptLogs())
service.EndFailedRelayAttempt(c, info)
cancel()
server.Close()
}
require.NoError(t, db.First(&user, user.Id).Error)
require.Equal(t, 9945, user.Quota)
require.NoError(t, db.First(&token, token.Id).Error)
require.Equal(t, 55, token.UsedQuota)
var logs int64
require.NoError(t, db.Model(&model.Log{}).Count(&logs).Error)
require.EqualValues(t, 3, logs)
}
+36 -14
View File
@@ -12,8 +12,10 @@ import (
"github.com/QuantumNous/new-api/common"
"github.com/QuantumNous/new-api/constant"
"github.com/QuantumNous/new-api/dto"
"github.com/QuantumNous/new-api/logger"
"github.com/QuantumNous/new-api/model"
"github.com/QuantumNous/new-api/pkg/billingexpr"
"github.com/QuantumNous/new-api/pkg/jsplugin"
"github.com/QuantumNous/new-api/relay/channel"
"github.com/QuantumNous/new-api/relay/channel/task/taskcommon"
relaycommon "github.com/QuantumNous/new-api/relay/common"
@@ -248,25 +250,29 @@ func RelayTaskSubmit(c *gin.Context, info *relaycommon.RelayInfo) (*TaskSubmitRe
info.OriginModelName = modelName
var priceData types.PriceData
var err error
useTiered := billing_setting.GetBillingMode(modelName) == billing_setting.BillingModeTieredExpr
var exprStr string
var exists bool
if useTiered {
exprStr, exists = billing_setting.GetBillingExpr(modelName)
} else if info.IsModelMapped {
if billing_setting.GetBillingMode(info.UpstreamModelName) == billing_setting.BillingModeTieredExpr {
if tailExpr, tailOK := billing_setting.GetBillingExpr(info.UpstreamModelName); tailOK && strings.TrimSpace(tailExpr) != "" {
exprStr = tailExpr
exists = true
useTiered = true
}
}
pluginKey := c.GetString("task_plugin_key")
pinnedValue, _ := c.Get(jsplugin.ContextKeyPinnedPlugin)
pinnedPlugin, _ := pinnedValue.(jsplugin.PinnedPlugin)
if pinnedPlugin.Plugin != nil {
pluginKey = pinnedPlugin.Plugin.Meta.Key
}
exprStr, exists := billing_setting.ResolveTaskBillingExpr(pluginKey, modelName, info.UpstreamModelName)
useTiered := exists || billing_setting.GetBillingMode(modelName) == billing_setting.BillingModeTieredExpr
if useTiered {
provider, supported := adaptor.(channel.TaskUsageFactsProvider)
if billingexpr.UsesFixedPricing(exprStr) {
return nil, service.TaskErrorWrapper(fmt.Errorf("fixed pricing is not supported for task usage expressions"), "model_price_error", http.StatusBadRequest)
}
if !exists || !supported {
return nil, service.TaskErrorWrapper(fmt.Errorf("task model %s has no usage expression or meter", modelName), "model_price_error", http.StatusBadRequest)
}
sharedModel := pinnedPlugin.Generation.SharedModel(modelName) || pinnedPlugin.Generation.SharedModel(info.UpstreamModelName)
if sharedModel && pinnedPlugin.Plugin != nil {
schema, _ := pinnedPlugin.Plugin.Meta.UsageForModel(info.UpstreamModelName)
if !billing_setting.TaskExprCompatible(exprStr, schema) {
return nil, service.TaskErrorWrapper(fmt.Errorf("task model %s pricing is not configured for plugin %s", modelName, pluginKey), "model_price_error", http.StatusBadRequest)
}
}
var facts map[string]any
if validatedProvider, ok := adaptor.(channel.TaskValidatedUsageFactsProvider); ok {
facts, err = validatedProvider.ExtractUsageFactsValidated(c, info)
@@ -364,7 +370,21 @@ func RelayTaskSubmit(c *gin.Context, info *relaycommon.RelayInfo) (*TaskSubmitRe
// 11. 提交后计费调整:让适配器根据上游实际返回调整 OtherRatios
finalQuota := info.PriceData.Quota
if info.TieredBillingSnapshot == nil {
if parsed.Immediate != nil && parsed.Immediate.Status == model.TaskStatusFailure {
finalQuota = 0
} else if snap := info.TieredBillingSnapshot; snap != nil {
if parsed.Immediate != nil && parsed.Immediate.Status == model.TaskStatusSuccess && len(parsed.Immediate.UsageFacts) > 0 {
settlement, facts, err := service.EvaluateTaskCompletionUsage(snap, parsed.Immediate.UsageFacts)
if err != nil {
logger.LogWarn(c, fmt.Sprintf("task immediate usage settlement failed; retaining reserved quota: %v", err))
} else {
finalQuota = settlement.ActualQuotaAfterGroup
snap.UsageFacts = facts
snap.EstimatedTier = settlement.MatchedTier
noteTaskQuotaClamp(info, settlement.Clamp)
}
}
} else {
if adjustedRatios := adaptor.AdjustBillingOnSubmit(info, parsed.TaskData); len(adjustedRatios) > 0 {
if adjustedQuota, ok := recalcQuotaFromRatios(info, adjustedRatios); ok {
// 基于调整后的 ratios 重新计算 quota
@@ -375,6 +395,8 @@ func RelayTaskSubmit(c *gin.Context, info *relaycommon.RelayInfo) (*TaskSubmitRe
}
}
info.PriceData.Quota = finalQuota
return &TaskSubmitResult{
UpstreamTaskID: parsed.UpstreamTaskID,
TaskData: parsed.TaskData,
+171 -2
View File
@@ -3,10 +3,12 @@ package relay
import (
"net/http"
"net/http/httptest"
"strings"
"testing"
"github.com/QuantumNous/new-api/common"
"github.com/QuantumNous/new-api/constant"
"github.com/QuantumNous/new-api/dto"
"github.com/QuantumNous/new-api/model"
"github.com/QuantumNous/new-api/pkg/billingexpr"
pluginruntime "github.com/QuantumNous/new-api/pkg/jsplugin"
@@ -15,10 +17,82 @@ import (
"github.com/QuantumNous/new-api/setting/billing_setting"
"github.com/QuantumNous/new-api/setting/config"
"github.com/gin-gonic/gin"
"github.com/glebarez/sqlite"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
"gorm.io/gorm"
)
func setupRelayChannelDB(t *testing.T) *gorm.DB {
t.Helper()
previousDB := model.DB
previousType := common.MainDatabaseType()
previousCache := common.MemoryCacheEnabled
database, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{})
require.NoError(t, err)
sqlDB, err := database.DB()
require.NoError(t, err)
sqlDB.SetMaxOpenConns(1)
require.NoError(t, database.AutoMigrate(&model.Channel{}))
model.DB = database
common.SetMainDatabaseType(common.DatabaseTypeSQLite)
common.MemoryCacheEnabled = false
t.Cleanup(func() {
model.DB = previousDB
common.SetMainDatabaseType(previousType)
common.MemoryCacheEnabled = previousCache
require.NoError(t, sqlDB.Close())
})
return database
}
func TestApplyChannelPinPreservesOriginTasksAndRetryMode(t *testing.T) {
database := setupRelayChannelDB(t)
channel := &model.Channel{Name: "origin-channel", Key: "sk-test", Status: common.ChannelStatusEnabled, Type: constant.ChannelTypeDoubaoVideo}
require.NoError(t, database.Create(channel).Error)
originTask := &model.Task{
TaskID: "task-lock", ChannelId: channel.Id, Action: "text_to_video", Status: model.TaskStatusSuccess,
PrivateData: model.TaskPrivateData{UpstreamTaskID: "upstream-task-lock"},
Data: []byte(`{"id":"upstream-task-lock"}`),
}
for _, tc := range []struct {
name string
tokenPin bool
apply func(*gin.Context, *relaycommon.RelayInfo) *dto.TaskError
}{
{name: "origin affinity", apply: ApplyOriginTaskAffinity},
{name: "same channel retry", apply: ApplyChannelPin},
{name: "token pin suppresses channel lock", tokenPin: true, apply: ApplyChannelPin},
} {
t.Run(tc.name, func(t *testing.T) {
c, _ := gin.CreateTestContext(httptest.NewRecorder())
common.SetContextKey(c, constant.ContextKeyOriginTasks, []*model.Task{originTask})
constraints := service.GetChannelConstraints(c)
constraints.AddPin(dto.ChannelPin{ChannelId: channel.Id, Source: dto.PinSourceOriginTask, Rank: dto.PinRankOriginTask, RetryMode: dto.PinRetrySameChannel})
if tc.tokenPin {
constraints.AddPin(dto.ChannelPin{ChannelId: channel.Id, Source: dto.PinSourceToken, Rank: dto.PinRankToken, RetryMode: dto.PinRetrySingleAttempt})
}
info := &relaycommon.RelayInfo{TaskRelayInfo: &relaycommon.TaskRelayInfo{}}
require.Nil(t, tc.apply(c, info))
if tc.tokenPin {
assert.Nil(t, info.LockedChannel)
} else {
locked, ok := info.LockedChannel.(*model.Channel)
require.True(t, ok)
require.NotNil(t, locked)
assert.Equal(t, channel.Id, locked.Id)
}
require.Len(t, info.OriginTasks, 1)
assert.Equal(t, "task-lock", info.OriginTasks[0].TaskID)
assert.Equal(t, "upstream-task-lock", info.OriginTasks[0].UpstreamTaskID)
assert.Equal(t, "text_to_video", info.OriginTasks[0].Action)
assert.Equal(t, string(model.TaskStatusSuccess), info.OriginTasks[0].Status)
assert.Equal(t, []byte(originTask.Data), info.OriginTasks[0].Data)
})
}
}
func TestTaskModel2DtoNormalizesLegacyAction(t *testing.T) {
task := &model.Task{Action: "firstTailGenerate"}
@@ -153,6 +227,7 @@ func TestRelayTaskSubmitAliasBillingIdentityAndExprFallback(t *testing.T) {
const mapping = `{"alias-model":"declared-model"}`
const aliasExpr = `tier("alias", 2)`
const tailExpr = `tier("tail", 3)`
const legacyExpr = `tier("legacy", u("old_units") * 2)`
tests := []struct {
name string
@@ -160,6 +235,7 @@ func TestRelayTaskSubmitAliasBillingIdentityAndExprFallback(t *testing.T) {
exprs map[string]string
wantTiered bool
wantExpr string
source string
}{
{
name: "alias own tiered wins",
@@ -175,6 +251,16 @@ func TestRelayTaskSubmitAliasBillingIdentityAndExprFallback(t *testing.T) {
wantTiered: true,
wantExpr: tailExpr,
},
{
name: "stored expression keeps running after schema narrows",
modes: map[string]string{"declared-model": "tiered_expr"},
exprs: map[string]string{"declared-model": legacyExpr},
wantTiered: true,
wantExpr: legacyExpr,
source: strings.Replace(billingFallbackPlugin, `fetchMode:"per_task"`, `fetchMode:"per_task", usageSchema:{old_units:{type:"number",unit:"count"}}, usageProfiles:[{models:["declared-model"],schema:{seconds:{type:"number",unit:"second"}}}]`, 1) + `
export function extractUsage(){return {old_units:2};}
`,
},
{
name: "neither tiered uses ordinary pricing",
wantTiered: false,
@@ -204,7 +290,11 @@ func TestRelayTaskSubmitAliasBillingIdentityAndExprFallback(t *testing.T) {
c.Set("group", "default")
info.UserGroup = "default"
info.UsingGroup = "default"
pinMappingOrderPlugin(t, c, billingFallbackPlugin)
source := testCase.source
if source == "" {
source = billingFallbackPlugin
}
pinMappingOrderPlugin(t, c, source)
info.OriginModelName = "alias-model"
_, taskErr := RelayTaskSubmit(c, info)
@@ -218,11 +308,14 @@ func TestRelayTaskSubmitAliasBillingIdentityAndExprFallback(t *testing.T) {
assert.Equal(t, "declared-model", task.Properties.UpstreamModelName)
if testCase.wantTiered {
require.NotNil(t, info.TieredBillingSnapshot)
require.NotNil(t, info.TieredBillingSnapshot, "submission error: %+v", taskErr)
assert.Equal(t, "alias-model", info.TieredBillingSnapshot.ModelName)
assert.Equal(t, testCase.wantExpr, info.TieredBillingSnapshot.ExprString)
assert.Equal(t, billingexpr.ExprHashString(testCase.wantExpr), info.TieredBillingSnapshot.ExprHash)
assert.NotEqual(t, "model_price_error", taskErr.Code)
if testCase.wantExpr == legacyExpr {
assert.Equal(t, 4*common.QuotaPerUnit, info.TieredBillingSnapshot.EstimatedQuotaBeforeGroup)
}
} else {
assert.Nil(t, info.TieredBillingSnapshot)
assert.Equal(t, "model_price_error", taskErr.Code)
@@ -230,3 +323,79 @@ func TestRelayTaskSubmitAliasBillingIdentityAndExprFallback(t *testing.T) {
})
}
}
func TestSharedTaskBillingExpressionSelectionAndFrozenSettlement(t *testing.T) {
const baseExpr = `tier("base", u("seconds") * 2)`
const alphaExpr = `tier("alpha", u("seconds") * 3)`
const betaExpr = `tier("beta", u("credits") * 5)`
const aliasExpr = `tier("alias", u("credits") * 7)`
for _, tc := range []struct {
name, plugin, model, mapping, modelExpr, mode, wantExpr string
variants map[string]string
wantPriceError bool
}{
{name: "executing plugin override", plugin: "billing-beta", model: "declared-model", modelExpr: baseExpr, mode: "tiered_expr", variants: map[string]string{"billing-alpha::declared-model": alphaExpr, "billing-beta::declared-model": betaExpr}, wantExpr: betaExpr},
{name: "override ignores model mode", plugin: "billing-beta", model: "declared-model", modelExpr: baseExpr, mode: "ratio", variants: map[string]string{"billing-beta::declared-model": betaExpr}, wantExpr: betaExpr},
{name: "model expression fallback", plugin: "billing-alpha", model: "declared-model", modelExpr: baseExpr, mode: "tiered_expr", wantExpr: baseExpr},
{name: "alias override precedes mapped override", plugin: "billing-beta", model: "alias-model", mapping: `{"alias-model":"declared-model"}`, modelExpr: baseExpr, mode: "tiered_expr", variants: map[string]string{"billing-beta::declared-model": betaExpr, "billing-beta::alias-model": aliasExpr}, wantExpr: aliasExpr},
{name: "mapped override precedes model fallback", plugin: "billing-beta", model: "alias-model", mapping: `{"alias-model":"declared-model"}`, modelExpr: baseExpr, mode: "tiered_expr", variants: map[string]string{"billing-beta::declared-model": betaExpr}, wantExpr: betaExpr},
{name: "unconfigured plugin cannot use another schema", plugin: "billing-beta", model: "declared-model", modelExpr: baseExpr, mode: "tiered_expr", wantPriceError: true},
{name: "missing usage in skipped branch remains incompatible", plugin: "billing-beta", model: "declared-model", modelExpr: `true ? tier("free", 0) : tier("missing", u("seconds"))`, mode: "tiered_expr", wantPriceError: true},
{name: "fixed pricing is still rejected", plugin: "billing-beta", model: "declared-model", variants: map[string]string{"billing-beta::declared-model": `tier("fixed", fixed(1))`}, wantPriceError: true},
} {
t.Run(tc.name, func(t *testing.T) {
saveBillingConfig(t)
registry := pluginruntime.NewRegistry()
for _, spec := range []struct{ key, field, unit string }{{"billing-alpha", "seconds", "second"}, {"billing-beta", "credits", "credit"}} {
source := strings.ReplaceAll(billingFallbackPlugin, "bill-fallback", spec.key)
source = strings.Replace(source, `fetchMode:"per_task"`, `fetchMode:"per_task",usageSchema:{`+spec.field+`:{type:"number",unit:"`+spec.unit+`"}}`, 1)
source += `export function extractUsage(){return {` + spec.field + `:2};}`
_, err := registry.Register(source, pluginruntime.Options{})
require.NoError(t, err)
}
variants := tc.variants
if variants == nil {
variants = map[string]string{}
}
rawVariants, err := common.Marshal(variants)
require.NoError(t, err)
modes, err := common.Marshal(map[string]string{"declared-model": tc.mode})
require.NoError(t, err)
expressions, err := common.Marshal(map[string]string{"declared-model": tc.modelExpr})
require.NoError(t, err)
require.NoError(t, config.GlobalConfig.LoadFromDB(map[string]string{
billing_setting.PluginBillingExprOption: string(rawVariants), "billing_setting.billing_mode": string(modes), "billing_setting.billing_expr": string(expressions),
}))
c, info := newTaskSubmitContext(t, tc.model, tc.mapping)
c.Set("group", "default")
c.Set("task_plugin_key", tc.plugin)
info.UserGroup = "default"
info.UsingGroup = "default"
info.OriginModelName = tc.model
plugin, ok := registry.Generation().Get(tc.plugin)
require.True(t, ok)
c.Set(pluginruntime.ContextKeyPinnedPlugin, pluginruntime.PinnedPlugin{Generation: registry.Generation(), Plugin: plugin})
_, taskErr := RelayTaskSubmit(c, info)
require.NotNil(t, taskErr) // This fixture stops at reservation, before upstream submission.
if tc.wantPriceError {
assert.Equal(t, "model_price_error", taskErr.Code)
assert.Nil(t, info.TieredBillingSnapshot)
return
}
require.NotNil(t, info.TieredBillingSnapshot, "submission error: %+v", taskErr)
assert.Equal(t, tc.wantExpr, info.TieredBillingSnapshot.ExprString)
assert.NotEqual(t, "model_price_error", taskErr.Code)
require.NoError(t, config.GlobalConfig.LoadFromDB(map[string]string{billing_setting.PluginBillingExprOption: `{}`, "billing_setting.billing_expr": `{}`}))
field := "seconds"
if tc.plugin == "billing-beta" {
field = "credits"
}
result, usage, err := service.EvaluateTaskCompletionUsage(info.TieredBillingSnapshot, map[string]any{field: float64(4)})
require.NoError(t, err)
assert.Equal(t, float64(4), usage[field])
assert.Equal(t, float64(2), info.TieredBillingSnapshot.UsageFacts[field])
assert.Equal(t, 2*info.TieredBillingSnapshot.EstimatedQuotaAfterGroup, result.ActualQuotaAfterGroup)
assert.Equal(t, tc.wantExpr, info.TieredBillingSnapshot.ExprString)
})
}
}
+87
View File
@@ -0,0 +1,87 @@
package relay
import (
"errors"
"fmt"
"net/http"
"strings"
"github.com/QuantumNous/new-api/constant"
"github.com/QuantumNous/new-api/logger"
relaycommon "github.com/QuantumNous/new-api/relay/common"
"github.com/QuantumNous/new-api/relay/helper"
"github.com/QuantumNous/new-api/relaykit/dto"
"github.com/QuantumNous/new-api/relaykit/types"
"github.com/QuantumNous/new-api/service"
"github.com/QuantumNous/new-api/setting"
"github.com/gin-gonic/gin"
"github.com/samber/lo"
)
// PrepareRequestBilling estimates and reserves one request's charge. Transports
// provide the current request body through BodyStorage or BillingRequestInput;
// channel retries retain the resulting billing session and pricing snapshot.
func PrepareRequestBilling(c *gin.Context, info *relaycommon.RelayInfo) *types.NewAPIError {
durable := info.OriginModelName == "gpt-6-astra" && c.Request.Method == http.MethodPost && c.Request.URL.Path == "/v1/responses"
needSensitiveCheck := setting.ShouldCheckPromptSensitive()
meta := &types.TokenCountMeta{TokenType: types.TokenTypeTokenizer}
if info.Request != nil && (needSensitiveCheck || constant.CountToken || durable) {
meta = info.Request.GetTokenCountMeta()
} else {
// Avoid building CombineText when only the pricing quantities are needed.
switch request := info.Request.(type) {
case *dto.GeneralOpenAIRequest:
meta.MaxTokens = int(max(lo.FromPtr(request.MaxTokens), lo.FromPtr(request.MaxCompletionTokens)))
case *dto.OpenAIResponsesRequest:
meta.MaxTokens = int(lo.FromPtr(request.MaxOutputTokens))
case *dto.ClaudeRequest:
meta.MaxTokens = int(lo.FromPtr(request.MaxTokens))
case *dto.ImageRequest:
meta = request.GetTokenCountMeta()
}
}
if needSensitiveCheck && meta != nil {
if contains, words := service.CheckSensitiveText(meta.CombineText); contains {
service.RequestPolicy(c).AddEvent(service.PolicyEvent{ErrorCode: string(types.ErrorCodeSensitiveWordsDetected), ErrorSource: "local", Decision: service.PolicyDecision{Action: "stop", Reason: "local_rejection", Source: "global"}, Health: "unchanged"})
message := fmt.Sprintf("user sensitive words detected: %s", strings.Join(words, ", "))
logger.LogWarn(c, message)
return types.NewError(errors.New(message), types.ErrorCodeSensitiveWordsDetected)
}
}
tokens, err := service.EstimateRequestToken(c, meta, info)
if durable { tokens,err=service.CountRequestToken(c,meta,info) }
if err != nil {
return types.NewError(err, types.ErrorCodeCountTokenFailed)
}
info.SetEstimatePromptTokens(tokens)
priceData, err := helper.ModelPriceHelper(c, info, tokens, meta)
if err != nil {
return types.NewError(err, types.ErrorCodeModelPriceError, types.ErrOptionWithStatusCode(http.StatusBadRequest))
}
if durable {
info.DurableAttemptBilling = true
return nil
}
if priceData.FreeModel {
logger.LogInfo(c, fmt.Sprintf("模型 %s 免费,跳过预扣费", info.OriginModelName))
return nil
}
return service.PreConsumeBilling(c, priceData.QuotaToPreConsume, info)
}
// RefundFailedRequestBilling applies the common final-failure policy after all
// eligible attempts have ended. A settled BillingSession never refunds again.
func RefundFailedRequestBilling(c *gin.Context, info *relaycommon.RelayInfo, apiErr *types.NewAPIError) *types.NewAPIError {
if apiErr == nil {
return nil
}
apiErr = service.NormalizeViolationFeeError(apiErr)
if info.Billing != nil {
info.Billing.Refund(c)
}
service.ChargeViolationFeeIfNeeded(c, info, apiErr)
return apiErr
}
+447
View File
@@ -0,0 +1,447 @@
package relay
import (
"context"
"fmt"
"io"
"net/http"
"net/http/httptest"
"strings"
"testing"
"time"
"github.com/QuantumNous/new-api/common"
"github.com/QuantumNous/new-api/constant"
"github.com/QuantumNous/new-api/relay/channel"
"github.com/QuantumNous/new-api/relay/channel/claude"
"github.com/QuantumNous/new-api/relay/channel/gemini"
"github.com/QuantumNous/new-api/relay/channel/openai"
relaycommon "github.com/QuantumNous/new-api/relay/common"
relayconstant "github.com/QuantumNous/new-api/relay/constant"
"github.com/QuantumNous/new-api/relaykit/dto"
"github.com/QuantumNous/new-api/relaykit/types"
"github.com/QuantumNous/new-api/service"
"github.com/gin-gonic/gin"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
)
func TestResponseGuardBeforeDelivery(t *testing.T) {
for _, tc := range []struct{ name, content, body, reason string }{
{"astra", "application/json", `{"model":"gpt-6-astra","output":[]}`, "accepted"},
{"luna", "application/json", `{"model":"gpt-5.6-luna","usage":{"input_tokens":8,"output_tokens":2}}`, "model_mismatch"},
{"missing", "application/json", `{"output":[]}`, "missing_model"},
{"not a prefix match", "application/json", `{"model":"gpt-6-astra-other"}`, "model_mismatch"},
{"sse astra", "text/event-stream", "event: response.created\ndata: {\"type\":\"response.created\",\"response\":{\"model\":\"gpt-6-astra\"}}\n\ndata: [DONE]\n\n", "accepted"},
{"sse luna", "text/event-stream", "data: {\"type\":\"response.created\",\"response\":{\"model\":\"gpt-5.6-luna\"}}\n\n", "model_mismatch"},
{"sse no model", "text/event-stream", "data: {\"type\":\"response.output_text.delta\",\"delta\":\"hidden\"}\n\ndata: [DONE]\n\n", "missing_model"},
{"sse buffer bound", "text/event-stream", ":" + strings.Repeat("x", 1<<20), "confirmation_buffer_limit"},
} {
t.Run(tc.name, func(t *testing.T) {
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
w.Header().Set("Content-Type", tc.content)
w.Header().Set("X-Codex-Turn-State", "must-not-leak")
_, _ = io.WriteString(w, tc.body)
}))
defer server.Close()
writer := httptest.NewRecorder()
c, _ := gin.CreateTestContext(writer)
c.Request = httptest.NewRequest("POST", "/v1/responses", strings.NewReader("{}"))
info := &relaycommon.RelayInfo{AstraAutoSwitch: true, ResponseGuard: &relaycommon.ResponseModelGuard{}, ChannelMeta: &relaycommon.ChannelMeta{ChannelId: 1}, IsStream: strings.Contains(tc.content, "event-stream")}
req, err := http.NewRequest("POST", server.URL, strings.NewReader("{}"))
require.NoError(t, err)
resp, err := channel.DoRequest(c, req, info)
assert.False(t, c.Writer.Written())
assert.Empty(t, writer.Header().Get("X-Codex-Turn-State"))
assert.Equal(t, tc.reason, info.ResponseGuard.Last().Reason)
if tc.reason == "accepted" {
require.NoError(t, err)
defer resp.Body.Close()
body, err := io.ReadAll(resp.Body)
require.NoError(t, err)
assert.Equal(t, tc.body, string(body))
} else {
require.Error(t, err)
assert.Nil(t, resp)
assert.Equal(t, http.StatusConflict, info.ResponseGuard.Last().APIError().StatusCode)
}
})
}
}
func TestResponseGuardAcceptsModelAfterFormerTenSecondCutoff(t *testing.T) {
previous := constant.StreamingTimeout
constant.StreamingTimeout = 30
t.Cleanup(func() { constant.StreamingTimeout = previous })
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
select {
case <-time.After(11 * time.Second):
w.Header().Set("Content-Type", "application/json")
_, _ = io.WriteString(w, `{"model":"gpt-6-astra","output":[]}`)
case <-r.Context().Done():
}
}))
defer server.Close()
c, _ := gin.CreateTestContext(httptest.NewRecorder())
c.Request = httptest.NewRequest("POST", "/v1/responses", nil)
info := &relaycommon.RelayInfo{AstraAutoSwitch: true, ResponseGuard: &relaycommon.ResponseModelGuard{}, ChannelMeta: &relaycommon.ChannelMeta{ChannelId: 1}}
req, err := http.NewRequest("POST", server.URL, strings.NewReader("{}"))
require.NoError(t, err)
resp, err := channel.DoRequest(c, req, info)
require.NoError(t, err)
defer resp.Body.Close()
assert.Equal(t, "accepted", info.ResponseGuard.Last().Reason)
assert.False(t, c.Writer.Written())
}
func TestResponseGuardOnlyEnablesForAstraRequests(t *testing.T) {
tests := []struct {
name string
path string
model string
enabled bool
want bool
}{
{name: "astra enabled", path: "/v1/responses", model: "gpt-6-astra", enabled: true, want: true},
{name: "direct luna unaffected", path: "/v1/responses", model: "gpt-5.6-luna", enabled: true},
{name: "token opt out", path: "/v1/responses", model: "gpt-6-astra"},
{name: "other endpoint", path: "/v1/chat/completions", model: "gpt-6-astra", enabled: true},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
info := &relaycommon.RelayInfo{
OriginModelName: tt.model,
AstraAutoSwitch: tt.enabled,
ChannelMeta: &relaycommon.ChannelMeta{
ChannelSetting: dto.ChannelSettings{ResponsesAstraModelGuard: common.GetPointer(!tt.enabled)},
},
}
assert.Equal(t, tt.want, shouldEnableAstraResponseGuard(tt.path, info))
})
}
}
func TestResponseGuardTokenOptOutAllowsMismatchDespiteChannelSetting(t *testing.T) {
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
w.Header().Set("Content-Type", "application/json")
_, _ = io.WriteString(w, `{"model":"gpt-5.6-luna","output":[]}`)
}))
defer server.Close()
c, _ := gin.CreateTestContext(httptest.NewRecorder())
c.Request = httptest.NewRequest("POST", "/v1/responses", strings.NewReader("{}"))
guard := &relaycommon.ResponseModelGuard{}
_, previous := guard.Begin(context.Background(), 1, 1)
previous.Close()
previous.Reason = "http_error"
info := &relaycommon.RelayInfo{
ResponseGuard: guard,
UpstreamRequests: 1,
ChannelMeta: &relaycommon.ChannelMeta{
ChannelId: 2,
ChannelSetting: dto.ChannelSettings{ResponsesAstraModelGuard: common.GetPointer(true)},
},
}
req, err := http.NewRequest("POST", server.URL, strings.NewReader("{}"))
require.NoError(t, err)
resp, err := channel.DoRequest(c, req, info)
require.NoError(t, err)
defer resp.Body.Close()
body, err := io.ReadAll(resp.Body)
require.NoError(t, err)
assert.Contains(t, string(body), `"model":"gpt-5.6-luna"`)
assert.Len(t, guard.Attempts, 1)
assert.Equal(t, "http_error", previous.Reason)
}
func TestResponseGuardPublicErrorPrefersObservedMismatch(t *testing.T) {
guard := &relaycommon.ResponseModelGuard{}
_, mismatch := guard.Begin(context.Background(), 1, 1)
mismatch.Returned = "gpt-5.6-luna"
require.Error(t, mismatch.Reject("model_mismatch"))
t.Cleanup(mismatch.Close)
err := guard.APIError()
require.NotNil(t, err)
assert.Equal(t, relaycommon.ResponseGuardErrorCode, err.GetErrorCode())
assert.Equal(t, http.StatusConflict, err.StatusCode)
assert.Contains(t, err.Error(), "gpt-6-astra")
assert.Contains(t, err.Error(), "gpt-5.6-luna")
}
func TestResponseGuardLateMismatchAndCancellation(t *testing.T) {
for _, cancelAfterAccept := range []bool{false, true} {
t.Run(fmt.Sprint(cancelAfterAccept), func(t *testing.T) {
ctx, cancel := context.WithCancel(context.Background())
defer cancel()
guard := &relaycommon.ResponseModelGuard{}
_, attempt := guard.Begin(ctx, 1, 1)
defer attempt.Close()
first := "data: {\"type\":\"response.created\",\"response\":{\"id\":\"resp_ok\",\"model\":\"gpt-6-astra\"}}\n\n"
model := "gpt-5.6-luna"
if cancelAfterAccept {
model = "gpt-6-astra"
}
last := fmt.Sprintf("data: {\"type\":\"response.completed\",\"response\":{\"model\":%q,\"usage\":{\"input_tokens\":8,\"output_tokens\":2}}}\n\n", model)
resp := &http.Response{StatusCode: 200, Header: http.Header{"Content-Type": []string{"text/event-stream"}}, Body: io.NopCloser(strings.NewReader(first + last))}
require.NoError(t, attempt.Confirm(resp, ctx))
defer resp.Body.Close()
if cancelAfterAccept {
cancel()
}
body, err := io.ReadAll(resp.Body)
if cancelAfterAccept {
require.NoError(t, err)
assert.Equal(t, first+last, string(body))
} else {
require.Error(t, err)
assert.Equal(t, first, string(body))
assert.Equal(t, "model_mismatch", attempt.Reason)
}
require.NotNil(t, attempt.Usage)
assert.Equal(t, 8, attempt.Usage.Input)
})
}
}
func TestResponseGuardRejectsCancelledClientAndExhaustedBudget(t *testing.T) {
for _, exhausted := range []bool{false, true} {
t.Run(fmt.Sprint(exhausted), func(t *testing.T) {
c, _ := gin.CreateTestContext(httptest.NewRecorder())
ctx, cancel := context.WithCancel(context.Background())
defer cancel()
c.Request = httptest.NewRequest("POST", "/v1/responses", strings.NewReader("{}")).WithContext(ctx)
info := &relaycommon.RelayInfo{AstraAutoSwitch: true, ResponseGuard: &relaycommon.ResponseModelGuard{}, ChannelMeta: &relaycommon.ChannelMeta{}}
if exhausted {
info.UpstreamRequests = 4
} else {
cancel()
}
req, err := http.NewRequest("POST", "http://127.0.0.1:1", strings.NewReader("{}"))
require.NoError(t, err)
resp, err := channel.DoRequest(c, req, info)
require.Error(t, err)
assert.Nil(t, resp)
assert.Empty(t, info.ResponseGuard.Attempts)
})
}
}
func TestResponseModelComparisonAndLog(t *testing.T) {
for _, tc := range []struct {
name string
models []string
returned string
mismatch bool
}{
{name: "absent", models: []string{"", " "}},
{name: "requested model", models: []string{"requested"}, returned: "requested"},
{name: "mapped model", models: []string{"mapped"}, returned: "mapped"},
{name: "different model", models: []string{"other"}, returned: "other", mismatch: true},
{name: "case differs", models: []string{"Requested"}, returned: "Requested"},
{name: "requested prefix", models: []string{"requested-2026-09-01"}, returned: "requested-2026-09-01"},
{name: "mapped case differs", models: []string{"MAPPED"}, returned: "MAPPED"},
{name: "mapped prefix", models: []string{"mapped-2026-09-01"}, returned: "mapped-2026-09-01"},
{name: "reverse prefix still warns", models: []string{"request"}, returned: "request", mismatch: true},
{name: "substring still warns", models: []string{"other-requested"}, returned: "other-requested", mismatch: true},
{name: "compatible difference survives matching frames", models: []string{"mapped", "Requested", "", "requested"}, returned: "Requested"},
{name: "warning supersedes compatible difference", models: []string{"Requested", "other"}, returned: "other", mismatch: true},
{name: "compatible difference cannot erase warning", models: []string{"other", "Requested"}, returned: "other", mismatch: true},
{name: "mismatch survives later frames", models: []string{"mapped", "other", "", "requested"}, returned: "other", mismatch: true},
} {
t.Run(tc.name, func(t *testing.T) {
c, _ := gin.CreateTestContext(httptest.NewRecorder())
info := &relaycommon.RelayInfo{
OriginModelName: "requested",
ChannelMeta: &relaycommon.ChannelMeta{UpstreamModelName: "mapped", IsModelMapped: true},
}
for _, name := range tc.models {
info.ObserveResponseModel(name)
}
other := service.GenerateTextOtherInfo(c, info, 1, 1, 1, 0, 0, 0, 1)
var stored struct {
ResponseModel *relaycommon.ResponseModel `json:"response_model"`
}
require.NoError(t, common.UnmarshalJsonStr(other.JSONString(), &stored))
if tc.returned == "" {
assert.Nil(t, stored.ResponseModel)
return
}
require.NotNil(t, stored.ResponseModel)
assert.Equal(t, &relaycommon.ResponseModel{
RequestedModel: "requested", UpstreamModel: "mapped", ReturnedModel: tc.returned, Mismatch: tc.mismatch,
}, stored.ResponseModel)
assert.Equal(t, "mapped", info.UpstreamModelName)
})
}
}
func TestResponseModelLogOmitsUnchangedModel(t *testing.T) {
for _, tc := range []struct {
name string
upstream string
returned string
mapped bool
record bool
}{
{name: "same model", upstream: "requested", returned: "requested"},
{name: "no upstream name", returned: "requested"},
{name: "mapped model", upstream: "mapped", returned: "mapped", mapped: true, record: true},
{name: "mapped response echoes request", upstream: "mapped", returned: "requested", mapped: true, record: true},
{name: "prefix difference", upstream: "requested", returned: "requested-2026-09-01", record: true},
{name: "case difference", upstream: "requested", returned: "REQUESTED", record: true},
{name: "mismatch", upstream: "requested", returned: "other", record: true},
} {
t.Run(tc.name, func(t *testing.T) {
c, _ := gin.CreateTestContext(httptest.NewRecorder())
info := &relaycommon.RelayInfo{
OriginModelName: "requested",
ChannelMeta: &relaycommon.ChannelMeta{UpstreamModelName: tc.upstream, IsModelMapped: tc.mapped},
}
info.ObserveResponseModel(tc.returned)
other := service.GenerateTextOtherInfo(c, info, 1, 1, 1, 0, 0, 0, 1)
require.NotNil(t, info.ResponseModel)
assert.Equal(t, float64(1), other.Snapshot()["model_ratio"])
if tc.record {
assert.Equal(t, *info.ResponseModel, other.Snapshot()["response_model"])
} else {
assert.NotContains(t, other.Snapshot(), "response_model")
}
})
}
}
func TestResponseModelEmptyExpectedNamesDoNotMatchEveryPrefix(t *testing.T) {
for _, tc := range []struct{ requested, upstream string }{
{}, {upstream: "mapped"}, {requested: "requested"},
} {
t.Run(tc.requested+"/"+tc.upstream, func(t *testing.T) {
info := &relaycommon.RelayInfo{
OriginModelName: tc.requested,
ChannelMeta: &relaycommon.ChannelMeta{UpstreamModelName: tc.upstream},
}
info.ObserveResponseModel("other")
require.NotNil(t, info.ResponseModel)
assert.True(t, info.ResponseModel.Mismatch)
})
}
}
func TestResponseModelRetryResetsObservation(t *testing.T) {
c, _ := gin.CreateTestContext(httptest.NewRecorder())
common.SetContextKey(c, constant.ContextKeyOriginalModel, "requested")
info := &relaycommon.RelayInfo{OriginModelName: "requested"}
info.InitChannelMeta(c)
info.ObserveResponseModel("other")
require.True(t, info.ResponseModel.Mismatch)
info.InitChannelMeta(c)
assert.Nil(t, info.ResponseModel)
info.ObserveResponseModel("requested")
assert.False(t, info.ResponseModel.Mismatch)
}
func TestResponseModelHandlersCaptureBeforeConversion(t *testing.T) {
gin.SetMode(gin.TestMode)
previousTimeout := constant.StreamingTimeout
constant.StreamingTimeout = 30
t.Cleanup(func() { constant.StreamingTimeout = previousTimeout })
chat := `{"id":"chat_1","model":"returned","choices":[{"index":0,"message":{"role":"assistant","content":"hello"},"finish_reason":"stop"}],"usage":{"prompt_tokens":2,"completion_tokens":3,"total_tokens":5}}`
chatStream := "data: " + `{"id":"chat_1","model":"returned","choices":[{"index":0,"delta":{"role":"assistant","content":"hello"}}]}` + "\n\ndata: " + `{"id":"chat_1","choices":[],"usage":{"prompt_tokens":2,"completion_tokens":3,"total_tokens":5}}` + "\n\ndata: [DONE]\n\n"
responses := `{"id":"resp_1","model":"returned","status":"completed","output":[{"type":"message","role":"assistant","content":[{"type":"output_text","text":"hello"}]}],"usage":{"input_tokens":2,"output_tokens":3,"total_tokens":5}}`
responsesStream := "data: " + `{"type":"response.created","response":{"id":"resp_1","model":"returned"}}` + "\n\ndata: " + `{"type":"response.completed","response":{"id":"resp_1","status":"completed","output":[],"usage":{"input_tokens":2,"output_tokens":3,"total_tokens":5}}}` + "\n\n"
claudeBody := `{"id":"msg_1","type":"message","model":"returned","role":"assistant","content":[{"type":"text","text":"hello"}],"stop_reason":"end_turn","usage":{"input_tokens":2,"output_tokens":3}}`
claudeStream := "data: " + `{"type":"message_start","message":{"id":"msg_1","type":"message","model":"returned","role":"assistant","content":[],"usage":{"input_tokens":2,"output_tokens":0}}}` + "\n\ndata: " + `{"type":"message_delta","delta":{"stop_reason":"end_turn"},"usage":{"output_tokens":3}}` + "\n\ndata: " + `{"type":"message_stop"}` + "\n\n"
geminiBody := `{"modelVersion":"returned","candidates":[{"content":{"role":"model","parts":[{"text":"hello"}]},"finishReason":"STOP"}],"usageMetadata":{"promptTokenCount":2,"candidatesTokenCount":3,"totalTokenCount":5}}`
geminiStream := "data: " + geminiBody + "\n\n"
for _, tc := range []struct {
name string
body string
format types.RelayFormat
stream bool
handler func(*gin.Context, *relaycommon.RelayInfo, *http.Response) (*dto.Usage, *types.NewAPIError)
}{
{"chat", chat, types.RelayFormatOpenAI, false, openai.OpenaiHandler},
{"chat stream", chatStream, types.RelayFormatOpenAI, true, openai.OaiStreamHandler},
{"responses", responses, types.RelayFormatOpenAIResponses, false, openai.OaiResponsesHandler},
{"responses stream", responsesStream, types.RelayFormatOpenAIResponses, true, openai.OaiResponsesStreamHandler},
{"responses to chat", responses, types.RelayFormatOpenAI, false, openai.OaiResponsesToChatHandler},
{"responses to chat stream", responsesStream, types.RelayFormatOpenAI, true, openai.OaiResponsesToChatStreamHandler},
{"responses buffered to chat", responsesStream, types.RelayFormatOpenAI, false, openai.OaiResponsesToChatBufferedStreamHandler},
{"chat to responses", chat, types.RelayFormatOpenAIResponses, false, openai.OaiChatToResponsesHandler},
{"chat to responses stream", chatStream, types.RelayFormatOpenAIResponses, true, openai.OaiChatToResponsesStreamHandler},
{"claude", claudeBody, types.RelayFormatClaude, false, func(c *gin.Context, info *relaycommon.RelayInfo, resp *http.Response) (*dto.Usage, *types.NewAPIError) {
return claude.ClaudeHandler(c, resp, info)
}},
{"claude stream", claudeStream, types.RelayFormatClaude, true, func(c *gin.Context, info *relaycommon.RelayInfo, resp *http.Response) (*dto.Usage, *types.NewAPIError) {
return claude.ClaudeStreamHandler(c, resp, info)
}},
{"claude to responses stream", claudeStream, types.RelayFormatOpenAIResponses, true, func(c *gin.Context, info *relaycommon.RelayInfo, resp *http.Response) (*dto.Usage, *types.NewAPIError) {
return claude.ClaudeResponsesStreamHandler(c, resp, info)
}},
{"gemini", geminiBody, types.RelayFormatGemini, false, gemini.GeminiTextGenerationHandler},
{"gemini stream", geminiStream, types.RelayFormatGemini, true, gemini.GeminiTextGenerationStreamHandler},
{"gemini to chat", geminiBody, types.RelayFormatOpenAI, false, gemini.GeminiChatHandler},
{"gemini to responses", geminiBody, types.RelayFormatOpenAIResponses, false, gemini.GeminiResponsesHandler},
{"gemini to responses stream", geminiStream, types.RelayFormatOpenAIResponses, true, gemini.GeminiResponsesStreamHandler},
} {
t.Run(tc.name, func(t *testing.T) {
c, _ := gin.CreateTestContext(httptest.NewRecorder())
c.Request = httptest.NewRequest(http.MethodPost, "/v1/responses", nil)
info := &relaycommon.RelayInfo{
OriginModelName: "requested", RelayFormat: tc.format, RelayMode: relayconstant.RelayModeChatCompletions,
IsStream: tc.stream, DisablePing: true, StartTime: time.Now(),
ChannelMeta: &relaycommon.ChannelMeta{UpstreamModelName: "mapped"},
ThinkingContentInfo: relaycommon.ThinkingContentInfo{IsFirstThinkingContent: true},
}
usage, apiErr := tc.handler(c, info, &http.Response{StatusCode: http.StatusOK, Header: make(http.Header), Body: io.NopCloser(strings.NewReader(tc.body))})
require.Nil(t, apiErr)
require.NotNil(t, usage)
assert.Equal(t, 5, usage.TotalTokens)
require.NotNil(t, info.ResponseModel)
assert.Equal(t, &relaycommon.ResponseModel{RequestedModel: "requested", UpstreamModel: "mapped", ReturnedModel: "returned", Mismatch: true}, info.ResponseModel)
other := service.GenerateTextOtherInfo(c, info, 1, 1, 1, 0, 0, 0, 1)
assert.Equal(t, *info.ResponseModel, other.Snapshot()["response_model"])
})
}
}
func TestResponseModelSharedResponsesAccumulator(t *testing.T) {
// WebSocket turns and HTTP SSE use this same accumulator. A model-free
// terminal must not erase the model declared before an interrupted stream.
info := &relaycommon.RelayInfo{OriginModelName: "requested", ChannelMeta: &relaycommon.ChannelMeta{UpstreamModelName: "mapped"}}
accumulator := service.NewResponsesUsageAccumulator(info)
accumulator.Observe(&dto.ResponsesStreamResponse{Type: "response.created", Response: &dto.OpenAIResponsesResponse{Model: "returned"}})
accumulator.Observe(&dto.ResponsesStreamResponse{Type: "response.failed", Response: &dto.OpenAIResponsesResponse{Usage: &dto.Usage{InputTokens: 2, OutputTokens: 3}}})
assert.Equal(t, 5, accumulator.Finish().TotalTokens)
require.NotNil(t, info.ResponseModel)
assert.True(t, info.ResponseModel.Mismatch)
assert.Equal(t, "returned", info.ResponseModel.ReturnedModel)
}
func TestResponseModelDoesNotRecordSynthesizedModel(t *testing.T) {
c, _ := gin.CreateTestContext(httptest.NewRecorder())
c.Request = httptest.NewRequest(http.MethodPost, "/v1/chat/completions", nil)
info := &relaycommon.RelayInfo{
OriginModelName: "requested", RelayFormat: types.RelayFormatOpenAI,
ChannelMeta: &relaycommon.ChannelMeta{UpstreamModelName: "mapped"},
}
_, apiErr := openai.OaiResponsesToChatBufferedStreamHandler(c, info, &http.Response{
StatusCode: http.StatusOK, Header: make(http.Header), Body: io.NopCloser(strings.NewReader("data: [DONE]\n\n")),
})
require.Nil(t, apiErr)
assert.Nil(t, info.ResponseModel)
}
func TestAstraDetectionCanBeDisabledIndependently(t *testing.T) {
for _, disabled := range []bool{false, true} {
info := &relaycommon.RelayInfo{OriginModelName: "gpt-6-astra", AstraDetectionDisabled: disabled, AstraAutoSwitch: true, ChannelMeta: &relaycommon.ChannelMeta{}}
info.ObserveResponseModel("gpt-5.6-luna")
if disabled {
require.Nil(t, info.ResponseModel)
require.False(t, shouldEnableAstraResponseGuard("/v1/responses", info))
} else {
require.NotNil(t, info.ResponseModel)
require.True(t, info.ResponseModel.Mismatch)
require.True(t, shouldEnableAstraResponseGuard("/v1/responses", info))
}
}
}
+93 -79
View File
@@ -1,36 +1,40 @@
package relay
import (
"encoding/json"
"fmt"
"io"
"net/http"
"strings"
"github.com/QuantumNous/new-api/common"
"github.com/QuantumNous/new-api/logger"
relaycommon "github.com/QuantumNous/new-api/relay/common"
relayconstant "github.com/QuantumNous/new-api/relay/constant"
"github.com/QuantumNous/new-api/relay/helper"
"github.com/QuantumNous/new-api/relaykit/dto"
"github.com/QuantumNous/new-api/relaykit/types"
"github.com/QuantumNous/new-api/service"
"github.com/QuantumNous/new-api/setting/model_setting"
"github.com/gin-gonic/gin"
)
func ResponsesHelper(c *gin.Context, info *relaycommon.RelayInfo) (newAPIError *types.NewAPIError) {
info.InitChannelMeta(c)
if info.RelayMode == relayconstant.RelayModeResponsesCompact &&
!common.SupportsResponsesCompact(info.ChannelType, info.ApiType) {
return types.NewErrorWithStatusCode(
fmt.Errorf("unsupported endpoint %q for api type %d", "/v1/responses/compact", info.ApiType),
types.ErrorCodeInvalidRequest,
http.StatusBadRequest,
types.ErrOptionWithSkipRetry(),
)
}
// Independent cursors retain the payload metadata used by net/http redirects
// and transport retries without closing the shared request storage.
type responsesAttemptBody struct {
common.ReplayableBody
reader io.Reader
}
func (b responsesAttemptBody) Read(p []byte) (int, error) { return b.reader.Read(p) }
func shouldEnableAstraResponseGuard(path string, info *relaycommon.RelayInfo) bool {
return path == "/v1/responses" &&
info.OriginModelName == "gpt-6-astra" &&
info.AstraAutoSwitch && !info.AstraDetectionDisabled &&
info.ResponseGuard == nil
}
func ResponsesHelper(c *gin.Context, info *relaycommon.RelayInfo) (newAPIError *types.NewAPIError) {
var responsesReq *dto.OpenAIResponsesRequest
switch req := info.Request.(type) {
case *dto.OpenAIResponsesRequest:
@@ -61,77 +65,64 @@ func ResponsesHelper(c *gin.Context, info *relaycommon.RelayInfo) (newAPIError *
)
}
request, err := common.DeepCopy(responsesReq)
if err != nil {
return types.NewError(fmt.Errorf("failed to copy request to GeneralOpenAIRequest: %w", err), types.ErrorCodeInvalidRequest, types.ErrOptionWithSkipRetry())
adaptor, requestBody, closer, apiErr := PrepareResponsesRequest(c, info, responsesReq)
if apiErr != nil {
return apiErr
}
err = helper.ModelMappedHelper(c, info, request)
if err != nil {
return types.NewError(err, types.ErrorCodeChannelModelMappedError, types.ErrOptionWithSkipRetry())
}
if err := helper.ApplyReasoningModelSuffix(c, info, request); err != nil {
return newConvertRequestFailedError(c, info, err)
}
adaptor := GetAdaptor(info.ApiType)
if adaptor == nil {
return types.NewError(fmt.Errorf("invalid api type: %d", info.ApiType), types.ErrorCodeInvalidApiType, types.ErrOptionWithSkipRetry())
}
adaptor.Init(info)
var requestBody io.Reader
if model_setting.GetGlobalSettings().PassThroughRequestEnabled || info.ChannelSetting.PassThroughBodyEnabled {
storage, err := common.GetBodyStorage(c)
if err != nil {
return types.NewError(err, types.ErrorCodeReadRequestBodyFailed, types.ErrOptionWithSkipRetry())
}
requestBody = common.NewReplayableBodyReader(storage)
} else {
convertedRequest, err := adaptor.ConvertOpenAIResponsesRequest(c, info, *request)
if err != nil {
return newConvertRequestFailedError(c, info, err)
}
relaycommon.AppendRequestConversionFromRequest(info, convertedRequest)
jsonData, err := common.Marshal(convertedRequest)
if err != nil {
return types.NewError(err, types.ErrorCodeConvertRequestFailed, types.ErrOptionWithSkipRetry())
}
// remove disabled fields for OpenAI Responses API
jsonData, err = relaycommon.RemoveDisabledFields(jsonData, info.ChannelOtherSettings, info.ChannelSetting.PassThroughBodyEnabled)
if err != nil {
return types.NewError(err, types.ErrorCodeConvertRequestFailed, types.ErrOptionWithSkipRetry())
}
// apply param override
if len(info.ParamOverride) > 0 {
jsonData, err = relaycommon.ApplyParamOverrideWithRelayInfo(jsonData, info)
if err != nil {
return newAPIErrorFromParamOverride(err)
defer closer.Close()
// Missing-usage estimation follows the actual prepared payload, including
// parameter overrides, rather than the client's pre-transformation request.
if info.DurableAttemptBilling {
if cursor, err := requestBody.NewReader(); err == nil {
var sent dto.OpenAIResponsesRequest
decodeErr := json.NewDecoder(cursor).Decode(&sent)
_ = cursor.Close()
if decodeErr == nil {
if tokens, err := service.CountRequestToken(c, sent.GetTokenCountMeta(), info); err == nil {
info.SetEstimatePromptTokens(tokens)
}
}
}
logger.LogDebug(c, "requestBody: %s", jsonData)
body, closer, err := relaycommon.NewOutboundJSONBody(jsonData)
if err != nil {
return types.NewError(err, types.ErrorCodeConvertRequestFailed, types.ErrOptionWithSkipRetry())
}
defer closer.Close()
jsonData = nil
requestBody = body
}
if shouldEnableAstraResponseGuard(c.Request.URL.Path, info) {
info.ResponseGuard = &relaycommon.ResponseModelGuard{Deadline: info.GuardDeadline}
}
var httpResp *http.Response
resp, err := adaptor.DoRequest(c, info, requestBody)
guardEnabled := info.ResponseGuard != nil && info.AstraAutoSwitch
body, err := requestBody.NewReader()
if err != nil {
return types.NewOpenAIError(err, types.ErrorCodeDoRequestFailed, http.StatusInternalServerError)
return types.NewOpenAIError(err, types.ErrorCodeReadRequestBodyFailed, http.StatusInternalServerError)
}
resp, requestErr := adaptor.DoRequest(c, info, responsesAttemptBody{ReplayableBody: requestBody, reader: body})
_ = body.Close()
if requestErr != nil {
apiErr := types.NewError(requestErr, types.ErrorCodeDoRequestFailed)
guard := info.ResponseGuard
if guard == nil || apiErr.GetErrorCode() != relaycommon.ResponseGuardErrorCode {
return apiErr
}
attempt := guard.Last()
if attempt != nil {
service.RequestPolicy(c).AddEvent(service.PolicyEvent{ChannelID: info.ChannelId, ErrorCode: string(relaycommon.ResponseGuardErrorCode), ErrorSource: "model_guard", Health: "unchanged", Decision: service.PolicyDecision{Action: "reject", Reason: attempt.Reason, Source: "model_guard"}})
}
return guard.APIError()
}
var httpResp *http.Response
if resp != nil {
httpResp = resp.(*http.Response)
}
if guardEnabled {
defer func() {
if attempt := info.ResponseGuard.Last(); attempt != nil {
attempt.Finish(c.Writer.Written())
}
}()
}
statusCodeMappingStr := c.GetString("status_code_mapping")
if resp != nil {
httpResp = resp.(*http.Response)
if httpResp != nil {
if httpResp.StatusCode != http.StatusOK {
newAPIError = service.RelayErrorHandler(c.Request.Context(), httpResp, false)
// reset status code 重置状态码
@@ -141,6 +132,23 @@ func ResponsesHelper(c *gin.Context, info *relaycommon.RelayInfo) (newAPIError *
}
usage, newAPIError := adaptor.DoResponse(c, httpResp, info)
if guardEnabled {
if attempt := info.ResponseGuard.Last(); attempt != nil {
attempt.Finish(c.Writer.Written())
if guardErr := attempt.APIError(); guardErr != nil {
if info.StreamStatus != nil {
info.StreamStatus.MarkFailed(string(relaycommon.ResponseGuardErrorCode), "api_error", http.StatusConflict)
}
if c.Writer.Written() && c.Request.Context().Err() == nil {
message := common.MessageWithRequestId(attempt.PublicErrorMessage(), c.GetString(common.RequestIdKey))
c.SSEvent("response.failed", gin.H{"type": "response.failed", "sequence_number": attempt.Sequence + 1, "response": gin.H{"id": attempt.ResponseID, "object": "response", "status": "failed", "error": gin.H{"code": relaycommon.ResponseGuardErrorCode, "message": message}}})
c.Writer.Flush()
c.Set("response_guard_error_sent", true)
}
return guardErr
}
}
}
if newAPIError != nil {
// reset status code 重置状态码
service.ResetStatusCode(newAPIError, statusCodeMappingStr)
@@ -165,10 +173,16 @@ func ResponsesHelper(c *gin.Context, info *relaycommon.RelayInfo) (newAPIError *
return nil
}
if strings.HasPrefix(info.OriginModelName, "gpt-4o-audio") {
service.PostAudioConsumeQuota(c, info, usageDto, "")
} else {
service.PostTextConsumeQuota(c, info, usageDto, nil)
}
ConsumeResponsesQuota(c, info, usageDto)
return nil
}
// ConsumeResponsesQuota applies the same settlement dispatch to HTTP and
// WebSocket Responses usage. Compact requests keep their separate repricing.
func ConsumeResponsesQuota(c *gin.Context, info *relaycommon.RelayInfo, usage *dto.Usage) {
if strings.HasPrefix(info.OriginModelName, "gpt-4o-audio") {
service.PostAudioConsumeQuota(c, info, usage, "")
return
}
service.PostTextConsumeQuota(c, info, usage, nil)
}
+87
View File
@@ -0,0 +1,87 @@
package relay
import (
"fmt"
"io"
"net/http"
"github.com/QuantumNous/new-api/common"
"github.com/QuantumNous/new-api/logger"
relaychannel "github.com/QuantumNous/new-api/relay/channel"
relaycommon "github.com/QuantumNous/new-api/relay/common"
relayconstant "github.com/QuantumNous/new-api/relay/constant"
"github.com/QuantumNous/new-api/relay/helper"
"github.com/QuantumNous/new-api/relaykit/dto"
"github.com/QuantumNous/new-api/relaykit/types"
"github.com/QuantumNous/new-api/setting/model_setting"
"github.com/gin-gonic/gin"
)
// PrepareResponsesRequest applies the same model, conversion and channel rules
// for HTTP and WebSocket requests. The caller closes closer after the attempt;
// passthrough bodies remain owned by the incoming request's BodyStorage.
// The returned adaptor retains route/conversion state for DoRequest/DoResponse.
func PrepareResponsesRequest(c *gin.Context, info *relaycommon.RelayInfo, req *dto.OpenAIResponsesRequest) (relaychannel.Adaptor, common.ReplayableBody, io.Closer, *types.NewAPIError) {
info.InitChannelMeta(c)
if info.RelayMode == relayconstant.RelayModeResponsesCompact &&
!common.SupportsResponsesCompact(info.ChannelType, info.ApiType) {
return nil, nil, nil, types.NewErrorWithStatusCode(
fmt.Errorf("unsupported endpoint %q for api type %d", "/v1/responses/compact", info.ApiType),
types.ErrorCodeInvalidRequest,
http.StatusBadRequest,
types.ErrOptionWithSkipRetry(),
)
}
request, err := common.DeepCopy(req)
if err != nil {
return nil, nil, nil, types.NewError(fmt.Errorf("failed to copy responses request: %w", err), types.ErrorCodeInvalidRequest, types.ErrOptionWithSkipRetry())
}
if err := helper.ModelMappedHelper(c, info, request); err != nil {
return nil, nil, nil, types.NewError(err, types.ErrorCodeChannelModelMappedError, types.ErrOptionWithSkipRetry())
}
if err := helper.ApplyReasoningModelSuffix(c, info, request); err != nil {
return nil, nil, nil, newConvertRequestFailedError(c, info, err)
}
adaptor := GetAdaptor(info.ApiType)
if adaptor == nil {
return nil, nil, nil, types.NewError(fmt.Errorf("invalid api type: %d", info.ApiType), types.ErrorCodeInvalidApiType, types.ErrOptionWithSkipRetry())
}
adaptor.Init(info)
if model_setting.GetGlobalSettings().PassThroughRequestEnabled || info.ChannelSetting.PassThroughBodyEnabled {
storage, err := common.GetBodyStorage(c)
if err != nil {
return nil, nil, nil, types.NewError(err, types.ErrorCodeReadRequestBodyFailed, types.ErrOptionWithSkipRetry())
}
body := common.NewReplayableBodyReader(storage)
return adaptor, body, io.NopCloser(body), nil
}
convertedRequest, err := adaptor.ConvertOpenAIResponsesRequest(c, info, *request)
if err != nil {
return nil, nil, nil, newConvertRequestFailedError(c, info, err)
}
relaycommon.AppendRequestConversionFromRequest(info, convertedRequest)
jsonData, err := common.Marshal(convertedRequest)
if err != nil {
return nil, nil, nil, types.NewError(err, types.ErrorCodeConvertRequestFailed, types.ErrOptionWithSkipRetry())
}
jsonData, err = relaycommon.RemoveDisabledFields(jsonData, info.ChannelOtherSettings, info.ChannelSetting.PassThroughBodyEnabled)
if err != nil {
return nil, nil, nil, types.NewError(err, types.ErrorCodeConvertRequestFailed, types.ErrOptionWithSkipRetry())
}
if len(info.ParamOverride) > 0 {
jsonData, err = relaycommon.ApplyParamOverrideWithRelayInfo(jsonData, info)
if err != nil {
return nil, nil, nil, newAPIErrorFromParamOverride(err)
}
}
logger.LogDebug(c, "requestBody: %s", jsonData)
body, closer, err := relaycommon.NewOutboundJSONBody(jsonData)
if err != nil {
return nil, nil, nil, types.NewError(err, types.ErrorCodeConvertRequestFailed, types.ErrOptionWithSkipRetry())
}
return adaptor, body, closer, nil
}
+939
View File
@@ -0,0 +1,939 @@
package relay
import (
"bytes"
"context"
"errors"
"fmt"
"io"
"net/http"
"reflect"
"slices"
"strings"
"sync"
"time"
"github.com/QuantumNous/new-api/common"
appconstant "github.com/QuantumNous/new-api/constant"
appdto "github.com/QuantumNous/new-api/dto"
"github.com/QuantumNous/new-api/i18n"
"github.com/QuantumNous/new-api/middleware"
appmodel "github.com/QuantumNous/new-api/model"
perfmetrics "github.com/QuantumNous/new-api/pkg/perf_metrics"
"github.com/QuantumNous/new-api/pkg/wsmanager"
relaychannel "github.com/QuantumNous/new-api/relay/channel"
relaycommon "github.com/QuantumNous/new-api/relay/common"
"github.com/QuantumNous/new-api/relay/helper"
"github.com/QuantumNous/new-api/relaykit/dto"
"github.com/QuantumNous/new-api/relaykit/types"
"github.com/QuantumNous/new-api/service"
"github.com/gin-gonic/gin"
"github.com/gorilla/websocket"
)
const responsesWSEventTypeResponseCreate = "response.create"
const responsesWSWriteTimeout = 30 * time.Second
// ResponsesWSRequestRunner executes the existing authentication and rate-limit
// middleware around one complete request, without a second HTTP connection.
type ResponsesWSRequestRunner func(*http.Request, string, func(*gin.Context) *types.NewAPIError) *types.NewAPIError
type responsesWSCreateEvent struct {
Type string `json:"type"`
EventID string `json:"event_id,omitempty"`
StreamID common.RawMessage `json:"stream_id,omitempty"`
Generate common.RawMessage `json:"generate,omitempty"`
Request common.RawMessage `json:"response,omitempty"`
}
type responsesWSCreateRequest struct {
Request dto.OpenAIResponsesRequest
Body []byte
Generate common.RawMessage
StreamID string
}
type responsesWSErrorEvent struct {
Type string `json:"type"`
Status int `json:"status"`
EventID string `json:"event_id,omitempty"`
StreamID string `json:"stream_id,omitempty"`
ResponseID string `json:"response_id,omitempty"`
Error *types.OpenAIError `json:"error"`
}
type responsesWSMessage struct {
kind int
body []byte
err error
}
// responsesWSControl is a client control event (response.cancel) whose
// envelope the read loop already parsed.
type responsesWSControl struct {
body []byte
eventID, streamID string
}
// Only the request worker reads or changes billing state. Socket readers pass
// bounded messages to it; cancellation never performs an independent refund.
type responsesWSCallState struct {
inbox chan responsesWSMessage
controls chan responsesWSControl
done chan struct{}
terminal *responsesWSMessage
closeAfter bool
}
type responsesWSSession struct {
ctx context.Context
cancel context.CancelFunc
client *websocket.Conn
runner ResponsesWSRequestRunner
request *http.Request
requestID string
nextEventIndex int
workers sync.WaitGroup
clientWriteMu sync.Mutex
targetWriteMu sync.Mutex
connectionMu sync.Mutex
target *websocket.Conn
unregister func()
stateMu sync.Mutex
current *responsesWSCallState
// These fields belong to the serial request worker and describe the actual
// established connection. Per-request token/user data is never stored here.
lastResponseID string
lockedModel string
lockedChannelID int
lockedGroup string
lockedKey string
lockedKeyIndex int
lockedContext map[appconstant.ContextKey]any
lockedRoute dto.AdvancedCustomRoute
}
func ResponsesWebSocketHelper(c *gin.Context, client *websocket.Conn, runner ResponsesWSRequestRunner) *types.NewAPIError {
ctx, cancel := context.WithCancel(c.Request.Context())
s := &responsesWSSession{ctx: ctx, cancel: cancel, client: client, runner: runner,
request: c.Request.Clone(ctx), requestID: c.GetString(common.RequestIdKey)}
if s.requestID == "" {
s.requestID = common.NewRequestId()
}
maxMB := appconstant.MaxRequestBodyMB
if maxMB <= 0 {
maxMB = 128
}
client.SetReadLimit(int64(maxMB) << 20)
defer func() {
s.shutdown()
s.workers.Wait()
}()
for {
_, message, err := client.ReadMessage()
if err != nil {
return nil
}
envelope, streamID, err := parseResponsesWSEnvelope(message)
eventType := envelope.Type
if err != nil {
s.sendError(envelope.EventID, streamID, newResponsesWSInvalidRequestError(err))
continue
}
if eventType != responsesWSEventTypeResponseCreate {
// Controls are owned by the active request too. In particular a cancel
// arriving during authentication must not precede its upstream create.
state := s.getCurrent()
if eventType != "response.cancel" || state == nil {
s.sendError(envelope.EventID, streamID, newResponsesWSInvalidRequestError(fmt.Errorf("unsupported websocket event %q", eventType)))
continue
}
select {
case state.controls <- responsesWSControl{body: message, eventID: envelope.EventID, streamID: streamID}:
case <-state.done:
case <-s.ctx.Done():
return nil
default:
s.sendError(envelope.EventID, streamID, newResponsesWSInvalidRequestError(errors.New("a response control event is already pending")))
}
continue
}
state := &responsesWSCallState{inbox: make(chan responsesWSMessage), controls: make(chan responsesWSControl, 1), done: make(chan struct{})}
if !s.tryReserveCurrent(state) {
s.sendError(envelope.EventID, streamID, types.NewErrorWithStatusCode(errors.New("another response.create is already in progress on this websocket connection"), types.ErrorCodeInvalidRequest, http.StatusConflict, types.ErrOptionWithSkipRetry()))
continue
}
requestID := fmt.Sprintf("%s-ws-%d", s.requestID, s.nextEventIndex)
s.nextEventIndex++
s.workers.Go(func() { s.runRequest(state, message, envelope, streamID, requestID) })
}
}
func (s *responsesWSSession) runRequest(state *responsesWSCallState, message []byte, envelope responsesWSCreateEvent, streamID string, requestID string) {
create, parseErr := normalizeResponsesWSCreateEvent(message, envelope, streamID)
var apiErr *types.NewAPIError
defer func() {
if recovered := recover(); recovered != nil {
apiErr = types.NewError(fmt.Errorf("responses websocket request panic: %v", recovered), types.ErrorCodeBadResponse, types.ErrOptionWithSkipRetry())
state.closeAfter = true
}
// Finish the middleware count before publishing the terminal event. Hold
// admission while writing it so an immediate next create cannot race
// the current request's release; socket close needs neither lock.
outgoing := state.terminal
if apiErr != nil {
if body, err := buildResponsesWSErrorPayload(envelope.EventID, streamID, apiErr); err == nil {
outgoing = &responsesWSMessage{kind: websocket.TextMessage, body: body}
}
}
s.clientWriteMu.Lock()
s.stateMu.Lock()
if outgoing != nil {
if err := s.client.SetWriteDeadline(time.Now().Add(responsesWSWriteTimeout)); err != nil {
state.closeAfter = true
} else if err := s.client.WriteMessage(outgoing.kind, outgoing.body); err != nil {
state.closeAfter = true
}
}
s.current = nil
close(state.done)
s.stateMu.Unlock()
s.clientWriteMu.Unlock()
if state.closeAfter {
s.shutdown()
}
}()
request := s.request.Clone(s.ctx)
request.Body = io.NopCloser(bytes.NewReader(message))
request.ContentLength = int64(len(message))
request.Header.Set("Content-Type", "application/json")
apiErr = s.runner(request, requestID, func(c *gin.Context) *types.NewAPIError {
if parseErr != nil {
return newResponsesWSInvalidRequestError(parseErr)
}
c.Request.Body = io.NopCloser(bytes.NewReader(create.Body))
c.Request.ContentLength = int64(len(create.Body))
return s.runCall(c, state, create)
})
if apiErr != nil && (apiErr.StatusCode == http.StatusUnauthorized || apiErr.StatusCode == http.StatusForbidden) {
state.closeAfter = true
}
}
func (s *responsesWSSession) runCall(c *gin.Context, state *responsesWSCallState, create responsesWSCreateRequest) (apiErr *types.NewAPIError) {
policy := service.RequestPolicy(c)
modelName := create.Request.Model
started := time.Now()
var info *relaycommon.RelayInfo
billingPrepared := false
defer func() {
if recovered := recover(); recovered != nil {
apiErr = types.NewError(fmt.Errorf("responses websocket call panic: %v", recovered), types.ErrorCodeBadResponse, types.ErrOptionWithSkipRetry())
state.closeAfter = true
}
if info == nil && modelName != "" {
info = &relaycommon.RelayInfo{OriginModelName: modelName, UsingGroup: common.GetContextKeyString(c, appconstant.ContextKeyUsingGroup), StartTime: started}
}
perfmetrics.RecordRelayResult(c.Request.Context(), info, apiErr)
// Settlement already marks the request policy successful, and nothing
// reads a termination decision after this point on the WebSocket path,
// so neither policy record belongs here.
if info != nil && billingPrepared {
apiErr = RefundFailedRequestBilling(c, info, apiErr)
}
}()
if modelName == "" {
return newResponsesWSInvalidRequestError(errors.New("model is required"))
}
if s.lockedModel != "" && modelName != s.lockedModel {
return newResponsesWSInvalidRequestError(fmt.Errorf("responses websocket connection is locked to model %q", s.lockedModel))
}
if apiErr = checkResponsesWSModelAccess(c, modelName); apiErr != nil {
return apiErr
}
common.SetContextKey(c, appconstant.ContextKeyOriginalModel, modelName)
common.SetContextKey(c, appconstant.ContextKeyRequestStartTime, time.Now())
service.GetChannelConstraints(c).AddFilter(appdto.ChannelFilter{Kind: appdto.FilterRequestPath, RequestPath: c.Request.URL.Path})
if s.lockedChannelID != 0 {
if apiErr = s.restoreConnectionContext(c, modelName); apiErr != nil {
return apiErr
}
info = relaycommon.GenRelayInfoResponses(c, &create.Request)
info.IsStream = true
common.SetContextKey(c, appconstant.ContextKeyIsStream, true)
if apiErr = PrepareRequestBilling(c, info); apiErr != nil {
return apiErr
}
billingPrepared = true
var payload []byte
payload, apiErr = buildResponsesWSCreatePayload(c, info, create.Request, create.Generate, create.StreamID)
if apiErr != nil {
return apiErr
}
if err := s.writeTarget(websocket.TextMessage, payload); err != nil {
state.closeAfter = true
return types.NewError(err, types.ErrorCodeBadResponse, types.ErrOptionWithSkipRetry())
}
} else {
retry := &service.RetryParam{Ctx: c, TokenGroup: common.GetContextKeyString(c, appconstant.ContextKeyUsingGroup), ModelName: modelName, RequestPath: c.Request.URL.Path, Retry: common.GetPointer(0)}
for ; retry.GetRetry() <= common.RetryTimes; retry.IncreaseRetry() {
var channel *appmodel.Channel
channel, apiErr = selectResponsesWSChannel(c, modelName, retry)
if apiErr != nil {
return apiErr
}
service.AppendUsedChannel(c, channel.Id)
if info == nil {
info = relaycommon.GenRelayInfoResponses(c, &create.Request)
info.IsStream = true
common.SetContextKey(c, appconstant.ContextKeyIsStream, true)
if apiErr = PrepareRequestBilling(c, info); apiErr != nil {
return apiErr
}
billingPrepared = true
} else {
info.PriceData.GroupRatioInfo = helper.HandleGroupRatio(c, info)
if apiErr = service.PrepareTieredBillingForSelectedGroup(c, info); apiErr != nil {
return apiErr
}
}
info.RetryIndex = retry.GetRetry()
policy.BeginAttempt(channel, info.UsingGroup)
var payload []byte
payload, apiErr = buildResponsesWSCreatePayload(c, info, create.Request, create.Generate, create.StreamID)
if apiErr != nil {
return apiErr
}
adaptor := GetAdaptor(info.ApiType)
adaptor.Init(info)
target, dialErr := relaychannel.DoWssRequest(adaptor, c, info, nil)
if dialErr != nil {
apiErr = service.NormalizeViolationFeeError(types.NewError(dialErr, types.ErrorCodeDoRequestFailed))
service.ResetStatusCode(apiErr, c.GetString("status_code_mapping"))
info.LastError = apiErr
decision := service.DecideRelayRetry(c, apiErr, common.RetryTimes-retry.GetRetry())
service.RecordPolicyFailure(c, channel.Id, apiErr, decision)
service.ProcessChannelError(c, *types.NewChannelError(channel.Id, channel.Type, channel.Name, channel.ChannelInfo.IsMultiKey, info.ApiKey, channel.GetAutoBan()), apiErr, info)
if decision.Action == "retry" {
continue
}
return apiErr
}
if !s.setTarget(target) {
return types.NewError(context.Canceled, types.ErrorCodeBadResponse, types.ErrOptionWithSkipRetry())
}
if err := s.writeTarget(websocket.TextMessage, payload); err != nil {
state.closeAfter = true
return types.NewError(err, types.ErrorCodeBadResponse, types.ErrOptionWithSkipRetry())
}
s.lockedModel, s.lockedChannelID, s.lockedGroup = modelName, channel.Id, info.UsingGroup
s.lockedKey, s.lockedKeyIndex = info.ApiKey, info.ChannelMultiKeyIndex
s.lockedRoute, _ = channel.GetOtherSettings().AdvancedCustom.MatchPathForModel(c.Request.URL.Path, modelName)
s.lockedContext = make(map[appconstant.ContextKey]any)
for _, key := range []appconstant.ContextKey{
appconstant.ContextKeyChannelId, appconstant.ContextKeyChannelName, appconstant.ContextKeyChannelType,
appconstant.ContextKeyChannelCreateTime, appconstant.ContextKeyChannelSetting, appconstant.ContextKeyChannelOtherSetting,
appconstant.ContextKeyChannelParamOverride, appconstant.ContextKeyChannelHeaderOverride, appconstant.ContextKeyChannelOrganization,
appconstant.ContextKeyChannelAutoBan, appconstant.ContextKeyChannelModelMapping, appconstant.ContextKeyChannelStatusCodeMapping,
appconstant.ContextKeyChannelIsMultiKey, appconstant.ContextKeyChannelMultiKeyIndex, appconstant.ContextKeyChannelKey, appconstant.ContextKeyChannelBaseUrl,
} {
if value, ok := c.Get(string(key)); ok {
s.lockedContext[key] = value
}
}
s.registerChannelClose(channel.Id)
s.startTargetReader(target)
apiErr = nil
break
}
if apiErr != nil {
return apiErr
}
}
accumulator := service.NewResponsesUsageAccumulator(info)
info.StreamStatus = relaycommon.NewStreamStatus()
info.StreamStatus.RequireTerminal()
common.SetContextKey(c, appconstant.ContextKeyResponseStreamStatus, info.StreamStatus)
timeout := time.Duration(appconstant.StreamingTimeout) * time.Second
if timeout <= 0 {
timeout = 300 * time.Second
}
idle := time.NewTimer(timeout)
defer idle.Stop()
accepted := false
var pendingControl []byte
var sentControl []byte
var responseID string
for {
select {
case incoming := <-state.inbox:
idle.Reset(timeout)
if incoming.err != nil {
info.StreamStatus.SetEndReason(relaycommon.StreamEndReasonScannerErr, incoming.err)
state.closeAfter = true
ConsumeResponsesQuota(c, info, accumulator.Finish())
return nil
}
info.SetFirstResponseTime()
var event struct {
dto.ResponsesStreamResponse
StreamID string `json:"stream_id"`
}
if err := common.Unmarshal(incoming.body, &event); err != nil {
info.StreamStatus.RecordError("invalid upstream websocket event")
} else {
if event.Type != "error" && event.StreamID != "" && event.StreamID != create.StreamID {
if err := s.writeClient(incoming.kind, incoming.body); err != nil {
s.shutdown()
}
continue
}
// A repeated terminal from the previous response must never finish
// a subsequent request on this persistent connection.
if event.Response != nil && event.Response.ID != "" && event.Response.ID == s.lastResponseID {
continue
}
if event.Type == "error" {
var rejection responsesWSErrorEvent
_ = common.Unmarshal(incoming.body, &rejection)
if rejection.ResponseID != "" && rejection.ResponseID == s.lastResponseID {
continue
}
terminal, ambiguous, controlError := responsesWSErrorEndsRequest(rejection, create.StreamID, responseID, sentControl)
if !terminal {
if err := s.writeClient(incoming.kind, incoming.body); err != nil {
s.shutdown()
}
// Only a control error in this stream resolves its pending control.
if controlError {
sentControl = nil
}
continue
}
if accepted {
if rejection.Error != nil {
code := ""
if rejection.Error.Code != nil {
code = fmt.Sprint(rejection.Error.Code)
}
info.StreamStatus.MarkFailed(code, rejection.Error.Type, rejection.Status)
}
accumulator.Observe(&event.ResponsesStreamResponse)
info.StreamStatus.SetEndReason(relaycommon.StreamEndReasonDone, nil)
s.lastResponseID = responseID
state.terminal, state.closeAfter = &incoming, ambiguous
ConsumeResponsesQuota(c, info, accumulator.Finish())
return nil
}
if rejection.Error == nil {
// An error frame without an error object still describes the
// rejected request; keep the client-facing type stable.
rejection.Error = &types.OpenAIError{Type: "invalid_request_error", Message: event.Message, Code: event.Code}
}
rejected := types.WithOpenAIError(*rejection.Error, rejection.Status, types.ErrOptionWithSkipRetry())
if rejected.StatusCode < 400 || rejected.StatusCode > 599 {
rejected.StatusCode = http.StatusBadRequest
}
return rejected
}
if strings.HasPrefix(event.Type, "response.") {
if !accepted {
// Like HTTP, bind the session only once upstream accepted the request.
service.RecordChannelAffinity(c, s.lockedChannelID)
}
accepted = true
if event.Response != nil && event.Response.ID != "" {
responseID = event.Response.ID
}
}
accumulator.Observe(&event.ResponsesStreamResponse)
}
switch event.Type {
case "response.completed", "response.done", "response.incomplete", "response.failed", "response.cancelled", "response.canceled":
if event.Response != nil {
s.lastResponseID = event.Response.ID
}
info.StreamStatus.SetEndReason(relaycommon.StreamEndReasonDone, nil)
state.terminal = &incoming
ConsumeResponsesQuota(c, info, accumulator.Finish())
return nil
}
if err := s.writeClient(incoming.kind, incoming.body); err != nil {
s.shutdown()
}
if accepted && pendingControl != nil {
if err := s.writeTarget(websocket.TextMessage, pendingControl); err != nil {
s.shutdown()
}
sentControl = pendingControl
pendingControl = nil
}
case control := <-state.controls:
if pendingControl != nil || sentControl != nil {
s.sendError(control.eventID, control.streamID, newResponsesWSInvalidRequestError(errors.New("a response control event is already pending")))
continue
}
if !accepted {
pendingControl = control.body
continue
}
if err := s.writeTarget(websocket.TextMessage, control.body); err != nil {
s.shutdown()
}
sentControl = control.body
case <-idle.C:
info.StreamStatus.SetEndReason(relaycommon.StreamEndReasonTimeout, context.DeadlineExceeded)
state.closeAfter = true
ConsumeResponsesQuota(c, info, accumulator.Finish())
return nil
case <-s.ctx.Done():
info.StreamStatus.SetEndReason(relaycommon.StreamEndReasonClientGone, s.ctx.Err())
ConsumeResponsesQuota(c, info, accumulator.Finish())
return nil
}
}
}
// A single generation is active, but a cancel can fail independently. Correlate
// available IDs first; uncorrelated control/server errors require a bounded
// connection close because continuing could leave the generation stuck forever.
func responsesWSErrorEndsRequest(event responsesWSErrorEvent, streamID, responseID string, control []byte) (terminal, ambiguous, controlError bool) {
if len(control) > 0 {
var pending struct {
EventID string `json:"event_id"`
StreamID string `json:"stream_id"`
ResponseID string `json:"response_id"`
}
_ = common.Unmarshal(control, &pending)
if event.EventID != "" && event.EventID == pending.EventID {
return false, false, true
}
if event.ResponseID != "" && event.ResponseID == pending.ResponseID && event.ResponseID != responseID {
return false, false, true
}
if (event.StreamID == "" || event.StreamID == pending.StreamID) && event.Error != nil && event.Error.Type == "invalid_request_error" {
switch event.Error.Code {
case "response_not_found", "response_not_active", "response_already_completed":
return false, false, true
}
}
}
if event.StreamID != "" && event.StreamID != streamID || responseID != "" && event.ResponseID != "" && event.ResponseID != responseID {
return false, false, false
}
return true, len(control) > 0 && event.ResponseID == "", false
}
func (s *responsesWSSession) restoreConnectionContext(c *gin.Context, model string) *types.NewAPIError {
// The channel type is not rechecked: a locked channel keeps its type, and
// the FilterResponsesWebSocket rules were applied when it was selected.
channel, err := appmodel.CacheGetChannel(s.lockedChannelID)
if err != nil || channel == nil || channel.Status != common.ChannelStatusEnabled || !channel.GetSetting().ResponsesWebSocketEnabled {
return types.NewErrorWithStatusCode(errors.New("Responses WebSocket is disabled for this channel"), types.ErrorCode(appdto.FilterResponsesWebSocket), http.StatusForbidden, types.ErrOptionWithSkipRetry())
}
keyEnabled := channel.Key == s.lockedKey
if channel.ChannelInfo.IsMultiKey {
keys := channel.GetKeys()
status := channel.ChannelInfo.MultiKeyStatusList[s.lockedKeyIndex]
keyEnabled = s.lockedKeyIndex >= 0 && s.lockedKeyIndex < len(keys) && keys[s.lockedKeyIndex] == s.lockedKey && (status == 0 || status == common.ChannelStatusEnabled)
}
if !keyEnabled {
return types.NewErrorWithStatusCode(errors.New("the upstream connection credential is no longer enabled"), types.ErrorCodeAccessDenied, http.StatusForbidden, types.ErrOptionWithSkipRetry())
}
// Changes that alter the physical upstream connection require a new
// handshake. Request-level settings can be refreshed without rotating keys.
if channel.GetBaseURL() != s.lockedContext[appconstant.ContextKeyChannelBaseUrl] ||
!reflect.DeepEqual(channel.GetHeaderOverride(), s.lockedContext[appconstant.ContextKeyChannelHeaderOverride]) {
return types.NewErrorWithStatusCode(errors.New("upstream connection settings changed; reconnect required"), types.ErrorCodeAccessDenied, http.StatusForbidden, types.ErrOptionWithSkipRetry())
}
if previous, ok := s.lockedContext[appconstant.ContextKeyChannelSetting].(dto.ChannelSettings); ok && previous.Proxy != channel.GetSetting().Proxy {
return types.NewErrorWithStatusCode(errors.New("upstream proxy changed; reconnect required"), types.ErrorCodeAccessDenied, http.StatusForbidden, types.ErrOptionWithSkipRetry())
}
// An advanced custom connection was dialed through the route matched for
// this model. Only the fields that shape the upstream handshake force a
// reconnect: target path, protocol conversion and credential placement.
// Other channel types match the zero route on both sides.
route, _ := channel.GetOtherSettings().AdvancedCustom.MatchPathForModel(c.Request.URL.Path, model)
if strings.TrimSpace(route.UpstreamPath) != strings.TrimSpace(s.lockedRoute.UpstreamPath) ||
route.IsNative() != s.lockedRoute.IsNative() ||
!reflect.DeepEqual(route.Auth, s.lockedRoute.Auth) {
return types.NewErrorWithStatusCode(errors.New("upstream route changed; reconnect required"), types.ErrorCodeAccessDenied, http.StatusForbidden, types.ErrOptionWithSkipRetry())
}
if pin, found, _ := service.GetChannelConstraints(c).ResolvedPin(); found {
if pin.ChannelId != s.lockedChannelID {
return types.NewErrorWithStatusCode(errors.New("channel pin changed; reconnect required"), types.ErrorCodeAccessDenied, http.StatusForbidden, types.ErrOptionWithSkipRetry())
}
} else {
group := common.GetContextKeyString(c, appconstant.ContextKeyUsingGroup)
if group == "auto" {
if !slices.Contains(service.GetRequestAutoGroups(c, common.GetContextKeyString(c, appconstant.ContextKeyUserGroup)), s.lockedGroup) {
return types.NewErrorWithStatusCode(errors.New("the connection group is no longer allowed"), types.ErrorCodeAccessDenied, http.StatusForbidden, types.ErrOptionWithSkipRetry())
}
group = s.lockedGroup
common.SetContextKey(c, appconstant.ContextKeyAutoGroup, group)
}
if !appmodel.IsChannelEnabledForGroupModel(group, model, s.lockedChannelID) {
return types.NewErrorWithStatusCode(errors.New("the connection channel is no longer allowed for this group and model"), types.ErrorCodeAccessDenied, http.StatusForbidden, types.ErrOptionWithSkipRetry())
}
}
for key, value := range s.lockedContext {
c.Set(string(key), value)
}
common.SetContextKey(c, appconstant.ContextKeyChannelSetting, channel.GetSetting())
common.SetContextKey(c, appconstant.ContextKeyChannelOtherSetting, channel.GetOtherSettings())
common.SetContextKey(c, appconstant.ContextKeyChannelParamOverride, channel.GetParamOverride())
common.SetContextKey(c, appconstant.ContextKeyChannelModelMapping, channel.GetModelMapping())
common.SetContextKey(c, appconstant.ContextKeyChannelStatusCodeMapping, channel.GetStatusCodeMapping())
policy := service.RequestPolicy(c)
policy.BeginAttempt(channel, s.lockedGroup)
policy.AddEvent(service.PolicyEvent{ChannelID: channel.Id, Decision: service.PolicyDecision{Action: "select", Reason: "pinned_channel", Source: "channel_constraint"}})
return nil
}
func buildResponsesWSCreatePayload(c *gin.Context, info *relaycommon.RelayInfo, req dto.OpenAIResponsesRequest, generate common.RawMessage, streamID string) ([]byte, *types.NewAPIError) {
_, body, closer, apiErr := PrepareResponsesRequest(c, info, &req)
if apiErr != nil {
return nil, apiErr
}
defer closer.Close()
jsonData, err := io.ReadAll(body)
if err != nil {
return nil, types.NewError(err, types.ErrorCodeReadRequestBodyFailed, types.ErrOptionWithSkipRetry())
}
event, err := buildResponsesWSCreateEvent(jsonData, generate, streamID)
if err != nil {
return nil, types.NewError(err, types.ErrorCodeConvertRequestFailed, types.ErrOptionWithSkipRetry())
}
return event, nil
}
func (s *responsesWSSession) startTargetReader(target *websocket.Conn) {
limit := int64(helper.DefaultMaxScannerBufferSize)
if appconstant.StreamScannerMaxBufferMB > 0 {
limit = int64(appconstant.StreamScannerMaxBufferMB) << 20
}
target.SetReadLimit(limit)
s.workers.Go(func() {
for {
kind, body, err := target.ReadMessage()
incoming := responsesWSMessage{kind: kind, body: body, err: err}
if state := s.getCurrent(); state != nil {
select {
case state.inbox <- incoming:
case <-state.done:
if err == nil {
if writeErr := s.writeClient(kind, body); writeErr != nil {
s.shutdown()
return
}
} else {
s.shutdown()
}
case <-s.ctx.Done():
return
}
} else if err == nil {
if writeErr := s.writeClient(kind, body); writeErr != nil {
s.shutdown()
return
}
} else {
s.shutdown()
}
if err != nil {
return
}
}
})
}
func (s *responsesWSSession) getCurrent() *responsesWSCallState {
s.stateMu.Lock()
defer s.stateMu.Unlock()
return s.current
}
func (s *responsesWSSession) tryReserveCurrent(state *responsesWSCallState) bool {
s.stateMu.Lock()
defer s.stateMu.Unlock()
if s.current != nil {
return false
}
s.current = state
return true
}
func (s *responsesWSSession) getTarget() *websocket.Conn {
s.connectionMu.Lock()
defer s.connectionMu.Unlock()
return s.target
}
func (s *responsesWSSession) setTarget(target *websocket.Conn) bool {
s.connectionMu.Lock()
defer s.connectionMu.Unlock()
if s.ctx.Err() != nil {
_ = target.Close()
return false
}
s.target = target
return true
}
func (s *responsesWSSession) writeTarget(kind int, message []byte) error {
s.targetWriteMu.Lock()
defer s.targetWriteMu.Unlock()
target := s.getTarget()
if target == nil {
return errors.New("responses websocket upstream is not connected")
}
if err := target.SetWriteDeadline(time.Now().Add(responsesWSWriteTimeout)); err != nil {
return err
}
return target.WriteMessage(kind, message)
}
func (s *responsesWSSession) writeClient(kind int, message []byte) error {
s.clientWriteMu.Lock()
defer s.clientWriteMu.Unlock()
if err := s.client.SetWriteDeadline(time.Now().Add(responsesWSWriteTimeout)); err != nil {
return err
}
return s.client.WriteMessage(kind, message)
}
func (s *responsesWSSession) sendError(eventID, streamID string, apiErr *types.NewAPIError) {
if apiErr == nil {
return
}
payload, err := buildResponsesWSErrorPayload(eventID, streamID, apiErr)
if err == nil {
_ = s.writeClient(websocket.TextMessage, payload)
}
}
func (s *responsesWSSession) closeTarget() {
s.connectionMu.Lock()
target, unregister := s.target, s.unregister
s.target, s.unregister = nil, nil
s.connectionMu.Unlock()
if unregister != nil {
unregister()
}
if target != nil {
_ = target.Close()
}
}
func (s *responsesWSSession) shutdown() {
s.cancel()
s.closeTarget()
_ = s.client.Close()
}
func (s *responsesWSSession) registerChannelClose(channelID int) {
unregister := wsmanager.Register(channelID, wsmanager.KindResponses, s.closeForPolicy)
s.connectionMu.Lock()
if s.ctx.Err() != nil {
s.connectionMu.Unlock()
unregister()
return
}
s.unregister = unregister
s.connectionMu.Unlock()
}
// closeForPolicy tells both peers why the connection ends, like the realtime
// relay does; WriteControl is safe alongside a blocked data writer.
func (s *responsesWSSession) closeForPolicy(reason string) {
closeMessage := websocket.FormatCloseMessage(websocket.ClosePolicyViolation, reason)
deadline := time.Now().Add(time.Second)
_ = s.client.WriteControl(websocket.CloseMessage, closeMessage, deadline)
if target := s.getTarget(); target != nil {
_ = target.WriteControl(websocket.CloseMessage, closeMessage, deadline)
}
s.shutdown()
}
// Stream identity belongs to the WebSocket envelope. For the legacy wrapped
// input, the top-level field takes precedence over response.stream_id.
func parseResponsesWSEnvelope(message []byte) (responsesWSCreateEvent, string, error) {
var event responsesWSCreateEvent
if err := common.Unmarshal(message, &event); err != nil {
return event, "", fmt.Errorf("invalid websocket event json: %w", err)
}
streamRaw := event.StreamID
if len(streamRaw) == 0 && len(event.Request) > 0 {
var wrapped struct {
StreamID common.RawMessage `json:"stream_id"`
}
if err := common.Unmarshal(event.Request, &wrapped); err == nil {
streamRaw = wrapped.StreamID
}
}
var streamID string
if len(streamRaw) > 0 {
if err := common.Unmarshal(streamRaw, &streamID); err != nil || len(streamID) < 1 || len(streamID) > 256 {
return event, "", errors.New("stream_id must contain 1-256 ASCII letters, digits, underscores, hyphens, or periods")
}
for _, char := range streamID {
if !(char >= 'a' && char <= 'z' || char >= 'A' && char <= 'Z' || char >= '0' && char <= '9' || char == '_' || char == '-' || char == '.') {
return event, "", errors.New("stream_id must contain only ASCII letters, digits, underscores, hyphens, or periods")
}
}
}
if strings.TrimSpace(event.Type) == "" {
return event, streamID, errors.New("websocket event type is required")
}
return event, streamID, nil
}
func newResponsesWSInvalidRequestError(err error) *types.NewAPIError {
return types.NewErrorWithStatusCode(err, types.ErrorCodeInvalidRequest, http.StatusBadRequest, types.ErrOptionWithSkipRetry())
}
// normalizeResponsesWSCreateEvent builds the HTTP-shaped request body from a
// response.create whose envelope the read loop already parsed and validated.
// The message is decoded once more as a raw map (the body source) and once as
// the typed request; the envelope is not parsed again.
func normalizeResponsesWSCreateEvent(message []byte, event responsesWSCreateEvent, streamID string) (responsesWSCreateRequest, error) {
create := responsesWSCreateRequest{StreamID: streamID, Generate: event.Generate}
// Decode the wrapped request on its own so outer fields cannot enter the HTTP body.
source := message
if len(event.Request) > 0 {
source = event.Request
}
var raw map[string]common.RawMessage
if err := common.Unmarshal(source, &raw); err != nil {
return create, err
}
if len(event.Request) > 0 {
if len(create.Generate) == 0 {
create.Generate = raw["generate"]
}
} else {
for _, key := range []string{"type", "event_id", "background", "stream", "stream_options"} {
delete(raw, key)
}
}
delete(raw, "generate")
delete(raw, "stream_id")
payload, err := common.Marshal(raw)
if err != nil {
return create, err
}
if err := common.Unmarshal(payload, &create.Request); err != nil {
return create, err
}
if helper.ExceedsMaxTokensLimit(create.Request.MaxOutputTokens) {
return create, errors.New("max_output_tokens is invalid")
}
create.Request.Stream = nil
create.Request.StreamOptions = nil
create.Body = payload
return create, nil
}
func buildResponsesWSCreateEvent(jsonData []byte, generate common.RawMessage, streamID string) ([]byte, error) {
var event map[string]common.RawMessage
if err := common.Unmarshal(jsonData, &event); err != nil {
return nil, err
}
typeData, err := common.Marshal(responsesWSEventTypeResponseCreate)
if err != nil {
return nil, err
}
event["type"] = typeData
// The normalized body never carries stream_id; a channel parameter
// override may inject one, and it must not shadow the envelope's value.
delete(event, "stream_id")
if streamID != "" {
streamData, err := common.Marshal(streamID)
if err != nil {
return nil, err
}
event["stream_id"] = streamData
}
delete(event, "event_id")
delete(event, "background")
delete(event, "stream")
delete(event, "stream_options")
if len(generate) > 0 {
event["generate"] = generate
}
return common.Marshal(event)
}
func buildResponsesWSErrorPayload(eventID, streamID string, apiErr *types.NewAPIError) ([]byte, error) {
if apiErr == nil {
return nil, errors.New("api error is nil")
}
status := apiErr.StatusCode
if status == 0 {
status = http.StatusInternalServerError
}
openaiErr := apiErr.ToOpenAIError()
return common.Marshal(&responsesWSErrorEvent{
Type: "error",
Status: status,
EventID: eventID,
StreamID: streamID,
Error: &openaiErr,
})
}
// checkResponsesWSModelAccess applies the token model limit with the same name
// matching as the HTTP distributor. Unlike HTTP, it also runs for requests
// pinned to a channel: the check is stricter there on purpose, because a
// persistent connection keeps serving the model after the pin was resolved.
func checkResponsesWSModelAccess(c *gin.Context, modelName string) *types.NewAPIError {
if !common.GetContextKeyBool(c, appconstant.ContextKeyTokenModelLimitEnabled) {
return nil
}
raw, ok := common.GetContextKey(c, appconstant.ContextKeyTokenModelLimit)
if !ok {
return types.NewErrorWithStatusCode(errors.New("token has no model access"), types.ErrorCodeAccessDenied, http.StatusForbidden, types.ErrOptionWithSkipRetry())
}
tokenModelLimit, ok := raw.(map[string]bool)
if !ok {
tokenModelLimit = map[string]bool{}
}
if !middleware.TokenModelLimitAllows(tokenModelLimit, modelName) {
return types.NewErrorWithStatusCode(fmt.Errorf("token is not allowed to use model %s", modelName), types.ErrorCodeAccessDenied, http.StatusForbidden, types.ErrOptionWithSkipRetry())
}
return nil
}
// selectResponsesWSChannel narrows the shared HTTP selection to channels that
// speak the Responses WebSocket protocol and renders its outcome as a
// non-retryable NewAPIError.
func selectResponsesWSChannel(c *gin.Context, modelName string, retryParam *service.RetryParam) (*appmodel.Channel, *types.NewAPIError) {
constraints := service.GetChannelConstraints(c)
if !slices.ContainsFunc(constraints.Filters, func(filter appdto.ChannelFilter) bool {
return filter.Kind == appdto.FilterResponsesWebSocket
}) {
constraints.AddFilter(appdto.ChannelFilter{Kind: appdto.FilterResponsesWebSocket})
}
channel, _, selectErr := service.SelectChannelForRequest(c, modelName, retryParam)
if selectErr != nil {
message := selectErr.Message
if selectErr.MessageID != "" {
message = i18n.T(c, selectErr.MessageID, selectErr.Params)
}
code := selectErr.Code
if code == "" {
code = types.ErrorCodeGetChannelFailed
}
return nil, types.NewErrorWithStatusCode(errors.New(message), code, selectErr.StatusCode, types.ErrOptionWithSkipRetry())
}
if err := middleware.SetupContextForSelectedChannel(c, channel, modelName); err != nil {
return nil, err
}
return channel, nil
}
+682
View File
@@ -0,0 +1,682 @@
package relay
import (
"context"
"errors"
"fmt"
"net/http"
"net/http/httptest"
"strings"
"sync/atomic"
"testing"
"time"
"github.com/QuantumNous/new-api/common"
"github.com/QuantumNous/new-api/constant"
appdto "github.com/QuantumNous/new-api/dto"
"github.com/QuantumNous/new-api/i18n"
"github.com/QuantumNous/new-api/model"
relaycommon "github.com/QuantumNous/new-api/relay/common"
"github.com/QuantumNous/new-api/relaykit/dto"
"github.com/QuantumNous/new-api/relaykit/types"
"github.com/QuantumNous/new-api/service"
"github.com/QuantumNous/new-api/setting/operation_setting"
"github.com/gin-gonic/gin"
"github.com/gorilla/websocket"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
)
// normalizeResponsesWSTestMessage runs the read-loop envelope parse followed by
// request normalization, exactly as the session does for one response.create.
func normalizeResponsesWSTestMessage(message []byte) (responsesWSCreateRequest, string, error) {
envelope, streamID, err := parseResponsesWSEnvelope(message)
if err != nil {
return responsesWSCreateRequest{StreamID: streamID}, envelope.EventID, err
}
create, err := normalizeResponsesWSCreateEvent(message, envelope, streamID)
return create, envelope.EventID, err
}
func TestNormalizeResponsesWSMaxOutputTokens(t *testing.T) {
for _, tc := range []struct {
value string
valid bool
}{
{value: "0", valid: true},
{value: "1073741823", valid: true},
{value: "1073741824"},
{value: "18446744073686646784"},
{value: "-1"},
} {
for _, wrapped := range []bool{false, true} {
t.Run(fmt.Sprintf("%s/wrapped=%t", tc.value, wrapped), func(t *testing.T) {
fields := `"model":"gpt-5.1","input":"hi","max_output_tokens":` + tc.value
payload := `{"type":"response.create",` + fields + `}`
if wrapped {
payload = `{"type":"response.create","response":{` + fields + `}}`
}
create, _, err := normalizeResponsesWSTestMessage([]byte(payload))
if !tc.valid {
require.Error(t, err)
assert.Equal(t, http.StatusBadRequest, newResponsesWSInvalidRequestError(err).StatusCode)
return
}
require.NoError(t, err)
require.NotNil(t, create.Request.MaxOutputTokens)
assert.Equal(t, tc.value, fmt.Sprint(*create.Request.MaxOutputTokens))
})
}
}
}
func TestSelectResponsesWSChannelHonorsPinsAndFilters(t *testing.T) {
require.NoError(t, i18n.Init())
database := setupRelayChannelDB(t)
enabled := &model.Channel{Name: "enabled", Key: "sk-test", Status: common.ChannelStatusEnabled, Type: constant.ChannelTypeOpenAI}
enabled.SetSetting(dto.ChannelSettings{ResponsesWebSocketEnabled: true})
wsDisabled := &model.Channel{Name: "ws-disabled", Key: "sk-test", Status: common.ChannelStatusEnabled, Type: constant.ChannelTypeOpenAI}
disabled := &model.Channel{Name: "disabled", Key: "sk-test", Status: common.ChannelStatusManuallyDisabled, Type: constant.ChannelTypeOpenAI}
filtered := &model.Channel{Name: "filtered", Key: "sk-test", Status: common.ChannelStatusEnabled, Type: constant.ChannelTypeAdvancedCustom}
for _, channel := range []*model.Channel{enabled, disabled, filtered, wsDisabled} {
require.NoError(t, database.Create(channel).Error)
}
for _, tc := range []struct {
name string
channelID int
status int
}{
{name: "token pin overrides origin pin", channelID: enabled.Id},
{name: "disabled pin rejects", channelID: disabled.Id, status: http.StatusForbidden},
{name: "pin cannot bypass websocket switch", channelID: wsDisabled.Id, status: http.StatusBadRequest},
{name: "pin cannot bypass path filter", channelID: filtered.Id, status: http.StatusBadRequest},
{name: "missing pin rejects", channelID: 99999, status: http.StatusBadRequest},
} {
t.Run(tc.name, func(t *testing.T) {
c, _ := gin.CreateTestContext(httptest.NewRecorder())
c.Request = httptest.NewRequest(http.MethodGet, "/v1/responses", nil)
constraints := service.GetChannelConstraints(c)
constraints.AddPin(appdto.ChannelPin{ChannelId: disabled.Id, Source: appdto.PinSourceOriginTask, Rank: appdto.PinRankOriginTask, RetryMode: appdto.PinRetrySameChannel})
constraints.AddPin(appdto.ChannelPin{ChannelId: tc.channelID, Source: appdto.PinSourceToken, Rank: appdto.PinRankToken, RetryMode: appdto.PinRetrySingleAttempt})
constraints.AddFilter(appdto.ChannelFilter{Kind: appdto.FilterRequestPath, RequestPath: c.Request.URL.Path})
channel, apiErr := selectResponsesWSChannel(c, "gpt-5.1", &service.RetryParam{Ctx: c, ModelName: "gpt-5.1", TokenGroup: "default"})
if tc.status != 0 {
require.NotNil(t, apiErr)
assert.Equal(t, tc.status, apiErr.StatusCode)
assert.Nil(t, channel)
assert.False(t, service.ShouldRetryRelayError(c, apiErr, 2))
return
}
require.Nil(t, apiErr)
require.NotNil(t, channel)
assert.Equal(t, enabled.Id, channel.Id)
assert.Equal(t, enabled.Id, common.GetContextKeyInt(c, constant.ContextKeyChannelId))
assert.False(t, service.ShouldRetryRelayError(c, types.NewErrorWithStatusCode(errors.New("upstream failed"), types.ErrorCodeDoRequestFailed, 503), 2))
})
}
}
// The WebSocket relay must admit the same model names as the HTTP distributor
// when a token restricts models (reasoning suffixes and @modifiers included).
func TestCheckResponsesWSModelAccessMatchesHTTPTokenLimits(t *testing.T) {
for _, tc := range []struct {
model string
status int
}{
{model: "gpt-5.1"},
{model: "gpt-5.1-high"},
{model: "gpt-5.1@thinking:on"},
{model: "gpt-4o", status: http.StatusForbidden},
} {
t.Run(tc.model, func(t *testing.T) {
c, _ := gin.CreateTestContext(httptest.NewRecorder())
c.Request = httptest.NewRequest(http.MethodGet, "/v1/responses", nil)
common.SetContextKey(c, constant.ContextKeyTokenModelLimitEnabled, true)
common.SetContextKey(c, constant.ContextKeyTokenModelLimit, map[string]bool{"gpt-5.1": true})
apiErr := checkResponsesWSModelAccess(c, tc.model)
if tc.status == 0 {
assert.Nil(t, apiErr)
return
}
require.NotNil(t, apiErr)
assert.Equal(t, tc.status, apiErr.StatusCode)
assert.False(t, service.ShouldRetryRelayError(c, apiErr, 2))
})
}
}
func TestSelectResponsesWSChannelAcceptsNativeResponsesChannelTypes(t *testing.T) {
require.NoError(t, i18n.Init())
database := setupRelayChannelDB(t)
nativeRoute := &dto.AdvancedCustomConfig{Routes: []dto.AdvancedCustomRoute{{IncomingPath: "/v1/responses", UpstreamPath: "/v1/responses"}}}
noneRoute := &dto.AdvancedCustomConfig{Routes: []dto.AdvancedCustomRoute{{IncomingPath: "/v1/responses", UpstreamPath: "/v1/responses", Converter: "none"}}}
convertedRoute := &dto.AdvancedCustomConfig{Routes: []dto.AdvancedCustomRoute{{IncomingPath: "/v1/responses", UpstreamPath: "/v1/chat/completions", Converter: "openai_responses_to_openai_chat_completions"}}}
chatRoute := &dto.AdvancedCustomConfig{Routes: []dto.AdvancedCustomRoute{{IncomingPath: "/v1/chat/completions", UpstreamPath: "/v1/chat/completions"}}}
for _, tc := range []struct {
name string
channelType int
advanced *dto.AdvancedCustomConfig
rejectedBy appdto.ChannelFilterKind
}{
{name: "new api", channelType: constant.ChannelTypeNewAPI},
{name: "sub2api", channelType: constant.ChannelTypeSub2API},
{name: "advanced custom native responses route", channelType: constant.ChannelTypeAdvancedCustom, advanced: nativeRoute},
{name: "advanced custom explicit none converter", channelType: constant.ChannelTypeAdvancedCustom, advanced: noneRoute},
{name: "advanced custom converter route", channelType: constant.ChannelTypeAdvancedCustom, advanced: convertedRoute, rejectedBy: appdto.FilterResponsesWebSocket},
{name: "advanced custom without responses route", channelType: constant.ChannelTypeAdvancedCustom, advanced: chatRoute, rejectedBy: appdto.FilterRequestPath},
{name: "anthropic", channelType: constant.ChannelTypeAnthropic, rejectedBy: appdto.FilterResponsesWebSocket},
} {
t.Run(tc.name, func(t *testing.T) {
channel := &model.Channel{Name: tc.name, Key: "sk-test", Status: common.ChannelStatusEnabled, Type: tc.channelType}
channel.SetSetting(dto.ChannelSettings{ResponsesWebSocketEnabled: true})
channel.SetOtherSettings(dto.ChannelOtherSettings{AdvancedCustom: tc.advanced})
require.NoError(t, database.Create(channel).Error)
c, _ := gin.CreateTestContext(httptest.NewRecorder())
c.Request = httptest.NewRequest(http.MethodGet, "/v1/responses", nil)
constraints := service.GetChannelConstraints(c)
constraints.AddPin(appdto.ChannelPin{ChannelId: channel.Id, Source: appdto.PinSourceToken, Rank: appdto.PinRankToken, RetryMode: appdto.PinRetrySingleAttempt})
constraints.AddFilter(appdto.ChannelFilter{Kind: appdto.FilterRequestPath, RequestPath: c.Request.URL.Path})
selected, apiErr := selectResponsesWSChannel(c, "ws-model", &service.RetryParam{Ctx: c, ModelName: "ws-model", TokenGroup: "default"})
if tc.rejectedBy != "" {
require.NotNil(t, apiErr)
assert.Nil(t, selected)
assert.Equal(t, http.StatusBadRequest, apiErr.StatusCode)
assert.Equal(t, types.ErrorCode(tc.rejectedBy), apiErr.GetErrorCode())
return
}
require.Nil(t, apiErr)
require.NotNil(t, selected)
assert.Equal(t, channel.Id, selected.Id)
})
}
}
func TestRestoreConnectionContextRejectsChangedAdvancedCustomRoute(t *testing.T) {
database := setupRelayChannelDB(t)
baseURL := "http://upstream.example"
channel := &model.Channel{Name: "advanced", Key: "sk-test", Status: common.ChannelStatusEnabled, Type: constant.ChannelTypeAdvancedCustom, BaseURL: &baseURL}
channel.SetSetting(dto.ChannelSettings{ResponsesWebSocketEnabled: true})
channel.SetOtherSettings(dto.ChannelOtherSettings{AdvancedCustom: &dto.AdvancedCustomConfig{Routes: []dto.AdvancedCustomRoute{{IncomingPath: "/v1/responses", UpstreamPath: "/v1/responses"}}}})
require.NoError(t, database.Create(channel).Error)
c, _ := gin.CreateTestContext(httptest.NewRecorder())
c.Request = httptest.NewRequest(http.MethodGet, "/v1/responses", nil)
service.GetChannelConstraints(c).AddPin(appdto.ChannelPin{ChannelId: channel.Id, Source: appdto.PinSourceToken, Rank: appdto.PinRankToken, RetryMode: appdto.PinRetrySingleAttempt})
route, ok := channel.GetOtherSettings().AdvancedCustom.MatchPathForModel("/v1/responses", "ws-model")
require.True(t, ok)
session := &responsesWSSession{lockedChannelID: channel.Id, lockedModel: "ws-model", lockedKey: "sk-test", lockedRoute: route,
lockedContext: map[constant.ContextKey]any{constant.ContextKeyChannelBaseUrl: channel.GetBaseURL(), constant.ContextKeyChannelHeaderOverride: channel.GetHeaderOverride()}}
require.Nil(t, session.restoreConnectionContext(c, "ws-model"))
// Request-level edits (model list, explicit none converter) keep the connection.
channel.SetOtherSettings(dto.ChannelOtherSettings{AdvancedCustom: &dto.AdvancedCustomConfig{Routes: []dto.AdvancedCustomRoute{{IncomingPath: "/v1/responses", UpstreamPath: "/v1/responses", Converter: "none", Models: []string{"ws-model", "other-model"}}}}})
require.NoError(t, database.Model(channel).Update("settings", channel.OtherSettings).Error)
require.Nil(t, session.restoreConnectionContext(c, "ws-model"))
channel.SetOtherSettings(dto.ChannelOtherSettings{AdvancedCustom: &dto.AdvancedCustomConfig{Routes: []dto.AdvancedCustomRoute{{IncomingPath: "/v1/responses", UpstreamPath: "/v2/responses"}}}})
require.NoError(t, database.Model(channel).Update("settings", channel.OtherSettings).Error)
apiErr := session.restoreConnectionContext(c, "ws-model")
require.NotNil(t, apiErr)
assert.Equal(t, http.StatusForbidden, apiErr.StatusCode)
assert.ErrorContains(t, apiErr, "upstream route changed")
}
func TestResponsesWSChannelRoutingRequiresExplicitOptIn(t *testing.T) {
require.NoError(t, i18n.Init())
database := setupRelayChannelDB(t)
require.NoError(t, database.AutoMigrate(&model.Ability{}))
legacy := &model.Channel{Name: "legacy-http", Key: "sk-test", Type: constant.ChannelTypeOpenAI, Status: common.ChannelStatusEnabled, Group: "default", Models: "ws-model", Priority: common.GetPointer(int64(10))}
enabled := &model.Channel{Name: "websocket", Key: "sk-test", Type: constant.ChannelTypeCodex, Status: common.ChannelStatusEnabled, Group: "default", Models: "ws-model", Priority: common.GetPointer(int64(0))}
unsupported := &model.Channel{Name: "unsupported", Key: "sk-test", Type: constant.ChannelTypeAnthropic, Status: common.ChannelStatusEnabled, Group: "default", Models: "ws-model", Priority: common.GetPointer(int64(5))}
enabled.SetSetting(dto.ChannelSettings{ResponsesWebSocketEnabled: true})
unsupported.SetSetting(dto.ChannelSettings{ResponsesWebSocketEnabled: true})
for _, channel := range []*model.Channel{legacy, enabled, unsupported} {
require.NoError(t, database.Create(channel).Error)
require.NoError(t, database.Create(&model.Ability{ChannelId: channel.Id, Model: "ws-model", Group: "default", Enabled: true, Priority: channel.Priority}).Error)
}
previousCache := common.MemoryCacheEnabled
t.Cleanup(func() {
defer func() { common.MemoryCacheEnabled = previousCache }()
ids := []int{legacy.Id, enabled.Id, unsupported.Id}
require.NoError(t, database.Where("channel_id IN ?", ids).Delete(&model.Ability{}).Error)
require.NoError(t, database.Where("id IN ?", ids).Delete(&model.Channel{}).Error)
model.InitChannelCache()
})
common.MemoryCacheEnabled = true
model.InitChannelCache()
c, _ := gin.CreateTestContext(httptest.NewRecorder())
c.Request = httptest.NewRequest(http.MethodGet, "/v1/responses", nil)
params := &service.RetryParam{Ctx: c, ModelName: "ws-model", TokenGroup: "default"}
channel, apiErr := selectResponsesWSChannel(c, "ws-model", params)
require.Nil(t, apiErr)
require.NotNil(t, channel)
assert.Equal(t, enabled.Id, channel.Id)
httpChannel, err := model.GetRandomSatisfiedChannel("default", "ws-model", 0, nil)
require.NoError(t, err)
require.NotNil(t, httpChannel)
assert.Equal(t, legacy.Id, httpChannel.Id)
// Disabling the saved setting takes effect for the next create on an existing session.
enabled.SetSetting(dto.ChannelSettings{ResponsesWebSocketEnabled: false})
require.NoError(t, database.Model(enabled).Update("setting", enabled.Setting).Error)
model.InitChannelCache()
session := &responsesWSSession{lockedChannelID: enabled.Id, lockedModel: "ws-model"}
apiErr = session.restoreConnectionContext(c, "ws-model")
require.NotNil(t, apiErr)
assert.Equal(t, http.StatusForbidden, apiErr.StatusCode)
channel, apiErr = selectResponsesWSChannel(c, "ws-model", params)
require.NotNil(t, apiErr)
assert.Nil(t, channel)
}
// Session affinity follows the HTTP distributor: a strict binding to an
// unusable channel fails the request, a prefer binding falls back, and an
// unusable binding is dropped from the cache either way.
func TestSelectResponsesWSChannelHonorsStrictSessionBinding(t *testing.T) {
require.NoError(t, i18n.Init())
database := setupRelayChannelDB(t)
require.NoError(t, database.AutoMigrate(&model.Ability{}))
bound := &model.Channel{Name: "bound", Key: "sk-test", Type: constant.ChannelTypeOpenAI, Status: common.ChannelStatusAutoDisabled, Group: "default", Models: "ws-model", Priority: common.GetPointer(int64(10))}
fallback := &model.Channel{Name: "fallback", Key: "sk-test", Type: constant.ChannelTypeOpenAI, Status: common.ChannelStatusEnabled, Group: "default", Models: "ws-model", Priority: common.GetPointer(int64(0))}
for _, channel := range []*model.Channel{bound, fallback} {
channel.SetSetting(dto.ChannelSettings{ResponsesWebSocketEnabled: true})
require.NoError(t, database.Create(channel).Error)
require.NoError(t, database.Create(&model.Ability{ChannelId: channel.Id, Model: "ws-model", Group: "default", Enabled: channel.Status == common.ChannelStatusEnabled, Priority: channel.Priority}).Error)
}
previousCache := common.MemoryCacheEnabled
t.Cleanup(func() {
defer func() { common.MemoryCacheEnabled = previousCache }()
ids := []int{bound.Id, fallback.Id}
require.NoError(t, database.Where("channel_id IN ?", ids).Delete(&model.Ability{}).Error)
require.NoError(t, database.Where("id IN ?", ids).Delete(&model.Channel{}).Error)
model.InitChannelCache()
})
common.MemoryCacheEnabled = true
model.InitChannelCache()
affinity := operation_setting.GetChannelAffinitySetting()
previousAffinity := *affinity
t.Cleanup(func() { *affinity = previousAffinity })
snapshot, err := model.BuildRequestPolicy(map[string]string{
"channel_affinity_setting.enabled": "true",
"channel_affinity_setting.session_mode": "strict",
"channel_affinity_setting.rules": `[{"name":"session","model_regex":[".*"],"key_sources":[{"type":"request_header","key":"X-Session"}],"session_mode":"inherit"}]`,
})
require.NoError(t, err)
*affinity = snapshot.Affinity
newSessionContext := func() *gin.Context {
c, _ := gin.CreateTestContext(httptest.NewRecorder())
c.Request = httptest.NewRequest(http.MethodGet, "/v1/responses", nil)
c.Request.Header.Set("X-Session", t.Name())
return c
}
seed := newSessionContext()
_, found := service.GetPreferredChannelByAffinity(seed, "ws-model", "default")
require.False(t, found)
seed.Set("channel_id", bound.Id)
service.RecordChannelAffinity(seed, bound.Id)
t.Cleanup(func() { service.ClearCurrentChannelAffinityCache(seed) })
strict := newSessionContext()
channel, apiErr := selectResponsesWSChannel(strict, "ws-model", &service.RetryParam{Ctx: strict, ModelName: "ws-model", TokenGroup: "default", Retry: common.GetPointer(0)})
require.NotNil(t, apiErr)
assert.Nil(t, channel)
assert.Equal(t, http.StatusServiceUnavailable, apiErr.StatusCode)
assert.False(t, service.ShouldRetryRelayError(strict, apiErr, 2))
_, found = service.GetPreferredChannelByAffinity(seed, "ws-model", "default")
assert.False(t, found, "an unusable binding is cleared unless keep_on_channel_disabled is set")
affinity.SessionMode = "prefer"
service.RecordChannelAffinity(seed, bound.Id)
prefer := newSessionContext()
channel, apiErr = selectResponsesWSChannel(prefer, "ws-model", &service.RetryParam{Ctx: prefer, ModelName: "ws-model", TokenGroup: "default", Retry: common.GetPointer(0)})
require.Nil(t, apiErr)
require.NotNil(t, channel)
assert.Equal(t, fallback.Id, channel.Id)
}
func TestNormalizeResponsesWSCreateEventWrapper(t *testing.T) {
message := []byte(`{
"type": "response.create",
"event_id": "evt_1",
"generate": false,
"response": {
"model": "gpt-5.3-codex-spark",
"input": "hi",
"store": false,
"stream": true,
"stream_options": {"include_usage": true}
}
}`)
create, eventID, err := normalizeResponsesWSTestMessage(message)
if err != nil {
t.Fatalf("normalizeResponsesWSTestMessage() error = %v", err)
}
req := create.Request
if eventID != "evt_1" {
t.Fatalf("eventID = %q, want evt_1", eventID)
}
if req.Model != "gpt-5.3-codex-spark" {
t.Fatalf("model = %q", req.Model)
}
if strings.TrimSpace(string(create.Generate)) != "false" {
t.Fatalf("generate = %s, want false", create.Generate)
}
if req.Stream != nil {
t.Fatalf("stream = %v, want nil", req.Stream)
}
if req.StreamOptions != nil {
t.Fatalf("stream_options = %#v, want nil", req.StreamOptions)
}
if strings.TrimSpace(string(req.Store)) != "false" {
t.Fatalf("store = %s, want false", req.Store)
}
}
func TestNormalizeResponsesWSCreateEventFlat(t *testing.T) {
message := []byte(`{
"type": "response.create",
"event_id": "evt_2",
"model": "gpt-5.3-codex-spark",
"input": "hi",
"generate": false,
"stream": true,
"background": true,
"stream_options": {"include_usage": true}
}`)
create, eventID, err := normalizeResponsesWSTestMessage(message)
if err != nil {
t.Fatalf("normalizeResponsesWSTestMessage() error = %v", err)
}
req := create.Request
if eventID != "evt_2" {
t.Fatalf("eventID = %q, want evt_2", eventID)
}
if req.Model != "gpt-5.3-codex-spark" {
t.Fatalf("model = %q", req.Model)
}
if strings.TrimSpace(string(create.Generate)) != "false" {
t.Fatalf("generate = %s, want false", create.Generate)
}
if req.Stream != nil {
t.Fatalf("stream = %v, want nil", req.Stream)
}
if req.StreamOptions != nil {
t.Fatalf("stream_options = %#v, want nil", req.StreamOptions)
}
}
func TestBuildResponsesWSCreateEventIsFlat(t *testing.T) {
payload := []byte(`{
"model": "gpt-5.3-codex-spark",
"input": "hi",
"store": false,
"event_id": "evt_upstream",
"stream": true,
"background": true,
"stream_options": {"include_usage": true}
}`)
got, err := buildResponsesWSCreateEvent(payload, common.RawMessage(`false`), "")
if err != nil {
t.Fatalf("buildResponsesWSCreateEvent() error = %v", err)
}
var data map[string]any
if err := common.Unmarshal(got, &data); err != nil {
t.Fatalf("unmarshal result: %v", err)
}
if data["type"] != responsesWSEventTypeResponseCreate {
t.Fatalf("type = %#v", data["type"])
}
if data["model"] != "gpt-5.3-codex-spark" || data["input"] != "hi" || data["store"] != false {
t.Fatalf("unexpected flat event fields: %s", got)
}
if data["generate"] != false {
t.Fatalf("generate = %#v, want false", data["generate"])
}
for _, key := range []string{"response", "event_id", "stream", "background", "stream_options"} {
if _, ok := data[key]; ok {
t.Fatalf("field %q should not be present in upstream event: %s", key, got)
}
}
}
func TestHTTPResponsesRequestOmitsWebSocketMetadata(t *testing.T) {
var req dto.OpenAIResponsesRequest
require.NoError(t, common.Unmarshal([]byte(`{"model":"gpt-5.3-codex-spark","input":"hi","generate":false,"stream_id":"planner"}`), &req))
got, err := common.Marshal(req)
require.NoError(t, err)
var data map[string]any
require.NoError(t, common.Unmarshal(got, &data))
assert.NotContains(t, data, "stream_id")
assert.NotContains(t, data, "generate")
}
func TestBuildResponsesWSErrorPayloadIncludesStatus(t *testing.T) {
payload, err := buildResponsesWSErrorPayload("evt_err", "", types.NewErrorWithStatusCode(
errors.New("model is required"),
types.ErrorCodeInvalidRequest,
http.StatusBadRequest,
types.ErrOptionWithSkipRetry(),
))
if err != nil {
t.Fatalf("buildResponsesWSErrorPayload() error = %v", err)
}
var data struct {
Type string `json:"type"`
Status int `json:"status"`
EventID string `json:"event_id"`
Error *types.OpenAIError `json:"error"`
}
if err := common.Unmarshal(payload, &data); err != nil {
t.Fatalf("unmarshal result: %v", err)
}
if data.Type != "error" || data.Status != http.StatusBadRequest || data.EventID != "evt_err" {
t.Fatalf("unexpected error event: %s", payload)
}
if data.Error == nil || data.Error.Code != string(types.ErrorCodeInvalidRequest) {
t.Fatalf("unexpected error body: %#v", data.Error)
}
}
func TestResponsesWSInvalidRequestErrorUsesBadRequestStatus(t *testing.T) {
payload, err := buildResponsesWSErrorPayload("", "", newResponsesWSInvalidRequestError(errors.New("bad event")))
if err != nil {
t.Fatalf("buildResponsesWSErrorPayload() error = %v", err)
}
var data struct {
Status int `json:"status"`
}
if err := common.Unmarshal(payload, &data); err != nil {
t.Fatalf("unmarshal result: %v", err)
}
if data.Status != http.StatusBadRequest {
t.Fatalf("status = %d, want %d", data.Status, http.StatusBadRequest)
}
}
func newTestResponsesWSTarget(t *testing.T) (*websocket.Conn, func()) {
t.Helper()
target, _, cleanup := newTestWebSocketPair(t)
return target, cleanup
}
func newTestWebSocketPair(t *testing.T) (*websocket.Conn, *websocket.Conn, func()) {
t.Helper()
upgrader := websocket.Upgrader{}
serverConnCh := make(chan *websocket.Conn, 1)
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
conn, err := upgrader.Upgrade(w, r, nil)
if err != nil {
t.Errorf("upgrade websocket: %v", err)
return
}
serverConnCh <- conn
}))
targetURL := "ws" + strings.TrimPrefix(server.URL, "http")
target, _, err := websocket.DefaultDialer.Dial(targetURL, nil)
if err != nil {
server.Close()
t.Fatalf("dial websocket: %v", err)
}
serverConn := <-serverConnCh
cleanup := func() {
_ = target.Close()
_ = serverConn.Close()
server.Close()
}
return target, serverConn, cleanup
}
func TestResponsesWSMessageSizeLimit(t *testing.T) {
previous := constant.MaxRequestBodyMB
constant.MaxRequestBodyMB = 1
t.Cleanup(func() { constant.MaxRequestBodyMB = previous })
client, server, cleanup := newTestWebSocketPair(t)
defer cleanup()
c, _ := gin.CreateTestContext(httptest.NewRecorder())
c.Request = httptest.NewRequest(http.MethodGet, "/v1/responses", nil)
var admitted atomic.Bool
done := make(chan struct{})
go func() {
defer close(done)
ResponsesWebSocketHelper(c, server, func(*http.Request, string, func(*gin.Context) *types.NewAPIError) *types.NewAPIError {
admitted.Store(true)
return nil
})
}()
_ = client.WriteMessage(websocket.TextMessage, []byte(strings.Repeat("x", (1<<20)+1)))
require.NoError(t, client.SetReadDeadline(time.Now().Add(5*time.Second)))
_, _, err := client.ReadMessage()
assert.True(t, websocket.IsCloseError(err, websocket.CloseMessageTooBig), "oversized message must be rejected before admission: %v", err)
select {
case <-done:
case <-time.After(5 * time.Second):
t.Fatal("websocket request did not exit after oversized message")
}
assert.False(t, admitted.Load())
}
func TestResponsesWSShutdownInterruptsBusyWriter(t *testing.T) {
client, server, cleanupClient := newTestWebSocketPair(t)
defer cleanupClient()
target, peer, cleanupTarget := newTestWebSocketPair(t)
defer cleanupTarget()
ctx, cancel := context.WithCancel(context.Background())
s := &responsesWSSession{ctx: ctx, cancel: cancel, client: server, target: target}
// A blocked network writer owns this lock. Closing the connection must
// remain possible so that writer can be interrupted.
s.targetWriteMu.Lock()
defer s.targetWriteMu.Unlock()
done := make(chan struct{})
go func() { s.shutdown(); close(done) }()
select {
case <-done:
case <-time.After(5 * time.Second):
t.Fatal("shutdown waited for the network writer")
}
require.NoError(t, peer.SetReadDeadline(time.Now().Add(time.Second)))
_, _, err := peer.ReadMessage()
assert.Error(t, err)
require.NoError(t, client.SetReadDeadline(time.Now().Add(time.Second)))
_, _, err = client.ReadMessage()
assert.Error(t, err)
assert.Nil(t, s.getTarget())
}
func TestResponsesWSPassthroughPreservesRawPricingParameters(t *testing.T) {
create, _, err := normalizeResponsesWSTestMessage([]byte(`{"type":"response.create","generate":false,"response":{"model":"gpt-5.1","input":"hi","vendor":{"tier":"premium"},"stream":true}}`))
require.NoError(t, err)
c, _ := gin.CreateTestContext(httptest.NewRecorder())
c.Request = httptest.NewRequest(http.MethodPost, "/v1/responses", strings.NewReader(string(create.Body)))
c.Request.Header.Set("Content-Type", "application/json")
common.SetContextKey(c, constant.ContextKeyOriginalModel, create.Request.Model)
common.SetContextKey(c, constant.ContextKeyChannelType, constant.ChannelTypeOpenAI)
common.SetContextKey(c, constant.ContextKeyChannelSetting, dto.ChannelSettings{PassThroughBodyEnabled: true})
info := relaycommon.GenRelayInfoResponses(c, &create.Request)
payload, apiErr := buildResponsesWSCreatePayload(c, info, create.Request, create.Generate, create.StreamID)
require.Nil(t, apiErr)
assert.JSONEq(t, `{"type":"response.create","generate":false,"model":"gpt-5.1","input":"hi","vendor":{"tier":"premium"}}`, string(payload))
storage, err := common.GetBodyStorage(c)
require.NoError(t, err)
require.NoError(t, storage.Close())
}
func TestResponsesWSStreamIdentity(t *testing.T) {
for _, tc := range []struct {
name, fields, want string
invalid bool
}{
{name: "default", fields: `"model":"gpt-4o"`},
{name: "named", fields: `"model":"gpt-4o","stream_id":"planner.A-1_2"`, want: "planner.A-1_2"},
{name: "wrapped", fields: `"response":{"model":"gpt-4o","stream_id":"wrapped"}`, want: "wrapped"},
{name: "top-level-wins", fields: `"stream_id":"outer","response":{"model":"gpt-4o","stream_id":"inner"}`, want: "outer"},
{name: "maximum", fields: `"model":"gpt-4o","stream_id":"` + strings.Repeat("a", 256) + `"`, want: strings.Repeat("a", 256)},
{name: "too-long", fields: `"stream_id":"` + strings.Repeat("a", 257) + `"`, invalid: true},
{name: "empty", fields: `"stream_id":""`, invalid: true},
{name: "null", fields: `"stream_id":null`, invalid: true},
{name: "number", fields: `"stream_id":123`, invalid: true},
{name: "unicode", fields: `"stream_id":"计划"`, invalid: true},
{name: "spaces", fields: `"stream_id":"a b"`, invalid: true},
{name: "invalid-top-level", fields: `"stream_id":"","response":{"stream_id":"inner"}`, invalid: true},
} {
t.Run(tc.name, func(t *testing.T) {
create, eventID, err := normalizeResponsesWSTestMessage([]byte(`{"type":"response.create","event_id":"test",` + tc.fields + `}`))
assert.Equal(t, "test", eventID)
if tc.invalid {
require.ErrorContains(t, err, "stream_id")
return
}
require.NoError(t, err)
assert.Equal(t, tc.want, create.StreamID)
assert.NotContains(t, string(create.Body), "stream_id")
for _, passthrough := range []bool{false, true} {
c, _ := gin.CreateTestContext(httptest.NewRecorder())
c.Request = httptest.NewRequest(http.MethodPost, "/v1/responses", strings.NewReader(string(create.Body)))
common.SetContextKey(c, constant.ContextKeyOriginalModel, create.Request.Model)
common.SetContextKey(c, constant.ContextKeyChannelType, constant.ChannelTypeOpenAI)
common.SetContextKey(c, constant.ContextKeyChannelSetting, dto.ChannelSettings{PassThroughBodyEnabled: passthrough})
info := relaycommon.GenRelayInfoResponses(c, &create.Request)
payload, apiErr := buildResponsesWSCreatePayload(c, info, create.Request, create.Generate, create.StreamID)
require.Nil(t, apiErr)
var event map[string]any
require.NoError(t, common.Unmarshal(payload, &event))
if tc.want == "" {
assert.NotContains(t, event, "stream_id")
} else {
assert.Equal(t, tc.want, event["stream_id"])
}
if storage, err := common.GetBodyStorage(c); err == nil {
require.NoError(t, storage.Close())
}
}
})
}
}
func TestResponsesWSErrorAttribution(t *testing.T) {
for _, tc := range []struct {
name, payload, control string
terminal, ambiguous, controlError bool
}{
{name: "active-server-failure", payload: `{"status":500,"error":{"type":"server_error"}}`, terminal: true},
{name: "other-stream", payload: `{"stream_id":"other","status":500}`},
{name: "previous-response", payload: `{"response_id":"previous","status":500}`},
{name: "cancel-event-id", payload: `{"event_id":"cancel","status":500}`, control: `{"event_id":"cancel"}`, controlError: true},
{name: "cancel-target", payload: `{"response_id":"missing","status":400}`, control: `{"response_id":"missing"}`, controlError: true},
{name: "cancel-code", payload: `{"error":{"type":"invalid_request_error","code":"response_not_found"}}`, control: `{"response_id":"missing"}`, controlError: true},
{name: "ambiguous-after-cancel", payload: `{"status":500}`, control: `{"response_id":"active"}`, terminal: true, ambiguous: true},
{name: "explicit-generation-failure", payload: `{"response_id":"active","status":500}`, control: `{"response_id":"active"}`, terminal: true},
} {
t.Run(tc.name, func(t *testing.T) {
var event responsesWSErrorEvent
require.NoError(t, common.Unmarshal([]byte(tc.payload), &event))
terminal, ambiguous, controlError := responsesWSErrorEndsRequest(event, "planner", "active", []byte(tc.control))
assert.Equal(t, tc.terminal, terminal)
assert.Equal(t, tc.ambiguous, ambiguous)
assert.Equal(t, tc.controlError, controlError)
})
}
}
+15
View File
@@ -2,7 +2,10 @@ package relay
import (
"fmt"
"sync"
"time"
"github.com/QuantumNous/new-api/pkg/wsmanager"
relaycommon "github.com/QuantumNous/new-api/relay/common"
"github.com/QuantumNous/new-api/relaykit/dto"
"github.com/QuantumNous/new-api/relaykit/types"
@@ -33,6 +36,18 @@ func WssHelper(c *gin.Context, info *relaycommon.RelayInfo) (newAPIError *types.
if resp != nil {
info.TargetWs = resp.(*websocket.Conn)
defer info.TargetWs.Close()
var closeOnce sync.Once
unregister := wsmanager.Register(info.ChannelId, wsmanager.KindRealtime, func(reason string) {
closeOnce.Do(func() {
deadline := time.Now().Add(time.Second)
closeMessage := websocket.FormatCloseMessage(websocket.ClosePolicyViolation, reason)
_ = info.ClientWs.WriteControl(websocket.CloseMessage, closeMessage, deadline)
_ = info.TargetWs.WriteControl(websocket.CloseMessage, closeMessage, deadline)
_ = info.ClientWs.Close()
_ = info.TargetWs.Close()
})
})
defer unregister()
}
usage, newAPIError := adaptor.DoResponse(c, nil, info)