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")
}