feat(task): replace built-in task adaptors with a sandboxed JS plugin system (#7076)

setting/billing_setting/tiered_billing.go

@@ -2,8 +2,14 @@ package billing_setting
 
 import (
 	"fmt"
+	"math"
+	"sort"
 
+	"github.com/QuantumNous/new-api/common"
+	"github.com/QuantumNous/new-api/relaykit/dto"
 	"github.com/QuantumNous/new-api/pkg/billingexpr"
+	"github.com/QuantumNous/new-api/pkg/jsplugin"
+	relaycommon "github.com/QuantumNous/new-api/relay/common"
 	"github.com/QuantumNous/new-api/setting/config"
 	"github.com/samber/lo"
 )
@@ -13,6 +19,7 @@ const (
 	BillingModeTieredExpr = "tiered_expr"
 	BillingModeField      = "billing_mode"
 	BillingExprField      = "billing_expr"
+	maxTaskExprSmokeTests = 64
 )
 
 // BillingSetting is managed by config.GlobalConfig.Register.
@@ -75,13 +82,167 @@ func SmokeTestExpr(exprStr string) error {
 }
 
 func smokeTestExpr(exprStr string) error {
+	if _, err := billingexpr.CompileFromCache(exprStr); err != nil {
+		return err
+	}
+	usageKeys := billingexpr.UsedUsageKeys(exprStr)
+	if len(usageKeys) > 0 {
+		sortedKeys := make([]string, 0, len(usageKeys))
+		for key := range usageKeys {
+			sortedKeys = append(sortedKeys, key)
+		}
+		sort.Strings(sortedKeys)
+		return fmt.Errorf("expression references usage keys %v but the model has no task plugin usage schema", sortedKeys)
+	}
+
 	vectors := []billingexpr.TokenParams{
 		{P: 0, C: 0, Len: 0},
 		{P: 1000, C: 1000, Len: 1000},
 		{P: 100000, C: 100000, Len: 100000},
 		{P: 1000000, C: 1000000, Len: 1000000},
 	}
-	requests := []billingexpr.RequestInput{
+
+	for _, v := range vectors {
+		for _, request := range billingExprSmokeRequests() {
+			result, _, err := billingexpr.RunExprWithRequest(exprStr, v, request)
+			if err != nil {
+				return fmt.Errorf("vector {p=%g, c=%g}: run failed: %w", v.P, v.C, err)
+			}
+			if math.IsNaN(result) || math.IsInf(result, 0) || result < 0 {
+				return fmt.Errorf("vector {p=%g, c=%g}: result must be finite and non-negative, got %f", v.P, v.C, result)
+			}
+		}
+	}
+	return nil
+}
+
+// SmokeTestTaskExpr validates a task usage expression against the usage facts
+// declared by its plugin. Literal u() keys must be declared; dynamic calls are
+// still exercised by the generated runtime vectors when possible.
+func SmokeTestTaskExpr(exprStr string, schema map[string]jsplugin.UsageFieldSchema) error {
+	if _, err := billingexpr.CompileFromCache(exprStr); err != nil {
+		return err
+	}
+	for key := range billingexpr.UsedUsageKeys(exprStr) {
+		if _, declared := schema[key]; !declared {
+			return fmt.Errorf("usage key %q is not declared by the task plugin", key)
+		}
+	}
+
+	for _, usage := range taskUsageSmokeVectors(schema) {
+		for _, request := range billingExprSmokeRequests() {
+			request.Usage = usage
+			result, _, err := billingexpr.RunExprWithRequest(exprStr, billingexpr.TokenParams{}, request)
+			if err != nil {
+				return fmt.Errorf("usage vector %v: run failed: %w", usage, err)
+			}
+			if math.IsNaN(result) || math.IsInf(result, 0) || result < 0 {
+				return fmt.Errorf("usage vector %v: result must be finite and non-negative, got %f", usage, result)
+			}
+		}
+	}
+	return nil
+}
+
+type usageSmokeDimension struct {
+	name   string
+	values []any
+}
+
+func taskUsageSmokeVectors(schema map[string]jsplugin.UsageFieldSchema) []map[string]any {
+	names := make([]string, 0, len(schema))
+	for name := range schema {
+		names = append(names, name)
+	}
+	sort.Strings(names)
+
+	dimensions := make([]usageSmokeDimension, 0, len(names))
+	for _, name := range names {
+		field := schema[name]
+		if len(field.Enum) > 0 {
+			values := make([]any, len(field.Enum))
+			for index, value := range field.Enum {
+				values[index] = value
+			}
+			dimensions = append(dimensions, usageSmokeDimension{name: name, values: values})
+			continue
+		}
+		if field.Type == "boolean" {
+			dimensions = append(dimensions, usageSmokeDimension{name: name, values: []any{false, true}})
+			continue
+		}
+		limit := relaycommon.MaxTaskDurationSeconds
+		if field.Unit == "count" {
+			limit = dto.MaxImageN
+		}
+		if field.Unit == "token" || field.Unit == "credit" {
+			limit = common.MaxQuota
+		}
+		dimensions = append(dimensions, usageSmokeDimension{
+			name:   name,
+			values: []any{float64(0), float64(1), float64(limit)},
+		})
+	}
+
+	if usageSmokeCombinationCount(dimensions, maxTaskExprSmokeTests) > maxTaskExprSmokeTests {
+		for index := range dimensions {
+			field := schema[dimensions[index].name]
+			if len(field.Enum) <= 2 {
+				continue
+			}
+			dimensions[index].values = []any{field.Enum[0], field.Enum[len(field.Enum)-1]}
+		}
+	}
+
+	vectors := make([]map[string]any, 0, maxTaskExprSmokeTests)
+	var appendVectors func(int, map[string]any)
+	appendVectors = func(index int, current map[string]any) {
+		if len(vectors) >= maxTaskExprSmokeTests {
+			return
+		}
+		if index == len(dimensions) {
+			vector := make(map[string]any, len(current))
+			for key, value := range current {
+				vector[key] = value
+			}
+			vectors = append(vectors, vector)
+			return
+		}
+		for _, value := range dimensions[index].values {
+			current[dimensions[index].name] = value
+			appendVectors(index+1, current)
+		}
+		delete(current, dimensions[index].name)
+	}
+	appendVectors(0, make(map[string]any, len(dimensions)))
+
+	combinationCount := usageSmokeCombinationCount(dimensions, maxTaskExprSmokeTests)
+	if combinationCount > maxTaskExprSmokeTests && len(vectors) > 0 {
+		last := make(map[string]any, len(dimensions))
+		for _, dimension := range dimensions {
+			last[dimension.name] = dimension.values[len(dimension.values)-1]
+		}
+		vectors[len(vectors)-1] = last
+	}
+	return vectors
+}
+
+func usageSmokeCombinationCount(dimensions []usageSmokeDimension, stopAfter int) int {
+	count := 1
+	for _, dimension := range dimensions {
+		if len(dimension.values) == 0 {
+			return 0
+		}
+		if count > stopAfter/len(dimension.values) {
+			return stopAfter + 1
+		}
+		count *= len(dimension.values)
+	}
+	return count
+}
+
+func billingExprSmokeRequests() []billingexpr.RequestInput {
+	return []billingexpr.RequestInput{
 		{},
 		{
 			Headers: map[string]string{
@@ -90,17 +251,4 @@ func smokeTestExpr(exprStr string) error {
 			Body: []byte(`{"service_tier":"fast","stream_options":{"include_usage":true},"messages":[1,2,3,4,5,6,7,8,9,10,11,12,13,14,15,16,17,18,19,20,21]}`),
 		},
 	}
-
-	for _, v := range vectors {
-		for _, request := range requests {
-			result, _, err := billingexpr.RunExprWithRequest(exprStr, v, request)
-			if err != nil {
-				return fmt.Errorf("vector {p=%g, c=%g}: run failed: %w", v.P, v.C, err)
-			}
-			if result < 0 {
-				return fmt.Errorf("vector {p=%g, c=%g}: result %f < 0", v.P, v.C, result)
-			}
-		}
-	}
-	return nil
 }

setting/billing_setting/tiered_billing_test.go

@@ -0,0 +1,105 @@
+package billing_setting
+
+import (
+	"fmt"
+	"testing"
+
+	"github.com/QuantumNous/new-api/common"
+	"github.com/QuantumNous/new-api/relaykit/dto"
+	"github.com/QuantumNous/new-api/pkg/jsplugin"
+	relaycommon "github.com/QuantumNous/new-api/relay/common"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+)
+
+func TestSmokeTestTaskExprValidatesDeclaredUsageVectors(t *testing.T) {
+	videoSchema := map[string]jsplugin.UsageFieldSchema{
+		"seconds": {Type: "number", Unit: "second"},
+		"mode":    {Enum: []string{"std", "pro"}},
+		"quality": {Enum: []string{"sd", "hd"}},
+	}
+
+	tests := []struct {
+		name          string
+		schema        map[string]jsplugin.UsageFieldSchema
+		expression    string
+		expectedError string
+	}{
+		{
+			name:       "declared numeric and enum facts",
+			schema:     videoSchema,
+			expression: `u("mode") == "pro" ? tier("pro", u("seconds") * 0.8) : tier("std", u("seconds") * 0.4)`,
+		},
+		{
+			name:          "undeclared literal key",
+			schema:        videoSchema,
+			expression:    `tier("base", u("clips") * 0.1)`,
+			expectedError: `usage key "clips" is not declared`,
+		},
+		{
+			name:          "negative duration boundary",
+			schema:        videoSchema,
+			expression:    fmt.Sprintf(`u("seconds") == %d ? -1 : 0`, relaycommon.MaxTaskDurationSeconds),
+			expectedError: "result must be finite and non-negative",
+		},
+		{
+			name:          "negative count boundary",
+			schema:        map[string]jsplugin.UsageFieldSchema{"clips": {Type: "number", Unit: "count"}},
+			expression:    fmt.Sprintf(`u("clips") == %d ? -1 : 0`, dto.MaxImageN),
+			expectedError: "result must be finite and non-negative",
+		},
+		{
+			name:          "negative token boundary",
+			schema:        map[string]jsplugin.UsageFieldSchema{"tokens": {Type: "number", Unit: "token"}},
+			expression:    fmt.Sprintf(`u("tokens") == %d ? -1 : 0`, common.MaxQuota),
+			expectedError: "result must be finite and non-negative",
+		},
+		{
+			name:          "negative credit boundary",
+			schema:        map[string]jsplugin.UsageFieldSchema{"units": {Type: "number", Unit: "credit"}},
+			expression:    fmt.Sprintf(`u("units") == %d ? -1 : 0`, common.MaxQuota),
+			expectedError: "result must be finite and non-negative",
+		},
+		{
+			name:          "negative enum combination",
+			schema:        videoSchema,
+			expression:    `u("mode") == "pro" && u("quality") == "hd" ? -1 : 0`,
+			expectedError: "result must be finite and non-negative",
+		},
+	}
+
+	for _, testCase := range tests {
+		t.Run(testCase.name, func(t *testing.T) {
+			err := SmokeTestTaskExpr(testCase.expression, testCase.schema)
+			if testCase.expectedError == "" {
+				require.NoError(t, err)
+				return
+			}
+			require.ErrorContains(t, err, testCase.expectedError)
+		})
+	}
+}
+
+func TestSmokeTestTaskExprCapsOversizedEnumProductsAtLastCombination(t *testing.T) {
+	schema := make(map[string]jsplugin.UsageFieldSchema, 7)
+	condition := ""
+	for index := 0; index < 7; index++ {
+		schema[fmt.Sprintf("enum_%d", index)] = jsplugin.UsageFieldSchema{Enum: []string{"first", "middle", "last"}}
+		if condition != "" {
+			condition += " && "
+		}
+		condition += fmt.Sprintf(`u("enum_%d") == "last"`, index)
+	}
+
+	err := SmokeTestTaskExpr(condition+" ? -1 : 0", schema)
+	require.ErrorContains(t, err, "result must be finite and non-negative")
+}
+
+func TestSmokeTestExprRejectsTaskUsageWithoutSchema(t *testing.T) {
+	err := SmokeTestExpr(`u("mode") == "std" ? 1 : 2`)
+	require.Error(t, err)
+	assert.ErrorContains(t, err, "mode")
+	assert.ErrorContains(t, err, "no task plugin usage schema")
+
+	require.NoError(t, SmokeTestExpr(`tier("base", p * 2 + c * 8)`))
+}