mirror of
https://gh-proxy.org/https://github.com/klinsmaya/new-api.git
synced 2026-11-08 01:17:09 +08:00
CaIon
55e2ae54d1
A task retry that moved to another plugin or channel kept the previous attempt's usage expression. A retry from an expression-priced plugin to one priced per request charged, and later settled, at the first plugin's price. Each attempt now prices itself from scratch. A retry to a more expensive channel was topped up only after the upstream had accepted the task. A user who could not cover the difference went into arrears, or the task was dropped while the upstream was already running it. The retry now raises the reservation before the request is sent, with the same balance check as the first pre-consume, and a channel the balance cannot cover is never called. A wallet top-up taken before submission replaced the amount held by the first pre-consume instead of adding to it, so a request that then failed refunded only the top-up. Image requests whose quantity was raised before submission were affected too. The full amount is now refunded.
654 lines
30 KiB
Go
654 lines
30 KiB
Go
package relay
|
|
|
|
import (
|
|
"encoding/base64"
|
|
"net/http"
|
|
"net/http/httptest"
|
|
"strconv"
|
|
"strings"
|
|
"testing"
|
|
|
|
"github.com/QuantumNous/new-api/common"
|
|
"github.com/QuantumNous/new-api/constant"
|
|
"github.com/QuantumNous/new-api/dto"
|
|
"github.com/QuantumNous/new-api/i18n"
|
|
"github.com/QuantumNous/new-api/model"
|
|
"github.com/QuantumNous/new-api/pkg/billingexpr"
|
|
pluginruntime "github.com/QuantumNous/new-api/pkg/jsplugin"
|
|
relaycommon "github.com/QuantumNous/new-api/relay/common"
|
|
"github.com/QuantumNous/new-api/service"
|
|
"github.com/QuantumNous/new-api/setting/billing_setting"
|
|
"github.com/QuantumNous/new-api/setting/config"
|
|
"github.com/QuantumNous/new-api/setting/ratio_setting"
|
|
"github.com/gin-gonic/gin"
|
|
"github.com/glebarez/sqlite"
|
|
"github.com/stretchr/testify/assert"
|
|
"github.com/stretchr/testify/require"
|
|
"gorm.io/gorm"
|
|
)
|
|
|
|
func setupRelayChannelDB(t *testing.T) *gorm.DB {
|
|
t.Helper()
|
|
previousDB := model.DB
|
|
previousType := common.MainDatabaseType()
|
|
previousCache := common.MemoryCacheEnabled
|
|
database, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{})
|
|
require.NoError(t, err)
|
|
sqlDB, err := database.DB()
|
|
require.NoError(t, err)
|
|
sqlDB.SetMaxOpenConns(1)
|
|
require.NoError(t, database.AutoMigrate(&model.Channel{}))
|
|
model.DB = database
|
|
common.SetMainDatabaseType(common.DatabaseTypeSQLite)
|
|
common.MemoryCacheEnabled = false
|
|
t.Cleanup(func() {
|
|
model.DB = previousDB
|
|
common.SetMainDatabaseType(previousType)
|
|
common.MemoryCacheEnabled = previousCache
|
|
require.NoError(t, sqlDB.Close())
|
|
})
|
|
return database
|
|
}
|
|
|
|
func TestRealtimeFetchAppliesChannelHeaderOverrides(t *testing.T) {
|
|
database := setupRelayChannelDB(t)
|
|
service.InitHttpClient()
|
|
requests := make(chan *http.Request, 1)
|
|
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
requests <- r.Clone(r.Context())
|
|
_, _ = w.Write([]byte(`{"name":"operations/task-1","done":false}`))
|
|
}))
|
|
t.Cleanup(server.Close)
|
|
headerOverride := `{"X-API-Key":"{api_key}"}`
|
|
channel := &model.Channel{
|
|
Name: "realtime-header-override", Type: constant.ChannelTypeGemini,
|
|
Key: "channel-key", BaseURL: &server.URL, HeaderOverride: &headerOverride,
|
|
}
|
|
require.NoError(t, database.Create(channel).Error)
|
|
task := &model.Task{
|
|
TaskID: "public-task", ChannelId: channel.Id,
|
|
Status: model.TaskStatusInProgress, Progress: "50%",
|
|
PrivateData: model.TaskPrivateData{
|
|
Key: "saved-key", UpstreamTaskID: base64.RawURLEncoding.EncodeToString([]byte("operations/task-1")),
|
|
},
|
|
}
|
|
|
|
tryRealtimeFetch(task, true)
|
|
|
|
require.Len(t, requests, 1)
|
|
request := <-requests
|
|
assert.Equal(t, "/v1beta/operations/task-1", request.URL.Path)
|
|
assert.Equal(t, "saved-key", request.Header.Get("X-API-Key"))
|
|
assert.Equal(t, "saved-key", request.Header.Get("X-Goog-Api-Key"))
|
|
}
|
|
|
|
func TestRealtimeFetchRefundsTaskItFindsFailed(t *testing.T) {
|
|
database := setupRelayChannelDB(t)
|
|
require.NoError(t, database.AutoMigrate(&model.Task{}, &model.User{}, &model.Log{}))
|
|
previousLogDB := model.LOG_DB
|
|
previousFactory := service.GetTaskAdaptorFunc
|
|
previousRedis := common.RedisEnabled
|
|
model.LOG_DB = database
|
|
common.RedisEnabled = false
|
|
service.GetTaskAdaptorFunc = func(platform constant.TaskPlatform) service.TaskPollingAdaptor {
|
|
if adaptor := GetTaskAdaptor(platform); adaptor != nil {
|
|
return adaptor
|
|
}
|
|
return nil
|
|
}
|
|
t.Cleanup(func() {
|
|
model.LOG_DB = previousLogDB
|
|
service.GetTaskAdaptorFunc = previousFactory
|
|
common.RedisEnabled = previousRedis
|
|
})
|
|
service.InitHttpClient()
|
|
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
_, _ = w.Write([]byte(`{"name":"operations/task-1","done":true,"error":{"code":3,"message":"blocked by safety"}}`))
|
|
}))
|
|
t.Cleanup(server.Close)
|
|
channel := &model.Channel{Name: "realtime-refund", Type: constant.ChannelTypeGemini, Key: "channel-key", BaseURL: &server.URL}
|
|
require.NoError(t, database.Create(channel).Error)
|
|
require.NoError(t, database.Create(&model.User{Id: 7, Username: "realtime-owner", Quota: 100}).Error)
|
|
task := &model.Task{
|
|
TaskID: "public-task", Platform: constant.TaskPlatform(strconv.Itoa(constant.ChannelTypeGemini)),
|
|
UserId: 7, ChannelId: channel.Id, Quota: 1000,
|
|
Status: model.TaskStatusInProgress, Progress: "50%",
|
|
PrivateData: model.TaskPrivateData{
|
|
Key: "channel-key", UpstreamTaskID: base64.RawURLEncoding.EncodeToString([]byte("operations/task-1")),
|
|
},
|
|
}
|
|
require.NoError(t, database.Create(task).Error)
|
|
|
|
tryRealtimeFetch(task, true)
|
|
|
|
var stored model.Task
|
|
require.NoError(t, database.First(&stored, task.ID).Error)
|
|
assert.Equal(t, model.TaskStatus(model.TaskStatusFailure), stored.Status)
|
|
assert.Equal(t, "blocked by safety", stored.FailReason)
|
|
assert.Zero(t, stored.Quota, "the failed task is refunded")
|
|
var user model.User
|
|
require.NoError(t, database.First(&user, 7).Error)
|
|
assert.Equal(t, 1100, user.Quota)
|
|
}
|
|
|
|
func TestApplyChannelPinPreservesOriginTasksAndRetryMode(t *testing.T) {
|
|
database := setupRelayChannelDB(t)
|
|
channel := &model.Channel{Name: "origin-channel", Key: "sk-test", Status: common.ChannelStatusEnabled, Type: constant.ChannelTypeDoubaoVideo}
|
|
require.NoError(t, database.Create(channel).Error)
|
|
originTask := &model.Task{
|
|
TaskID: "task-lock", ChannelId: channel.Id, Action: "text_to_video", Status: model.TaskStatusSuccess,
|
|
PrivateData: model.TaskPrivateData{UpstreamTaskID: "upstream-task-lock"},
|
|
Data: []byte(`{"id":"upstream-task-lock"}`),
|
|
}
|
|
|
|
for _, tc := range []struct {
|
|
name string
|
|
tokenPin bool
|
|
apply func(*gin.Context, *relaycommon.RelayInfo) *dto.TaskError
|
|
}{
|
|
{name: "origin affinity", apply: ApplyOriginTaskAffinity},
|
|
{name: "same channel retry", apply: ApplyChannelPin},
|
|
{name: "token pin suppresses channel lock", tokenPin: true, apply: ApplyChannelPin},
|
|
} {
|
|
t.Run(tc.name, func(t *testing.T) {
|
|
c, _ := gin.CreateTestContext(httptest.NewRecorder())
|
|
common.SetContextKey(c, constant.ContextKeyOriginTasks, []*model.Task{originTask})
|
|
constraints := service.GetChannelConstraints(c)
|
|
constraints.AddPin(dto.ChannelPin{ChannelId: channel.Id, Source: dto.PinSourceOriginTask, Rank: dto.PinRankOriginTask, RetryMode: dto.PinRetrySameChannel})
|
|
if tc.tokenPin {
|
|
constraints.AddPin(dto.ChannelPin{ChannelId: channel.Id, Source: dto.PinSourceToken, Rank: dto.PinRankToken, RetryMode: dto.PinRetrySingleAttempt})
|
|
}
|
|
info := &relaycommon.RelayInfo{TaskRelayInfo: &relaycommon.TaskRelayInfo{}}
|
|
require.Nil(t, tc.apply(c, info))
|
|
if tc.tokenPin {
|
|
assert.Nil(t, info.LockedChannel)
|
|
} else {
|
|
locked, ok := info.LockedChannel.(*model.Channel)
|
|
require.True(t, ok)
|
|
require.NotNil(t, locked)
|
|
assert.Equal(t, channel.Id, locked.Id)
|
|
}
|
|
require.Len(t, info.OriginTasks, 1)
|
|
assert.Equal(t, "task-lock", info.OriginTasks[0].TaskID)
|
|
assert.Equal(t, "upstream-task-lock", info.OriginTasks[0].UpstreamTaskID)
|
|
assert.Equal(t, "text_to_video", info.OriginTasks[0].Action)
|
|
assert.Equal(t, string(model.TaskStatusSuccess), info.OriginTasks[0].Status)
|
|
assert.Equal(t, []byte(originTask.Data), info.OriginTasks[0].Data)
|
|
})
|
|
}
|
|
}
|
|
|
|
func TestTaskModel2DtoNormalizesLegacyAction(t *testing.T) {
|
|
task := &model.Task{Action: "firstTailGenerate"}
|
|
|
|
dtoTask := TaskModel2Dto(task)
|
|
|
|
assert.Equal(t, constant.TaskActionFirstTailToVideo, dtoTask.Action)
|
|
assert.Equal(t, "firstTailGenerate", task.Action)
|
|
}
|
|
|
|
const mappingOrderSubmitPlugin = `
|
|
export const meta = {apiVersion:1,key:"maporder",name:"Map Order",version:"1.0.0",author:{name:"Test"},models:["declared-model"],fetchMode:"per_task"};
|
|
export function buildSubmitRequest(ctx) {
|
|
return {url: ctx.baseUrl+"/submit", method:"POST", body:{upstreamModel: ctx.upstreamModel, model: ctx.model}, action:"text_to_video"};
|
|
}
|
|
export function parseSubmitResponse(){return {taskId:"1"};}
|
|
export function buildQueryRequest(){return {url:"https://provider.example"};}
|
|
export function parseTaskResult(){return {status:"SUCCESS"};}
|
|
`
|
|
|
|
const mappingOrderRewritePlugin = `
|
|
export const meta = {apiVersion:1,key:"maporder-rw",name:"Map Order RW",version:"1.0.0",author:{name:"Test"},models:["declared-model"],fetchMode:"per_task"};
|
|
export function buildSubmitRequest(ctx) {
|
|
return {url: ctx.baseUrl+"/submit", method:"POST", body:{upstreamModel: ctx.upstreamModel}, rewriteModel:"rewritten"};
|
|
}
|
|
export function parseSubmitResponse(){return {taskId:"1"};}
|
|
export function buildQueryRequest(){return {url:"https://provider.example"};}
|
|
export function parseTaskResult(){return {status:"SUCCESS"};}
|
|
`
|
|
|
|
func pinMappingOrderPlugin(t *testing.T, c *gin.Context, source string) {
|
|
t.Helper()
|
|
plugin, err := pluginruntime.NewRegistry().Register(source, pluginruntime.Options{})
|
|
require.NoError(t, err)
|
|
c.Set(pluginruntime.ContextKeyPinnedPlugin, pluginruntime.PinnedPlugin{Plugin: plugin})
|
|
}
|
|
|
|
func newTaskSubmitContext(t *testing.T, originalModel, mapping string) (*gin.Context, *relaycommon.RelayInfo) {
|
|
t.Helper()
|
|
require.NoError(t, i18n.Init())
|
|
gin.SetMode(gin.TestMode)
|
|
recorder := httptest.NewRecorder()
|
|
c, _ := gin.CreateTestContext(recorder)
|
|
c.Request = httptest.NewRequest(http.MethodPost, "/v1/videos", nil)
|
|
common.SetContextKey(c, constant.ContextKeyOriginalModel, originalModel)
|
|
common.SetContextKey(c, constant.ContextKeyChannelBaseUrl, "https://provider.example")
|
|
if mapping != "" {
|
|
c.Set("model_mapping", mapping)
|
|
}
|
|
c.Set("task_request", map[string]any{"prompt": "p"})
|
|
return c, &relaycommon.RelayInfo{TaskRelayInfo: &relaycommon.TaskRelayInfo{}}
|
|
}
|
|
|
|
func TestRelayTaskSubmitMapsBeforeValidateWhenOriginSet(t *testing.T) {
|
|
const mapping = `{"alias-model":"mid-model","mid-model":"declared-model"}`
|
|
|
|
c, info := newTaskSubmitContext(t, "alias-model", mapping)
|
|
pinMappingOrderPlugin(t, c, mappingOrderSubmitPlugin)
|
|
info.OriginModelName = "alias-model"
|
|
|
|
_, taskErr := RelayTaskSubmit(c, info)
|
|
require.NotNil(t, taskErr)
|
|
assert.Equal(t, "model_price_error", taskErr.Code)
|
|
assert.Equal(t, "alias-model", info.OriginModelName)
|
|
assert.Equal(t, "declared-model", info.UpstreamModelName)
|
|
assert.True(t, info.IsModelMapped)
|
|
}
|
|
|
|
func TestRelayTaskSubmitDeclaredNameWithoutMappingIsUnchanged(t *testing.T) {
|
|
c, info := newTaskSubmitContext(t, "declared-model", "")
|
|
pinMappingOrderPlugin(t, c, mappingOrderSubmitPlugin)
|
|
info.OriginModelName = "declared-model"
|
|
|
|
_, taskErr := RelayTaskSubmit(c, info)
|
|
require.NotNil(t, taskErr)
|
|
assert.Equal(t, "model_price_error", taskErr.Code)
|
|
assert.Equal(t, "declared-model", info.OriginModelName)
|
|
assert.Equal(t, "declared-model", info.UpstreamModelName)
|
|
assert.False(t, info.IsModelMapped)
|
|
}
|
|
|
|
func TestRelayTaskSubmitDoesNotApplyMappingTwice(t *testing.T) {
|
|
c, info := newTaskSubmitContext(t, "alias-model", `{"alias-model":"declared-model"}`)
|
|
pinMappingOrderPlugin(t, c, mappingOrderRewritePlugin)
|
|
info.OriginModelName = "alias-model"
|
|
|
|
_, taskErr := RelayTaskSubmit(c, info)
|
|
require.NotNil(t, taskErr)
|
|
assert.Equal(t, "model_price_error", taskErr.Code)
|
|
assert.Equal(t, "rewritten", info.UpstreamModelName, "late mapping would overwrite rewriteModel with the chain tail")
|
|
assert.Equal(t, "alias-model", info.OriginModelName)
|
|
}
|
|
|
|
func TestRelayTaskSubmitEmptyOriginKeepsLateMapping(t *testing.T) {
|
|
plugin, err := pluginruntime.NewRegistry().Register(mappingOrderSubmitPlugin, pluginruntime.Options{})
|
|
require.NoError(t, err)
|
|
synthesized := service.CoverTaskActionToModelName(constant.TaskPlatform(plugin.Meta.Key), "text_to_video")
|
|
c, info := newTaskSubmitContext(t, "pre-validate-upstream",
|
|
`{"pre-validate-upstream":"should-not-apply-early","`+synthesized+`":"legacy-tail"}`)
|
|
c.Set(pluginruntime.ContextKeyPinnedPlugin, pluginruntime.PinnedPlugin{Plugin: plugin})
|
|
info.OriginModelName = ""
|
|
|
|
_, taskErr := RelayTaskSubmit(c, info)
|
|
require.NotNil(t, taskErr)
|
|
assert.Equal(t, "model_price_error", taskErr.Code)
|
|
assert.Equal(t, synthesized, info.OriginModelName)
|
|
assert.Equal(t, "legacy-tail", info.UpstreamModelName)
|
|
assert.True(t, info.IsModelMapped)
|
|
}
|
|
|
|
const billingFallbackPlugin = `
|
|
export const meta = {apiVersion:1,key:"bill-fallback",name:"Bill Fallback",version:"1.0.0",author:{name:"Test"},models:["declared-model"],fetchMode:"per_task"};
|
|
export function buildSubmitRequest(ctx) {
|
|
return {url: ctx.baseUrl+"/submit", method:"POST", body:{upstreamModel: ctx.upstreamModel, model: ctx.model}, action:"text_to_video"};
|
|
}
|
|
export function parseSubmitResponse(){return {taskId:"1"};}
|
|
export function buildQueryRequest(){return {url:"https://provider.example"};}
|
|
export function parseTaskResult(){return {status:"SUCCESS"};}
|
|
`
|
|
|
|
func saveBillingConfig(t *testing.T) {
|
|
t.Helper()
|
|
saved := map[string]string{}
|
|
require.NoError(t, config.GlobalConfig.SaveToDB(func(key, value string) error {
|
|
saved[key] = value
|
|
return nil
|
|
}))
|
|
t.Cleanup(func() {
|
|
require.NoError(t, config.GlobalConfig.LoadFromDB(saved))
|
|
})
|
|
}
|
|
|
|
func TestRelayTaskSubmitAliasBillingIdentityAndExprFallback(t *testing.T) {
|
|
const mapping = `{"alias-model":"declared-model"}`
|
|
const aliasExpr = `tier("alias", 2)`
|
|
const tailExpr = `tier("tail", 3)`
|
|
const legacyExpr = `tier("legacy", u("old_units") * 2)`
|
|
|
|
tests := []struct {
|
|
name string
|
|
modes map[string]string
|
|
exprs map[string]string
|
|
wantTiered bool
|
|
wantExpr string
|
|
source string
|
|
}{
|
|
{
|
|
name: "alias own tiered wins",
|
|
modes: map[string]string{"alias-model": "tiered_expr", "declared-model": "tiered_expr"},
|
|
exprs: map[string]string{"alias-model": aliasExpr, "declared-model": tailExpr},
|
|
wantTiered: true,
|
|
wantExpr: aliasExpr,
|
|
},
|
|
{
|
|
name: "fallback uses tail expr",
|
|
modes: map[string]string{"declared-model": "tiered_expr"},
|
|
exprs: map[string]string{"declared-model": tailExpr},
|
|
wantTiered: true,
|
|
wantExpr: tailExpr,
|
|
},
|
|
{
|
|
name: "stored expression keeps running after schema narrows",
|
|
modes: map[string]string{"declared-model": "tiered_expr"},
|
|
exprs: map[string]string{"declared-model": legacyExpr},
|
|
wantTiered: true,
|
|
wantExpr: legacyExpr,
|
|
source: strings.Replace(billingFallbackPlugin, `fetchMode:"per_task"`, `fetchMode:"per_task", usageSchema:{old_units:{type:"number",unit:"count"}}, usageProfiles:[{models:["declared-model"],schema:{seconds:{type:"number",unit:"second"}}}]`, 1) + `
|
|
export function extractUsage(){return {old_units:2};}
|
|
`,
|
|
},
|
|
{
|
|
name: "neither tiered uses ordinary pricing",
|
|
wantTiered: false,
|
|
},
|
|
}
|
|
for _, testCase := range tests {
|
|
t.Run(testCase.name, func(t *testing.T) {
|
|
saveBillingConfig(t)
|
|
if len(testCase.modes) > 0 {
|
|
modeJSON, marshalErr := common.Marshal(testCase.modes)
|
|
require.NoError(t, marshalErr)
|
|
exprJSON, marshalErr := common.Marshal(testCase.exprs)
|
|
require.NoError(t, marshalErr)
|
|
require.NoError(t, config.GlobalConfig.LoadFromDB(map[string]string{
|
|
"billing_setting.billing_mode": string(modeJSON),
|
|
"billing_setting.billing_expr": string(exprJSON),
|
|
}))
|
|
if testCase.wantExpr == aliasExpr {
|
|
require.Equal(t, billing_setting.BillingModeTieredExpr, billing_setting.GetBillingMode("alias-model"))
|
|
} else {
|
|
require.Equal(t, billing_setting.BillingModeRatio, billing_setting.GetBillingMode("alias-model"))
|
|
require.Equal(t, billing_setting.BillingModeTieredExpr, billing_setting.GetBillingMode("declared-model"))
|
|
}
|
|
}
|
|
|
|
c, info := newTaskSubmitContext(t, "alias-model", mapping)
|
|
c.Set("group", "default")
|
|
info.UserGroup = "default"
|
|
info.UsingGroup = "default"
|
|
source := testCase.source
|
|
if source == "" {
|
|
source = billingFallbackPlugin
|
|
}
|
|
pinMappingOrderPlugin(t, c, source)
|
|
info.OriginModelName = "alias-model"
|
|
|
|
_, taskErr := RelayTaskSubmit(c, info)
|
|
require.NotNil(t, taskErr)
|
|
assert.Equal(t, "alias-model", info.OriginModelName)
|
|
assert.Equal(t, "declared-model", info.UpstreamModelName)
|
|
assert.True(t, info.IsModelMapped)
|
|
|
|
task := model.InitTask(constant.TaskPlatform("bill-fallback"), info)
|
|
assert.Equal(t, "alias-model", task.Properties.OriginModelName)
|
|
assert.Equal(t, "declared-model", task.Properties.UpstreamModelName)
|
|
|
|
if testCase.wantTiered {
|
|
require.NotNil(t, info.TieredBillingSnapshot, "submission error: %+v", taskErr)
|
|
assert.Equal(t, "alias-model", info.TieredBillingSnapshot.ModelName)
|
|
assert.Equal(t, testCase.wantExpr, info.TieredBillingSnapshot.ExprString)
|
|
assert.Equal(t, billingexpr.ExprHashString(testCase.wantExpr), info.TieredBillingSnapshot.ExprHash)
|
|
assert.NotEqual(t, "model_price_error", taskErr.Code)
|
|
if testCase.wantExpr == legacyExpr {
|
|
assert.Equal(t, 4*common.QuotaPerUnit, info.TieredBillingSnapshot.EstimatedQuotaBeforeGroup)
|
|
}
|
|
} else {
|
|
assert.Nil(t, info.TieredBillingSnapshot)
|
|
assert.Equal(t, "model_price_error", taskErr.Code)
|
|
}
|
|
})
|
|
}
|
|
}
|
|
|
|
func TestSharedTaskBillingExpressionSelectionAndFrozenSettlement(t *testing.T) {
|
|
const baseExpr = `tier("base", u("seconds") * 2)`
|
|
const alphaExpr = `tier("alpha", u("seconds") * 3)`
|
|
const betaExpr = `tier("beta", u("credits") * 5)`
|
|
const aliasExpr = `tier("alias", u("credits") * 7)`
|
|
for _, tc := range []struct {
|
|
name, plugin, model, mapping, modelExpr, mode, wantExpr string
|
|
variants map[string]string
|
|
wantPriceError, profiled bool
|
|
}{
|
|
{name: "executing plugin override", plugin: "billing-beta", model: "declared-model", modelExpr: baseExpr, mode: "tiered_expr", variants: map[string]string{"billing-alpha::declared-model": alphaExpr, "billing-beta::declared-model": betaExpr}, wantExpr: betaExpr},
|
|
{name: "override ignores model mode", plugin: "billing-beta", model: "declared-model", modelExpr: baseExpr, mode: "ratio", variants: map[string]string{"billing-beta::declared-model": betaExpr}, wantExpr: betaExpr},
|
|
{name: "model expression fallback", plugin: "billing-alpha", model: "declared-model", modelExpr: baseExpr, mode: "tiered_expr", wantExpr: baseExpr},
|
|
{name: "alias override precedes mapped override", plugin: "billing-beta", model: "alias-model", mapping: `{"alias-model":"declared-model"}`, modelExpr: baseExpr, mode: "tiered_expr", variants: map[string]string{"billing-beta::declared-model": betaExpr, "billing-beta::alias-model": aliasExpr}, wantExpr: aliasExpr},
|
|
{name: "mapped override precedes model fallback", plugin: "billing-beta", model: "alias-model", mapping: `{"alias-model":"declared-model"}`, modelExpr: baseExpr, mode: "tiered_expr", variants: map[string]string{"billing-beta::declared-model": betaExpr}, wantExpr: betaExpr},
|
|
{name: "unconfigured plugin cannot use another schema", plugin: "billing-beta", model: "declared-model", modelExpr: baseExpr, mode: "tiered_expr", wantPriceError: true},
|
|
{name: "missing usage in skipped branch remains incompatible", plugin: "billing-beta", model: "declared-model", modelExpr: `true ? tier("free", 0) : tier("missing", u("seconds"))`, mode: "tiered_expr", wantPriceError: true},
|
|
{name: "fixed pricing is still rejected", plugin: "billing-beta", model: "declared-model", variants: map[string]string{"billing-beta::declared-model": `tier("fixed", fixed(1))`}, wantPriceError: true},
|
|
{name: "endpoint mapping keeps the declared profile", plugin: "billing-beta", model: "declared-model", mapping: `{"declared-model":"ep-endpoint"}`, modelExpr: baseExpr, mode: "tiered_expr", variants: map[string]string{"billing-beta::declared-model": betaExpr}, wantExpr: betaExpr, profiled: true},
|
|
} {
|
|
t.Run(tc.name, func(t *testing.T) {
|
|
saveBillingConfig(t)
|
|
registry := pluginruntime.NewRegistry()
|
|
for _, spec := range []struct{ key, field, unit string }{{"billing-alpha", "seconds", "second"}, {"billing-beta", "credits", "credit"}} {
|
|
source := strings.ReplaceAll(billingFallbackPlugin, "bill-fallback", spec.key)
|
|
schema := `usageSchema:{` + spec.field + `:{type:"number",unit:"` + spec.unit + `"}}`
|
|
if tc.profiled && spec.key == "billing-beta" {
|
|
// The endpoint ID is undeclared, so only the declared model's profile carries this schema.
|
|
schema = `usageSchema:{seconds:{type:"number",unit:"second"}},usageProfiles:[{models:["declared-model"],schema:{` + spec.field + `:{type:"number",unit:"` + spec.unit + `"}}}]`
|
|
}
|
|
source = strings.Replace(source, `fetchMode:"per_task"`, `fetchMode:"per_task",`+schema, 1)
|
|
source += `export function extractUsage(){return {` + spec.field + `:2};}`
|
|
_, err := registry.Register(source, pluginruntime.Options{})
|
|
require.NoError(t, err)
|
|
}
|
|
variants := tc.variants
|
|
if variants == nil {
|
|
variants = map[string]string{}
|
|
}
|
|
rawVariants, err := common.Marshal(variants)
|
|
require.NoError(t, err)
|
|
modes, err := common.Marshal(map[string]string{"declared-model": tc.mode})
|
|
require.NoError(t, err)
|
|
expressions, err := common.Marshal(map[string]string{"declared-model": tc.modelExpr})
|
|
require.NoError(t, err)
|
|
require.NoError(t, config.GlobalConfig.LoadFromDB(map[string]string{
|
|
billing_setting.PluginBillingExprOption: string(rawVariants), "billing_setting.billing_mode": string(modes), "billing_setting.billing_expr": string(expressions),
|
|
}))
|
|
c, info := newTaskSubmitContext(t, tc.model, tc.mapping)
|
|
c.Set("group", "default")
|
|
c.Set("task_plugin_key", tc.plugin)
|
|
info.UserGroup = "default"
|
|
info.UsingGroup = "default"
|
|
info.OriginModelName = tc.model
|
|
plugin, ok := registry.Generation().Get(tc.plugin)
|
|
require.True(t, ok)
|
|
c.Set(pluginruntime.ContextKeyPinnedPlugin, pluginruntime.PinnedPlugin{Generation: registry.Generation(), Plugin: plugin})
|
|
_, taskErr := RelayTaskSubmit(c, info)
|
|
require.NotNil(t, taskErr) // This fixture stops at reservation, before upstream submission.
|
|
if tc.wantPriceError {
|
|
assert.Equal(t, "model_price_error", taskErr.Code)
|
|
assert.Nil(t, info.TieredBillingSnapshot)
|
|
return
|
|
}
|
|
require.NotNil(t, info.TieredBillingSnapshot, "submission error: %+v", taskErr)
|
|
assert.Equal(t, tc.wantExpr, info.TieredBillingSnapshot.ExprString)
|
|
assert.NotEqual(t, "model_price_error", taskErr.Code)
|
|
require.NoError(t, config.GlobalConfig.LoadFromDB(map[string]string{billing_setting.PluginBillingExprOption: `{}`, "billing_setting.billing_expr": `{}`}))
|
|
field := "seconds"
|
|
if tc.plugin == "billing-beta" {
|
|
field = "credits"
|
|
}
|
|
result, usage, err := service.EvaluateTaskCompletionUsage(info.TieredBillingSnapshot, map[string]any{field: float64(4)})
|
|
require.NoError(t, err)
|
|
assert.Equal(t, float64(4), usage[field])
|
|
assert.Equal(t, float64(2), info.TieredBillingSnapshot.UsageFacts[field])
|
|
assert.Equal(t, 2*info.TieredBillingSnapshot.EstimatedQuotaAfterGroup, result.ActualQuotaAfterGroup)
|
|
assert.Equal(t, tc.wantExpr, info.TieredBillingSnapshot.ExprString)
|
|
})
|
|
}
|
|
}
|
|
|
|
// Issue #7478: task APIs answer 201 Created or 202 Accepted on submission.
|
|
// Every 2xx must reach parseSubmitResponse; only other statuses are upstream
|
|
// failures that keep the upstream status and body.
|
|
func TestRelayTaskSubmitAcceptsAnySuccessfulUpstreamStatus(t *testing.T) {
|
|
service.InitHttpClient()
|
|
const source = `
|
|
export const meta = {apiVersion:1,key:"status-echo",name:"Status Echo",version:"1.0.0",author:{name:"Test"},models:["declared-model"],fetchMode:"per_task"};
|
|
export function buildSubmitRequest(ctx) { return {url: ctx.baseUrl+"/submit", method:"POST", body:{model: ctx.model}, action:"text_to_video"}; }
|
|
export function parseSubmitResponse(ctx, response) { return {taskId: response.body.id, taskData: {status: response.statusCode}}; }
|
|
export function buildQueryRequest(ctx) { return {url: ctx.baseUrl+"/query"}; }
|
|
export function parseTaskResult() { return {status:"SUCCESS"}; }
|
|
`
|
|
for _, tc := range []struct {
|
|
status int
|
|
wantCode string
|
|
}{
|
|
{status: http.StatusOK},
|
|
{status: http.StatusCreated},
|
|
{status: http.StatusAccepted},
|
|
{status: http.StatusBadRequest, wantCode: "fail_to_fetch_task"},
|
|
{status: http.StatusBadGateway, wantCode: "fail_to_fetch_task"},
|
|
} {
|
|
t.Run(strconv.Itoa(tc.status), func(t *testing.T) {
|
|
saveBillingConfig(t)
|
|
require.NoError(t, config.GlobalConfig.LoadFromDB(map[string]string{
|
|
"billing_setting.billing_mode": `{"declared-model":"tiered_expr"}`,
|
|
"billing_setting.billing_expr": `{"declared-model":"tier(\"flat\", 3)"}`,
|
|
}))
|
|
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
|
|
w.Header().Set("Content-Type", "application/json")
|
|
w.WriteHeader(tc.status)
|
|
_, _ = w.Write([]byte(`{"id":"job-42","message":"upstream body"}`))
|
|
}))
|
|
defer server.Close()
|
|
|
|
c, info := newTaskSubmitContext(t, "declared-model", "")
|
|
common.SetContextKey(c, constant.ContextKeyChannelBaseUrl, server.URL)
|
|
c.Set("group", "default")
|
|
info.UserGroup, info.UsingGroup = "default", "default"
|
|
info.OriginModelName = "declared-model"
|
|
// An existing reservation stands in for the wallet so the fixture reaches the upstream call.
|
|
info.Billing = &imageReservation{limit: 1 << 30}
|
|
pinMappingOrderPlugin(t, c, source)
|
|
|
|
result, taskErr := RelayTaskSubmit(c, info)
|
|
if tc.wantCode != "" {
|
|
require.NotNil(t, taskErr)
|
|
assert.Equal(t, tc.wantCode, taskErr.Code)
|
|
assert.Equal(t, tc.status, taskErr.StatusCode)
|
|
assert.Contains(t, taskErr.Message, "upstream body")
|
|
assert.Nil(t, result)
|
|
return
|
|
}
|
|
require.Nil(t, taskErr, "submission error: %+v", taskErr)
|
|
require.NotNil(t, result)
|
|
assert.Equal(t, "job-42", result.UpstreamTaskID)
|
|
assert.JSONEq(t, `{"status":`+strconv.Itoa(tc.status)+`}`, string(result.TaskData))
|
|
})
|
|
}
|
|
}
|
|
|
|
// A retry prices the plugin it lands on from scratch and holds that price
|
|
// before the upstream sees the request. Plugin A fails upstream on the first
|
|
// attempt; the retry runs on plugin B, pinned the way
|
|
// SetupContextForSelectedChannel pins another candidate plugin's channel.
|
|
func TestRelayTaskSubmitRetryRepricesAndReservesBeforeUpstream(t *testing.T) {
|
|
service.InitHttpClient()
|
|
const source = `
|
|
export const meta = {apiVersion:1,key:"KEY",name:"KEY",version:"1.0.0",author:{name:"Test"},models:["declared-model"],fetchMode:"per_task",usageSchema:{seconds:{type:"number",unit:"second",description:"Video generation unit price"}}};
|
|
export function buildSubmitRequest(ctx) { return {url: ctx.baseUrl+"/submit", method:"POST", body:{model: ctx.model}, action:"text_to_video"}; }
|
|
export function parseSubmitResponse(ctx, response) { return {taskId: response.body.id}; }
|
|
export function buildQueryRequest(ctx) { return {url: ctx.baseUrl+"/query"}; }
|
|
export function parseTaskResult() { return {status:"SUCCESS"}; }
|
|
export function extractUsage() { return {seconds:2}; }
|
|
`
|
|
pricier := map[string]string{"retry-a::declared-model": `tier("a", u("seconds") * 0.05)`, "retry-b::declared-model": `tier("b", u("seconds") * 0.4)`}
|
|
for _, tc := range []struct {
|
|
name, perCallPrice, wantCode string
|
|
expressions map[string]string
|
|
wallet, wantQuota int
|
|
wantSubmitted []string
|
|
}{
|
|
{name: "per-call plugin drops the previous expression", expressions: map[string]string{"retry-a::declared-model": `tier("a", u("seconds") * 0.02)`}, perCallPrice: `{"declared-model":0.04}`, wallet: 1 << 30, wantQuota: 40000, wantSubmitted: []string{"a", "b"}},
|
|
{name: "pricier plugin is held before submission", expressions: pricier, perCallPrice: `{}`, wallet: 1 << 30, wantQuota: 400000, wantSubmitted: []string{"a", "b"}},
|
|
{name: "unaffordable plugin is never submitted", expressions: pricier, perCallPrice: `{}`, wallet: 100000, wantCode: "insufficient_user_quota", wantSubmitted: []string{"a"}},
|
|
} {
|
|
t.Run(tc.name, func(t *testing.T) {
|
|
saveBillingConfig(t)
|
|
savedPrices := ratio_setting.ModelPrice2JSONString()
|
|
t.Cleanup(func() { require.NoError(t, ratio_setting.UpdateModelPriceByJSONString(savedPrices)) })
|
|
require.NoError(t, ratio_setting.UpdateModelPriceByJSONString(tc.perCallPrice))
|
|
expressions, err := common.Marshal(tc.expressions)
|
|
require.NoError(t, err)
|
|
require.NoError(t, config.GlobalConfig.LoadFromDB(map[string]string{
|
|
billing_setting.PluginBillingExprOption: string(expressions),
|
|
"billing_setting.billing_mode": `{}`,
|
|
"billing_setting.billing_expr": `{}`,
|
|
"group_ratio_setting.group_ratio": `{"default":1}`,
|
|
}))
|
|
registry := pluginruntime.NewRegistry()
|
|
for _, key := range []string{"retry-a", "retry-b"} {
|
|
_, err := registry.Register(strings.ReplaceAll(source, "KEY", key), pluginruntime.Options{})
|
|
require.NoError(t, err)
|
|
}
|
|
// The fake stands in for the session the first attempt created; its
|
|
// limit is the most the wallet can hold.
|
|
billing := &imageReservation{limit: tc.wallet}
|
|
var submitted []string
|
|
var heldAtSubmit []int
|
|
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
attempt, _, _ := strings.Cut(strings.TrimPrefix(r.URL.Path, "/"), "/")
|
|
submitted = append(submitted, attempt)
|
|
heldAtSubmit = append(heldAtSubmit, billing.held)
|
|
w.Header().Set("Content-Type", "application/json")
|
|
if attempt == "a" {
|
|
w.WriteHeader(http.StatusInternalServerError)
|
|
}
|
|
_, _ = w.Write([]byte(`{"id":"job-` + attempt + `"}`))
|
|
}))
|
|
defer server.Close()
|
|
|
|
c, info := newTaskSubmitContext(t, "declared-model", "")
|
|
c.Set("group", "default")
|
|
info.UserGroup, info.UsingGroup, info.OriginModelName = "default", "default", "declared-model"
|
|
info.Billing = billing
|
|
var result *TaskSubmitResult
|
|
var taskErr *dto.TaskError
|
|
for _, attempt := range []string{"a", "b"} {
|
|
plugin, ok := registry.Generation().Get("retry-" + attempt)
|
|
require.True(t, ok)
|
|
c.Set(pluginruntime.ContextKeyPinnedPlugin, pluginruntime.PinnedPlugin{Generation: registry.Generation(), Plugin: plugin})
|
|
c.Set("task_plugin_key", plugin.Meta.Key)
|
|
common.SetContextKey(c, constant.ContextKeyChannelBaseUrl, server.URL+"/"+attempt)
|
|
result, taskErr = RelayTaskSubmit(c, info)
|
|
if attempt == "a" {
|
|
require.NotNil(t, taskErr)
|
|
require.Equal(t, http.StatusInternalServerError, taskErr.StatusCode, "the first attempt fails upstream and is retried")
|
|
}
|
|
}
|
|
|
|
require.Equal(t, tc.wantSubmitted, submitted)
|
|
if tc.wantCode != "" {
|
|
require.NotNil(t, taskErr)
|
|
assert.Equal(t, tc.wantCode, taskErr.Code)
|
|
assert.Equal(t, http.StatusForbidden, taskErr.StatusCode)
|
|
return
|
|
}
|
|
require.Nil(t, taskErr, "submission error: %+v", taskErr)
|
|
assert.Equal(t, tc.wantQuota, result.Quota)
|
|
assert.Equal(t, tc.wantQuota, heldAtSubmit[1], "plugin B must not see the request before its price is held")
|
|
if tc.perCallPrice != `{}` {
|
|
assert.Nil(t, info.TieredBillingSnapshot, "a per-call plugin must not settle with the previous plugin's expression")
|
|
return
|
|
}
|
|
require.NotNil(t, info.TieredBillingSnapshot)
|
|
assert.Equal(t, tc.expressions["retry-b::declared-model"], info.TieredBillingSnapshot.ExprString)
|
|
})
|
|
}
|
|
}
|