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