From 6c22550ea325d4fea0e0ece52412cb1e4449291c Mon Sep 17 00:00:00 2001 From: CaIon Date: Sun, 30 Aug 2026 19:13:51 +0800 Subject: [PATCH] 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. --- controller/billing_option_test.go | 103 ++++++++++ controller/option.go | 6 + controller/plugin_protocol_test.go | 26 +++ docs/plugin-api/v1.md | 2 + middleware/task_plugin.go | 85 +++++++- middleware/task_plugin_model.go | 45 ++++ middleware/task_plugin_test.go | 8 +- model/channel_cache.go | 4 +- model/pricing.go | 15 +- model/pricing_usage_schema_test.go | 70 +++++++ model/task_model_alias.go | 212 +++++++++++++++++++ pkg/jsplugin/model_fold.go | 26 +++ pkg/jsplugin/registry.go | 7 +- pkg/jsplugin/routing.go | 91 ++++++--- plugins/tasks/doubao/plugin.js | 2 +- plugins/tasks/hailuo/plugin.js | 3 +- plugins/tasks/jimeng/plugin.js | 5 +- plugins/tasks/kling/plugin.js | 4 +- plugins/tasks/sunoapi/plugin.js | 2 +- plugins/tasks/vidu/plugin.js | 2 +- relay/channel/task/jsplugin/adaptor_test.go | 54 +++++ relay/relay_task.go | 41 +++- relay/relay_task_test.go | 214 ++++++++++++++++++++ 23 files changed, 968 insertions(+), 59 deletions(-) create mode 100644 middleware/task_plugin_model.go create mode 100644 model/task_model_alias.go create mode 100644 pkg/jsplugin/model_fold.go diff --git a/controller/billing_option_test.go b/controller/billing_option_test.go index c02334d52a..cf6b37ea53 100644 --- a/controller/billing_option_test.go +++ b/controller/billing_option_test.go @@ -7,10 +7,14 @@ import ( "testing" "github.com/QuantumNous/new-api/common" + "github.com/QuantumNous/new-api/model" "github.com/QuantumNous/new-api/pkg/jsplugin" + "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 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(), "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") +} diff --git a/controller/option.go b/controller/option.go index 70a5f19218..1feb4a818d 100644 --- a/controller/option.go +++ b/controller/option.go @@ -352,6 +352,12 @@ func UpdateOption(c *gin.Context) { expression := expressions[modelName] if plugin, ok := generation.GetByModel(modelName); ok { 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 { err = billing_setting.SmokeTestExpr(expression) } diff --git a/controller/plugin_protocol_test.go b/controller/plugin_protocol_test.go index a6a5fdc1bd..edc64034aa 100644 --- a/controller/plugin_protocol_test.go +++ b/controller/plugin_protocol_test.go @@ -1088,6 +1088,7 @@ func TestRetrieveTaskPluginResponsePendingSkipsRenderFinal(t *testing.T) { require.NoError(t, common.Unmarshal(recorder.Body.Bytes(), &response)) assert.Equal(t, "resp_retrieve_pending", response["id"]) assert.Equal(t, "in_progress", response["status"]) + assert.Equal(t, "video-model", response["model"]) assert.Equal(t, true, response["background"]) assert.Nil(t, response["completed_at"]) 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"]) } +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) { logs := make([]string, 0, 1) pinned := compilePluginProtocolRetrieveEndpoint(t, "retrieve-success", ` diff --git a/docs/plugin-api/v1.md b/docs/plugin-api/v1.md index 814513b06f..62fba917a1 100644 --- a/docs/plugin-api/v1.md +++ b/docs/plugin-api/v1.md @@ -138,3 +138,5 @@ Protocol media uses host-injected `ctx.artifacts[key].url`. Provider URLs from ` ## 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. + +`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. diff --git a/middleware/task_plugin.go b/middleware/task_plugin.go index 3ca33ab732..7d271c7723 100644 --- a/middleware/task_plugin.go +++ b/middleware/task_plugin.go @@ -332,18 +332,50 @@ func PinTaskPluginEndpoint() gin.HandlerFunc { c.Next() return } - c.Set(contextKeyTaskPluginEndpointModel, *modelRequest) claimedModel := modelRequest.Model if strings.TrimSpace(claimedModel) == "" { + c.Set(contextKeyTaskPluginEndpointModel, *modelRequest) c.Next() 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 { + c.Set(contextKeyTaskPluginEndpointModel, *modelRequest) c.Next() 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 { candidates = []pluginruntime.ProtocolBinding{binding} } @@ -386,12 +418,13 @@ func PinTaskPluginEndpoint() gin.HandlerFunc { pin := pluginruntime.PinnedPlugin{Generation: generation, Plugin: binding.Plugin} pinnedEndpoint := pluginruntime.PinnedEndpoint{ - Generation: generation, - Plugin: binding.Plugin, - Protocol: binding.Protocol, - Operation: binding.Operation, - Model: claimedModel, - Candidates: candidates, + Generation: generation, + Plugin: binding.Plugin, + Protocol: binding.Protocol, + Operation: binding.Operation, + Model: pinModel, + MappedModel: mappedModel, + Candidates: candidates, } c.Set(pluginruntime.ContextKeyPinnedPlugin, pin) c.Set(pluginruntime.ContextKeyPinnedEndpoint, pinnedEndpoint) @@ -403,7 +436,7 @@ func PinTaskPluginEndpoint() gin.HandlerFunc { binding.Plugin.Meta.Version, binding.Operation.Methods[0], binding.Protocol, - claimedModel, + pinModel, ) c.Next() } @@ -515,6 +548,13 @@ func PrepareTaskPluginEndpoint() gin.HandlerFunc { abortWithOpenAiMessage(c, status, err.Error()) 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) bodyKind, _ := bodyObject["kind"].(string) allowedBody := false @@ -561,6 +601,7 @@ func PrepareTaskPluginEndpoint() gin.HandlerFunc { Protocol: pinned.Protocol, Operation: pinned.Operation.Name, Model: pinned.Model, + UpstreamModel: pinned.MappedModel, Stream: stream, } c.Set(pluginruntime.ContextKeyProtocolRequest, protocolContext) @@ -624,7 +665,8 @@ func PrepareTaskPluginEndpoint() gin.HandlerFunc { return } modelOwned := slices.Contains(pinned.Plugin.Meta.Models, resolvedModel) - if !modelOwned || resolvedModel != pinned.Model { + mappedPin := pinned.MappedModel != "" + if resolvedModel != pinned.Model || (!modelOwned && !mappedPin) { logger.LogWarn( 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", @@ -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"}}) 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("resolved_task_model", modelName) c.Set("expected_task_plugin_key", pluginKey) diff --git a/middleware/task_plugin_model.go b/middleware/task_plugin_model.go new file mode 100644 index 0000000000..e7bcba7539 --- /dev/null +++ b/middleware/task_plugin_model.go @@ -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 +} diff --git a/middleware/task_plugin_test.go b/middleware/task_plugin_test.go index c2e8cbfc14..c7e43e7de8 100644 --- a/middleware/task_plugin_test.go +++ b/middleware/task_plugin_test.go @@ -630,9 +630,11 @@ func TestTaskPluginEndpointMissPreservesOrdinaryRequestBody(t *testing.T) { func(c *gin.Context) { _, pinned := c.Get(jsplugin.ContextKeyPinnedEndpoint) assert.False(t, pinned) - var body map[string]any - require.NoError(t, common.UnmarshalBodyReusable(c, &body)) - assert.Equal(t, "ordinary-model", body["model"]) + storage, storageErr := common.GetBodyStorage(c) + require.NoError(t, storageErr) + raw, bytesErr := storage.Bytes() + require.NoError(t, bytesErr) + assert.Equal(t, []byte(`{"model":"ordinary-model","input":"hello"}`), raw) c.Status(http.StatusNoContent) }, ) diff --git a/model/channel_cache.go b/model/channel_cache.go index 97b80cac44..a992c1961c 100644 --- a/model/channel_cache.go +++ b/model/channel_cache.go @@ -11,8 +11,8 @@ import ( "github.com/QuantumNous/new-api/common" "github.com/QuantumNous/new-api/constant" - "github.com/QuantumNous/new-api/logger" "github.com/QuantumNous/new-api/dto" + "github.com/QuantumNous/new-api/logger" kitdto "github.com/QuantumNous/new-api/relaykit/dto" "github.com/QuantumNous/new-api/setting/ratio_setting" ) @@ -27,6 +27,7 @@ var channelSyncLock sync.RWMutex func InitChannelCache() { if !common.MemoryCacheEnabled { InvalidatePricingCache() + rebuildTaskAliasView() return } newChannelId2channel := make(map[int]*Channel) @@ -101,6 +102,7 @@ func InitChannelCache() { // loadPricingAdvancedCustomConfigs. channelSyncLock MUST be released before // invalidating the pricing cache, otherwise the reversed order deadlocks. InvalidatePricingCache() + rebuildTaskAliasView() common.SysLog("channels synced from database") } diff --git a/model/pricing.go b/model/pricing.go index 9d9f5c50e3..4cec75f3e1 100644 --- a/model/pricing.go +++ b/model/pricing.go @@ -410,8 +410,21 @@ func updatePricing() { pricing.BillingMode = billingMode 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)) for key, field := range plugin.Meta.UsageSchema { field.Enum = append([]string(nil), field.Enum...) diff --git a/model/pricing_usage_schema_test.go b/model/pricing_usage_schema_test.go index 593bc1df55..b947b9749c 100644 --- a/model/pricing_usage_schema_test.go +++ b/model/pricing_usage_schema_test.go @@ -8,6 +8,7 @@ import ( "github.com/QuantumNous/new-api/constant" "github.com/QuantumNous/new-api/pkg/jsplugin" "github.com/QuantumNous/new-api/relaykit/dto" + "github.com/QuantumNous/new-api/setting/config" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" ) @@ -60,6 +61,75 @@ func TestPricingCarriesTaskUsageSchemaAndRefreshesWithPluginGeneration(t *testin 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 { result := make(map[string]Pricing, len(pricings)) for _, pricing := range pricings { diff --git a/model/task_model_alias.go b/model/task_model_alias.go new file mode 100644 index 0000000000..c975386ede --- /dev/null +++ b/model/task_model_alias.go @@ -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 + } +} diff --git a/pkg/jsplugin/model_fold.go b/pkg/jsplugin/model_fold.go new file mode 100644 index 0000000000..5e4cb037d1 --- /dev/null +++ b/pkg/jsplugin/model_fold.go @@ -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) +} diff --git a/pkg/jsplugin/registry.go b/pkg/jsplugin/registry.go index e16b7041b4..1ad50e54db 100644 --- a/pkg/jsplugin/registry.go +++ b/pkg/jsplugin/registry.go @@ -1144,13 +1144,16 @@ func normalizeV1Meta(meta *Meta) error { seenChannelTypes[channelType] = struct{}{} } models := make(map[string]struct{}, len(meta.Models)) + seenFold := make(map[string]struct{}, len(meta.Models)) for _, model := range meta.Models { if strings.TrimSpace(model) == "" || strings.TrimSpace(model) != model { return fmt.Errorf("plugin meta models must contain non-empty canonical names") } - if _, exists := models[model]; exists { - return fmt.Errorf("plugin meta models must be unique") + folded := asciiFold(model) + if _, exists := seenFold[folded]; exists { + return fmt.Errorf("plugin meta models must be unique case-insensitively") } + seenFold[folded] = struct{}{} models[model] = struct{}{} } hosts := make(map[string]struct{}, len(meta.AllowedHosts)) diff --git a/pkg/jsplugin/routing.go b/pkg/jsplugin/routing.go index 3faac1b1fe..03845e7e1b 100644 --- a/pkg/jsplugin/routing.go +++ b/pkg/jsplugin/routing.go @@ -219,12 +219,13 @@ type PinnedRoute struct { // distribution may rebind it to another candidate from the same generation // when multiple legacy providers expose the same model. type PinnedEndpoint struct { - Generation *RoutingGeneration - Plugin *LoadedPlugin - Protocol string - Operation HostProtocolOperation - Model string - Candidates []ProtocolBinding + Generation *RoutingGeneration + Plugin *LoadedPlugin + Protocol string + Operation HostProtocolOperation + Model string + MappedModel string + Candidates []ProtocolBinding } // RouteRequestContext is the canonical request view exposed to declarative @@ -296,7 +297,11 @@ type ProtocolRequestContext struct { Protocol string `json:"protocol"` Operation string `json:"operation"` Model string `json:"model"` - Stream bool `json:"stream"` + // 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"` } func (p ProtocolRequestContext) JSValue() map[string]any { @@ -304,6 +309,9 @@ func (p ProtocolRequestContext) JSValue() map[string]any { value["protocol"] = p.Protocol value["operation"] = p.Operation value["model"] = p.Model + if p.UpstreamModel != "" { + value["upstreamModel"] = p.UpstreamModel + } value["stream"] = p.Stream return value } @@ -318,15 +326,16 @@ type RoutingGeneration struct { Number uint64 PublishedAt time.Time - byKey map[string]*LoadedPlugin - byModel map[string]*LoadedPlugin - byChannelType map[int]*LoadedPlugin - routeIndex map[string]RouteBinding - protocolIndex map[string][]ProtocolBinding - plugins []*LoadedPlugin - routes []RouteBinding - runtime http.Handler - retainCurrent map[string]struct{} + byKey map[string]*LoadedPlugin + byModel map[string]*LoadedPlugin + canonicalModelByFold map[string]string + byChannelType map[int]*LoadedPlugin + routeIndex map[string]RouteBinding + protocolIndex map[string][]ProtocolBinding + plugins []*LoadedPlugin + routes []RouteBinding + runtime http.Handler + retainCurrent map[string]struct{} } var ( @@ -409,6 +418,20 @@ func (g *RoutingGeneration) GetByModel(model string) (*LoadedPlugin, bool) { 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 // an incoming concrete URL; runtime matching is delegated to Gin. 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 { return fmt.Errorf("plugin %s models must contain non-empty canonical names", subject) } - if _, duplicate := seen[model]; duplicate { - return fmt.Errorf("plugin %s models must be unique", subject) + folded := asciiFold(model) + 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 } @@ -836,14 +860,15 @@ func buildRoutingGenerationFromPlugins(effective map[string]*LoadedPlugin, numbe sort.Strings(keys) generation := &RoutingGeneration{ - Number: number, - PublishedAt: time.Now(), - byKey: make(map[string]*LoadedPlugin, len(effective)), - byModel: make(map[string]*LoadedPlugin), - byChannelType: make(map[int]*LoadedPlugin), - routeIndex: make(map[string]RouteBinding), - protocolIndex: make(map[string][]ProtocolBinding), - plugins: make([]*LoadedPlugin, 0, len(effective)), + Number: number, + PublishedAt: time.Now(), + byKey: make(map[string]*LoadedPlugin, len(effective)), + byModel: make(map[string]*LoadedPlugin), + canonicalModelByFold: make(map[string]string), + byChannelType: make(map[int]*LoadedPlugin), + routeIndex: make(map[string]RouteBinding), + protocolIndex: make(map[string][]ProtocolBinding), + plugins: make([]*LoadedPlugin, 0, len(effective)), } for _, key := range keys { plugin := effective[key] @@ -853,6 +878,18 @@ func buildRoutingGenerationFromPlugins(effective map[string]*LoadedPlugin, numbe if _, exists := generation.byModel[model]; !exists { 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 { diff --git a/plugins/tasks/doubao/plugin.js b/plugins/tasks/doubao/plugin.js index 07afceaaf4..e6dd7e9ed7 100644 --- a/plugins/tasks/doubao/plugin.js +++ b/plugins/tasks/doubao/plugin.js @@ -279,7 +279,7 @@ export function extractUsage(ctx) { const req = ctx.requestBody || {}; const metadata = req.metadata || {}; 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 }; } let seconds = Number(req.seconds || req.duration || metadata.duration || 0); diff --git a/plugins/tasks/hailuo/plugin.js b/plugins/tasks/hailuo/plugin.js index e54b72438c..b26397609e 100644 --- a/plugins/tasks/hailuo/plugin.js +++ b/plugins/tasks/hailuo/plugin.js @@ -388,7 +388,8 @@ protocols.openai_video = { } const hasImage = hasHailuoImage(req, hasInputReferenceFile); 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 { kind: "submit", model: ctx.model, diff --git a/plugins/tasks/jimeng/plugin.js b/plugins/tasks/jimeng/plugin.js index 5aaa9ff874..f7714b1670 100644 --- a/plugins/tasks/jimeng/plugin.js +++ b/plugins/tasks/jimeng/plugin.js @@ -533,7 +533,10 @@ protocols.openai_video = { } const seconds = req.seconds === undefined ? req.duration : req.seconds; 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 { kind: "submit", diff --git a/plugins/tasks/kling/plugin.js b/plugins/tasks/kling/plugin.js index 74a9c11d8e..c61ca925d2 100644 --- a/plugins/tasks/kling/plugin.js +++ b/plugins/tasks/kling/plugin.js @@ -357,7 +357,7 @@ export const protocols = { if (!prompt && images.length === 0) throw new Error("input is required"); const metadata = Object.assign({}, req.metadata || {}); 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]; const requestBody = { model: model, prompt: prompt, metadata: metadata }; if (images.length) requestBody.image = images[0]; @@ -443,7 +443,7 @@ export const protocols = { const image = trimmed(req.input_reference || req.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 || {}; req.mode = resolveKlingMode(model, req.mode || metadata.mode); const hasImage = hasKlingImage(req, hasInputReferenceFile); diff --git a/plugins/tasks/sunoapi/plugin.js b/plugins/tasks/sunoapi/plugin.js index 240009d5e7..3749b2c08f 100644 --- a/plugins/tasks/sunoapi/plugin.js +++ b/plugins/tasks/sunoapi/plugin.js @@ -128,7 +128,7 @@ export function parseSubmitResponse(ctx, resp) { export function extractUsage(ctx) { 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"); return { clips: action === "lyrics" ? 1 : 2, action: action }; } diff --git a/plugins/tasks/vidu/plugin.js b/plugins/tasks/vidu/plugin.js index 99fffbdf18..3b529f66a5 100644 --- a/plugins/tasks/vidu/plugin.js +++ b/plugins/tasks/vidu/plugin.js @@ -417,7 +417,7 @@ protocols.openai_video = { if (req.seconds !== undefined) req.seconds = Number(req.seconds); 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; if (seconds !== undefined) req.duration = Number(seconds); else req.duration = defaultDuration(model); diff --git a/relay/channel/task/jsplugin/adaptor_test.go b/relay/channel/task/jsplugin/adaptor_test.go index f7a48df304..136844f640 100644 --- a/relay/channel/task/jsplugin/adaptor_test.go +++ b/relay/channel/task/jsplugin/adaptor_test.go @@ -19,6 +19,7 @@ import ( pluginruntime "github.com/QuantumNous/new-api/pkg/jsplugin" "github.com/QuantumNous/new-api/relay/channel" relaycommon "github.com/QuantumNous/new-api/relay/common" + "github.com/QuantumNous/new-api/relay/helper" "github.com/QuantumNous/new-api/relaykit/dto" "github.com/QuantumNous/new-api/service" "github.com/gin-gonic/gin" @@ -1154,3 +1155,56 @@ func TestTaskAdaptorBatchBridge(t *testing.T) { assert.Equal(t, "40%", pending.TaskInfo.Progress) 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"]) +} diff --git a/relay/relay_task.go b/relay/relay_task.go index 63270b5def..ed39570606 100644 --- a/relay/relay_task.go +++ b/relay/relay_task.go @@ -212,6 +212,19 @@ func RelayTaskSubmit(c *gin.Context, info *relaycommon.RelayInfo) (*TaskSubmitRe info.PublicTaskID = model.GenerateTaskID() } 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 { return nil, taskErr } @@ -222,19 +235,33 @@ func RelayTaskSubmit(c *gin.Context, info *relaycommon.RelayInfo) (*TaskSubmitRe modelName = service.CoverTaskActionToModelName(platform, info.Action) } - // 2.5 应用渠道的模型映射（与同步任务对齐） - info.OriginModelName = modelName - info.UpstreamModelName = modelName - if err := helper.ModelMappedHelper(c, info, nil); err != nil { - return nil, service.TaskErrorWrapperLocal(err, "model_mapping_failed", http.StatusBadRequest) + if !mappedBeforeValidate { + info.OriginModelName = modelName + info.UpstreamModelName = modelName + if err := helper.ModelMappedHelper(c, info, nil); err != nil { + return nil, service.TaskErrorWrapperLocal(err, "model_mapping_failed", http.StatusBadRequest) + } } // 4. 价格计算：基础模型价格 info.OriginModelName = modelName var priceData types.PriceData var err error - if billing_setting.GetBillingMode(modelName) == billing_setting.BillingModeTieredExpr { - exprStr, exists := billing_setting.GetBillingExpr(modelName) + useTiered := billing_setting.GetBillingMode(modelName) == billing_setting.BillingModeTieredExpr + var exprStr string + var exists bool + if useTiered { + exprStr, exists = billing_setting.GetBillingExpr(modelName) + } else if info.IsModelMapped { + if billing_setting.GetBillingMode(info.UpstreamModelName) == billing_setting.BillingModeTieredExpr { + if tailExpr, tailOK := billing_setting.GetBillingExpr(info.UpstreamModelName); tailOK && strings.TrimSpace(tailExpr) != "" { + exprStr = tailExpr + exists = true + useTiered = true + } + } + } + if useTiered { provider, supported := adaptor.(channel.TaskUsageFactsProvider) if !exists || !supported { return nil, service.TaskErrorWrapper(fmt.Errorf("task model %s has no usage expression or meter", modelName), "model_price_error", http.StatusBadRequest) diff --git a/relay/relay_task_test.go b/relay/relay_task_test.go index 39e0fd1945..8bd046d7a4 100644 --- a/relay/relay_task_test.go +++ b/relay/relay_task_test.go @@ -1,11 +1,22 @@ package relay import ( + "net/http" + "net/http/httptest" "testing" + "github.com/QuantumNous/new-api/common" "github.com/QuantumNous/new-api/constant" "github.com/QuantumNous/new-api/model" + "github.com/QuantumNous/new-api/pkg/billingexpr" + pluginruntime "github.com/QuantumNous/new-api/pkg/jsplugin" + 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/require" ) func TestTaskModel2DtoNormalizesLegacyAction(t *testing.T) { @@ -16,3 +27,206 @@ func TestTaskModel2DtoNormalizesLegacyAction(t *testing.T) { 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() + 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) + } + }) + } +}