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
1 parent
aa6c3f50f4
commit
ceddfd3d6b
828 files changed
+122253
-13702
No files matched your search
+171
-2
@@ -3,10 +3,12 @@ package relay
|
||||
import (
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
"github.com/QuantumNous/new-api/common"
|
||||
"github.com/QuantumNous/new-api/constant"
|
||||
"github.com/QuantumNous/new-api/dto"
|
||||
"github.com/QuantumNous/new-api/model"
|
||||
"github.com/QuantumNous/new-api/pkg/billingexpr"
|
||||
pluginruntime "github.com/QuantumNous/new-api/pkg/jsplugin"
|
||||
@@ -15,10 +17,82 @@ import (
|
||||
"github.com/QuantumNous/new-api/setting/billing_setting"
|
||||
"github.com/QuantumNous/new-api/setting/config"
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/glebarez/sqlite"
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/require"
|
||||
"gorm.io/gorm"
|
||||
)
|
||||
|
||||
func setupRelayChannelDB(t *testing.T) *gorm.DB {
|
||||
t.Helper()
|
||||
previousDB := model.DB
|
||||
previousType := common.MainDatabaseType()
|
||||
previousCache := common.MemoryCacheEnabled
|
||||
database, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{})
|
||||
require.NoError(t, err)
|
||||
sqlDB, err := database.DB()
|
||||
require.NoError(t, err)
|
||||
sqlDB.SetMaxOpenConns(1)
|
||||
require.NoError(t, database.AutoMigrate(&model.Channel{}))
|
||||
model.DB = database
|
||||
common.SetMainDatabaseType(common.DatabaseTypeSQLite)
|
||||
common.MemoryCacheEnabled = false
|
||||
t.Cleanup(func() {
|
||||
model.DB = previousDB
|
||||
common.SetMainDatabaseType(previousType)
|
||||
common.MemoryCacheEnabled = previousCache
|
||||
require.NoError(t, sqlDB.Close())
|
||||
})
|
||||
return database
|
||||
}
|
||||
|
||||
func TestApplyChannelPinPreservesOriginTasksAndRetryMode(t *testing.T) {
|
||||
database := setupRelayChannelDB(t)
|
||||
channel := &model.Channel{Name: "origin-channel", Key: "sk-test", Status: common.ChannelStatusEnabled, Type: constant.ChannelTypeDoubaoVideo}
|
||||
require.NoError(t, database.Create(channel).Error)
|
||||
originTask := &model.Task{
|
||||
TaskID: "task-lock", ChannelId: channel.Id, Action: "text_to_video", Status: model.TaskStatusSuccess,
|
||||
PrivateData: model.TaskPrivateData{UpstreamTaskID: "upstream-task-lock"},
|
||||
Data: []byte(`{"id":"upstream-task-lock"}`),
|
||||
}
|
||||
|
||||
for _, tc := range []struct {
|
||||
name string
|
||||
tokenPin bool
|
||||
apply func(*gin.Context, *relaycommon.RelayInfo) *dto.TaskError
|
||||
}{
|
||||
{name: "origin affinity", apply: ApplyOriginTaskAffinity},
|
||||
{name: "same channel retry", apply: ApplyChannelPin},
|
||||
{name: "token pin suppresses channel lock", tokenPin: true, apply: ApplyChannelPin},
|
||||
} {
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
c, _ := gin.CreateTestContext(httptest.NewRecorder())
|
||||
common.SetContextKey(c, constant.ContextKeyOriginTasks, []*model.Task{originTask})
|
||||
constraints := service.GetChannelConstraints(c)
|
||||
constraints.AddPin(dto.ChannelPin{ChannelId: channel.Id, Source: dto.PinSourceOriginTask, Rank: dto.PinRankOriginTask, RetryMode: dto.PinRetrySameChannel})
|
||||
if tc.tokenPin {
|
||||
constraints.AddPin(dto.ChannelPin{ChannelId: channel.Id, Source: dto.PinSourceToken, Rank: dto.PinRankToken, RetryMode: dto.PinRetrySingleAttempt})
|
||||
}
|
||||
info := &relaycommon.RelayInfo{TaskRelayInfo: &relaycommon.TaskRelayInfo{}}
|
||||
require.Nil(t, tc.apply(c, info))
|
||||
if tc.tokenPin {
|
||||
assert.Nil(t, info.LockedChannel)
|
||||
} else {
|
||||
locked, ok := info.LockedChannel.(*model.Channel)
|
||||
require.True(t, ok)
|
||||
require.NotNil(t, locked)
|
||||
assert.Equal(t, channel.Id, locked.Id)
|
||||
}
|
||||
require.Len(t, info.OriginTasks, 1)
|
||||
assert.Equal(t, "task-lock", info.OriginTasks[0].TaskID)
|
||||
assert.Equal(t, "upstream-task-lock", info.OriginTasks[0].UpstreamTaskID)
|
||||
assert.Equal(t, "text_to_video", info.OriginTasks[0].Action)
|
||||
assert.Equal(t, string(model.TaskStatusSuccess), info.OriginTasks[0].Status)
|
||||
assert.Equal(t, []byte(originTask.Data), info.OriginTasks[0].Data)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestTaskModel2DtoNormalizesLegacyAction(t *testing.T) {
|
||||
task := &model.Task{Action: "firstTailGenerate"}
|
||||
|
||||
@@ -153,6 +227,7 @@ func TestRelayTaskSubmitAliasBillingIdentityAndExprFallback(t *testing.T) {
|
||||
const mapping = `{"alias-model":"declared-model"}`
|
||||
const aliasExpr = `tier("alias", 2)`
|
||||
const tailExpr = `tier("tail", 3)`
|
||||
const legacyExpr = `tier("legacy", u("old_units") * 2)`
|
||||
|
||||
tests := []struct {
|
||||
name string
|
||||
@@ -160,6 +235,7 @@ func TestRelayTaskSubmitAliasBillingIdentityAndExprFallback(t *testing.T) {
|
||||
exprs map[string]string
|
||||
wantTiered bool
|
||||
wantExpr string
|
||||
source string
|
||||
}{
|
||||
{
|
||||
name: "alias own tiered wins",
|
||||
@@ -175,6 +251,16 @@ func TestRelayTaskSubmitAliasBillingIdentityAndExprFallback(t *testing.T) {
|
||||
wantTiered: true,
|
||||
wantExpr: tailExpr,
|
||||
},
|
||||
{
|
||||
name: "stored expression keeps running after schema narrows",
|
||||
modes: map[string]string{"declared-model": "tiered_expr"},
|
||||
exprs: map[string]string{"declared-model": legacyExpr},
|
||||
wantTiered: true,
|
||||
wantExpr: legacyExpr,
|
||||
source: strings.Replace(billingFallbackPlugin, `fetchMode:"per_task"`, `fetchMode:"per_task", usageSchema:{old_units:{type:"number",unit:"count"}}, usageProfiles:[{models:["declared-model"],schema:{seconds:{type:"number",unit:"second"}}}]`, 1) + `
|
||||
export function extractUsage(){return {old_units:2};}
|
||||
`,
|
||||
},
|
||||
{
|
||||
name: "neither tiered uses ordinary pricing",
|
||||
wantTiered: false,
|
||||
@@ -204,7 +290,11 @@ func TestRelayTaskSubmitAliasBillingIdentityAndExprFallback(t *testing.T) {
|
||||
c.Set("group", "default")
|
||||
info.UserGroup = "default"
|
||||
info.UsingGroup = "default"
|
||||
pinMappingOrderPlugin(t, c, billingFallbackPlugin)
|
||||
source := testCase.source
|
||||
if source == "" {
|
||||
source = billingFallbackPlugin
|
||||
}
|
||||
pinMappingOrderPlugin(t, c, source)
|
||||
info.OriginModelName = "alias-model"
|
||||
|
||||
_, taskErr := RelayTaskSubmit(c, info)
|
||||
@@ -218,11 +308,14 @@ func TestRelayTaskSubmitAliasBillingIdentityAndExprFallback(t *testing.T) {
|
||||
assert.Equal(t, "declared-model", task.Properties.UpstreamModelName)
|
||||
|
||||
if testCase.wantTiered {
|
||||
require.NotNil(t, info.TieredBillingSnapshot)
|
||||
require.NotNil(t, info.TieredBillingSnapshot, "submission error: %+v", taskErr)
|
||||
assert.Equal(t, "alias-model", info.TieredBillingSnapshot.ModelName)
|
||||
assert.Equal(t, testCase.wantExpr, info.TieredBillingSnapshot.ExprString)
|
||||
assert.Equal(t, billingexpr.ExprHashString(testCase.wantExpr), info.TieredBillingSnapshot.ExprHash)
|
||||
assert.NotEqual(t, "model_price_error", taskErr.Code)
|
||||
if testCase.wantExpr == legacyExpr {
|
||||
assert.Equal(t, 4*common.QuotaPerUnit, info.TieredBillingSnapshot.EstimatedQuotaBeforeGroup)
|
||||
}
|
||||
} else {
|
||||
assert.Nil(t, info.TieredBillingSnapshot)
|
||||
assert.Equal(t, "model_price_error", taskErr.Code)
|
||||
@@ -230,3 +323,79 @@ func TestRelayTaskSubmitAliasBillingIdentityAndExprFallback(t *testing.T) {
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestSharedTaskBillingExpressionSelectionAndFrozenSettlement(t *testing.T) {
|
||||
const baseExpr = `tier("base", u("seconds") * 2)`
|
||||
const alphaExpr = `tier("alpha", u("seconds") * 3)`
|
||||
const betaExpr = `tier("beta", u("credits") * 5)`
|
||||
const aliasExpr = `tier("alias", u("credits") * 7)`
|
||||
for _, tc := range []struct {
|
||||
name, plugin, model, mapping, modelExpr, mode, wantExpr string
|
||||
variants map[string]string
|
||||
wantPriceError bool
|
||||
}{
|
||||
{name: "executing plugin override", plugin: "billing-beta", model: "declared-model", modelExpr: baseExpr, mode: "tiered_expr", variants: map[string]string{"billing-alpha::declared-model": alphaExpr, "billing-beta::declared-model": betaExpr}, wantExpr: betaExpr},
|
||||
{name: "override ignores model mode", plugin: "billing-beta", model: "declared-model", modelExpr: baseExpr, mode: "ratio", variants: map[string]string{"billing-beta::declared-model": betaExpr}, wantExpr: betaExpr},
|
||||
{name: "model expression fallback", plugin: "billing-alpha", model: "declared-model", modelExpr: baseExpr, mode: "tiered_expr", wantExpr: baseExpr},
|
||||
{name: "alias override precedes mapped override", plugin: "billing-beta", model: "alias-model", mapping: `{"alias-model":"declared-model"}`, modelExpr: baseExpr, mode: "tiered_expr", variants: map[string]string{"billing-beta::declared-model": betaExpr, "billing-beta::alias-model": aliasExpr}, wantExpr: aliasExpr},
|
||||
{name: "mapped override precedes model fallback", plugin: "billing-beta", model: "alias-model", mapping: `{"alias-model":"declared-model"}`, modelExpr: baseExpr, mode: "tiered_expr", variants: map[string]string{"billing-beta::declared-model": betaExpr}, wantExpr: betaExpr},
|
||||
{name: "unconfigured plugin cannot use another schema", plugin: "billing-beta", model: "declared-model", modelExpr: baseExpr, mode: "tiered_expr", wantPriceError: true},
|
||||
{name: "missing usage in skipped branch remains incompatible", plugin: "billing-beta", model: "declared-model", modelExpr: `true ? tier("free", 0) : tier("missing", u("seconds"))`, mode: "tiered_expr", wantPriceError: true},
|
||||
{name: "fixed pricing is still rejected", plugin: "billing-beta", model: "declared-model", variants: map[string]string{"billing-beta::declared-model": `tier("fixed", fixed(1))`}, wantPriceError: true},
|
||||
} {
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
saveBillingConfig(t)
|
||||
registry := pluginruntime.NewRegistry()
|
||||
for _, spec := range []struct{ key, field, unit string }{{"billing-alpha", "seconds", "second"}, {"billing-beta", "credits", "credit"}} {
|
||||
source := strings.ReplaceAll(billingFallbackPlugin, "bill-fallback", spec.key)
|
||||
source = strings.Replace(source, `fetchMode:"per_task"`, `fetchMode:"per_task",usageSchema:{`+spec.field+`:{type:"number",unit:"`+spec.unit+`"}}`, 1)
|
||||
source += `export function extractUsage(){return {` + spec.field + `:2};}`
|
||||
_, err := registry.Register(source, pluginruntime.Options{})
|
||||
require.NoError(t, err)
|
||||
}
|
||||
variants := tc.variants
|
||||
if variants == nil {
|
||||
variants = map[string]string{}
|
||||
}
|
||||
rawVariants, err := common.Marshal(variants)
|
||||
require.NoError(t, err)
|
||||
modes, err := common.Marshal(map[string]string{"declared-model": tc.mode})
|
||||
require.NoError(t, err)
|
||||
expressions, err := common.Marshal(map[string]string{"declared-model": tc.modelExpr})
|
||||
require.NoError(t, err)
|
||||
require.NoError(t, config.GlobalConfig.LoadFromDB(map[string]string{
|
||||
billing_setting.PluginBillingExprOption: string(rawVariants), "billing_setting.billing_mode": string(modes), "billing_setting.billing_expr": string(expressions),
|
||||
}))
|
||||
c, info := newTaskSubmitContext(t, tc.model, tc.mapping)
|
||||
c.Set("group", "default")
|
||||
c.Set("task_plugin_key", tc.plugin)
|
||||
info.UserGroup = "default"
|
||||
info.UsingGroup = "default"
|
||||
info.OriginModelName = tc.model
|
||||
plugin, ok := registry.Generation().Get(tc.plugin)
|
||||
require.True(t, ok)
|
||||
c.Set(pluginruntime.ContextKeyPinnedPlugin, pluginruntime.PinnedPlugin{Generation: registry.Generation(), Plugin: plugin})
|
||||
_, taskErr := RelayTaskSubmit(c, info)
|
||||
require.NotNil(t, taskErr) // This fixture stops at reservation, before upstream submission.
|
||||
if tc.wantPriceError {
|
||||
assert.Equal(t, "model_price_error", taskErr.Code)
|
||||
assert.Nil(t, info.TieredBillingSnapshot)
|
||||
return
|
||||
}
|
||||
require.NotNil(t, info.TieredBillingSnapshot, "submission error: %+v", taskErr)
|
||||
assert.Equal(t, tc.wantExpr, info.TieredBillingSnapshot.ExprString)
|
||||
assert.NotEqual(t, "model_price_error", taskErr.Code)
|
||||
require.NoError(t, config.GlobalConfig.LoadFromDB(map[string]string{billing_setting.PluginBillingExprOption: `{}`, "billing_setting.billing_expr": `{}`}))
|
||||
field := "seconds"
|
||||
if tc.plugin == "billing-beta" {
|
||||
field = "credits"
|
||||
}
|
||||
result, usage, err := service.EvaluateTaskCompletionUsage(info.TieredBillingSnapshot, map[string]any{field: float64(4)})
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, float64(4), usage[field])
|
||||
assert.Equal(t, float64(2), info.TieredBillingSnapshot.UsageFacts[field])
|
||||
assert.Equal(t, 2*info.TieredBillingSnapshot.EstimatedQuotaAfterGroup, result.ActualQuotaAfterGroup)
|
||||
assert.Equal(t, tc.wantExpr, info.TieredBillingSnapshot.ExprString)
|
||||
})
|
||||
}
|
||||
}
|
||||
Reference in new issue
Block a user