feat(task): resolve channel-mapped aliases and case variants for plugin models

Channel model_mapping keys exposed in a channel's model list now act as
first-class aliases for task-plugin models across the whole line:

- Derived alias view (model/task_model_alias.go): built from enabled
  channels' model_mapping, chain-following with cycle detection, declared
  names always win, cross-plugin conflicts dropped. Rebuilt on channel
  cache refresh, registry generation change, and a 60s TTL.
- Request path: PinTaskPluginEndpoint resolves declared-name case folds
  and mapping aliases before endpoint lookup (never rewriting the body
  until the endpoint is claimed), pins with MappedModel, and the decode
  contract accepts alias echoes without loosening model ownership for
  normal pins. Legacy /v1/tasks submit folds case variants the same way.
  Fixes aliases on POST /v1/responses silently falling through to the
  main relay against task channels.
- Mapping order: ModelMappedHelper now runs before the plugin submit
  hook builds and caches the upstream body, so channel model_mapping
  actually reaches the upstream request. Plugins receive the mapped
  name as ctx.upstreamModel in both decode and submit contexts.
- Billing: identity stays the origin name; when the alias has no tiered
  expression, the selected channel's mapping tail expression applies.
  Pricing page and billing-expr smoke tests resolve aliases to the
  owning plugin's usage schema.
- Case folding: ASCII-only fold with exact-match priority; same-plugin
  and cross-plugin fold collisions rejected at registration.
- Plugins: model-keyed rate tables, req_key derivation, and combo
  validation in doubao/kling/jimeng/hailuo/vidu/sunoapi now key on
  ctx.upstreamModel || ctx.model; render/echo paths keep ctx.model.
This commit is contained in:
CaIon
2026-08-30 19:13:51 +08:00
parent 918427d8ab
commit 6c22550ea3
23 changed files with 968 additions and 59 deletions
+103
View File
@@ -7,10 +7,14 @@ import (
"testing" "testing"
"github.com/QuantumNous/new-api/common" "github.com/QuantumNous/new-api/common"
"github.com/QuantumNous/new-api/model"
"github.com/QuantumNous/new-api/pkg/jsplugin" "github.com/QuantumNous/new-api/pkg/jsplugin"
"github.com/QuantumNous/new-api/setting/config"
"github.com/gin-gonic/gin" "github.com/gin-gonic/gin"
"github.com/glebarez/sqlite"
"github.com/stretchr/testify/assert" "github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require" "github.com/stretchr/testify/require"
"gorm.io/gorm"
) )
func TestUpdateOptionRejectsInvalidTaskBillingExpressions(t *testing.T) { func TestUpdateOptionRejectsInvalidTaskBillingExpressions(t *testing.T) {
@@ -98,3 +102,102 @@ func TestUpdateOptionRejectsUsageExpressionWithoutTaskPlugin(t *testing.T) {
assert.Contains(t, recorder.Body.String(), "mode") assert.Contains(t, recorder.Body.String(), "mode")
assert.Contains(t, recorder.Body.String(), "no task plugin usage schema") assert.Contains(t, recorder.Body.String(), "no task plugin usage schema")
} }
func setupBillingAliasOptionDB(t *testing.T) {
t.Helper()
previousDB := model.DB
previousLogDB := model.LOG_DB
previousType := common.MainDatabaseType()
previousCache := common.MemoryCacheEnabled
previousMap := common.OptionMap
previousRedis := common.RedisEnabled
database, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{})
require.NoError(t, err)
require.NoError(t, database.AutoMigrate(&model.Channel{}, &model.Option{}, &model.Log{}, &model.User{}))
model.DB = database
model.LOG_DB = database
common.SetMainDatabaseType(common.DatabaseTypeSQLite)
common.MemoryCacheEnabled = false
common.RedisEnabled = false
common.OptionMap = map[string]string{}
t.Cleanup(func() {
model.DB = previousDB
model.LOG_DB = previousLogDB
common.SetMainDatabaseType(previousType)
common.MemoryCacheEnabled = previousCache
common.OptionMap = previousMap
common.RedisEnabled = previousRedis
model.InitChannelCache()
})
}
func TestUpdateOptionAliasBillingExprUsesPluginSchema(t *testing.T) {
setupBillingAliasOptionDB(t)
const pluginKey = "billing-alias-probe"
source := `
export const meta = {
apiVersion: 1, key: "billing-alias-probe", name: "Billing Alias Probe", version: "1.0.0", author: {name: "Test"},
models: ["declared-model"], fetchMode: "per_task",
usageSchema: {seconds: {type: "number", unit: "second"}}
};
export function buildSubmitRequest() { return {}; }
export function parseSubmitResponse() { return {}; }
export function buildQueryRequest() { return {}; }
export function parseTaskResult() { return {}; }
`
_, err := jsplugin.DefaultRegistry.Register(source, jsplugin.Options{})
require.NoError(t, err)
t.Cleanup(func() { jsplugin.DefaultRegistry.Unregister(pluginKey) })
mapping := `{"alias-model":"declared-model"}`
require.NoError(t, model.DB.Create(&model.Channel{
Id: 1,
Type: 54,
Key: "key-1",
Status: common.ChannelStatusEnabled,
Name: "ch-1",
Group: "default",
Models: "alias-model,declared-model",
ModelMapping: &mapping,
}).Error)
model.InitChannelCache()
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))
})
putExpr := func(modelName, expression string) *httptest.ResponseRecorder {
t.Helper()
expressions, marshalErr := common.Marshal(map[string]string{modelName: expression})
require.NoError(t, marshalErr)
body, marshalErr := common.Marshal(OptionUpdateRequest{
Key: "billing_setting.billing_expr",
Value: string(expressions),
})
require.NoError(t, marshalErr)
recorder := httptest.NewRecorder()
context, _ := gin.CreateTestContext(recorder)
context.Request = httptest.NewRequest(http.MethodPut, "/api/option/", strings.NewReader(string(body)))
UpdateOption(context)
return recorder
}
accepted := putExpr("alias-model", `u("seconds")`)
assert.Equal(t, http.StatusOK, accepted.Code)
assert.Contains(t, accepted.Body.String(), `"success":true`)
rejectedKey := putExpr("alias-model", `u("clips")`)
assert.Equal(t, http.StatusOK, rejectedKey.Code)
assert.Contains(t, rejectedKey.Body.String(), `"success":false`)
assert.Contains(t, rejectedKey.Body.String(), `usage key \"clips\" is not declared`)
unresolvable := putExpr("unknown-alias-model", `u("seconds")`)
assert.Equal(t, http.StatusOK, unresolvable.Code)
assert.Contains(t, unresolvable.Body.String(), `"success":false`)
assert.Contains(t, unresolvable.Body.String(), "no task plugin usage schema")
}
+6
View File
@@ -352,6 +352,12 @@ func UpdateOption(c *gin.Context) {
expression := expressions[modelName] expression := expressions[modelName]
if plugin, ok := generation.GetByModel(modelName); ok { if plugin, ok := generation.GetByModel(modelName); ok {
err = billing_setting.SmokeTestTaskExpr(expression, plugin.Meta.UsageSchema) err = billing_setting.SmokeTestTaskExpr(expression, plugin.Meta.UsageSchema)
} else if target, resolved := model.ResolveTaskModelAlias(generation, modelName); resolved {
if plugin, ok := generation.Get(target.PluginKey); ok {
err = billing_setting.SmokeTestTaskExpr(expression, plugin.Meta.UsageSchema)
} else {
err = billing_setting.SmokeTestExpr(expression)
}
} else { } else {
err = billing_setting.SmokeTestExpr(expression) err = billing_setting.SmokeTestExpr(expression)
} }
+26
View File
@@ -1088,6 +1088,7 @@ func TestRetrieveTaskPluginResponsePendingSkipsRenderFinal(t *testing.T) {
require.NoError(t, common.Unmarshal(recorder.Body.Bytes(), &response)) require.NoError(t, common.Unmarshal(recorder.Body.Bytes(), &response))
assert.Equal(t, "resp_retrieve_pending", response["id"]) assert.Equal(t, "resp_retrieve_pending", response["id"])
assert.Equal(t, "in_progress", response["status"]) assert.Equal(t, "in_progress", response["status"])
assert.Equal(t, "video-model", response["model"])
assert.Equal(t, true, response["background"]) assert.Equal(t, true, response["background"])
assert.Nil(t, response["completed_at"]) assert.Nil(t, response["completed_at"])
assert.Empty(t, response["output"]) assert.Empty(t, response["output"])
@@ -1096,6 +1097,31 @@ func TestRetrieveTaskPluginResponsePendingSkipsRenderFinal(t *testing.T) {
assert.Equal(t, "/v1/responses/resp_retrieve_pending", metadata["retrieval_path"]) assert.Equal(t, "/v1/responses/resp_retrieve_pending", metadata["retrieval_path"])
} }
func TestRetrieveTaskPluginResponseEchoesOriginModelName(t *testing.T) {
pinned := compilePluginProtocolRetrieveEndpoint(t, "retrieve-alias-echo", `
export const protocols = {openai_responses: {
renderEvents: function() { throw new Error("pending retrieve called renderEvents"); },
renderFinal: function() { throw new Error("pending retrieve called renderFinal"); }
}};
`, pluginruntime.Options{})
c, recorder := newPluginProtocolRetrieveContext("resp_retrieve_alias")
deps := pluginProtocolRetrieveDeps(pinned, &model.Task{
TaskID: "task_retrieve_alias",
Platform: constant.TaskPlatform(pinned.Plugin.Meta.Key),
UserId: 71,
Status: model.TaskStatusInProgress,
Properties: model.Properties{OriginModelName: "alias-model"},
CreatedAt: 1_710_000_000,
}, true, nil)
retrieveTaskPluginResponse(c, deps)
assert.Equal(t, http.StatusOK, recorder.Code)
var response map[string]any
require.NoError(t, common.Unmarshal(recorder.Body.Bytes(), &response))
assert.Equal(t, "alias-model", response["model"])
}
func TestRetrieveTaskPluginResponseSuccessRendersFinal(t *testing.T) { func TestRetrieveTaskPluginResponseSuccessRendersFinal(t *testing.T) {
logs := make([]string, 0, 1) logs := make([]string, 0, 1)
pinned := compilePluginProtocolRetrieveEndpoint(t, "retrieve-success", ` pinned := compilePluginProtocolRetrieveEndpoint(t, "retrieve-success", `
+2
View File
@@ -138,3 +138,5 @@ Protocol media uses host-injected `ctx.artifacts[key].url`. Provider URLs from `
## Persisted data and driver hooks ## Persisted data and driver hooks
The persisted field remains `task.data`; there is no `task.raw` alias. Driver hooks (`buildSubmitRequest`, `parseSubmitResponse`, query/result, usage, artifact, and content hooks) stay flat and must not branch on the client path or protocol. The persisted field remains `task.data`; there is no `task.raw` alias. Driver hooks (`buildSubmitRequest`, `parseSubmitResponse`, query/result, usage, artifact, and content hooks) stay flat and must not branch on the client path or protocol.
`ctx.model` is the billing and display identity (the origin name the client sent, including a channel-mapping alias). `ctx.upstreamModel` is the machine identity after channel `model_mapping`. Rate tables and model-keyed usage facts must use `ctx.upstreamModel || ctx.model`. Decode and render hooks that echo the client model must keep `ctx.model`. `buildSubmitRequest` must not set descriptor top-level `model` on a mapped pin; the host requires the plugin to echo the alias verbatim.
+69 -6
View File
@@ -332,18 +332,50 @@ func PinTaskPluginEndpoint() gin.HandlerFunc {
c.Next() c.Next()
return return
} }
c.Set(contextKeyTaskPluginEndpointModel, *modelRequest)
claimedModel := modelRequest.Model claimedModel := modelRequest.Model
if strings.TrimSpace(claimedModel) == "" { if strings.TrimSpace(claimedModel) == "" {
c.Set(contextKeyTaskPluginEndpointModel, *modelRequest)
c.Next() c.Next()
return return
} }
binding, found := generation.LookupEndpoint(c.Request.Method, c.Request.URL.Path, claimedModel) lookupModel := claimedModel
pinModel := claimedModel
mappedModel := ""
rewriteTo := ""
if declared, ok := generation.CanonicalModel(claimedModel); ok {
lookupModel = declared
pinModel = declared
if claimedModel != declared {
rewriteTo = declared
}
} else if target, ok := model.ResolveTaskModelAlias(generation, claimedModel); ok {
if target.Declared == "" {
c.Set(contextKeyTaskPluginEndpointModel, *modelRequest)
c.Next()
return
}
lookupModel = target.Declared
pinModel = target.Alias
mappedModel = target.Declared
if claimedModel != target.Alias {
rewriteTo = target.Alias
}
}
binding, found := generation.LookupEndpoint(c.Request.Method, c.Request.URL.Path, lookupModel)
if !found || binding.Plugin == nil { if !found || binding.Plugin == nil {
c.Set(contextKeyTaskPluginEndpointModel, *modelRequest)
c.Next() c.Next()
return return
} }
candidates := generation.LookupEndpointCandidates(c.Request.Method, c.Request.URL.Path, claimedModel) if rewriteTo != "" {
if rewriteErr := rewriteTaskPluginJSONModel(c, rewriteTo); rewriteErr != nil {
abortWithOpenAiMessage(c, http.StatusBadRequest, "Invalid task protocol request")
return
}
}
modelRequest.Model = pinModel
c.Set(contextKeyTaskPluginEndpointModel, *modelRequest)
candidates := generation.LookupEndpointCandidates(c.Request.Method, c.Request.URL.Path, lookupModel)
if len(candidates) == 0 { if len(candidates) == 0 {
candidates = []pluginruntime.ProtocolBinding{binding} candidates = []pluginruntime.ProtocolBinding{binding}
} }
@@ -390,7 +422,8 @@ func PinTaskPluginEndpoint() gin.HandlerFunc {
Plugin: binding.Plugin, Plugin: binding.Plugin,
Protocol: binding.Protocol, Protocol: binding.Protocol,
Operation: binding.Operation, Operation: binding.Operation,
Model: claimedModel, Model: pinModel,
MappedModel: mappedModel,
Candidates: candidates, Candidates: candidates,
} }
c.Set(pluginruntime.ContextKeyPinnedPlugin, pin) c.Set(pluginruntime.ContextKeyPinnedPlugin, pin)
@@ -403,7 +436,7 @@ func PinTaskPluginEndpoint() gin.HandlerFunc {
binding.Plugin.Meta.Version, binding.Plugin.Meta.Version,
binding.Operation.Methods[0], binding.Operation.Methods[0],
binding.Protocol, binding.Protocol,
claimedModel, pinModel,
) )
c.Next() c.Next()
} }
@@ -515,6 +548,13 @@ func PrepareTaskPluginEndpoint() gin.HandlerFunc {
abortWithOpenAiMessage(c, status, err.Error()) abortWithOpenAiMessage(c, status, err.Error())
return return
} }
if body, ok := requestContext.Body.(map[string]any); ok {
if fields, ok := body["fields"].(map[string][]string); ok {
if values := fields["model"]; len(values) > 0 && values[0] != pinned.Model {
fields["model"][0] = pinned.Model
}
}
}
bodyObject, _ := requestContext.Body.(map[string]any) bodyObject, _ := requestContext.Body.(map[string]any)
bodyKind, _ := bodyObject["kind"].(string) bodyKind, _ := bodyObject["kind"].(string)
allowedBody := false allowedBody := false
@@ -561,6 +601,7 @@ func PrepareTaskPluginEndpoint() gin.HandlerFunc {
Protocol: pinned.Protocol, Protocol: pinned.Protocol,
Operation: pinned.Operation.Name, Operation: pinned.Operation.Name,
Model: pinned.Model, Model: pinned.Model,
UpstreamModel: pinned.MappedModel,
Stream: stream, Stream: stream,
} }
c.Set(pluginruntime.ContextKeyProtocolRequest, protocolContext) c.Set(pluginruntime.ContextKeyProtocolRequest, protocolContext)
@@ -624,7 +665,8 @@ func PrepareTaskPluginEndpoint() gin.HandlerFunc {
return return
} }
modelOwned := slices.Contains(pinned.Plugin.Meta.Models, resolvedModel) modelOwned := slices.Contains(pinned.Plugin.Meta.Models, resolvedModel)
if !modelOwned || resolvedModel != pinned.Model { mappedPin := pinned.MappedModel != ""
if resolvedModel != pinned.Model || (!modelOwned && !mappedPin) {
logger.LogWarn( logger.LogWarn(
c, c,
"task_plugin subsystem=endpoint event=prepare_rejected generation=%d plugin=%q stage=parse_request reason=resolved_model_not_owned claimed_model=%q resolved_model=%q", "task_plugin subsystem=endpoint event=prepare_rejected generation=%d plugin=%q stage=parse_request reason=resolved_model_not_owned claimed_model=%q resolved_model=%q",
@@ -1380,6 +1422,27 @@ func PrepareTaskPluginSubmit() gin.HandlerFunc {
c.AbortWithStatusJSON(http.StatusBadRequest, gin.H{"error": gin.H{"message": "model is required", "type": "invalid_request_error"}}) c.AbortWithStatusJSON(http.StatusBadRequest, gin.H{"error": gin.H{"message": "model is required", "type": "invalid_request_error"}})
return return
} }
exactOwned := slices.Contains(plugin.Meta.Models, modelName)
exactAlias := false
if target, resolved := model.ResolveTaskModelAlias(generation, modelName); resolved && target.Alias == modelName && target.PluginKey == plugin.Meta.Key {
exactAlias = true
}
if !exactOwned && !exactAlias {
folded := ""
if declared, ok := generation.CanonicalModel(modelName); ok && slices.Contains(plugin.Meta.Models, declared) && declared != modelName {
folded = declared
} else if target, resolved := model.ResolveTaskModelAlias(generation, modelName); resolved && target.PluginKey == plugin.Meta.Key && target.Alias != "" && target.Alias != modelName {
folded = target.Alias
}
if folded != "" {
if rewriteErr := rewriteTaskPluginJSONModel(c, folded); rewriteErr != nil {
c.AbortWithStatusJSON(http.StatusBadRequest, gin.H{"error": gin.H{"message": rewriteErr.Error(), "type": "invalid_request_error"}})
return
}
requestBody["model"] = folded
modelName = folded
}
}
c.Set("task_request", requestBody) c.Set("task_request", requestBody)
c.Set("resolved_task_model", modelName) c.Set("resolved_task_model", modelName)
c.Set("expected_task_plugin_key", pluginKey) c.Set("expected_task_plugin_key", pluginKey)
+45
View File
@@ -0,0 +1,45 @@
package middleware
import (
"io"
"mime"
"strings"
"github.com/QuantumNous/new-api/common"
"github.com/gin-gonic/gin"
"github.com/tidwall/sjson"
)
// rewriteTaskPluginJSONModel patches the top-level JSON "model" field and
// replaces BodyStorage. Non-JSON bodies are left untouched.
func rewriteTaskPluginJSONModel(c *gin.Context, spelling string) error {
mediaType, _, err := mime.ParseMediaType(c.GetHeader("Content-Type"))
if err != nil {
return nil
}
if mediaType != "application/json" && !strings.HasSuffix(mediaType, "+json") {
return nil
}
storage, err := common.GetBodyStorage(c)
if err != nil {
return err
}
raw, err := storage.Bytes()
if err != nil {
return err
}
patched, err := sjson.SetBytes(raw, "model", spelling)
if err != nil {
return err
}
newStorage, err := common.CreateBodyStorage(patched)
if err != nil {
return err
}
_ = storage.Close()
c.Set(common.KeyBodyStorage, newStorage)
c.Set(common.KeyRequestBody, nil)
c.Request.Body = io.NopCloser(newStorage)
c.Request.ContentLength = int64(len(patched))
return nil
}
+5 -3
View File
@@ -630,9 +630,11 @@ func TestTaskPluginEndpointMissPreservesOrdinaryRequestBody(t *testing.T) {
func(c *gin.Context) { func(c *gin.Context) {
_, pinned := c.Get(jsplugin.ContextKeyPinnedEndpoint) _, pinned := c.Get(jsplugin.ContextKeyPinnedEndpoint)
assert.False(t, pinned) assert.False(t, pinned)
var body map[string]any storage, storageErr := common.GetBodyStorage(c)
require.NoError(t, common.UnmarshalBodyReusable(c, &body)) require.NoError(t, storageErr)
assert.Equal(t, "ordinary-model", body["model"]) raw, bytesErr := storage.Bytes()
require.NoError(t, bytesErr)
assert.Equal(t, []byte(`{"model":"ordinary-model","input":"hello"}`), raw)
c.Status(http.StatusNoContent) c.Status(http.StatusNoContent)
}, },
) )
+3 -1
View File
@@ -11,8 +11,8 @@ import (
"github.com/QuantumNous/new-api/common" "github.com/QuantumNous/new-api/common"
"github.com/QuantumNous/new-api/constant" "github.com/QuantumNous/new-api/constant"
"github.com/QuantumNous/new-api/logger"
"github.com/QuantumNous/new-api/dto" "github.com/QuantumNous/new-api/dto"
"github.com/QuantumNous/new-api/logger"
kitdto "github.com/QuantumNous/new-api/relaykit/dto" kitdto "github.com/QuantumNous/new-api/relaykit/dto"
"github.com/QuantumNous/new-api/setting/ratio_setting" "github.com/QuantumNous/new-api/setting/ratio_setting"
) )
@@ -27,6 +27,7 @@ var channelSyncLock sync.RWMutex
func InitChannelCache() { func InitChannelCache() {
if !common.MemoryCacheEnabled { if !common.MemoryCacheEnabled {
InvalidatePricingCache() InvalidatePricingCache()
rebuildTaskAliasView()
return return
} }
newChannelId2channel := make(map[int]*Channel) newChannelId2channel := make(map[int]*Channel)
@@ -101,6 +102,7 @@ func InitChannelCache() {
// loadPricingAdvancedCustomConfigs. channelSyncLock MUST be released before // loadPricingAdvancedCustomConfigs. channelSyncLock MUST be released before
// invalidating the pricing cache, otherwise the reversed order deadlocks. // invalidating the pricing cache, otherwise the reversed order deadlocks.
InvalidatePricingCache() InvalidatePricingCache()
rebuildTaskAliasView()
common.SysLog("channels synced from database") common.SysLog("channels synced from database")
} }
+14 -1
View File
@@ -410,8 +410,21 @@ func updatePricing() {
pricing.BillingMode = billingMode pricing.BillingMode = billingMode
pricing.BillingExpr = expr pricing.BillingExpr = expr
} }
} else if target, resolved := ResolveTaskModelAlias(pluginGeneration, model); resolved && target.Declared != "" {
if tailMode := billing_setting.GetBillingMode(target.Declared); tailMode == "tiered_expr" {
if expr, ok := billing_setting.GetBillingExpr(target.Declared); ok && strings.TrimSpace(expr) != "" {
pricing.BillingMode = tailMode
pricing.BillingExpr = expr
} }
if plugin, ok := pluginGeneration.GetByModel(model); ok && len(plugin.Meta.UsageSchema) > 0 { }
}
plugin, ok := pluginGeneration.GetByModel(model)
if !ok {
if target, resolved := ResolveTaskModelAlias(pluginGeneration, model); resolved {
plugin, ok = pluginGeneration.Get(target.PluginKey)
}
}
if ok && plugin != nil && len(plugin.Meta.UsageSchema) > 0 {
pricing.BillingUsageSchema = make(map[string]jsplugin.UsageFieldSchema, len(plugin.Meta.UsageSchema)) pricing.BillingUsageSchema = make(map[string]jsplugin.UsageFieldSchema, len(plugin.Meta.UsageSchema))
for key, field := range plugin.Meta.UsageSchema { for key, field := range plugin.Meta.UsageSchema {
field.Enum = append([]string(nil), field.Enum...) field.Enum = append([]string(nil), field.Enum...)
+70
View File
@@ -8,6 +8,7 @@ import (
"github.com/QuantumNous/new-api/constant" "github.com/QuantumNous/new-api/constant"
"github.com/QuantumNous/new-api/pkg/jsplugin" "github.com/QuantumNous/new-api/pkg/jsplugin"
"github.com/QuantumNous/new-api/relaykit/dto" "github.com/QuantumNous/new-api/relaykit/dto"
"github.com/QuantumNous/new-api/setting/config"
"github.com/stretchr/testify/assert" "github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require" "github.com/stretchr/testify/require"
) )
@@ -60,6 +61,75 @@ func TestPricingCarriesTaskUsageSchemaAndRefreshesWithPluginGeneration(t *testin
assert.Equal(t, "count", refreshedPricing["pricing-usage-model"].BillingUsageSchema["clips"].Unit) assert.Equal(t, "count", refreshedPricing["pricing-usage-model"].BillingUsageSchema["clips"].Unit)
} }
func TestPricingAliasCarriesPluginUsageSchemaAndTailExpr(t *testing.T) {
resetPricingEndpointTestTables(t)
const pluginKey = "pricing-usage-probe"
source := pricingUsagePluginSource("1.0.0", `{
seconds: {type: "number", unit: "second", description: "Estimated duration."}
}`)
_, err := jsplugin.DefaultRegistry.Register(source, jsplugin.Options{})
require.NoError(t, err)
t.Cleanup(func() { jsplugin.DefaultRegistry.Unregister(pluginKey) })
mapping := `{"alias-model":"pricing-usage-model"}`
channel := &Channel{
Id: 910,
Type: constant.ChannelTypeTaskPlugin,
Key: "key-910",
Status: 1,
Name: "channel-910",
Models: "alias-model,pricing-usage-model",
ModelMapping: &mapping,
}
require.NoError(t, DB.Create(channel).Error)
insertPricingEndpointAbility(t, 910, "alias-model")
insertPricingEndpointAbility(t, 910, "pricing-usage-model")
InitChannelCache()
saved := map[string]string{}
require.NoError(t, config.GlobalConfig.SaveToDB(func(key, value string) error {
saved[key] = value
return nil
}))
t.Cleanup(func() {
require.NoError(t, config.GlobalConfig.LoadFromDB(saved))
})
require.NoError(t, config.GlobalConfig.LoadFromDB(map[string]string{
"billing_setting.billing_mode": `{"pricing-usage-model":"tiered_expr","alias-own-expr":"tiered_expr"}`,
"billing_setting.billing_expr": `{"pricing-usage-model":"u(\"seconds\")","alias-own-expr":"u(\"seconds\") * 2"}`,
}))
InvalidatePricingCache()
pricing := pricingByModel(GetPricing())
require.Contains(t, pricing, "alias-model")
require.Contains(t, pricing, "pricing-usage-model")
assert.Equal(t, "second", pricing["alias-model"].BillingUsageSchema["seconds"].Unit)
assert.Equal(t, "Estimated duration.", pricing["alias-model"].BillingUsageSchema["seconds"].Description["en"])
assert.Equal(t, "tiered_expr", pricing["alias-model"].BillingMode)
assert.Equal(t, `u("seconds")`, pricing["alias-model"].BillingExpr)
assert.Equal(t, "tiered_expr", pricing["pricing-usage-model"].BillingMode)
assert.Equal(t, `u("seconds")`, pricing["pricing-usage-model"].BillingExpr)
ownMapping := `{"alias-own-expr":"pricing-usage-model"}`
own := &Channel{
Id: 911,
Type: constant.ChannelTypeTaskPlugin,
Key: "key-911",
Status: 1,
Name: "channel-911",
Models: "alias-own-expr,pricing-usage-model",
ModelMapping: &ownMapping,
}
require.NoError(t, DB.Create(own).Error)
insertPricingEndpointAbility(t, 911, "alias-own-expr")
InitChannelCache()
InvalidatePricingCache()
refreshed := pricingByModel(GetPricing())
assert.Equal(t, `u("seconds") * 2`, refreshed["alias-own-expr"].BillingExpr)
assert.Equal(t, "second", refreshed["alias-own-expr"].BillingUsageSchema["seconds"].Unit)
}
func pricingByModel(pricings []Pricing) map[string]Pricing { func pricingByModel(pricings []Pricing) map[string]Pricing {
result := make(map[string]Pricing, len(pricings)) result := make(map[string]Pricing, len(pricings))
for _, pricing := range pricings { for _, pricing := range pricings {
+212
View File
@@ -0,0 +1,212 @@
package model
import (
"fmt"
"sync"
"sync/atomic"
"time"
"github.com/QuantumNous/new-api/common"
"github.com/QuantumNous/new-api/pkg/jsplugin"
)
// TaskAliasTarget is one mapping-derived alias after cross-channel aggregation.
// Declared is empty when the same plugin resolves the alias to more than one
// declared tail (display-only; pin is disabled).
type TaskAliasTarget struct {
Alias string
Declared string
PluginKey string
}
type taskAliasView struct {
generation uint64
expiresAt time.Time
byFold map[string]TaskAliasTarget
}
const taskAliasViewTTL = 60 * time.Second
var (
taskAliasViewPtr atomic.Pointer[taskAliasView]
taskAliasRebuildMu sync.Mutex
)
// ResolveTaskModelAlias returns the mapping-derived alias target for name
// (exact or ASCII-folded). g is the caller's routing generation: a Number
// mismatch or TTL expiry rebuilds the view against that generation.
func ResolveTaskModelAlias(g *jsplugin.RoutingGeneration, name string) (TaskAliasTarget, bool) {
if g == nil || name == "" {
return TaskAliasTarget{}, false
}
view := loadFreshTaskAliasView(g)
if view == nil {
return TaskAliasTarget{}, false
}
target, ok := view.byFold[jsplugin.ASCIIFold(name)]
return target, ok
}
func loadFreshTaskAliasView(g *jsplugin.RoutingGeneration) *taskAliasView {
view := taskAliasViewPtr.Load()
if taskAliasViewFresh(view, g.Number) {
return view
}
taskAliasRebuildMu.Lock()
defer taskAliasRebuildMu.Unlock()
view = taskAliasViewPtr.Load()
if taskAliasViewFresh(view, g.Number) {
return view
}
rebuilt := buildTaskAliasView(g)
taskAliasViewPtr.Store(rebuilt)
return rebuilt
}
func taskAliasViewFresh(view *taskAliasView, generation uint64) bool {
return view != nil && view.generation == generation && time.Now().Before(view.expiresAt)
}
func rebuildTaskAliasView() {
taskAliasRebuildMu.Lock()
defer taskAliasRebuildMu.Unlock()
taskAliasViewPtr.Store(buildTaskAliasView(jsplugin.DefaultRegistry.Generation()))
}
type taskAliasDraft struct {
spellings []string
byPlugin map[string]map[string]struct{}
}
func buildTaskAliasView(generation *jsplugin.RoutingGeneration) *taskAliasView {
genNum := uint64(0)
if generation != nil {
genNum = generation.Number
}
view := &taskAliasView{
generation: genNum,
expiresAt: time.Now().Add(taskAliasViewTTL),
byFold: make(map[string]TaskAliasTarget),
}
if DB == nil {
return view
}
var channels []Channel
err := DB.Select("id", "type", "models", "model_mapping").
Where("status = ?", common.ChannelStatusEnabled).
Find(&channels).Error
if err != nil {
common.SysError(fmt.Sprintf("rebuild task alias view: %s", err.Error()))
return view
}
drafts := make(map[string]*taskAliasDraft)
for i := range channels {
channel := &channels[i]
mappingJSON := channel.GetModelMapping()
if mappingJSON == "" || mappingJSON == "{}" {
continue
}
modelMap := make(map[string]string)
if err := common.UnmarshalJsonStr(mappingJSON, &modelMap); err != nil {
common.SysError(fmt.Sprintf("task alias view: channel %d model_mapping: %s", channel.Id, err.Error()))
continue
}
inModels := make(map[string]struct{})
for _, modelName := range channel.GetModels() {
inModels[modelName] = struct{}{}
}
for alias, mapped := range modelMap {
if mapped == "" {
continue
}
if _, exposed := inModels[alias]; !exposed {
continue
}
if _, declared := generation.CanonicalModel(alias); declared {
continue
}
tail, cyclic := followChannelModelMapping(modelMap, alias)
if cyclic {
common.SysError(fmt.Sprintf("task alias mapping cycle dropped: channel=%d key=%q", channel.Id, alias))
continue
}
declared, ok := generation.CanonicalModel(tail)
if !ok {
continue
}
plugin, ok := generation.GetByModel(declared)
if !ok {
continue
}
fold := jsplugin.ASCIIFold(alias)
draft := drafts[fold]
if draft == nil {
draft = &taskAliasDraft{byPlugin: make(map[string]map[string]struct{})}
drafts[fold] = draft
}
draft.spellings = append(draft.spellings, alias)
declareds := draft.byPlugin[plugin.Meta.Key]
if declareds == nil {
declareds = make(map[string]struct{})
draft.byPlugin[plugin.Meta.Key] = declareds
}
declareds[declared] = struct{}{}
}
}
for _, draft := range drafts {
alias := draft.spellings[0]
for _, spelling := range draft.spellings[1:] {
if spelling < alias {
alias = spelling
}
}
if len(draft.byPlugin) != 1 {
pluginKeys := make([]string, 0, len(draft.byPlugin))
for key := range draft.byPlugin {
pluginKeys = append(pluginKeys, key)
}
common.SysLog(fmt.Sprintf("task model alias %q dropped: maps to multiple plugins %v", alias, pluginKeys))
continue
}
var pluginKey, declared string
for key, declareds := range draft.byPlugin {
pluginKey = key
if len(declareds) == 1 {
for name := range declareds {
declared = name
}
}
}
view.byFold[jsplugin.ASCIIFold(alias)] = TaskAliasTarget{
Alias: alias,
Declared: declared,
PluginKey: pluginKey,
}
}
return view
}
// followChannelModelMapping walks one channel's mapping the same way
// ModelMappedHelper does: visited-set cycle detection, self-map stops at
// the current hop, a non-self cycle is reported to the caller.
func followChannelModelMapping(modelMap map[string]string, start string) (string, bool) {
current := start
visited := map[string]bool{current: true}
for {
mapped, exists := modelMap[current]
if !exists || mapped == "" {
return current, false
}
if visited[mapped] {
if mapped == current {
return current, false
}
return "", true
}
visited[mapped] = true
current = mapped
}
}
+26
View File
@@ -0,0 +1,26 @@
package jsplugin
// asciiFold maps only 'A''Z' onto 'a''z'. Every other byte is unchanged.
// Unicode case folding is intentionally not applied: U+212A (KELVIN SIGN)
// must not become 'k' and impersonate an ASCII model name.
func asciiFold(s string) string {
var buf []byte
for i := 0; i < len(s); i++ {
c := s[i]
if c >= 'A' && c <= 'Z' {
if buf == nil {
buf = []byte(s)
}
buf[i] = c + ('a' - 'A')
}
}
if buf == nil {
return s
}
return string(buf)
}
// ASCIIFold is the exported form of asciiFold for consumers outside this package.
func ASCIIFold(s string) string {
return asciiFold(s)
}
+5 -2
View File
@@ -1144,13 +1144,16 @@ func normalizeV1Meta(meta *Meta) error {
seenChannelTypes[channelType] = struct{}{} seenChannelTypes[channelType] = struct{}{}
} }
models := make(map[string]struct{}, len(meta.Models)) models := make(map[string]struct{}, len(meta.Models))
seenFold := make(map[string]struct{}, len(meta.Models))
for _, model := range meta.Models { for _, model := range meta.Models {
if strings.TrimSpace(model) == "" || strings.TrimSpace(model) != model { if strings.TrimSpace(model) == "" || strings.TrimSpace(model) != model {
return fmt.Errorf("plugin meta models must contain non-empty canonical names") return fmt.Errorf("plugin meta models must contain non-empty canonical names")
} }
if _, exists := models[model]; exists { folded := asciiFold(model)
return fmt.Errorf("plugin meta models must be unique") if _, exists := seenFold[folded]; exists {
return fmt.Errorf("plugin meta models must be unique case-insensitively")
} }
seenFold[folded] = struct{}{}
models[model] = struct{}{} models[model] = struct{}{}
} }
hosts := make(map[string]struct{}, len(meta.AllowedHosts)) hosts := make(map[string]struct{}, len(meta.AllowedHosts))
+40 -3
View File
@@ -224,6 +224,7 @@ type PinnedEndpoint struct {
Protocol string Protocol string
Operation HostProtocolOperation Operation HostProtocolOperation
Model string Model string
MappedModel string
Candidates []ProtocolBinding Candidates []ProtocolBinding
} }
@@ -296,6 +297,10 @@ type ProtocolRequestContext struct {
Protocol string `json:"protocol"` Protocol string `json:"protocol"`
Operation string `json:"operation"` Operation string `json:"operation"`
Model string `json:"model"` Model string `json:"model"`
// UpstreamModel is the declared machine identity when Model is a
// channel-mapping alias; empty otherwise. Decode hooks that key rate
// tables or request shaping by model must use it over Model.
UpstreamModel string `json:"upstreamModel,omitempty"`
Stream bool `json:"stream"` Stream bool `json:"stream"`
} }
@@ -304,6 +309,9 @@ func (p ProtocolRequestContext) JSValue() map[string]any {
value["protocol"] = p.Protocol value["protocol"] = p.Protocol
value["operation"] = p.Operation value["operation"] = p.Operation
value["model"] = p.Model value["model"] = p.Model
if p.UpstreamModel != "" {
value["upstreamModel"] = p.UpstreamModel
}
value["stream"] = p.Stream value["stream"] = p.Stream
return value return value
} }
@@ -320,6 +328,7 @@ type RoutingGeneration struct {
byKey map[string]*LoadedPlugin byKey map[string]*LoadedPlugin
byModel map[string]*LoadedPlugin byModel map[string]*LoadedPlugin
canonicalModelByFold map[string]string
byChannelType map[int]*LoadedPlugin byChannelType map[int]*LoadedPlugin
routeIndex map[string]RouteBinding routeIndex map[string]RouteBinding
protocolIndex map[string][]ProtocolBinding protocolIndex map[string][]ProtocolBinding
@@ -409,6 +418,20 @@ func (g *RoutingGeneration) GetByModel(model string) (*LoadedPlugin, bool) {
return plugin, ok return plugin, ok
} }
// CanonicalModel returns the declared spelling for model. An exact byModel
// hit wins and returns the input unchanged; otherwise the ASCII-folded
// index is consulted. Miss and nil-receiver return ("", false).
func (g *RoutingGeneration) CanonicalModel(model string) (string, bool) {
if g == nil || model == "" {
return "", false
}
if _, ok := g.byModel[model]; ok {
return model, true
}
declared, ok := g.canonicalModelByFold[asciiFold(model)]
return declared, ok
}
// LookupDeclaredRoute resolves a manifest path declaration. It does not match // LookupDeclaredRoute resolves a manifest path declaration. It does not match
// an incoming concrete URL; runtime matching is delegated to Gin. // an incoming concrete URL; runtime matching is delegated to Gin.
func (g *RoutingGeneration) LookupDeclaredRoute(method, path string) (RouteBinding, bool) { func (g *RoutingGeneration) LookupDeclaredRoute(method, path string) (RouteBinding, bool) {
@@ -721,10 +744,11 @@ func validateModelScope(models []string, subject string) error {
if strings.TrimSpace(model) == "" || strings.TrimSpace(model) != model { if strings.TrimSpace(model) == "" || strings.TrimSpace(model) != model {
return fmt.Errorf("plugin %s models must contain non-empty canonical names", subject) return fmt.Errorf("plugin %s models must contain non-empty canonical names", subject)
} }
if _, duplicate := seen[model]; duplicate { folded := asciiFold(model)
return fmt.Errorf("plugin %s models must be unique", subject) if _, duplicate := seen[folded]; duplicate {
return fmt.Errorf("plugin %s models must be unique case-insensitively", subject)
} }
seen[model] = struct{}{} seen[folded] = struct{}{}
} }
return nil return nil
} }
@@ -840,6 +864,7 @@ func buildRoutingGenerationFromPlugins(effective map[string]*LoadedPlugin, numbe
PublishedAt: time.Now(), PublishedAt: time.Now(),
byKey: make(map[string]*LoadedPlugin, len(effective)), byKey: make(map[string]*LoadedPlugin, len(effective)),
byModel: make(map[string]*LoadedPlugin), byModel: make(map[string]*LoadedPlugin),
canonicalModelByFold: make(map[string]string),
byChannelType: make(map[int]*LoadedPlugin), byChannelType: make(map[int]*LoadedPlugin),
routeIndex: make(map[string]RouteBinding), routeIndex: make(map[string]RouteBinding),
protocolIndex: make(map[string][]ProtocolBinding), protocolIndex: make(map[string][]ProtocolBinding),
@@ -853,6 +878,18 @@ func buildRoutingGenerationFromPlugins(effective map[string]*LoadedPlugin, numbe
if _, exists := generation.byModel[model]; !exists { if _, exists := generation.byModel[model]; !exists {
generation.byModel[model] = plugin generation.byModel[model] = plugin
} }
folded := asciiFold(model)
if existing, exists := generation.canonicalModelByFold[folded]; exists {
if existing != model {
otherKey := plugin.Meta.Key
if other, ok := generation.byModel[existing]; ok {
otherKey = other.Meta.Key
}
return nil, fmt.Errorf("plugin %s model %q conflicts with plugin %s model %q", plugin.Meta.Key, model, otherKey, existing)
}
continue
}
generation.canonicalModelByFold[folded] = model
} }
for _, channelType := range plugin.Meta.ChannelTypes { for _, channelType := range plugin.Meta.ChannelTypes {
+1 -1
View File
@@ -279,7 +279,7 @@ export function extractUsage(ctx) {
const req = ctx.requestBody || {}; const req = ctx.requestBody || {};
const metadata = req.metadata || {}; const metadata = req.metadata || {};
if (ctx.usagePurpose === "billing_ratios") { if (ctx.usagePurpose === "billing_ratios") {
const ratio = videoInputRatio(ctx.model, metadata.resolution, metadata.content); const ratio = videoInputRatio(ctx.upstreamModel || ctx.model, metadata.resolution, metadata.content);
return ratio === 1 ? null : { video_input_ratio: ratio }; return ratio === 1 ? null : { video_input_ratio: ratio };
} }
let seconds = Number(req.seconds || req.duration || metadata.duration || 0); let seconds = Number(req.seconds || req.duration || metadata.duration || 0);
+2 -1
View File
@@ -388,7 +388,8 @@ protocols.openai_video = {
} }
const hasImage = hasHailuoImage(req, hasInputReferenceFile); const hasImage = hasHailuoImage(req, hasInputReferenceFile);
const duration = req.duration === undefined ? undefined : Number(req.duration); const duration = req.duration === undefined ? undefined : Number(req.duration);
validateHailuoCombo(ctx.model, duration, outboundResolution(req, ctx.model), hasImage); const comboModel = ctx.upstreamModel || ctx.model;
validateHailuoCombo(comboModel, duration, outboundResolution(req, comboModel), hasImage);
return { return {
kind: "submit", kind: "submit",
model: ctx.model, model: ctx.model,
+4 -1
View File
@@ -533,7 +533,10 @@ protocols.openai_video = {
} }
const seconds = req.seconds === undefined ? req.duration : req.seconds; const seconds = req.seconds === undefined ? req.duration : req.seconds;
if (seconds !== undefined) { if (seconds !== undefined) {
req.duration = validateSecondsForReqKey(convertedReqKey(String(ctx.model || req.model || ""), decodeImageCount(req, hasInputReferenceFile)), seconds); req.duration = validateSecondsForReqKey(
convertedReqKey(String(ctx.upstreamModel || ctx.model || req.model || ""), decodeImageCount(req, hasInputReferenceFile)),
seconds
);
} }
return { return {
kind: "submit", kind: "submit",
+2 -2
View File
@@ -357,7 +357,7 @@ export const protocols = {
if (!prompt && images.length === 0) throw new Error("input is required"); if (!prompt && images.length === 0) throw new Error("input is required");
const metadata = Object.assign({}, req.metadata || {}); const metadata = Object.assign({}, req.metadata || {});
if (Object.prototype.hasOwnProperty.call(req, "mode")) metadata.mode = req.mode; if (Object.prototype.hasOwnProperty.call(req, "mode")) metadata.mode = req.mode;
metadata.mode = resolveKlingMode(model, metadata.mode); metadata.mode = resolveKlingMode(ctx.upstreamModel || model, metadata.mode);
if (images.length > 1 && !metadata.image_tail) metadata.image_tail = images[1]; if (images.length > 1 && !metadata.image_tail) metadata.image_tail = images[1];
const requestBody = { model: model, prompt: prompt, metadata: metadata }; const requestBody = { model: model, prompt: prompt, metadata: metadata };
if (images.length) requestBody.image = images[0]; if (images.length) requestBody.image = images[0];
@@ -443,7 +443,7 @@ export const protocols = {
const image = trimmed(req.input_reference || req.image); const image = trimmed(req.input_reference || req.image);
if (image) req.image = image; if (image) req.image = image;
} }
const model = ctx.model || req.model || "kling-v1"; const model = ctx.upstreamModel || ctx.model || req.model || "kling-v1";
const metadata = req.metadata || {}; const metadata = req.metadata || {};
req.mode = resolveKlingMode(model, req.mode || metadata.mode); req.mode = resolveKlingMode(model, req.mode || metadata.mode);
const hasImage = hasKlingImage(req, hasInputReferenceFile); const hasImage = hasKlingImage(req, hasInputReferenceFile);
+1 -1
View File
@@ -128,7 +128,7 @@ export function parseSubmitResponse(ctx, resp) {
export function extractUsage(ctx) { export function extractUsage(ctx) {
if (ctx.usagePurpose === "billing_ratios") return null; if (ctx.usagePurpose === "billing_ratios") return null;
const model = trimmed(ctx.model || (ctx.requestBody || {}).model).toLowerCase(); const model = trimmed(ctx.upstreamModel || ctx.model || (ctx.requestBody || {}).model).toLowerCase();
const action = actionName(ctx).toLowerCase() || (model === "suno_lyrics" ? "lyrics" : "music"); const action = actionName(ctx).toLowerCase() || (model === "suno_lyrics" ? "lyrics" : "music");
return { clips: action === "lyrics" ? 1 : 2, action: action }; return { clips: action === "lyrics" ? 1 : 2, action: action };
} }
+1 -1
View File
@@ -417,7 +417,7 @@ protocols.openai_video = {
if (req.seconds !== undefined) req.seconds = Number(req.seconds); if (req.seconds !== undefined) req.seconds = Number(req.seconds);
else if (req.duration !== undefined) req.seconds = Number(req.duration); else if (req.duration !== undefined) req.seconds = Number(req.duration);
} }
const model = ctx.model || req.model; const model = ctx.upstreamModel || ctx.model || req.model;
const seconds = req.seconds === undefined ? req.duration : req.seconds; const seconds = req.seconds === undefined ? req.duration : req.seconds;
if (seconds !== undefined) req.duration = Number(seconds); if (seconds !== undefined) req.duration = Number(seconds);
else req.duration = defaultDuration(model); else req.duration = defaultDuration(model);
@@ -19,6 +19,7 @@ import (
pluginruntime "github.com/QuantumNous/new-api/pkg/jsplugin" pluginruntime "github.com/QuantumNous/new-api/pkg/jsplugin"
"github.com/QuantumNous/new-api/relay/channel" "github.com/QuantumNous/new-api/relay/channel"
relaycommon "github.com/QuantumNous/new-api/relay/common" relaycommon "github.com/QuantumNous/new-api/relay/common"
"github.com/QuantumNous/new-api/relay/helper"
"github.com/QuantumNous/new-api/relaykit/dto" "github.com/QuantumNous/new-api/relaykit/dto"
"github.com/QuantumNous/new-api/service" "github.com/QuantumNous/new-api/service"
"github.com/gin-gonic/gin" "github.com/gin-gonic/gin"
@@ -1154,3 +1155,56 @@ func TestTaskAdaptorBatchBridge(t *testing.T) {
assert.Equal(t, "40%", pending.TaskInfo.Progress) assert.Equal(t, "40%", pending.TaskInfo.Progress)
assert.Empty(t, pending.TaskInfo.Url) assert.Empty(t, pending.TaskInfo.Url)
} }
const mappingOrderAdaptorPlugin = `
export const meta = {apiVersion:1,key:"map-order-adaptor",name:"Map Order Adaptor",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}};
}
export function parseSubmitResponse(){return {taskId:"1"};}
export function buildQueryRequest(){return {url:"https://provider.example"};}
export function parseTaskResult(){return {status:"SUCCESS"};}
`
func mappingOrderSubmitBody(t *testing.T, origin, mapping string) []byte {
t.Helper()
plugin, err := pluginruntime.NewRegistry().Register(mappingOrderAdaptorPlugin, pluginruntime.Options{})
require.NoError(t, err)
adaptor := New(plugin)
info := &relaycommon.RelayInfo{
ChannelMeta: &relaycommon.ChannelMeta{ChannelBaseUrl: "https://provider.example"},
TaskRelayInfo: &relaycommon.TaskRelayInfo{},
OriginModelName: origin,
}
adaptor.Init(info)
c, _ := gin.CreateTestContext(httptest.NewRecorder())
c.Request = httptest.NewRequest(http.MethodPost, "/v1/videos", nil)
if mapping != "" {
c.Set("model_mapping", mapping)
}
c.Set("task_request", map[string]any{"prompt": "p"})
info.UpstreamModelName = info.OriginModelName
require.NoError(t, helper.ModelMappedHelper(c, info, nil))
require.Nil(t, adaptor.ValidateRequestAndSetAction(c, info))
body, err := adaptor.BuildRequestBody(c, info)
require.NoError(t, err)
raw, err := io.ReadAll(body)
require.NoError(t, err)
return raw
}
func TestTaskAdaptorBuildSubmitReceivesMappedUpstreamModel(t *testing.T) {
gin.SetMode(gin.TestMode)
mapped := mappingOrderSubmitBody(t, "alias-model", `{"alias-model":"mid-model","mid-model":"declared-model"}`)
var decoded map[string]any
require.NoError(t, common.Unmarshal(mapped, &decoded))
assert.Equal(t, "declared-model", decoded["upstreamModel"])
assert.Equal(t, "alias-model", decoded["model"])
withoutMapping := mappingOrderSubmitBody(t, "declared-model", "")
emptyMapping := mappingOrderSubmitBody(t, "declared-model", "{}")
assert.Equal(t, withoutMapping, emptyMapping)
require.NoError(t, common.Unmarshal(withoutMapping, &decoded))
assert.Equal(t, "declared-model", decoded["upstreamModel"])
assert.Equal(t, "declared-model", decoded["model"])
}
+30 -3
View File
@@ -212,6 +212,19 @@ func RelayTaskSubmit(c *gin.Context, info *relaycommon.RelayInfo) (*TaskSubmitRe
info.PublicTaskID = model.GenerateTaskID() info.PublicTaskID = model.GenerateTaskID()
} }
adaptor.Init(info) adaptor.Init(info)
// Plugin submit hooks run during ValidateRequestAndSetAction and cache the
// upstream body. OriginModelName is already seeded on that line (protocol
// resolved_task_model, legacy submit, or GenRelayInfo original_model), so
// map before validation. The empty-name CoverTaskActionToModelName
// synthesis happens after validate and cannot move; skip the late block
// when early mapping ran so a chain is never applied twice.
mappedBeforeValidate := info.OriginModelName != ""
if mappedBeforeValidate {
info.UpstreamModelName = info.OriginModelName
if err := helper.ModelMappedHelper(c, info, nil); err != nil {
return nil, service.TaskErrorWrapperLocal(err, "model_mapping_failed", http.StatusBadRequest)
}
}
if taskErr := adaptor.ValidateRequestAndSetAction(c, info); taskErr != nil { if taskErr := adaptor.ValidateRequestAndSetAction(c, info); taskErr != nil {
return nil, taskErr return nil, taskErr
} }
@@ -222,19 +235,33 @@ func RelayTaskSubmit(c *gin.Context, info *relaycommon.RelayInfo) (*TaskSubmitRe
modelName = service.CoverTaskActionToModelName(platform, info.Action) modelName = service.CoverTaskActionToModelName(platform, info.Action)
} }
// 2.5 应用渠道的模型映射(与同步任务对齐) if !mappedBeforeValidate {
info.OriginModelName = modelName info.OriginModelName = modelName
info.UpstreamModelName = modelName info.UpstreamModelName = modelName
if err := helper.ModelMappedHelper(c, info, nil); err != nil { if err := helper.ModelMappedHelper(c, info, nil); err != nil {
return nil, service.TaskErrorWrapperLocal(err, "model_mapping_failed", http.StatusBadRequest) return nil, service.TaskErrorWrapperLocal(err, "model_mapping_failed", http.StatusBadRequest)
} }
}
// 4. 价格计算:基础模型价格 // 4. 价格计算:基础模型价格
info.OriginModelName = modelName info.OriginModelName = modelName
var priceData types.PriceData var priceData types.PriceData
var err error var err error
if billing_setting.GetBillingMode(modelName) == billing_setting.BillingModeTieredExpr { useTiered := billing_setting.GetBillingMode(modelName) == billing_setting.BillingModeTieredExpr
exprStr, exists := billing_setting.GetBillingExpr(modelName) var exprStr string
var exists bool
if useTiered {
exprStr, exists = billing_setting.GetBillingExpr(modelName)
} else if info.IsModelMapped {
if billing_setting.GetBillingMode(info.UpstreamModelName) == billing_setting.BillingModeTieredExpr {
if tailExpr, tailOK := billing_setting.GetBillingExpr(info.UpstreamModelName); tailOK && strings.TrimSpace(tailExpr) != "" {
exprStr = tailExpr
exists = true
useTiered = true
}
}
}
if useTiered {
provider, supported := adaptor.(channel.TaskUsageFactsProvider) provider, supported := adaptor.(channel.TaskUsageFactsProvider)
if !exists || !supported { if !exists || !supported {
return nil, service.TaskErrorWrapper(fmt.Errorf("task model %s has no usage expression or meter", modelName), "model_price_error", http.StatusBadRequest) return nil, service.TaskErrorWrapper(fmt.Errorf("task model %s has no usage expression or meter", modelName), "model_price_error", http.StatusBadRequest)
+214
View File
@@ -1,11 +1,22 @@
package relay package relay
import ( import (
"net/http"
"net/http/httptest"
"testing" "testing"
"github.com/QuantumNous/new-api/common"
"github.com/QuantumNous/new-api/constant" "github.com/QuantumNous/new-api/constant"
"github.com/QuantumNous/new-api/model" "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/gin-gonic/gin"
"github.com/stretchr/testify/assert" "github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
) )
func TestTaskModel2DtoNormalizesLegacyAction(t *testing.T) { func TestTaskModel2DtoNormalizesLegacyAction(t *testing.T) {
@@ -16,3 +27,206 @@ func TestTaskModel2DtoNormalizesLegacyAction(t *testing.T) {
assert.Equal(t, constant.TaskActionFirstTailToVideo, dtoTask.Action) assert.Equal(t, constant.TaskActionFirstTailToVideo, dtoTask.Action)
assert.Equal(t, "firstTailGenerate", task.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()
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)`
tests := []struct {
name string
modes map[string]string
exprs map[string]string
wantTiered bool
wantExpr 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: "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"
pinMappingOrderPlugin(t, c, billingFallbackPlugin)
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)
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)
} else {
assert.Nil(t, info.TieredBillingSnapshot)
assert.Equal(t, "model_price_error", taskErr.Code)
}
})
}
}