feat: sync deployed v1.0.0-rc.38-custom.42
Sync production source including routing, billing, Astra detection and public uptime. Update custom build instructions and license packaging.
This commit is contained in:
cneicy committed
1 parent
aa6c3f50f4
commit
ceddfd3d6b
828 files changed
+122253
-13702
No files matched your search
@@ -2,6 +2,7 @@ package billingexpr_test
|
||||
|
||||
import (
|
||||
"math"
|
||||
"os"
|
||||
"testing"
|
||||
|
||||
"github.com/QuantumNous/new-api/common"
|
||||
@@ -10,6 +11,94 @@ import (
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
func TestFixedPriceBranches(t *testing.T) {
|
||||
const expression = `(len <= 32000 ? tier("short", fixed(0.01)) : tier("long", p * 2 + c * 8)) * (param("fast") == true ? 2 : 1)`
|
||||
for _, tc := range []struct {
|
||||
name string
|
||||
params billingexpr.TokenParams
|
||||
cost float64
|
||||
tier string
|
||||
}{
|
||||
{"zero usage still charges once", billingexpr.TokenParams{}, 20000, "short"},
|
||||
{"fixed branch ignores token prices", billingexpr.TokenParams{P: 32000, C: 9000, Len: 32000}, 20000, "short"},
|
||||
{"token branch uses actual usage", billingexpr.TokenParams{P: 32001, C: 100, Len: 32001}, 129604, "long"},
|
||||
} {
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
cost, trace, err := billingexpr.RunExprWithRequest(expression, tc.params, billingexpr.RequestInput{Body: []byte(`{"fast":true}`)})
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, tc.cost, cost)
|
||||
assert.Equal(t, tc.tier, trace.MatchedTier)
|
||||
require.Len(t, trace.RequestRules, 1)
|
||||
assert.True(t, trace.RequestRules[0].Matched)
|
||||
})
|
||||
}
|
||||
cost, _, err := billingexpr.RunExpr(`tier("free", fixed(0))`, billingexpr.TokenParams{P: 1000})
|
||||
require.NoError(t, err)
|
||||
assert.Zero(t, cost)
|
||||
for _, count := range []int{-1, 0, 129} {
|
||||
_, _, err := billingexpr.RunExprWithRequest(`tier("image", fixed(0.04)) * image_count`, billingexpr.TokenParams{}, billingexpr.RequestInput{ImageCount: &count})
|
||||
require.ErrorContains(t, err, "image_count")
|
||||
}
|
||||
}
|
||||
|
||||
func TestFixedPriceRejectsInvalidLeavesIncludingUnselectedBranches(t *testing.T) {
|
||||
for _, expression := range []string{
|
||||
`true ? tier("ok", p) : tier("bad", fixed(-0.01))`,
|
||||
`tier("bad", fixed(p))`,
|
||||
`tier("bad", fixed(0.01 + 0.02))`,
|
||||
`tier("bad", fixed(1e308))`,
|
||||
`tier("bad", fixed(0.01) + p * 2)`,
|
||||
`tier("bad", p * fixed(0.01))`,
|
||||
`tier("a", fixed(0.01)) + tier("b", p * 2)`,
|
||||
`fixed(0.01)`,
|
||||
`tier("bad", fixed(0.01)) * p`,
|
||||
`tier("bad", fixed(0.01)) * (param("fast") == true ? -2 : 1)`,
|
||||
} {
|
||||
t.Run(expression, func(t *testing.T) {
|
||||
_, err := billingexpr.CompileFromCache(expression)
|
||||
require.Error(t, err)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestFixedPriceQuotaBoundariesAndUnsupportedTaskSnapshots(t *testing.T) {
|
||||
for _, tc := range []struct {
|
||||
name, expression string
|
||||
group float64
|
||||
want int
|
||||
clamped bool
|
||||
}{
|
||||
{"normal group", `tier("request", fixed(0.01))`, 1.5, 7500, false},
|
||||
{"free group", `tier("request", fixed(0.01))`, 0, 0, false},
|
||||
{"explicit free price", `tier("free", fixed(0))`, 1, 0, false},
|
||||
{"oversized price saturates instead of crediting", `tier("huge", fixed(1e20))`, 1, math.MaxInt32, true},
|
||||
} {
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
snap := &billingexpr.BillingSnapshot{ExprString: tc.expression, ExprHash: billingexpr.ExprHashString(tc.expression), GroupRatio: tc.group, QuotaPerUnit: 500000}
|
||||
result, err := billingexpr.ComputeTieredQuota(snap, billingexpr.TokenParams{})
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, tc.want, result.ActualQuotaAfterGroup)
|
||||
assert.Equal(t, billingexpr.BillingUnitRequest, result.BillingUnit)
|
||||
assert.Equal(t, tc.clamped, result.Clamp != nil)
|
||||
_, err = billingexpr.QuotaRoundStrict(result.ActualQuotaBeforeGroup * tc.group)
|
||||
if tc.clamped {
|
||||
require.Error(t, err)
|
||||
} else {
|
||||
require.NoError(t, err)
|
||||
}
|
||||
snap.TaskUsageBilling = true
|
||||
_, err = billingexpr.ComputeTieredQuota(snap, billingexpr.TokenParams{})
|
||||
require.ErrorContains(t, err, "task usage")
|
||||
})
|
||||
}
|
||||
const optimized = `true ? tier("tokens", p) : tier("request", fixed(0.01))`
|
||||
assert.True(t, billingexpr.UsesFixedPricing(optimized))
|
||||
_, trace, err := billingexpr.RunExpr(optimized, billingexpr.TokenParams{P: 10})
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, billingexpr.BillingUnitToken, trace.BillingUnit)
|
||||
assert.Nil(t, trace.FixedPrice)
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Claude-style: fixed tiers, input > 200K changes both input & output price
|
||||
// ---------------------------------------------------------------------------
|
||||
@@ -1063,3 +1152,47 @@ func BenchmarkExprRunCached(b *testing.B) {
|
||||
billingexpr.RunExpr(benchComplexExpr, params)
|
||||
}
|
||||
}
|
||||
|
||||
// Shared fixtures protect the frontend simulator against drift from the real engine.
|
||||
func TestFrontendSimulationContract(t *testing.T) {
|
||||
data, err := os.ReadFile("testdata/frontend_simulation.json")
|
||||
require.NoError(t, err)
|
||||
var cases []struct {
|
||||
Name string
|
||||
Expression string
|
||||
Tokens billingexpr.TokenParams
|
||||
Body map[string]any
|
||||
Headers map[string]string
|
||||
Usage map[string]any
|
||||
Cost float64
|
||||
Tier string
|
||||
Multipliers []float64
|
||||
Matched []bool
|
||||
BillingUnit billingexpr.BillingUnit
|
||||
FixedPrice *float64
|
||||
ImageCount *int
|
||||
}
|
||||
require.NoError(t, common.Unmarshal(data, &cases))
|
||||
for _, tc := range cases {
|
||||
t.Run(tc.Name, func(t *testing.T) {
|
||||
body, err := common.Marshal(tc.Body)
|
||||
require.NoError(t, err)
|
||||
cost, trace, err := billingexpr.RunExprWithRequest(tc.Expression, tc.Tokens, billingexpr.RequestInput{Body: body, Headers: tc.Headers, Usage: tc.Usage, ImageCount: tc.ImageCount})
|
||||
require.NoError(t, err)
|
||||
assert.InDelta(t, tc.Cost, cost, 1e-9)
|
||||
assert.Equal(t, tc.Tier, trace.MatchedTier)
|
||||
if tc.ImageCount != nil {
|
||||
assert.Equal(t, tc.ImageCount, trace.ImageCount)
|
||||
}
|
||||
if tc.BillingUnit != "" {
|
||||
assert.Equal(t, tc.BillingUnit, trace.BillingUnit)
|
||||
assert.Equal(t, tc.FixedPrice, trace.FixedPrice)
|
||||
}
|
||||
require.Len(t, trace.RequestRules, len(tc.Multipliers))
|
||||
for i, rule := range trace.RequestRules {
|
||||
assert.Equal(t, tc.Multipliers[i], rule.Multiplier)
|
||||
assert.Equal(t, tc.Matched[i], rule.Matched)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
+53
-28
@@ -8,6 +8,7 @@ import (
|
||||
|
||||
"github.com/expr-lang/expr"
|
||||
"github.com/expr-lang/expr/ast"
|
||||
"github.com/expr-lang/expr/parser"
|
||||
"github.com/expr-lang/expr/vm"
|
||||
)
|
||||
|
||||
@@ -114,6 +115,7 @@ type cachedEntry struct {
|
||||
usedUsageKeys map[string]bool
|
||||
requestRules []RequestRuleTrace
|
||||
version int
|
||||
fixedPricing bool
|
||||
}
|
||||
|
||||
var (
|
||||
@@ -123,33 +125,36 @@ var (
|
||||
|
||||
// compileEnvPrototypeV1 is the v1 type-checking prototype used at compile time.
|
||||
var compileEnvPrototypeV1 = map[string]any{
|
||||
"p": float64(0),
|
||||
"c": float64(0),
|
||||
"len": float64(0),
|
||||
"cr": float64(0),
|
||||
"cc": float64(0),
|
||||
"cc1h": float64(0),
|
||||
"img": float64(0),
|
||||
"img_o": float64(0),
|
||||
"ai": float64(0),
|
||||
"ao": float64(0),
|
||||
"tier": func(string, float64) float64 { return 0 },
|
||||
"_trace": func(int, bool, float64) float64 { return 1 },
|
||||
"_trace_int": func(int, bool, int) int { return 1 },
|
||||
"header": func(string) string { return "" },
|
||||
"param": func(string) any { return nil },
|
||||
"u": func(string) any { return nil },
|
||||
"has": func(any, string) bool { return false },
|
||||
"hour": func(string) int { return 0 },
|
||||
"minute": func(string) int { return 0 },
|
||||
"weekday": func(string) int { return 0 },
|
||||
"month": func(string) int { return 0 },
|
||||
"day": func(string) int { return 0 },
|
||||
"max": math.Max,
|
||||
"min": math.Min,
|
||||
"abs": math.Abs,
|
||||
"ceil": math.Ceil,
|
||||
"floor": math.Floor,
|
||||
"image_count": float64(1),
|
||||
"p": float64(0),
|
||||
"c": float64(0),
|
||||
"len": float64(0),
|
||||
"cr": float64(0),
|
||||
"cc": float64(0),
|
||||
"cc1h": float64(0),
|
||||
"img": float64(0),
|
||||
"img_cr": float64(0),
|
||||
"img_o": float64(0),
|
||||
"ai": float64(0),
|
||||
"ao": float64(0),
|
||||
"tier": func(string, float64) float64 { return 0 },
|
||||
"fixed": func(float64) float64 { return 0 },
|
||||
"_trace": func(int, bool, float64) float64 { return 1 },
|
||||
"_trace_int": func(int, bool, int) int { return 1 },
|
||||
"header": func(string) string { return "" },
|
||||
"param": func(string) any { return nil },
|
||||
"u": func(string) any { return nil },
|
||||
"has": func(any, string) bool { return false },
|
||||
"hour": func(string) int { return 0 },
|
||||
"minute": func(string) int { return 0 },
|
||||
"weekday": func(string) int { return 0 },
|
||||
"month": func(string) int { return 0 },
|
||||
"day": func(string) int { return 0 },
|
||||
"max": math.Max,
|
||||
"min": math.Min,
|
||||
"abs": math.Abs,
|
||||
"ceil": math.Ceil,
|
||||
"floor": math.Floor,
|
||||
}
|
||||
|
||||
func getCompileEnv(version int) map[string]any {
|
||||
@@ -188,6 +193,21 @@ func compileEntryFromCacheByHash(exprStr, hash string) (*cachedEntry, error) {
|
||||
cacheMu.RUnlock()
|
||||
|
||||
version, body := ParseExprVersion(exprStr)
|
||||
// Validate before optimization so unreachable fixed-price branches cannot
|
||||
// bypass validation or a host's unsupported-protocol checks.
|
||||
tree, err := parser.Parse(body)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("expr compile error: %w", err)
|
||||
}
|
||||
fixedPricing := ast.Find(tree.Node, func(node ast.Node) bool {
|
||||
identifier, ok := node.(*ast.IdentifierNode)
|
||||
return ok && identifier.Value == "fixed"
|
||||
}) != nil
|
||||
if fixedPricing {
|
||||
if err := validateFixedPricingTree(tree.Node); err != nil {
|
||||
return nil, fmt.Errorf("expr compile error: %w", err)
|
||||
}
|
||||
}
|
||||
patcher := &requestRulePatcher{}
|
||||
prog, err := expr.Compile(body, expr.Env(getCompileEnv(version)), expr.Patch(patcher), expr.AsFloat64())
|
||||
if patcher.restrictedIdentifier != "" {
|
||||
@@ -203,6 +223,7 @@ func compileEntryFromCacheByHash(exprStr, hash string) (*cachedEntry, error) {
|
||||
usedUsageKeys: extractUsedUsageKeys(prog),
|
||||
requestRules: patcher.requestRules,
|
||||
version: version,
|
||||
fixedPricing: fixedPricing,
|
||||
}
|
||||
cacheMu.Lock()
|
||||
if len(cache) >= maxCacheSize {
|
||||
@@ -271,10 +292,14 @@ func extractUsedUsageKeys(prog *vm.Program) map[string]bool {
|
||||
// UsedVars returns the set of identifier names referenced by an expression.
|
||||
// The result is cached alongside the compiled program. Returns nil for empty input.
|
||||
func UsedVars(exprStr string) map[string]bool {
|
||||
return UsedVarsByHash(exprStr, ExprHashString(exprStr))
|
||||
}
|
||||
|
||||
// UsedVarsByHash reuses the digest captured by the host's billing snapshot.
|
||||
func UsedVarsByHash(exprStr, hash string) map[string]bool {
|
||||
if exprStr == "" {
|
||||
return nil
|
||||
}
|
||||
hash := ExprHashString(exprStr)
|
||||
cacheMu.RLock()
|
||||
if entry, ok := cache[hash]; ok {
|
||||
cacheMu.RUnlock()
|
||||
|
||||
@@ -0,0 +1,111 @@
|
||||
package billingexpr
|
||||
|
||||
import (
|
||||
"errors"
|
||||
"math"
|
||||
|
||||
"github.com/expr-lang/expr/ast"
|
||||
"github.com/expr-lang/expr/parser"
|
||||
)
|
||||
|
||||
// DetectionCostBound proves a conservative cost ceiling rather than sampling
|
||||
// a few token counts. Both conditional branches are bounded, including prices
|
||||
// selected by time/request rules. Unsupported expressions fail closed.
|
||||
func DetectionCostBound(expression string, input, output float64) (float64, error) {
|
||||
_, body := ParseExprVersion(expression)
|
||||
tree, err := parser.Parse(body)
|
||||
if err != nil {
|
||||
return 0, err
|
||||
}
|
||||
var bound func(ast.Node) (float64, error)
|
||||
unsupported := errors.New("detection pricing cannot be safely bounded")
|
||||
bound = func(node ast.Node) (float64, error) {
|
||||
switch n := node.(type) {
|
||||
case *ast.IntegerNode:
|
||||
if n.Value >= 0 {
|
||||
return float64(n.Value), nil
|
||||
}
|
||||
case *ast.FloatNode:
|
||||
if n.Value >= 0 {
|
||||
return n.Value, nil
|
||||
}
|
||||
case *ast.IdentifierNode:
|
||||
switch n.Value {
|
||||
case "p", "len", "cr", "cc", "cc1h":
|
||||
return input, nil
|
||||
case "c":
|
||||
return output, nil
|
||||
case "img", "img_cr", "img_o", "ai", "ao":
|
||||
return 0, nil
|
||||
case "image_count":
|
||||
return 1, nil
|
||||
}
|
||||
case *ast.BinaryNode:
|
||||
a, e := bound(n.Left)
|
||||
if e != nil {
|
||||
return 0, e
|
||||
}
|
||||
b, e := bound(n.Right)
|
||||
if e != nil {
|
||||
return 0, e
|
||||
}
|
||||
switch n.Operator {
|
||||
case "+":
|
||||
return a + b, nil
|
||||
case "*":
|
||||
return a * b, nil
|
||||
case "/":
|
||||
switch n.Right.(type) {
|
||||
case *ast.IntegerNode, *ast.FloatNode:
|
||||
if b > 0 {
|
||||
return a / b, nil
|
||||
}
|
||||
}
|
||||
}
|
||||
case *ast.ConditionalNode:
|
||||
a, e := bound(n.Exp1)
|
||||
if e != nil {
|
||||
return 0, e
|
||||
}
|
||||
b, e := bound(n.Exp2)
|
||||
return max(a, b), e
|
||||
case *ast.CallNode:
|
||||
name, ok := n.Callee.(*ast.IdentifierNode)
|
||||
if !ok {
|
||||
return 0, unsupported
|
||||
}
|
||||
if name.Value == "tier" && len(n.Arguments) == 2 {
|
||||
return bound(n.Arguments[1])
|
||||
}
|
||||
if name.Value == "fixed" && len(n.Arguments) == 1 {
|
||||
v, e := bound(n.Arguments[0])
|
||||
return v * 1_000_000, e
|
||||
}
|
||||
if name.Value == "ceil" && len(n.Arguments) == 1 {
|
||||
v, e := bound(n.Arguments[0])
|
||||
return math.Ceil(v), e
|
||||
}
|
||||
if name.Value == "floor" && len(n.Arguments) == 1 {
|
||||
v, e := bound(n.Arguments[0])
|
||||
return math.Floor(v), e
|
||||
}
|
||||
if name.Value == "min" || name.Value == "max" {
|
||||
v := float64(0)
|
||||
for _, arg := range n.Arguments {
|
||||
x, e := bound(arg)
|
||||
if e != nil {
|
||||
return 0, e
|
||||
}
|
||||
v = max(v, x)
|
||||
}
|
||||
return v, nil
|
||||
}
|
||||
}
|
||||
return 0, unsupported
|
||||
}
|
||||
value, err := bound(tree.Node)
|
||||
if math.IsNaN(value) || math.IsInf(value, 0) {
|
||||
return 0, unsupported
|
||||
}
|
||||
return value, err
|
||||
}
|
||||
+193
-3
@@ -12,7 +12,7 @@ The expression is the billing contract between the administrator and the system.
|
||||
|
||||
2. **Variables are opt-in** — `p` (prompt) and `c` (completion) are the base. Cache (`cr`, `cc`, `cc1h`), image (`img`), and audio (`ai`, `ao`) variables are optional. If omitted, those tokens are included in `p`/`c` and priced at their rate. The system automatically detects which variables the expression uses (via AST introspection) and adjusts token normalization accordingly.
|
||||
|
||||
3. **Prices are real prices** — Expression coefficients are actual $/1M tokens prices as published by providers. No ratio conversion, no `/2` convention. `p * 2.5` means $2.50 per 1M prompt tokens.
|
||||
3. **Prices are real prices** — Token coefficients are actual $/1M tokens prices as published by providers. `p * 2.5` means $2.50 per 1M prompt tokens; `fixed(0.01)` means $0.01 per request. No ratio conversion or `/2` convention is required.
|
||||
|
||||
4. **Upstream-agnostic** — The expression doesn't need to know whether the upstream API is OpenAI-format (prompt_tokens includes cache) or Claude-format (input_tokens excludes cache). The system normalizes token counts before evaluation based on the upstream response format.
|
||||
|
||||
@@ -36,6 +36,7 @@ Powered by [expr-lang/expr](https://github.com/expr-lang/expr). Expressions are
|
||||
| `cc` | 缓存创建 token 数（Claude 5分钟 TTL / 通用） |
|
||||
| `cc1h` | 缓存创建 token 数 — 1小时 TTL（Claude 专用） |
|
||||
| `img` | 图片输入 token 数 |
|
||||
| `img_cr` | 图片缓存读取 token 数；只有表达式显式引用且上游提供有效缓存模态明细时才单独拆分 |
|
||||
| `ai` | 音频输入 token 数 |
|
||||
|
||||
**输出侧变量：**
|
||||
@@ -69,13 +70,97 @@ Powered by [expr-lang/expr](https://github.com/expr-lang/expr). Expressions are
|
||||
| `p * 3 + c * 15` | 500 | 没用 `ao`，音频输出包含在 `c` 里按 $15 计费 |
|
||||
| `p * 3 + c * 15 + ao * 50` | 400 | 用了 `ao`，音频 100 从 `c` 中扣除按 $50 计费 |
|
||||
|
||||
> **注意：** 这个自动排除仅针对 GPT/OpenAI 格式的 API（prompt_tokens 包含所有子类别）。Claude 格式的 API（input_tokens 本身就只包含纯文本）不做任何减法。系统根据上游返回格式自动判断，表达式作者无需关心。
|
||||
#### 图片缓存与兼容性
|
||||
|
||||
OpenAI 兼容图片接口可报告
|
||||
`usage.input_tokens_details.cached_tokens_details.{text_tokens,image_tokens,audio_tokens}`。
|
||||
统一 DTO 保留各字段的缺失与显式零值；复制、用量合并及计费转换不得丢失这些信息。
|
||||
当前公开 Images API 文档未明确缓存模态字段，兼容结构的合成测试不能替代真实上游缓存回包验证。
|
||||
|
||||
当表达式使用 `img_cr` 且明细有效时，`img_cr` 从 `cr` 和 `img` 中同时移出，
|
||||
再按表达式引用的变量从 `p` 中扣除各互不重叠的类别。`len` 保持原始输入总长度。
|
||||
例如总输入 1000、图片输入 600、缓存总量 300、图片缓存 200：
|
||||
`p * 5 + cr * 1.25 + img * 8 + img_cr * 2 + c * 30` 对应
|
||||
`p=300, cr=100, img=400, img_cr=200`；输出 100 时费用为 $0.008225。
|
||||
|
||||
缺少图片缓存明细时，`img_cr=0`，缓存总量继续进入 `cr`，沿用原有通用缓存计费，
|
||||
不推测图片缓存比例。不引用 `img_cr` 的既有表达式保持原有归一化行为。
|
||||
负数、超过总量或无法形成有效输入集合的明细会记录诊断并回退原有逻辑。
|
||||
Anthropic 的独立缓存计数语义不采用这一 OpenAI 图片拆分。
|
||||
|
||||
图片缓存表达式的消费日志保存 `image_cache_tokens` 与 `billing_tokens`，后者由结算结果携带，
|
||||
通过公共日志注入记录实际参与结算的计价用量，覆盖文本、audio 和 realtime 入口。
|
||||
求值失败回退预扣时不伪造实际用量；固定价格和未引用 `img_cr` 的表达式不新增这些字段。
|
||||
文本/图片日志的 `cache_tokens` 继续保留上游缓存总量。Images 的 `output_tokens` 是图片输出，
|
||||
内置图片价格使用 `c`，无需上游额外报告 `output_tokens_details.image_tokens`。
|
||||
|
||||
### 旧定价转换
|
||||
|
||||
旧倍率和旧按次模式仍可运行，管理界面标记为弃用，新建定价默认使用表达式。
|
||||
`POST /api/option/model_pricing/convert` 接收 `{model_name, pricing}`，其中 `pricing`
|
||||
为完整的旧定价草稿，缺少的键继承运行时默认值。接口返回表达式及本次转换使用的生效定价快照，或具体不支持原因，不写数据库。
|
||||
管理界面先用该快照显示左右对照预览，确认后才更新草稿；取消不改变原草稿。
|
||||
保存仍使用现有 `PATCH /api/option/model_pricing` 及 `expected_version` 冲突检查。
|
||||
|
||||
倍率转换以 `ModelRatio * 1000000 / QuotaPerUnit` 得到 USD/百万 tokens 的基础单价，
|
||||
保留输出、缓存读取、适用的缓存创建和图片输入的生效倍率及显式零值。
|
||||
常规按次价格转换为 `tier("request", fixed(price))`。旧价格字段继续保存，供切回旧模式使用。
|
||||
转换后的预扣及舍入采用表达式规则，不保证旧模式每笔舍入结果完全一致。
|
||||
迁移、已保存价格展示与草稿价格预览共用生效价格解析：输出使用模型族的硬编码兜底/强制倍率；
|
||||
缺少配置的缓存读取、缓存创建、图片输入分别使用运行时的 1、1.25、1 倍。
|
||||
图片倍率为 1 时，转换不生成独立 `img` 项，图片输入继续包含在 `p` 中按相同价格计费；
|
||||
图片倍率为其他值（包括显式 0）时才生成 `img` 项。
|
||||
缓存读取倍率为 1 时在价格预览中合并为输入价格；没有独立图片输入或缓存写入项时，转换同时省略 `cr`。
|
||||
存在这些独立项目时，缓存计数可能交叠，输入余量的零下限使直接合并不再等价，因此保留原 `cr` 计算项而不单列同价价格。
|
||||
基础输入价为 0 时可直接省略倍率为 1 的 `cr`；其他缓存读取倍率（包括显式 0）保留独立项。
|
||||
省略 `cr` 时，Anthropic 用量中单独报告的缓存读取数会加回 `p`，而 `len` 保持完整上下文长度。
|
||||
缓存创建按对外计费名和原始配置处理：名称包含 `claude`（不区分大小写）时保留 `cc` 和 `cc1h`，
|
||||
其中 1h 价格沿用旧引擎的 `CreateCacheRatio * 6 / 3.75`；其他名称仅在草稿已配置 `CreateCacheRatio`
|
||||
时生成通用 `cc`（包括显式 0），不生成 `cc1h`。没有配置的通用 1.25 倍兜底不产生新计费项。
|
||||
不按渠道映射或模型白名单扩展缓存类型，也不要求管理员另选类型。
|
||||
定价快照条目、草稿预览响应和转换结果附带只读 `cache_write_mode`（`none` / `standard` / `claude_ttl`），
|
||||
用于统一展示。`POST /api/option/model_pricing/preview` 的响应统一为
|
||||
`{success, message, data: {effective, cache_write_mode, billing_details}}`，调用者从 `data.effective`
|
||||
读取生效价格；转换接口和快照条目的 JSON 结构不变。元信息不保存、不参与表达式执行。
|
||||
未设置与显式 0 必须区分。草稿关闭某项覆盖后按保存后的回退规则解析，不能继续读取运行进程中尚未被保存替换的旧值。
|
||||
图片和普通音频规则均可自动迁移，具体规则如下；任务插件、视频、Realtime，以及依赖上游 cost
|
||||
反推缓存写入用量的 OpenRouter Claude 分支仍返回具体原因，不按渠道类型或未知端点一概拦截。
|
||||
|
||||
Gemini 音频输入使用旧结算实际采用的美元单价生成 `ai`，不乘普通输入倍率，因此普通输入价为零时
|
||||
仍保留独立音频费用。普通音频的 `ai` 单价为基础输入价乘 AudioRatio，`ao` 再乘 AudioCompletionRatio，
|
||||
缺省倍率按 1 处理，显式零保留。普通音频结算不使用缓存/图片倍率；模型同时配置这些价格时，
|
||||
转换生成音频请求与纯文本请求两个分支。音频分支使用 `max(len - ai, 0)` 计量普通输入，避免纯文本
|
||||
分支中引用的缓存变量改变音频费用。Gemini 文本结算仍保留缓存/音频交叠后的输入余量零下限。
|
||||
只读 `billing_details` 为预览提供实际音频单价、图片数量和请求倍率规则，不成为新的持久化价格来源。
|
||||
同一计费名在 Gemini/兼容音频入口存在相互冲突的单价，或活动图片渠道采用不同请求倍率时，
|
||||
转换返回该具体冲突，不能静默选择其中一条规则覆盖其他入口。
|
||||
|
||||
图片按次价格转换为 `tier("image", fixed(price)) * image_count`，再追加旧尺寸、质量及适用的
|
||||
prompt_extend 条件倍率。DALL·E 尺寸/质量规则按原始请求模型名生成，模型映射不会额外引入这些规则。
|
||||
OpenAI 已于 2026-05-12 下线 DALL·E 2/3；其校验、默认值和倍率保留在独立 legacy 文件，
|
||||
仅用于历史配置及兼容上游。转换从 DTO 的旧计费逻辑获取倍率，不维护第二份价格表。
|
||||
图片 token 定价不额外乘数量。所有金额和倍率固化在表达式中，不再叠加旧 OtherRatios。
|
||||
|
||||
### Image Quantity
|
||||
|
||||
`image_count` 是独立于 token 的计费数量。省略上下文时默认 1；提供的数量必须是 1 到
|
||||
`dto.MaxImageN`（128）的整数。固定价格语法允许它作为乘数，仍禁止用 token 乘固定价格。
|
||||
图片入口只解析一次 provider 计费标量并随请求携带。Ali 遵循 `parameters.n → n → 1` 的优先级：
|
||||
缺失或 `null` 的嵌套数量回退顶层，显式 `parameters.n=0` 返回 400；顶层 `n=0` 仍兼容为 1。
|
||||
Ali JSON 和 multipart 请求都明确发送该生效数量。每次渠道重试或参数覆盖之后，再校验最终上游数量并在发送前补足预留，
|
||||
图片钱包预留使用原子余额检查，不能通过增加数量形成欠费后继续提交。
|
||||
`estimated_image_count` 保存本次发送前的数量；JSON/multipart 图片参数上下文只保留计费需要的标量，不保留图片或提示词内容。
|
||||
结算使用独立的实际数量，不修改被冻结的 `param("n")`。非法 Ali usage 数量记录诊断，并回退有效图片列表数量；缺少有效实际数量时保持发送数量。正常结束
|
||||
的 SSE 可以减少数量，客户端提前断开不能减少应收数量。最终消费日志记录 `image_count`。
|
||||
|
||||
> **注意：** 自动扣除针对 GPT/OpenAI 格式的 API（prompt_tokens 包含子类别）。Claude 格式的 API 不重复扣除缓存；未独立计价的缓存读取加回输入。系统根据上游返回格式自动处理。
|
||||
|
||||
### Built-in Functions
|
||||
|
||||
| Function | Signature | Purpose |
|
||||
|----------|-----------|---------|
|
||||
| `tier` | `tier(name, value) → float64` | Records which pricing tier matched; must wrap the cost expression |
|
||||
| `fixed` | `fixed(amount) → float64` | A USD price per request, used only as the complete price in `tier(name, fixed(amount))` |
|
||||
| `param` | `param(path) → any` | Reads a JSON path from the request body (uses gjson) |
|
||||
| `header` | `header(key) → string` | Reads a request header value |
|
||||
| `has` | `has(source, substr) → bool` | Substring check |
|
||||
@@ -96,6 +181,11 @@ Powered by [expr-lang/expr](https://github.com/expr-lang/expr). Expressions are
|
||||
# Simple flat pricing
|
||||
tier("base", p * 2.5 + c * 15 + cr * 0.25)
|
||||
|
||||
# Conditional per-request pricing, with token pricing for the fallback
|
||||
len <= 32000
|
||||
? tier("short", fixed(0.01))
|
||||
: tier("long", p * 2 + c * 8)
|
||||
|
||||
# Multi-tier (Claude Sonnet style) — use len for tier conditions
|
||||
len <= 200000
|
||||
? tier("standard", p * 3 + c * 15 + cr * 0.3 + cc * 3.75 + cc1h * 6)
|
||||
@@ -108,6 +198,40 @@ tier("base", p * 2 + c * 8 + img * 2.5)
|
||||
tier("base", p * 0.43 + c * 3.06 + img * 0.78 + ai * 3.81 + ao * 15.11)
|
||||
```
|
||||
|
||||
### Fixed Request Prices
|
||||
|
||||
`tier("request", fixed(0.01))` replaces all token charges in the selected leaf
|
||||
with a $0.01 base price for one successful HTTP/SSE request. Other leaves may
|
||||
still use token pricing. Group ratios and request multipliers continue to apply;
|
||||
existing tool surcharges are calculated separately and added as before. A stream
|
||||
does not incur a fixed fee per chunk. Realtime and task usage expressions reject
|
||||
`fixed()` before upstream submission or reservation.
|
||||
|
||||
The amount must be a finite, non-negative numeric literal whose v1 scaled value
|
||||
is finite. Explicit `fixed(0)` is valid and stays free. Expressions using `fixed`
|
||||
must consist of a conditional pricing tree and the standard request multiplier
|
||||
factors described below, with each leaf wrapped in `tier()`. Adding fixed and
|
||||
token charges in one leaf, adding multiple tiers together, or multiplying a
|
||||
fixed price by tokens is rejected. Validation examines the original AST,
|
||||
including branches that optimization or short-circuiting would skip. Existing
|
||||
expressions without `fixed()` retain their original grammar and behavior.
|
||||
|
||||
`fixed(amount)` returns `amount * 1,000,000` internally, preserving v1's existing
|
||||
quota conversion and rounding. Pre-consume matches conditions using estimates;
|
||||
settlement reevaluates against actual or existing locally estimated usage and
|
||||
reconciles any change of branch through the normal billing session. Missing
|
||||
upstream usage retains the existing estimation path; an actual request-priced
|
||||
branch can charge even with zero tokens. Failures retain the normal refund
|
||||
policy. Settlement evaluation errors retain the reservation and its estimated
|
||||
billing unit.
|
||||
|
||||
Evaluation results expose `billing_unit` (`token` or `request`) and, for a
|
||||
request-priced leaf, `fixed_price` in USD before multipliers. Pre-consume snapshots
|
||||
carry `estimated_billing_unit` and `estimated_fixed_price`; consume logs append
|
||||
the actual values to `other`. An explicit zero price is present in these fields.
|
||||
Older snapshots and logs need no migration. The expression remains the sole
|
||||
pricing configuration.
|
||||
|
||||
### Request Rules (appended after `|||`)
|
||||
|
||||
Request-conditional multipliers are appended to the expression after a `|||` separator:
|
||||
@@ -149,6 +273,72 @@ The usage-log UI treats `request_rules` as the authoritative rule list and rende
|
||||
Task plugins can expose validated, provider-specific billing facts through
|
||||
`meta.usageSchema`. Expressions read those facts with `u("key")`. A literal key
|
||||
must be declared by the plugin schema before the expression can be saved.
|
||||
|
||||
Multiple plugins may serve the same model with different usage schemas. The
|
||||
executing channel identifies the plugin; clients do not select a billing
|
||||
provider. Task expression resolution uses this order:
|
||||
|
||||
1. `billing_setting.plugin_billing_expr["<pluginKey>::<clientModel>"]`
|
||||
2. The same plugin's override for the final mapped model, when present
|
||||
3. The client model's expression when its billing mode is `tiered_expr`
|
||||
4. The final mapped model's expression when the client model has no expression
|
||||
billing mode
|
||||
|
||||
Plugin overrides always use task-usage semantics, independently of the model's
|
||||
billing mode. The flat override map is a separate option and never contributes
|
||||
composite keys to model names or model-price synchronization exports. Existing
|
||||
prices are not rewritten. The billing snapshot freezes the selected expression;
|
||||
completion uses that snapshot even if provider pricing later changes.
|
||||
|
||||
A model-level expression save smoke-tests each candidate's selected schema,
|
||||
skipping candidates with their own override in the same complete draft.
|
||||
Provider expression saves validate the plugin/model binding and that plugin's
|
||||
schema; errors identify both model and plugin. For a shared model, an effective
|
||||
expression referencing an undeclared literal `u()` key is unconfigured for that
|
||||
provider, including references in branches skipped by the current request.
|
||||
Submission rejects it with 400 `model_price_error`; administration and public
|
||||
pricing expose the unconfigured state. Public model-level metadata remains the
|
||||
first plugin's view in ascending key order, with provider variants added
|
||||
separately.
|
||||
|
||||
Saved provider overrides remain visible when only one provider remains.
|
||||
Administration exposes unavailable plugin/model bindings as `stale` and permits
|
||||
preserving or removing their existing expressions. Changes to those expressions
|
||||
still require a valid binding; writes compare against the locked database
|
||||
snapshot rather than the process cache. No overrides are silently pruned.
|
||||
Public provider variants include their effective `billing_mode` so inherited
|
||||
ratio/per-call pricing is distinct from an unconfigured usage expression.
|
||||
|
||||
Plugins serving different model families can optionally declare
|
||||
`meta.usageProfiles: [{models: ["image-model"], schema: {...}, examples: [...]}]`.
|
||||
Each profile provides the complete usage schema and optional display examples
|
||||
for its declared models. It replaces, rather than merges with, the plugin's
|
||||
default `usageSchema` and `usageExamples`. A model may belong to only one
|
||||
profile, and profile model names must use the spelling declared in `meta.models`.
|
||||
An empty profile schema is allowed; a missing or null schema is not. Token-unit
|
||||
profiles require their own examples, using the same validation as the defaults.
|
||||
|
||||
The host resolves model aliases before selecting usage metadata. Model-level
|
||||
pricing and expression-save APIs use the declared target; an ambiguous alias
|
||||
with multiple targets in the same plugin retains the plugin defaults. Request
|
||||
and usage validation use the executing plugin and final upstream model. Polling
|
||||
selects metadata from each saved task's model, including in mixed-model batches.
|
||||
Unmatched models and plugins without profiles retain the plugin defaults.
|
||||
|
||||
Profiles do not introduce a request-body whitelist or remove legacy multiplier
|
||||
extensions. Numeric hook facts retain a consistent floating-point representation
|
||||
for expression evaluation, including fields outside the selected profile.
|
||||
Saving an expression checks literal usage keys against the selected schema;
|
||||
already stored expressions are not migrated. Single-plugin submissions retain
|
||||
the existing key-check behavior; shared models enforce the compatibility rule
|
||||
above. Existing runtime error and quota safety checks still apply.
|
||||
Completion continues to evaluate the frozen expression after overlaying measured
|
||||
facts; no schema snapshot or database migration is introduced.
|
||||
|
||||
This is an optional addition to plugin API version 1. New hosts accept existing
|
||||
plugins unchanged. Older hosts reject the new metadata field, so upgrade the
|
||||
host before installing a plugin that declares profiles.
|
||||
|
||||
Numeric facts are finite, non-negative values in the declared canonical unit
|
||||
(`second`, `count`, `token`, or `credit`); enum facts are exact strings from the
|
||||
declared value list. The schema description is display-only metadata and never
|
||||
@@ -273,7 +463,7 @@ After the upstream response returns with actual token usage:
|
||||
1. `BuildTieredTokenParams(usage, isClaudeUsageSemantic, usedVars)`:
|
||||
- Reads actual token counts from `dto.Usage`
|
||||
- For GPT-format APIs (prompt_tokens includes everything): subtracts sub-categories from P/C **only when** the expression uses their variables (detected via AST introspection of the compiled expression)
|
||||
- For Claude-format APIs (input_tokens is text-only): no adjustment needed
|
||||
- For Claude-format APIs: cache reads omitted from the expression are added to input; separately priced cache remains separate
|
||||
|
||||
2. `TryTieredSettle(relayInfo, params)`:
|
||||
- Uses the captured `BillingSnapshot`, whose group-dependent fields have been refreshed from the final selected group
|
||||
|
||||
@@ -0,0 +1,87 @@
|
||||
package billingexpr
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"math"
|
||||
|
||||
"github.com/expr-lang/expr/ast"
|
||||
)
|
||||
|
||||
// UsesFixedPricing includes unselected branches, even when compilation later
|
||||
// optimizes them away. Hosts use it to reject unsupported billing entrances.
|
||||
func UsesFixedPricing(expression string) bool {
|
||||
return UsesFixedPricingByHash(expression, ExprHashString(expression))
|
||||
}
|
||||
|
||||
// UsesFixedPricingByHash avoids hashing an expression already in a snapshot.
|
||||
func UsesFixedPricingByHash(expression, hash string) bool {
|
||||
entry, err := compileEntryFromCacheByHash(expression, hash)
|
||||
return err == nil && entry.fixedPricing
|
||||
}
|
||||
|
||||
func containsPricingMarker(node ast.Node) bool {
|
||||
return ast.Find(node, func(part ast.Node) bool {
|
||||
identifier, ok := part.(*ast.IdentifierNode)
|
||||
return ok && (identifier.Value == "tier" || identifier.Value == "fixed")
|
||||
}) != nil
|
||||
}
|
||||
|
||||
func isRequestPriceMultiplier(node ast.Node) bool {
|
||||
if identifier, ok := node.(*ast.IdentifierNode); ok && identifier.Value == "image_count" {
|
||||
return true
|
||||
}
|
||||
conditional, ok := node.(*ast.ConditionalNode)
|
||||
if !ok || !usesRequestProbe(conditional.Cond) || containsPricingMarker(conditional.Cond) {
|
||||
return false
|
||||
}
|
||||
multiplier, multiplierOK := requestRuleNumber(conditional.Exp1)
|
||||
fallback, fallbackOK := requestRuleNumber(conditional.Exp2)
|
||||
return multiplierOK && fallbackOK && fallback == 1 && multiplier >= 0 && !math.IsNaN(multiplier) && !math.IsInf(multiplier, 0)
|
||||
}
|
||||
|
||||
// validateFixedPricingTree enforces one pricing leaf per execution. Without
|
||||
// this invariant, adding two tiers or multiplying a fixed price by tokens
|
||||
// would make both the request charge and its billing-unit trace ambiguous.
|
||||
func validateFixedPricingTree(node ast.Node) error {
|
||||
switch part := node.(type) {
|
||||
case *ast.ConditionalNode:
|
||||
if containsPricingMarker(part.Cond) {
|
||||
break
|
||||
}
|
||||
if err := validateFixedPricingTree(part.Exp1); err != nil {
|
||||
return err
|
||||
}
|
||||
return validateFixedPricingTree(part.Exp2)
|
||||
case *ast.BinaryNode:
|
||||
if part.Operator != "*" {
|
||||
break
|
||||
}
|
||||
if isRequestPriceMultiplier(part.Right) {
|
||||
return validateFixedPricingTree(part.Left)
|
||||
}
|
||||
if isRequestPriceMultiplier(part.Left) {
|
||||
return validateFixedPricingTree(part.Right)
|
||||
}
|
||||
case *ast.CallNode:
|
||||
callee, ok := part.Callee.(*ast.IdentifierNode)
|
||||
if !ok || callee.Value != "tier" || len(part.Arguments) != 2 || containsPricingMarker(part.Arguments[0]) {
|
||||
break
|
||||
}
|
||||
price := part.Arguments[1]
|
||||
fixed, ok := price.(*ast.CallNode)
|
||||
if ok {
|
||||
function, direct := fixed.Callee.(*ast.IdentifierNode)
|
||||
if direct && function.Value == "fixed" && len(fixed.Arguments) == 1 {
|
||||
amount, literal := requestRuleNumber(fixed.Arguments[0])
|
||||
if literal && amount >= 0 && !math.IsNaN(amount) && !math.IsInf(amount*1_000_000, 0) {
|
||||
return nil
|
||||
}
|
||||
return fmt.Errorf("fixed price must be a finite, non-negative numeric literal with a finite v1 value")
|
||||
}
|
||||
}
|
||||
if !containsPricingMarker(price) {
|
||||
return nil
|
||||
}
|
||||
}
|
||||
return fmt.Errorf("fixed pricing requires tier(name, fixed(amount)) leaves, conditional tiers and request multipliers; token and fixed charges cannot be combined in one leaf")
|
||||
}
|
||||
+48
-18
@@ -6,6 +6,8 @@ import (
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"github.com/QuantumNous/new-api/relaykit/dto"
|
||||
|
||||
"github.com/expr-lang/expr"
|
||||
"github.com/expr-lang/expr/vm"
|
||||
"github.com/tidwall/gjson"
|
||||
@@ -30,7 +32,7 @@ func RunExprWithRequest(exprStr string, params TokenParams, request RequestInput
|
||||
if err != nil {
|
||||
return 0, TraceResult{}, err
|
||||
}
|
||||
return runProgram(entry.prog, entry.requestRules, params, request)
|
||||
return runProgram(entry.prog, entry.requestRules, entry.usedVars, params, request)
|
||||
}
|
||||
|
||||
// RunExprByHash is like RunExpr but accepts a pre-computed hash for the cache
|
||||
@@ -45,31 +47,59 @@ func RunExprByHashWithRequest(exprStr, hash string, params TokenParams, request
|
||||
if err != nil {
|
||||
return 0, TraceResult{}, err
|
||||
}
|
||||
return runProgram(entry.prog, entry.requestRules, params, request)
|
||||
return runProgram(entry.prog, entry.requestRules, entry.usedVars, params, request)
|
||||
}
|
||||
|
||||
func runProgram(prog *vm.Program, requestRules []RequestRuleTrace, params TokenParams, request RequestInput) (float64, TraceResult, error) {
|
||||
func runProgram(prog *vm.Program, requestRules []RequestRuleTrace, usedVars map[string]bool, params TokenParams, request RequestInput) (float64, TraceResult, error) {
|
||||
trace := TraceResult{
|
||||
BillingUnit: BillingUnitToken,
|
||||
RequestRules: append([]RequestRuleTrace(nil), requestRules...),
|
||||
}
|
||||
headers := normalizeHeaders(request.Headers)
|
||||
timeForRequest := func(tz string) time.Time {
|
||||
if request.At == nil {
|
||||
return timeInZone(tz)
|
||||
}
|
||||
loc, err := time.LoadLocation(strings.TrimSpace(tz))
|
||||
if err != nil {
|
||||
return request.At.UTC()
|
||||
}
|
||||
return request.At.In(loc)
|
||||
}
|
||||
imageCount := 1
|
||||
if usedVars["image_count"] {
|
||||
if request.ImageCount != nil {
|
||||
imageCount = *request.ImageCount
|
||||
}
|
||||
if imageCount < 1 || imageCount > dto.MaxImageN {
|
||||
return 0, trace, fmt.Errorf("image_count must be between 1 and %d", dto.MaxImageN)
|
||||
}
|
||||
trace.ImageCount = &imageCount
|
||||
}
|
||||
|
||||
env := map[string]any{
|
||||
"p": params.P,
|
||||
"c": params.C,
|
||||
"len": params.Len,
|
||||
"cr": params.CR,
|
||||
"cc": params.CC,
|
||||
"cc1h": params.CC1h,
|
||||
"img": params.Img,
|
||||
"img_o": params.ImgO,
|
||||
"ai": params.AI,
|
||||
"ao": params.AO,
|
||||
"image_count": float64(imageCount),
|
||||
"p": params.P,
|
||||
"c": params.C,
|
||||
"len": params.Len,
|
||||
"cr": params.CR,
|
||||
"cc": params.CC,
|
||||
"cc1h": params.CC1h,
|
||||
"img": params.Img,
|
||||
"img_cr": params.ImgCR,
|
||||
"img_o": params.ImgO,
|
||||
"ai": params.AI,
|
||||
"ao": params.AO,
|
||||
"tier": func(name string, value float64) float64 {
|
||||
trace.MatchedTier = name
|
||||
trace.Cost = value
|
||||
return value
|
||||
},
|
||||
"fixed": func(amount float64) float64 {
|
||||
trace.BillingUnit = BillingUnitRequest
|
||||
trace.FixedPrice = &amount
|
||||
return amount * 1_000_000
|
||||
},
|
||||
requestRuleTraceFunction: func(ruleIndex int, matched bool, multiplier float64) float64 {
|
||||
if matched && ruleIndex >= 0 && ruleIndex < len(trace.RequestRules) {
|
||||
trace.RequestRules[ruleIndex].Matched = true
|
||||
@@ -114,11 +144,11 @@ func runProgram(prog *vm.Program, requestRules []RequestRuleTrace, params TokenP
|
||||
}
|
||||
return strings.Contains(fmt.Sprint(source), substr)
|
||||
},
|
||||
"hour": func(tz string) int { return timeInZone(tz).Hour() },
|
||||
"minute": func(tz string) int { return timeInZone(tz).Minute() },
|
||||
"weekday": func(tz string) int { return int(timeInZone(tz).Weekday()) },
|
||||
"month": func(tz string) int { return int(timeInZone(tz).Month()) },
|
||||
"day": func(tz string) int { return timeInZone(tz).Day() },
|
||||
"hour": func(tz string) int { return timeForRequest(tz).Hour() },
|
||||
"minute": func(tz string) int { return timeForRequest(tz).Minute() },
|
||||
"weekday": func(tz string) int { return int(timeForRequest(tz).Weekday()) },
|
||||
"month": func(tz string) int { return int(timeForRequest(tz).Month()) },
|
||||
"day": func(tz string) int { return timeForRequest(tz).Day() },
|
||||
"max": math.Max,
|
||||
"min": math.Min,
|
||||
"abs": math.Abs,
|
||||
|
||||
@@ -1,6 +1,10 @@
|
||||
package billingexpr
|
||||
|
||||
import "github.com/QuantumNous/new-api/common"
|
||||
import (
|
||||
"fmt"
|
||||
|
||||
"github.com/QuantumNous/new-api/common"
|
||||
)
|
||||
|
||||
// quotaConversion converts raw expression output to quota based on the
|
||||
// expression version. This is the central dispatch point for future versions
|
||||
@@ -22,6 +26,9 @@ func ComputeTieredQuota(snap *BillingSnapshot, params TokenParams) (TieredResult
|
||||
}
|
||||
|
||||
func ComputeTieredQuotaWithRequest(snap *BillingSnapshot, params TokenParams, request RequestInput) (TieredResult, error) {
|
||||
if snap.TaskUsageBilling && UsesFixedPricingByHash(snap.ExprString, snap.ExprHash) {
|
||||
return TieredResult{}, fmt.Errorf("fixed pricing is not supported for task usage expressions")
|
||||
}
|
||||
cost, trace, err := RunExprByHashWithRequest(snap.ExprString, snap.ExprHash, params, request)
|
||||
if err != nil {
|
||||
return TieredResult{}, err
|
||||
@@ -31,12 +38,19 @@ func ComputeTieredQuotaWithRequest(snap *BillingSnapshot, params TokenParams, re
|
||||
afterGroup, clamp := common.QuotaRoundChecked(quotaBeforeGroup * snap.GroupRatio)
|
||||
crossed := trace.MatchedTier != snap.EstimatedTier
|
||||
|
||||
return TieredResult{
|
||||
result := TieredResult{
|
||||
ImageCount: trace.ImageCount,
|
||||
BillingUnit: trace.BillingUnit,
|
||||
FixedPrice: trace.FixedPrice,
|
||||
ActualQuotaBeforeGroup: quotaBeforeGroup,
|
||||
ActualQuotaAfterGroup: afterGroup,
|
||||
MatchedTier: trace.MatchedTier,
|
||||
RequestRules: trace.RequestRules,
|
||||
CrossedTier: crossed,
|
||||
Clamp: clamp,
|
||||
}, nil
|
||||
}
|
||||
if trace.BillingUnit == BillingUnitToken && UsedVarsByHash(snap.ExprString, snap.ExprHash)["img_cr"] {
|
||||
result.BillingTokens = ¶ms
|
||||
}
|
||||
return result, nil
|
||||
}
|
||||
+423
@@ -0,0 +1,423 @@
|
||||
[
|
||||
{
|
||||
"name": "versioned tokens and cache",
|
||||
"expression": "v1:tier(\"base\", p * 2 + c * 8 + cr * 0.2)",
|
||||
"tokens": {
|
||||
"p": 900,
|
||||
"c": 100,
|
||||
"cr": 100
|
||||
},
|
||||
"cost": 2620,
|
||||
"tier": "base"
|
||||
},
|
||||
{
|
||||
"name": "nil false and zero",
|
||||
"expression": "param(\"missing\") == nil && param(\"zero\") == 0 && param(\"enabled\") == false ? tier(\"ok\", 3.5) : tier(\"bad\", 99)",
|
||||
"body": {
|
||||
"zero": 0,
|
||||
"enabled": false
|
||||
},
|
||||
"cost": 3.5,
|
||||
"tier": "ok"
|
||||
},
|
||||
{
|
||||
"name": "normalized headers and two rules",
|
||||
"expression": "tier(\"base\", p * 2) * (header(\" X-Priority \") == \"high\" ? 2 : 1) * (param(\"service_tier\") == \"flex\" ? 0.5 : 1)",
|
||||
"tokens": {
|
||||
"p": 100
|
||||
},
|
||||
"headers": {
|
||||
" X-Priority ": " high "
|
||||
},
|
||||
"body": {
|
||||
"service_tier": "flex"
|
||||
},
|
||||
"cost": 200,
|
||||
"tier": "base",
|
||||
"multipliers": [
|
||||
2,
|
||||
0.5
|
||||
],
|
||||
"matched": [
|
||||
true,
|
||||
true
|
||||
]
|
||||
},
|
||||
{
|
||||
"name": "unmatched rule",
|
||||
"expression": "tier(\"base\", 10) * (param(\"priority\") == true ? 2 : 1)",
|
||||
"body": {},
|
||||
"cost": 10,
|
||||
"tier": "base",
|
||||
"multipliers": [
|
||||
2
|
||||
],
|
||||
"matched": [
|
||||
false
|
||||
]
|
||||
},
|
||||
{
|
||||
"name": "rules skipped by short circuit",
|
||||
"expression": "param(\"enabled\") == true ? tier(\"yes\", 10) * (header(\"tier\") == \"fast\" ? 3 : 1) : tier(\"no\", 2)",
|
||||
"body": {},
|
||||
"headers": {
|
||||
"tier": "fast"
|
||||
},
|
||||
"cost": 2,
|
||||
"tier": "no",
|
||||
"multipliers": [
|
||||
3
|
||||
],
|
||||
"matched": [
|
||||
false
|
||||
]
|
||||
},
|
||||
{
|
||||
"name": "noncanonical factor still evaluates",
|
||||
"expression": "tier(\"base\", 10) * (param(\"priority\") == true ? 2 : 1.5)",
|
||||
"body": {},
|
||||
"cost": 15,
|
||||
"tier": "base"
|
||||
},
|
||||
{
|
||||
"name": "path escapes array indices and length",
|
||||
"expression": "param(\"a\\\\.b.items.0.price\") + param(\"a\\\\.b.items.#\")",
|
||||
"body": {
|
||||
"a.b": {
|
||||
"items": [
|
||||
{
|
||||
"price": 4
|
||||
},
|
||||
{
|
||||
"price": 8
|
||||
}
|
||||
]
|
||||
}
|
||||
},
|
||||
"cost": 6,
|
||||
"tier": ""
|
||||
},
|
||||
{
|
||||
"name": "literal path to constructor",
|
||||
"expression": "param(\"constructor\") == nil ? 4 : 99",
|
||||
"body": {},
|
||||
"cost": 4,
|
||||
"tier": ""
|
||||
},
|
||||
{
|
||||
"name": "contains scalars and objects",
|
||||
"expression": "has(param(\"n\"), \"12\") && has(param(\"object\"), \"map[a:1 b:[true <nil>]]\") && !has(nil, \"x\") && !has(\"abc\", \"\") ? 8 : 99",
|
||||
"body": {
|
||||
"n": 123,
|
||||
"object": {
|
||||
"b": [
|
||||
true,
|
||||
null
|
||||
],
|
||||
"a": 1
|
||||
}
|
||||
},
|
||||
"cost": 8,
|
||||
"tier": ""
|
||||
},
|
||||
{
|
||||
"name": "math and numeric division",
|
||||
"expression": "min(10, max(2, 5)) + abs(-2) + ceil(1.1) + floor(2.9) + 5 / 2 + 7 % 3",
|
||||
"cost": 14.5,
|
||||
"tier": ""
|
||||
},
|
||||
{
|
||||
"name": "task token scale and request rule",
|
||||
"expression": "tier(\"base\", 0.2 + u(\"tokens\") * 10 / 1000000) * (header(\"tier\") == \"fast\" ? 2 : 1)",
|
||||
"usage": {
|
||||
"tokens": 500000
|
||||
},
|
||||
"headers": {
|
||||
"tier": "fast"
|
||||
},
|
||||
"cost": 10.4,
|
||||
"tier": "base",
|
||||
"multipliers": [
|
||||
2
|
||||
],
|
||||
"matched": [
|
||||
true
|
||||
]
|
||||
},
|
||||
{
|
||||
"name": "task enum conditions",
|
||||
"expression": "u(\"mode\") == \"pro\" ? tier(\"pro\", u(\"seconds\") * 0.8) : tier(\"std\", u(\"seconds\") * 0.4)",
|
||||
"usage": {
|
||||
"mode": "pro",
|
||||
"seconds": 5
|
||||
},
|
||||
"cost": 4,
|
||||
"tier": "pro"
|
||||
},
|
||||
{
|
||||
"name": "scientific rates and parentheses",
|
||||
"expression": "(tier(\"base\", (p * 2e-3 + c * 1E+1)))",
|
||||
"tokens": {
|
||||
"p": 1000,
|
||||
"c": 2
|
||||
},
|
||||
"cost": 22,
|
||||
"tier": "base"
|
||||
},
|
||||
{
|
||||
"name": "fixed request with zero usage",
|
||||
"expression": "tier(\"request\", fixed(0.01))",
|
||||
"cost": 10000,
|
||||
"tier": "request",
|
||||
"billingUnit": "request",
|
||||
"fixedPrice": 0.01
|
||||
},
|
||||
{
|
||||
"name": "explicit free request",
|
||||
"expression": "tier(\"free\", fixed(0))",
|
||||
"tokens": {
|
||||
"p": 1000
|
||||
},
|
||||
"cost": 0,
|
||||
"tier": "free",
|
||||
"billingUnit": "request",
|
||||
"fixedPrice": 0
|
||||
},
|
||||
{
|
||||
"name": "mixed expression token fallback",
|
||||
"expression": "len <= 32000 ? tier(\"short\", fixed(0.01)) : tier(\"long\", p * 2 + c * 8)",
|
||||
"tokens": {
|
||||
"p": 32001,
|
||||
"c": 100,
|
||||
"len": 32001
|
||||
},
|
||||
"cost": 64802,
|
||||
"tier": "long",
|
||||
"billingUnit": "token"
|
||||
},
|
||||
{
|
||||
"name": "fixed request with multiplier",
|
||||
"expression": "(len <= 32000 ? tier(\"short\", fixed(0.01)) : tier(\"long\", p * 2 + c * 8)) * (param(\"fast\") == true ? 2 : 1)",
|
||||
"tokens": {
|
||||
"len": 1
|
||||
},
|
||||
"body": {
|
||||
"fast": true
|
||||
},
|
||||
"cost": 20000,
|
||||
"tier": "short",
|
||||
"multipliers": [
|
||||
2
|
||||
],
|
||||
"matched": [
|
||||
true
|
||||
],
|
||||
"billingUnit": "request",
|
||||
"fixedPrice": 0.01
|
||||
},
|
||||
{
|
||||
"name": "fixed time condition",
|
||||
"expression": "hour(\"Asia/Shanghai\") >= 0 && hour(\"Asia/Shanghai\") < 24 ? tier(\"today\", fixed(0.02)) : tier(\"fallback\", p * 2)",
|
||||
"cost": 20000,
|
||||
"tier": "today",
|
||||
"billingUnit": "request",
|
||||
"fixedPrice": 0.02
|
||||
},
|
||||
{
|
||||
"name": "image count is distinct from requested n",
|
||||
"expression": "tier(\"image\", fixed(0.04)) * image_count",
|
||||
"imageCount": 2,
|
||||
"body": {
|
||||
"n": 3
|
||||
},
|
||||
"cost": 80000,
|
||||
"tier": "image",
|
||||
"billingUnit": "request",
|
||||
"fixedPrice": 0.04
|
||||
},
|
||||
{
|
||||
"name": "image count and quality adjustment",
|
||||
"expression": "tier(\"image\", fixed(0.04)) * image_count * (param(\"quality\") == \"hd\" ? 2 : 1)",
|
||||
"imageCount": 3,
|
||||
"body": {
|
||||
"quality": "hd"
|
||||
},
|
||||
"cost": 240000,
|
||||
"tier": "image",
|
||||
"billingUnit": "request",
|
||||
"fixedPrice": 0.04,
|
||||
"multipliers": [
|
||||
2
|
||||
],
|
||||
"matched": [
|
||||
true
|
||||
]
|
||||
},
|
||||
{
|
||||
"name": "free image count maximum",
|
||||
"expression": "tier(\"image\", fixed(0)) * image_count",
|
||||
"imageCount": 128,
|
||||
"cost": 0,
|
||||
"tier": "image",
|
||||
"billingUnit": "request",
|
||||
"fixedPrice": 0
|
||||
},
|
||||
{
|
||||
"name": "default image quantity",
|
||||
"expression": "tier(\"image\", fixed(0.04)) * image_count",
|
||||
"cost": 40000,
|
||||
"tier": "image",
|
||||
"billingUnit": "request",
|
||||
"fixedPrice": 0.04
|
||||
},
|
||||
{
|
||||
"name": "audio pricing ignores text cache discounts",
|
||||
"expression": "(ai > 0 || ao > 0) ? tier(\"audio\", max(len - ai, 0) * 2 + c * 6 + ai * 4 + ao * 4) : tier(\"base\", p * 2 + c * 6 + cr * 0.2)",
|
||||
"tokens": {
|
||||
"len": 1000,
|
||||
"ai": 200,
|
||||
"ao": 50,
|
||||
"c": 50
|
||||
},
|
||||
"cost": 2900,
|
||||
"tier": "audio",
|
||||
"displayTiers": [
|
||||
{
|
||||
"label": "audio",
|
||||
"conditions": [],
|
||||
"prices": {
|
||||
"p": 2,
|
||||
"c": 6,
|
||||
"ai": 4,
|
||||
"ao": 4
|
||||
},
|
||||
"conditionText": "Audio requests"
|
||||
},
|
||||
{
|
||||
"label": "base",
|
||||
"conditions": [],
|
||||
"prices": {
|
||||
"p": 2,
|
||||
"c": 6,
|
||||
"cr": 0.2
|
||||
},
|
||||
"conditionText": "Text-only requests"
|
||||
}
|
||||
],
|
||||
"conversion": {
|
||||
"model_name": "contract-audio",
|
||||
"pricing": {
|
||||
"ModelRatio": 1,
|
||||
"CompletionRatio": 3,
|
||||
"CacheRatio": 0.1,
|
||||
"AudioRatio": 2
|
||||
}
|
||||
}
|
||||
},
|
||||
{
|
||||
"name": "text requests retain configured cache discount",
|
||||
"expression": "(ai > 0 || ao > 0) ? tier(\"audio\", max(len - ai, 0) * 2 + c * 6 + ai * 4 + ao * 4) : tier(\"base\", p * 2 + c * 6 + cr * 0.2)",
|
||||
"tokens": {
|
||||
"p": 600,
|
||||
"c": 100,
|
||||
"cr": 400,
|
||||
"ai": 0,
|
||||
"ao": 0
|
||||
},
|
||||
"cost": 1880,
|
||||
"tier": "base",
|
||||
"displayTiers": [
|
||||
{
|
||||
"label": "audio",
|
||||
"conditions": [],
|
||||
"prices": {
|
||||
"p": 2,
|
||||
"c": 6,
|
||||
"ai": 4,
|
||||
"ao": 4
|
||||
},
|
||||
"conditionText": "Audio requests"
|
||||
},
|
||||
{
|
||||
"label": "base",
|
||||
"conditions": [],
|
||||
"prices": {
|
||||
"p": 2,
|
||||
"c": 6,
|
||||
"cr": 0.2
|
||||
},
|
||||
"conditionText": "Text-only requests"
|
||||
}
|
||||
]
|
||||
},
|
||||
{
|
||||
"name": "audio display survives redundant parentheses",
|
||||
"expression": "(ai > 0 || ao > 0) ? tier(\"audio\", max((len - ai), (0)) * 2 + c * 6 + ai * 4 + ao * 4) : tier(\"base\", p * 2 + c * 6 + cr * 0.2)",
|
||||
"tokens": {
|
||||
"len": 1000,
|
||||
"ai": 200,
|
||||
"ao": 50,
|
||||
"c": 50
|
||||
},
|
||||
"cost": 2900,
|
||||
"tier": "audio",
|
||||
"displayTiers": [
|
||||
{
|
||||
"label": "audio",
|
||||
"conditions": [],
|
||||
"prices": {
|
||||
"p": 2,
|
||||
"c": 6,
|
||||
"ai": 4,
|
||||
"ao": 4
|
||||
},
|
||||
"conditionText": "Audio requests"
|
||||
},
|
||||
{
|
||||
"label": "base",
|
||||
"conditions": [],
|
||||
"prices": {
|
||||
"p": 2,
|
||||
"c": 6,
|
||||
"cr": 0.2
|
||||
},
|
||||
"conditionText": "Text-only requests"
|
||||
}
|
||||
]
|
||||
},
|
||||
{
|
||||
"name": "audio display survives commuted operands",
|
||||
"expression": "(ao > 0 || ai > 0) ? tier(\"audio\", 2 * max(0, len - ai) + 6 * c + 4 * ai + ao * 4) : tier(\"base\", p * 2 + 6 * c + cr * 0.2)",
|
||||
"tokens": {
|
||||
"len": 1000,
|
||||
"ai": 200,
|
||||
"ao": 50,
|
||||
"c": 50
|
||||
},
|
||||
"cost": 2900,
|
||||
"tier": "audio",
|
||||
"displayTiers": [
|
||||
{
|
||||
"label": "audio",
|
||||
"conditions": [],
|
||||
"prices": {
|
||||
"p": 2,
|
||||
"c": 6,
|
||||
"ai": 4,
|
||||
"ao": 4
|
||||
},
|
||||
"conditionText": "Audio requests"
|
||||
},
|
||||
{
|
||||
"label": "base",
|
||||
"conditions": [],
|
||||
"prices": {
|
||||
"p": 2,
|
||||
"c": 6,
|
||||
"cr": 0.2
|
||||
},
|
||||
"conditionText": "Text-only requests"
|
||||
}
|
||||
]
|
||||
}
|
||||
]
|
||||
+36
-10
@@ -3,30 +3,36 @@ package billingexpr
|
||||
import (
|
||||
"crypto/sha256"
|
||||
"fmt"
|
||||
"time"
|
||||
|
||||
"github.com/QuantumNous/new-api/common"
|
||||
)
|
||||
|
||||
type RequestInput struct {
|
||||
At *time.Time `json:"at,omitempty"`
|
||||
Headers map[string]string
|
||||
Body []byte
|
||||
Usage map[string]any
|
||||
// ImageCount is a validated billing quantity, separate from the frozen
|
||||
// request's n. Settlement can replace it with the actual returned count.
|
||||
ImageCount *int
|
||||
}
|
||||
|
||||
// TokenParams holds all token dimensions passed into an Expr evaluation.
|
||||
// Fields beyond P and C are optional — when absent they default to 0,
|
||||
// which means cache-unaware expressions keep working unchanged.
|
||||
type TokenParams struct {
|
||||
P float64 // prompt tokens (text) — auto-excludes sub-categories priced separately
|
||||
C float64 // completion tokens (text) — auto-excludes sub-categories priced separately
|
||||
Len float64 // total input context length for tier conditions (non-Claude: raw prompt_tokens; Claude: text + cache read + cache creation)
|
||||
CR float64 // cache read (hit) tokens
|
||||
CC float64 // cache creation tokens (5-min TTL for Claude, generic for others)
|
||||
CC1h float64 // cache creation tokens — 1-hour TTL (Claude only)
|
||||
Img float64 // image input tokens
|
||||
ImgO float64 // image output tokens
|
||||
AI float64 // audio input tokens
|
||||
AO float64 // audio output tokens
|
||||
P float64 // prompt tokens (text) — auto-excludes sub-categories priced separately
|
||||
C float64 // completion tokens (text) — auto-excludes sub-categories priced separately
|
||||
Len float64 // total input context length for tier conditions (non-Claude: raw prompt_tokens; Claude: text + cache read + cache creation)
|
||||
CR float64 // cache read (hit) tokens
|
||||
CC float64 // cache creation tokens (5-min TTL for Claude, generic for others)
|
||||
CC1h float64 // cache creation tokens — 1-hour TTL (Claude only)
|
||||
Img float64 // image input tokens
|
||||
ImgCR float64 // image cache read tokens, separated only when explicitly priced
|
||||
ImgO float64 // image output tokens
|
||||
AI float64 // audio input tokens
|
||||
AO float64 // audio output tokens
|
||||
}
|
||||
|
||||
// RequestRuleTrace describes one request-dependent multiplier detected at compile time.
|
||||
@@ -36,8 +42,18 @@ type RequestRuleTrace struct {
|
||||
Matched bool `json:"matched"`
|
||||
}
|
||||
|
||||
type BillingUnit string
|
||||
|
||||
const (
|
||||
BillingUnitToken BillingUnit = "token"
|
||||
BillingUnitRequest BillingUnit = "request"
|
||||
)
|
||||
|
||||
// TraceResult holds side-channel info captured while an expression runs.
|
||||
type TraceResult struct {
|
||||
ImageCount *int `json:"image_count,omitempty"`
|
||||
BillingUnit BillingUnit `json:"billing_unit"`
|
||||
FixedPrice *float64 `json:"fixed_price,omitempty"`
|
||||
MatchedTier string `json:"matched_tier"`
|
||||
RequestRules []RequestRuleTrace `json:"request_rules,omitempty"`
|
||||
Cost float64 `json:"cost"`
|
||||
@@ -48,6 +64,7 @@ type TraceResult struct {
|
||||
// auto-group retry and settlement. It is fully serializable and contains no
|
||||
// compiled program pointers.
|
||||
type BillingSnapshot struct {
|
||||
EstimatedImageCount *int `json:"estimated_image_count,omitempty"`
|
||||
BillingMode string `json:"billing_mode"`
|
||||
ModelName string `json:"model_name"`
|
||||
ExprString string `json:"expr_string"`
|
||||
@@ -58,6 +75,8 @@ type BillingSnapshot struct {
|
||||
EstimatedQuotaBeforeGroup float64 `json:"estimated_quota_before_group"`
|
||||
EstimatedQuotaAfterGroup int `json:"estimated_quota_after_group"`
|
||||
EstimatedTier string `json:"estimated_tier"`
|
||||
EstimatedBillingUnit BillingUnit `json:"estimated_billing_unit,omitempty"`
|
||||
EstimatedFixedPrice *float64 `json:"estimated_fixed_price,omitempty"`
|
||||
QuotaPerUnit float64 `json:"quota_per_unit"`
|
||||
ExprVersion int `json:"expr_version"`
|
||||
TaskUsageBilling bool `json:"task_usage_billing,omitempty"`
|
||||
@@ -66,6 +85,13 @@ type BillingSnapshot struct {
|
||||
|
||||
// TieredResult holds everything needed after running tiered settlement.
|
||||
type TieredResult struct {
|
||||
// BillingTokens records the actual normalized inputs for successful token
|
||||
// billing that explicitly references img_cr. Logs serialize these through
|
||||
// the shared injection path, not as additional snapshot state.
|
||||
BillingTokens *TokenParams `json:"-"`
|
||||
ImageCount *int `json:"image_count,omitempty"`
|
||||
BillingUnit BillingUnit `json:"billing_unit"`
|
||||
FixedPrice *float64 `json:"fixed_price,omitempty"`
|
||||
ActualQuotaBeforeGroup float64 `json:"actual_quota_before_group"`
|
||||
ActualQuotaAfterGroup int `json:"actual_quota_after_group"`
|
||||
MatchedTier string `json:"matched_tier"`
|
||||
|
||||
+48
-24
@@ -7,7 +7,6 @@ import (
|
||||
"regexp"
|
||||
"slices"
|
||||
"strings"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
"github.com/QuantumNous/new-api/logger"
|
||||
@@ -117,7 +116,7 @@ type Engine struct {
|
||||
now func() time.Time
|
||||
log func(string)
|
||||
module *sobek.SourceTextModuleRecord
|
||||
pool sync.Pool
|
||||
pool chan *runtimeInstance
|
||||
semaphore chan struct{}
|
||||
}
|
||||
|
||||
@@ -172,13 +171,14 @@ func Compile(source string, options Options) (*Engine, error) {
|
||||
log: options.Log,
|
||||
module: module,
|
||||
semaphore: make(chan struct{}, concurrency),
|
||||
pool: make(chan *runtimeInstance, concurrency),
|
||||
}
|
||||
instance, err := engine.newRuntime(context.Background())
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
instance.logContext.context = nil
|
||||
engine.pool.Put(instance)
|
||||
engine.putRuntime(instance)
|
||||
return engine, nil
|
||||
}
|
||||
|
||||
@@ -201,14 +201,12 @@ func (e *Engine) Export(ctx context.Context, exportName string) (result any, err
|
||||
instance.runtime.ClearInterrupt()
|
||||
instance.logContext.context = nil
|
||||
if reusable {
|
||||
e.pool.Put(instance)
|
||||
e.putRuntime(instance)
|
||||
}
|
||||
}()
|
||||
timedOut := errors.New("plugin export timed out")
|
||||
timer := time.AfterFunc(e.timeout, func() { instance.runtime.Interrupt(timedOut) })
|
||||
stopContext := context.AfterFunc(ctx, func() { instance.runtime.Interrupt(ctx.Err()) })
|
||||
defer stopContext()
|
||||
defer timer.Stop()
|
||||
stopInterrupt := watchRuntimeContext(instance.runtime, ctx, e.timeout, timedOut)
|
||||
defer stopInterrupt()
|
||||
defer func() {
|
||||
if recovered := recover(); recovered != nil {
|
||||
reusable = false
|
||||
@@ -244,7 +242,7 @@ func (e *Engine) HasExport(ctx context.Context, exportName string) (bool, error)
|
||||
}
|
||||
defer func() {
|
||||
instance.logContext.context = nil
|
||||
e.pool.Put(instance)
|
||||
e.putRuntime(instance)
|
||||
}()
|
||||
value := instance.module.GetBindingValue(exportName)
|
||||
return value != nil && !sobek.IsUndefined(value), nil
|
||||
@@ -268,14 +266,12 @@ func (e *Engine) HasCallablePath(ctx context.Context, exportName string, members
|
||||
instance.runtime.ClearInterrupt()
|
||||
instance.logContext.context = nil
|
||||
if reusable {
|
||||
e.pool.Put(instance)
|
||||
e.putRuntime(instance)
|
||||
}
|
||||
}()
|
||||
timedOut := errors.New("plugin inspection timed out")
|
||||
timer := time.AfterFunc(e.timeout, func() { instance.runtime.Interrupt(timedOut) })
|
||||
stopContext := context.AfterFunc(ctx, func() { instance.runtime.Interrupt(ctx.Err()) })
|
||||
defer stopContext()
|
||||
defer timer.Stop()
|
||||
stopInterrupt := watchRuntimeContext(instance.runtime, ctx, e.timeout, timedOut)
|
||||
defer stopInterrupt()
|
||||
defer func() {
|
||||
if recovered := recover(); recovered != nil {
|
||||
reusable = false
|
||||
@@ -350,16 +346,14 @@ func (e *Engine) call(
|
||||
instance.runtime.ClearInterrupt()
|
||||
instance.logContext.context = nil
|
||||
if reusable {
|
||||
e.pool.Put(instance)
|
||||
e.putRuntime(instance)
|
||||
}
|
||||
}()
|
||||
|
||||
hookName := strings.Join(append([]string{exportName}, members...), ".")
|
||||
timedOut := errors.New("plugin call timed out")
|
||||
timer := time.AfterFunc(e.timeout, func() { instance.runtime.Interrupt(timedOut) })
|
||||
stopContext := context.AfterFunc(ctx, func() { instance.runtime.Interrupt(ctx.Err()) })
|
||||
defer stopContext()
|
||||
defer timer.Stop()
|
||||
stopInterrupt := watchRuntimeContext(instance.runtime, ctx, e.timeout, timedOut)
|
||||
defer stopInterrupt()
|
||||
defer func() {
|
||||
if recovered := recover(); recovered != nil {
|
||||
reusable = false
|
||||
@@ -459,13 +453,40 @@ func resolveExportPath(instance *runtimeInstance, exportName string, members []s
|
||||
return value, hookName, true
|
||||
}
|
||||
|
||||
// Idle runtimes are bounded by the execution limit and survive GC. Start with
|
||||
// one instance and grow only when concurrent work actually needs more.
|
||||
func (e *Engine) putRuntime(instance *runtimeInstance) {
|
||||
select {
|
||||
case e.pool <- instance:
|
||||
default:
|
||||
}
|
||||
}
|
||||
|
||||
func (e *Engine) getRuntime(ctx context.Context) (*runtimeInstance, error) {
|
||||
if pooled := e.pool.Get(); pooled != nil {
|
||||
instance := pooled.(*runtimeInstance)
|
||||
select {
|
||||
case instance := <-e.pool:
|
||||
instance.logContext.context = ctx
|
||||
return instance, nil
|
||||
default:
|
||||
return e.newRuntime(ctx)
|
||||
}
|
||||
}
|
||||
|
||||
// A timeout callback must finish before its runtime can be reused. Merely
|
||||
// stopping a timer does not wait for an already-started Interrupt call.
|
||||
func watchRuntimeContext(runtime *sobek.Runtime, ctx context.Context, timeout time.Duration, timeoutError error) func() {
|
||||
callContext, cancel := context.WithTimeoutCause(ctx, timeout, timeoutError)
|
||||
interrupted := make(chan struct{})
|
||||
stop := context.AfterFunc(callContext, func() {
|
||||
runtime.Interrupt(context.Cause(callContext))
|
||||
close(interrupted)
|
||||
})
|
||||
return func() {
|
||||
if !stop() {
|
||||
<-interrupted
|
||||
}
|
||||
cancel()
|
||||
}
|
||||
return e.newRuntime(ctx)
|
||||
}
|
||||
|
||||
func (e *Engine) newRuntime(ctx context.Context) (instance *runtimeInstance, err error) {
|
||||
@@ -483,8 +504,11 @@ func (e *Engine) newRuntime(ctx context.Context) (instance *runtimeInstance, err
|
||||
return nil, fmt.Errorf("inject plugin utils: %w", err)
|
||||
}
|
||||
timedOut := errors.New("plugin initialization timed out")
|
||||
timer := time.AfterFunc(e.timeout, func() { runtime.Interrupt(timedOut) })
|
||||
defer timer.Stop()
|
||||
stopInterrupt := watchRuntimeContext(runtime, ctx, e.timeout, timedOut)
|
||||
defer func() {
|
||||
stopInterrupt()
|
||||
runtime.ClearInterrupt()
|
||||
}()
|
||||
defer func() {
|
||||
if recovered := recover(); recovered != nil {
|
||||
if interrupted, ok := recovered.(*sobek.InterruptedError); ok {
|
||||
|
||||
+325
-2
@@ -4,13 +4,21 @@ import (
|
||||
"bytes"
|
||||
"context"
|
||||
"errors"
|
||||
"io"
|
||||
"math"
|
||||
"os"
|
||||
"os/exec"
|
||||
"runtime"
|
||||
"strings"
|
||||
"sync/atomic"
|
||||
"testing"
|
||||
"time"
|
||||
"unicode/utf8"
|
||||
|
||||
"github.com/QuantumNous/new-api/common"
|
||||
"github.com/QuantumNous/new-api/relaykit/relayconvert/kitutil"
|
||||
"github.com/gin-gonic/gin"
|
||||
jsoniter "github.com/json-iterator/go"
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
@@ -27,6 +35,7 @@ export function sign(ctx) {
|
||||
encoded: utils.base64(ctx.message),
|
||||
};
|
||||
}
|
||||
|
||||
export const meta = { apiVersion: 1, key: "mock" };
|
||||
`, Options{
|
||||
Key: "mock", Version: "1.0.0",
|
||||
@@ -49,6 +58,253 @@ export const meta = { apiVersion: 1, key: "mock" };
|
||||
assert.Equal(t, map[string]any{"apiVersion": int64(1), "key": "mock"}, meta)
|
||||
}
|
||||
|
||||
func TestEngineJSONCloneProducesIndependentNativeContainers(t *testing.T) {
|
||||
engine, err := Compile(`
|
||||
export function clone(input) {
|
||||
const result = utils.json.clone(input);
|
||||
result.items.push({text:"appended"});
|
||||
result.items[0].text = "changed";
|
||||
return {result, array:Array.isArray(result.items),
|
||||
prototype:Object.getPrototypeOf(result) === Object.prototype,
|
||||
ownProto:Object.prototype.hasOwnProperty.call(result,"__proto__"),
|
||||
available:utils.hasCapability("json-clone@1"), unavailable:utils.hasCapability("json-clone@2")};
|
||||
}
|
||||
export function invalid(kind) {
|
||||
let value;
|
||||
switch(kind) {
|
||||
case "cycle": value={}; value.self=value; break;
|
||||
case "function": value={run:function(){}}; break;
|
||||
case "undefined": value={missing:undefined}; break;
|
||||
case "nan": value=NaN; break;
|
||||
case "bigint": value=1n; break;
|
||||
case "date": value=new Date(); break;
|
||||
case "sparse": value=[,1]; break;
|
||||
case "size": value="<".repeat(200000); break;
|
||||
case "nodes": value=new Array(32768).fill(0); break;
|
||||
}
|
||||
return utils.json.clone(value);
|
||||
}`, Options{})
|
||||
require.NoError(t, err)
|
||||
input := map[string]any{"items": []any{map[string]any{"text": "original"}}, "__proto__": map[string]any{"marker": "data"}, "zero": math.Copysign(0, -1)}
|
||||
value, err := engine.Call(t.Context(), "clone", input)
|
||||
require.NoError(t, err)
|
||||
result := value.(map[string]any)
|
||||
assert.Equal(t, true, result["array"])
|
||||
assert.Equal(t, true, result["prototype"])
|
||||
assert.Equal(t, true, result["ownProto"])
|
||||
assert.Equal(t, true, result["available"])
|
||||
assert.Equal(t, false, result["unavailable"])
|
||||
cloned := result["result"].(map[string]any)
|
||||
assert.Equal(t, []any{map[string]any{"text": "changed"}, map[string]any{"text": "appended"}}, cloned["items"])
|
||||
assert.Equal(t, map[string]any{"marker": "data"}, cloned["__proto__"])
|
||||
assert.True(t, math.Signbit(cloned["zero"].(float64)))
|
||||
assert.Equal(t, []any{map[string]any{"text": "original"}}, input["items"])
|
||||
for _, kind := range []string{"cycle", "function", "undefined", "nan", "bigint", "date", "sparse", "size", "nodes"} {
|
||||
t.Run(kind, func(t *testing.T) {
|
||||
_, err := engine.Call(t.Context(), "invalid", kind)
|
||||
require.Error(t, err)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestJSONStateChangesAndEncodedLimits(t *testing.T) {
|
||||
const changes = `[
|
||||
[{"op":"set","path":[],"value":{"text":"<","parts":[],"nested":{"value":null}}}],
|
||||
[{"op":"appendText","path":["text"],"value":"&\n\"图像😀\u2028"},
|
||||
{"op":"append","path":["parts"],"value":{"text":"a","enabled":false}},
|
||||
{"op":"set","path":["nested","value"],"value":[1,2]}],
|
||||
[{"op":"appendText","path":["parts",0,"text"],"value":"b"},
|
||||
{"op":"set","path":["nested","value",0],"value":0},
|
||||
{"op":"set","path":["nested","a<b"],"value":{}}]
|
||||
]`
|
||||
var frames []any
|
||||
require.NoError(t, common.UnmarshalJsonStr(changes, &frames))
|
||||
const expected = `{"text":"<&\n\"图像😀\u2028","parts":[{"text":"ab","enabled":false}],"nested":{"value":[0,2],"a<b":{}}}`
|
||||
var want any
|
||||
require.NoError(t, common.UnmarshalJsonStr(expected, &want))
|
||||
encoded, err := common.Marshal(want)
|
||||
require.NoError(t, err)
|
||||
state := NewJSONState(len(encoded))
|
||||
for _, frame := range frames {
|
||||
require.NoError(t, state.Apply(t.Context(), frame))
|
||||
_, err = state.Value()
|
||||
require.NoError(t, err, "every intermediate frame must match the codec's actual byte size")
|
||||
}
|
||||
value, err := state.Value()
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, want, value)
|
||||
// One extra byte fails at the update boundary and invalidates publication.
|
||||
require.Error(t, state.Apply(t.Context(), []any{map[string]any{"op": "appendText", "path": []any{"text"}, "value": "x"}}))
|
||||
_, err = state.Value()
|
||||
require.Error(t, err)
|
||||
|
||||
t.Run("replacement releases byte and node budget", func(t *testing.T) {
|
||||
state := NewJSONState(16)
|
||||
for _, value := range []any{"abcdefghijklmn", map[string]any{}, []any{int64(0), false, nil}, []any(nil), map[string]any(nil)} {
|
||||
require.NoError(t, state.Apply(t.Context(), []any{map[string]any{"op": "set", "path": []any{}, "value": value}}))
|
||||
_, err := state.Value()
|
||||
require.NoError(t, err)
|
||||
}
|
||||
})
|
||||
t.Run("encoded limit includes HTML escaping", func(t *testing.T) {
|
||||
state := NewJSONState(13)
|
||||
require.NoError(t, state.Apply(t.Context(), []any{map[string]any{"op": "set", "path": []any{}, "value": "<"}}))
|
||||
require.Error(t, state.Apply(t.Context(), []any{map[string]any{"op": "appendText", "path": []any{}, "value": "<"}}))
|
||||
})
|
||||
t.Run("request cancellation", func(t *testing.T) {
|
||||
ctx, cancel := context.WithCancel(t.Context())
|
||||
cancel()
|
||||
state := NewJSONState(64)
|
||||
require.ErrorIs(t, state.Apply(ctx, []any{}), context.Canceled)
|
||||
})
|
||||
t.Run("bounded operation and result complexity", func(t *testing.T) {
|
||||
many := make([]any, 257)
|
||||
for index := range many {
|
||||
many[index] = map[string]any{"op": "set", "path": []any{}, "value": nil}
|
||||
}
|
||||
require.Error(t, NewJSONState(1024).Apply(t.Context(), many))
|
||||
var deep any = nil
|
||||
for range 33 {
|
||||
deep = []any{deep}
|
||||
}
|
||||
for _, value := range []any{deep, make([]any, 32768), math.Inf(1)} {
|
||||
state := NewJSONState(1 << 20)
|
||||
require.Error(t, state.Apply(t.Context(), []any{map[string]any{"op": "set", "path": []any{}, "value": value}}))
|
||||
}
|
||||
})
|
||||
for _, invalid := range []string{
|
||||
`[{"op":"remove","path":[],"value":null}]`,
|
||||
`[{"op":"set","path":[],"value":null,"extra":true}]`,
|
||||
`[{"op":"set","path":["missing","child"],"value":1}]`,
|
||||
`[{"op":"set","path":["parts",-1],"value":1}]`,
|
||||
`[{"op":"set","path":["parts",0.5],"value":1}]`,
|
||||
`[{"op":"set","path":["parts",100000000000000000000],"value":1}]`,
|
||||
`[{"op":"appendText","path":["parts"],"value":"x"}]`,
|
||||
`[{"op":"append","path":["text"],"value":"x"}]`,
|
||||
} {
|
||||
t.Run(invalid, func(t *testing.T) {
|
||||
state := NewJSONState(1024)
|
||||
require.NoError(t, state.Apply(t.Context(), frames[0]))
|
||||
var changes any
|
||||
require.NoError(t, common.UnmarshalJsonStr(invalid, &changes))
|
||||
require.Error(t, state.Apply(t.Context(), changes))
|
||||
_, err := state.Value()
|
||||
require.Error(t, err)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
type bytePreservingJSONCodec struct {
|
||||
jsoniter.API
|
||||
}
|
||||
|
||||
func (c bytePreservingJSONCodec) Decode(reader io.Reader, value any) error {
|
||||
return c.NewDecoder(reader).Decode(value)
|
||||
}
|
||||
|
||||
func TestJSONStateUTF8Normalization(t *testing.T) {
|
||||
const codecEnv = "JSPLUGIN_TEST_BYTE_PRESERVING_JSON"
|
||||
if os.Getenv(codecEnv) == "1" {
|
||||
// Only this test runs in the child process; the parent's codec and
|
||||
// parallel tests are unaffected. Disabling HTML escaping in jsoniter
|
||||
// also preserves invalid UTF-8, independently of Sonic's platform support.
|
||||
kitutil.SetCodec(bytePreservingJSONCodec{jsoniter.Config{SortMapKeys: true}.Froze()})
|
||||
encoded, err := common.Marshal("\xff\xfe")
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, []byte("\"\xff\xfe\""), encoded)
|
||||
var decoded any
|
||||
require.NoError(t, common.Unmarshal(encoded, &decoded))
|
||||
require.Equal(t, "\xff\xfe", decoded)
|
||||
} else {
|
||||
t.Run("byte preserving codec", func(t *testing.T) {
|
||||
executable, err := os.Executable()
|
||||
require.NoError(t, err)
|
||||
command := exec.CommandContext(t.Context(), executable, "-test.run=^TestJSONStateUTF8Normalization$", "-test.v")
|
||||
command.Env = append(os.Environ(), codecEnv+"=1")
|
||||
output, err := command.CombinedOutput()
|
||||
require.NoError(t, err, "%s", output)
|
||||
})
|
||||
}
|
||||
|
||||
const raw = "图像\xff\xfe😀<&\n\"\\\u2028"
|
||||
const normalized = "图像\ufffd\ufffd😀<&\n\"\\\u2028"
|
||||
for _, tc := range []struct {
|
||||
name string
|
||||
op string
|
||||
path []any
|
||||
initial any
|
||||
value any
|
||||
want any
|
||||
}{
|
||||
{"set", "set", []any{}, false, raw, normalized},
|
||||
{"append", "append", []any{}, []any{float64(0), false}, raw, []any{float64(0), false, normalized}},
|
||||
{"appendText", "appendText", []any{"a<b"}, map[string]any{"a<b": "前\xff"}, raw, map[string]any{"a<b": "前\ufffd" + normalized}},
|
||||
{"separate invalid bytes", "appendText", []any{}, "\xff", "\xfe", "\ufffd\ufffd"},
|
||||
} {
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
encoded, err := common.Marshal(tc.want)
|
||||
require.NoError(t, err)
|
||||
for _, shortBy := range []int{0, 1} {
|
||||
state := NewJSONState(len(encoded) - shortBy)
|
||||
err := state.Apply(t.Context(), []any{
|
||||
map[string]any{"op": "set", "path": []any{}, "value": tc.initial},
|
||||
map[string]any{"op": tc.op, "path": tc.path, "value": tc.value},
|
||||
})
|
||||
if shortBy == 0 {
|
||||
require.NoError(t, err, "normalized JSON must fit exactly")
|
||||
value, err := state.Value()
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, tc.want, value)
|
||||
continue
|
||||
}
|
||||
require.Error(t, err, "normalized JSON must exceed the limit by one byte")
|
||||
value, valueErr := state.Value()
|
||||
require.ErrorIs(t, valueErr, err)
|
||||
assert.Nil(t, value, "a failed batch must not publish its successful prefix")
|
||||
require.ErrorIs(t, state.Apply(t.Context(), []any{
|
||||
map[string]any{"op": "set", "path": []any{}, "value": nil},
|
||||
}), err, "a later replacement must not revive failed state")
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
t.Run("clone", func(t *testing.T) {
|
||||
engine, err := Compile(`export function clone(input) { return utils.json.clone(input); }`, Options{})
|
||||
require.NoError(t, err)
|
||||
input := map[string]any{"text": raw, "items": []any{float64(0), false, raw}}
|
||||
value, err := engine.Call(t.Context(), "clone", input)
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, map[string]any{"text": normalized, "items": []any{int64(0), false, normalized}}, value)
|
||||
assert.Equal(t, raw, input["text"])
|
||||
|
||||
encoded, err := common.Marshal(normalized)
|
||||
require.NoError(t, err)
|
||||
padding := strings.Repeat("x", MaxJSONToolBytes-len(encoded))
|
||||
value, err = engine.Call(t.Context(), "clone", raw+padding)
|
||||
require.NoError(t, err, "normalized clone must fit exactly")
|
||||
assert.Equal(t, normalized, strings.TrimSuffix(value.(string), padding))
|
||||
value, err = engine.Call(t.Context(), "clone", raw+padding+"x")
|
||||
require.Error(t, err, "normalized clone must exceed the limit by one byte")
|
||||
assert.Nil(t, value)
|
||||
})
|
||||
|
||||
t.Run("invalid keys and paths remain rejected", func(t *testing.T) {
|
||||
for _, change := range []any{
|
||||
map[string]any{"op": "set", "path": []any{}, "value": map[string]any{"\xff\xfe": raw}},
|
||||
map[string]any{"op": "set", "path": []any{"\xff\xfe"}, "value": raw},
|
||||
} {
|
||||
state := NewJSONState(128)
|
||||
require.Error(t, state.Apply(t.Context(), []any{
|
||||
map[string]any{"op": "set", "path": []any{}, "value": map[string]any{"\ufffd\ufffd": "original"}},
|
||||
change,
|
||||
}))
|
||||
value, err := state.Value()
|
||||
require.Error(t, err)
|
||||
assert.Nil(t, value)
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
func TestEngineConsoleLogUsesDebugLoggerAndRequestContext(t *testing.T) {
|
||||
previousDebug := common.DebugEnabled
|
||||
common.DebugEnabled = false
|
||||
@@ -160,16 +416,83 @@ func TestCompileIgnoresSourceMapDirectives(t *testing.T) {
|
||||
|
||||
func TestEngineInterruptsLongRunningHook(t *testing.T) {
|
||||
t.Parallel()
|
||||
engine, err := Compile(`export function run() { while (true) {} }`, Options{
|
||||
engine, err := Compile(`let calls = 0; export function run(loop) { calls++; if (loop) { while (true) {} } return calls; }`, Options{
|
||||
Key: "loop", Version: "1", Timeout: 20 * time.Millisecond,
|
||||
})
|
||||
require.NoError(t, err)
|
||||
|
||||
_, err = engine.Call(context.Background(), "run")
|
||||
_, err = engine.Call(context.Background(), "run", true)
|
||||
require.Error(t, err)
|
||||
assert.Contains(t, err.Error(), "timed out")
|
||||
var hookErr *HookError
|
||||
assert.False(t, errors.As(err, &hookErr), "timeouts must not be HookError")
|
||||
value, err := engine.Call(context.Background(), "run", false)
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, int64(1), value, "an interrupted runtime must be discarded")
|
||||
}
|
||||
|
||||
func TestEngineDoesNotReinitializeIdleRuntimeAfterGC(t *testing.T) {
|
||||
var initialized atomic.Int64
|
||||
engine, err := Compile(`console.log("initialized"); export function run(){return 42;}`, Options{
|
||||
Log: func(string) { initialized.Add(1) },
|
||||
})
|
||||
require.NoError(t, err)
|
||||
runtime.GC()
|
||||
runtime.GC()
|
||||
value, err := engine.Call(t.Context(), "run")
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, int64(42), value)
|
||||
assert.Equal(t, int64(1), initialized.Load())
|
||||
}
|
||||
|
||||
func TestEngineBoundsConcurrentCallsAndCancelsAnOccupiedRuntime(t *testing.T) {
|
||||
entered := make(chan struct{}, 2)
|
||||
release := make(chan struct{})
|
||||
defer close(release)
|
||||
engine, err := Compile(`export function run(loop) { utils.unixNow(); if (loop) { while (true) {} } return 42; }`, Options{
|
||||
Concurrency: 2,
|
||||
Now: func() time.Time {
|
||||
entered <- struct{}{}
|
||||
<-release
|
||||
return time.Unix(1, 0)
|
||||
},
|
||||
})
|
||||
require.NoError(t, err)
|
||||
ctx, cancel := context.WithCancel(t.Context())
|
||||
defer cancel()
|
||||
canceled := make(chan error, 1)
|
||||
healthy := make(chan error, 1)
|
||||
go func() { _, err := engine.Call(ctx, "run", true); canceled <- err }()
|
||||
go func() { _, err := engine.Call(t.Context(), "run", false); healthy <- err }()
|
||||
<-entered
|
||||
<-entered
|
||||
_, err = engine.CallPathWithAdmissionTimeout(t.Context(), time.Nanosecond, "run", nil, false)
|
||||
require.ErrorIs(t, err, ErrCallAdmissionTimeout)
|
||||
cancel()
|
||||
release <- struct{}{}
|
||||
release <- struct{}{}
|
||||
require.ErrorIs(t, <-canceled, context.Canceled)
|
||||
require.NoError(t, <-healthy)
|
||||
}
|
||||
|
||||
func TestEnginePreservesMutableExportAndMemberBindings(t *testing.T) {
|
||||
engine, err := Compile(`
|
||||
export let run = function(){return 1;};
|
||||
export const native = {render:function(){return 1;}};
|
||||
export function replace(){run=function(){return 2;};native.render=function(){return 3;};}
|
||||
`, Options{Concurrency: 1})
|
||||
require.NoError(t, err)
|
||||
value, err := engine.Call(t.Context(), "run")
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, int64(1), value)
|
||||
_, err = engine.Call(t.Context(), "replace")
|
||||
require.NoError(t, err)
|
||||
value, err = engine.Call(t.Context(), "run")
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, int64(2), value)
|
||||
value, err = engine.CallMember(t.Context(), "native", "render")
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, int64(3), value)
|
||||
}
|
||||
|
||||
func TestEngineHookErrorExtractsSanitizedJSMessage(t *testing.T) {
|
||||
|
||||
@@ -0,0 +1,428 @@
|
||||
package jsplugin
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"math"
|
||||
"strconv"
|
||||
"strings"
|
||||
"unicode/utf8"
|
||||
|
||||
"github.com/QuantumNous/new-api/common"
|
||||
"github.com/grafana/sobek"
|
||||
)
|
||||
|
||||
const (
|
||||
CapabilityJSONClone = "json-clone@1"
|
||||
CapabilitySubmitSSEDelta = "submit-sse-delta@1"
|
||||
MaxJSONToolBytes = 1 << 20
|
||||
maxJSONToolDepth = 32
|
||||
maxJSONToolNodes = 32768
|
||||
maxJSONChanges = 256
|
||||
)
|
||||
|
||||
// HasCapability describes host APIs independently of a plugin's mutable globals.
|
||||
func HasCapability(name string) bool {
|
||||
return name == CapabilityJSONClone || name == CapabilitySubmitSSEDelta
|
||||
}
|
||||
|
||||
// JSONState owns a request-local JSON result. Appended strings stay in Go and
|
||||
// are materialized only when the stream finishes, independently of JS runtimes.
|
||||
type JSONState struct {
|
||||
root *jsonStateNode
|
||||
limit int
|
||||
failed error
|
||||
}
|
||||
|
||||
type jsonStateNode struct {
|
||||
scalar any
|
||||
object map[string]*jsonStateNode
|
||||
array []*jsonStateNode
|
||||
text *strings.Builder
|
||||
bytes int
|
||||
nodes int
|
||||
}
|
||||
|
||||
type jsonStateBudget struct {
|
||||
ctx context.Context
|
||||
bytes int
|
||||
nodes int
|
||||
}
|
||||
|
||||
func (b *jsonStateBudget) spend(bytes, nodes int) error {
|
||||
if err := b.ctx.Err(); err != nil {
|
||||
return err
|
||||
}
|
||||
if bytes > b.bytes || nodes > b.nodes {
|
||||
return fmt.Errorf("JSON state exceeds size or node limit")
|
||||
}
|
||||
b.bytes -= bytes
|
||||
b.nodes -= nodes
|
||||
return nil
|
||||
}
|
||||
|
||||
func newJSONStateNode(value any, depth int, budget *jsonStateBudget) (*jsonStateNode, error) {
|
||||
if depth > maxJSONToolDepth {
|
||||
return nil, fmt.Errorf("JSON state exceeds depth limit")
|
||||
}
|
||||
if js, ok := value.(sobek.Value); ok {
|
||||
if sobek.IsUndefined(js) {
|
||||
return nil, fmt.Errorf("undefined is not a JSON value")
|
||||
}
|
||||
if object, ok := js.(*sobek.Object); ok {
|
||||
switch object.ClassName() {
|
||||
case "Object":
|
||||
keys := object.Keys()
|
||||
if len(keys) > budget.nodes {
|
||||
return nil, fmt.Errorf("JSON object exceeds node limit")
|
||||
}
|
||||
fields := make(map[string]any, len(keys))
|
||||
for _, key := range keys {
|
||||
fields[key] = object.Get(key)
|
||||
}
|
||||
value = fields
|
||||
case "Array":
|
||||
length := object.Get("length").ToInteger()
|
||||
if length < 0 || length > int64(budget.nodes) {
|
||||
return nil, fmt.Errorf("JSON array exceeds node limit")
|
||||
}
|
||||
items := make([]any, int(length))
|
||||
for index := range items {
|
||||
item := object.Get(strconv.Itoa(index))
|
||||
if item == nil {
|
||||
return nil, fmt.Errorf("sparse arrays are not JSON values")
|
||||
}
|
||||
items[index] = item
|
||||
}
|
||||
value = items
|
||||
default:
|
||||
return nil, fmt.Errorf("json.clone accepts only plain objects, arrays and JSON scalars")
|
||||
}
|
||||
} else {
|
||||
value = js.Export()
|
||||
}
|
||||
}
|
||||
beforeBytes, beforeNodes := budget.bytes, budget.nodes
|
||||
if err := budget.spend(0, 1); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
switch typed := value.(type) {
|
||||
case map[string]any:
|
||||
if typed == nil {
|
||||
value = nil
|
||||
}
|
||||
case []any:
|
||||
if typed == nil {
|
||||
value = nil
|
||||
}
|
||||
}
|
||||
node := &jsonStateNode{}
|
||||
switch typed := value.(type) {
|
||||
case map[string]any:
|
||||
if len(typed) > budget.nodes {
|
||||
return nil, fmt.Errorf("JSON object exceeds node limit")
|
||||
}
|
||||
if err := budget.spend(2, 0); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
node.object = make(map[string]*jsonStateNode, len(typed))
|
||||
for key, child := range typed {
|
||||
if !utf8.ValidString(key) || len(key) > budget.bytes {
|
||||
return nil, fmt.Errorf("invalid or oversized JSON object key")
|
||||
}
|
||||
encoded, err := common.Marshal(key)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
cost := len(encoded) + 1
|
||||
if len(node.object) != 0 {
|
||||
cost++
|
||||
}
|
||||
if err = budget.spend(cost, 0); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
node.object[key], err = newJSONStateNode(child, depth+1, budget)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
}
|
||||
case []any:
|
||||
if len(typed) > budget.nodes {
|
||||
return nil, fmt.Errorf("JSON array exceeds node limit")
|
||||
}
|
||||
if err := budget.spend(2+max(0, len(typed)-1), 0); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
node.array = make([]*jsonStateNode, 0, len(typed))
|
||||
for _, child := range typed {
|
||||
item, err := newJSONStateNode(child, depth+1, budget)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
node.array = append(node.array, item)
|
||||
}
|
||||
default:
|
||||
switch number := value.(type) {
|
||||
case int:
|
||||
value = float64(number)
|
||||
case int64:
|
||||
value = float64(number)
|
||||
}
|
||||
switch scalar := value.(type) {
|
||||
case nil, bool:
|
||||
case float64:
|
||||
if math.IsNaN(scalar) || math.IsInf(scalar, 0) {
|
||||
return nil, fmt.Errorf("JSON numbers must be finite")
|
||||
}
|
||||
case string:
|
||||
if len(scalar) > budget.bytes {
|
||||
return nil, fmt.Errorf("JSON string exceeds size limit")
|
||||
}
|
||||
if !utf8.ValidString(scalar) {
|
||||
value = string([]rune(scalar))
|
||||
}
|
||||
default:
|
||||
return nil, fmt.Errorf("value must contain only JSON objects, arrays and scalars")
|
||||
}
|
||||
encoded, err := common.Marshal(value)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if err = budget.spend(len(encoded), 0); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
node.scalar = value
|
||||
}
|
||||
node.bytes, node.nodes = beforeBytes-budget.bytes, beforeNodes-budget.nodes
|
||||
return node, nil
|
||||
}
|
||||
|
||||
func NewJSONState(limit int) *JSONState {
|
||||
return &JSONState{root: &jsonStateNode{bytes: 4, nodes: 1}, limit: max(0, min(limit, MaxJSONToolBytes))}
|
||||
}
|
||||
|
||||
// Apply accepts a bounded batch of set, append and appendText operations.
|
||||
// Failure invalidates the stream: callers must discard the state rather than
|
||||
// continue after an error. No partially updated result may be published.
|
||||
func (s *JSONState) Apply(ctx context.Context, changes any) (err error) {
|
||||
if s.failed != nil {
|
||||
return s.failed
|
||||
}
|
||||
defer func() {
|
||||
if err != nil {
|
||||
s.failed = err
|
||||
}
|
||||
}()
|
||||
if err = ctx.Err(); err != nil {
|
||||
return err
|
||||
}
|
||||
items, ok := changes.([]any)
|
||||
if !ok || len(items) > maxJSONChanges {
|
||||
return fmt.Errorf("changes must be an array of at most %d operations", maxJSONChanges)
|
||||
}
|
||||
for _, item := range items {
|
||||
if err := ctx.Err(); err != nil {
|
||||
return err
|
||||
}
|
||||
change, ok := item.(map[string]any)
|
||||
if !ok || len(change) != 3 {
|
||||
return fmt.Errorf("each change must contain only op, path and value")
|
||||
}
|
||||
op, _ := change["op"].(string)
|
||||
path, ok := change["path"].([]any)
|
||||
value, hasValue := change["value"]
|
||||
if !ok || !hasValue || len(path) > maxJSONToolDepth {
|
||||
return fmt.Errorf("change requires a JSON value and a path of at most %d segments", maxJSONToolDepth)
|
||||
}
|
||||
if op != "set" && op != "append" && op != "appendText" {
|
||||
return fmt.Errorf("unsupported JSON change operation %q", op)
|
||||
}
|
||||
ancestors := make([]*jsonStateNode, 0, len(path))
|
||||
target := s.root
|
||||
var key string
|
||||
var index int
|
||||
for offset, segment := range path {
|
||||
if target == nil {
|
||||
return fmt.Errorf("JSON change parent does not exist")
|
||||
}
|
||||
ancestors = append(ancestors, target)
|
||||
switch {
|
||||
case target.object != nil:
|
||||
key, ok = segment.(string)
|
||||
if !ok || !utf8.ValidString(key) || len(key) > s.limit {
|
||||
return fmt.Errorf("JSON object paths require valid string keys")
|
||||
}
|
||||
target = target.object[key]
|
||||
case target.array != nil:
|
||||
var number float64
|
||||
switch typed := segment.(type) {
|
||||
case int:
|
||||
number = float64(typed)
|
||||
case int64:
|
||||
number = float64(typed)
|
||||
case float64:
|
||||
number = typed
|
||||
default:
|
||||
return fmt.Errorf("JSON array paths require integer indices")
|
||||
}
|
||||
if math.IsNaN(number) || number < 0 || number >= float64(len(target.array)) || math.Trunc(number) != number {
|
||||
return fmt.Errorf("JSON array index is out of range")
|
||||
}
|
||||
index = int(number)
|
||||
target = target.array[index]
|
||||
default:
|
||||
return fmt.Errorf("JSON change path traverses a scalar")
|
||||
}
|
||||
if target == nil && (offset != len(path)-1 || op != "set") {
|
||||
return fmt.Errorf("JSON change target does not exist")
|
||||
}
|
||||
}
|
||||
addedBytes, addedNodes := 0, 0
|
||||
switch op {
|
||||
case "set":
|
||||
oldBytes, oldNodes := 0, 0
|
||||
if target != nil {
|
||||
oldBytes, oldNodes = target.bytes, target.nodes
|
||||
} else {
|
||||
encodedKey, err := common.Marshal(key)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
addedBytes = len(encodedKey) + 1
|
||||
if len(ancestors[len(ancestors)-1].object) != 0 {
|
||||
addedBytes++
|
||||
}
|
||||
}
|
||||
budget := &jsonStateBudget{ctx: ctx, bytes: s.limit - s.root.bytes + oldBytes - addedBytes, nodes: maxJSONToolNodes - s.root.nodes + oldNodes}
|
||||
next, err := newJSONStateNode(value, len(path), budget)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
addedBytes += next.bytes - oldBytes
|
||||
addedNodes = next.nodes - oldNodes
|
||||
if len(ancestors) == 0 {
|
||||
s.root = next
|
||||
} else if parent := ancestors[len(ancestors)-1]; parent.object != nil {
|
||||
parent.object[key] = next
|
||||
} else {
|
||||
parent.array[index] = next
|
||||
}
|
||||
case "append":
|
||||
if target.array == nil {
|
||||
return fmt.Errorf("append requires an array target")
|
||||
}
|
||||
if len(target.array) != 0 {
|
||||
addedBytes = 1
|
||||
}
|
||||
budget := &jsonStateBudget{ctx: ctx, bytes: s.limit - s.root.bytes - addedBytes, nodes: maxJSONToolNodes - s.root.nodes}
|
||||
next, err := newJSONStateNode(value, len(path)+1, budget)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
addedBytes += next.bytes
|
||||
addedNodes = next.nodes
|
||||
target.array = append(target.array, next)
|
||||
target.bytes += addedBytes
|
||||
target.nodes += addedNodes
|
||||
case "appendText":
|
||||
previous, isText := target.scalar.(string)
|
||||
text, ok := value.(string)
|
||||
if !isText || !ok || len(text) > s.limit-s.root.bytes {
|
||||
return fmt.Errorf("appendText requires strings within the state size limit")
|
||||
}
|
||||
if !utf8.ValidString(text) {
|
||||
text = string([]rune(text))
|
||||
}
|
||||
encoded, err := common.Marshal(text)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
addedBytes = len(encoded) - 2
|
||||
if addedBytes > s.limit-s.root.bytes {
|
||||
return fmt.Errorf("JSON state exceeds size limit")
|
||||
}
|
||||
if text != "" {
|
||||
if target.text == nil {
|
||||
target.text = &strings.Builder{}
|
||||
target.text.WriteString(previous)
|
||||
target.scalar = ""
|
||||
}
|
||||
target.text.WriteString(text)
|
||||
}
|
||||
target.bytes += addedBytes
|
||||
}
|
||||
for _, ancestor := range ancestors {
|
||||
ancestor.bytes += addedBytes
|
||||
ancestor.nodes += addedNodes
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// Value returns the final ordinary JSON tree and verifies the incremental byte
|
||||
// accounting against the configured codec before callers can persist it.
|
||||
func (s *JSONState) Value() (any, error) {
|
||||
if s.failed != nil {
|
||||
return nil, s.failed
|
||||
}
|
||||
value := s.root.value()
|
||||
encoded, err := common.Marshal(value)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if len(encoded) != s.root.bytes || len(encoded) > s.limit {
|
||||
return nil, fmt.Errorf("JSON state encoded size does not match its bounded representation")
|
||||
}
|
||||
return value, nil
|
||||
}
|
||||
|
||||
func (n *jsonStateNode) value() any {
|
||||
if n.object != nil {
|
||||
result := make(map[string]any, len(n.object))
|
||||
for key, child := range n.object {
|
||||
result[key] = child.value()
|
||||
}
|
||||
return result
|
||||
}
|
||||
if n.array != nil {
|
||||
result := make([]any, len(n.array))
|
||||
for index, child := range n.array {
|
||||
result[index] = child.value()
|
||||
}
|
||||
return result
|
||||
}
|
||||
if n.text != nil {
|
||||
return n.text.String()
|
||||
}
|
||||
return n.scalar
|
||||
}
|
||||
|
||||
// Native JS containers are essential here: wrapping Go slices in ToValue would
|
||||
// retain host-backed array mutation behavior instead of producing a JS clone.
|
||||
func (n *jsonStateNode) jsValue(runtime *sobek.Runtime) (sobek.Value, error) {
|
||||
if n.object != nil {
|
||||
result := runtime.NewObject()
|
||||
for key, child := range n.object {
|
||||
value, err := child.jsValue(runtime)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if err = result.DefineDataProperty(key, value, sobek.FLAG_TRUE, sobek.FLAG_TRUE, sobek.FLAG_TRUE); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
}
|
||||
return result, nil
|
||||
}
|
||||
if n.array != nil {
|
||||
items := make([]any, len(n.array))
|
||||
for index, child := range n.array {
|
||||
value, err := child.jsValue(runtime)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
items[index] = value
|
||||
}
|
||||
return runtime.NewArray(items...), nil
|
||||
}
|
||||
return runtime.ToValue(n.value()), nil
|
||||
}
|
||||
+232
-50
@@ -40,6 +40,11 @@ var pluginKeyPattern = regexp.MustCompile(`^[a-z0-9][a-z0-9_-]*$`)
|
||||
var pluginVersionPattern = regexp.MustCompile(`^(0|[1-9][0-9]*)\.(0|[1-9][0-9]*)\.(0|[1-9][0-9]*)(?:-[0-9A-Za-z.-]+)?(?:\+[0-9A-Za-z.-]+)?$`)
|
||||
var localeTagPattern = regexp.MustCompile(`^[a-zA-Z]{2,3}(-[a-zA-Z0-9]{2,8})*$`)
|
||||
|
||||
// ValidPluginKey checks the canonical identifier accepted by plugin manifests.
|
||||
func ValidPluginKey(key string) bool {
|
||||
return len(key) <= 30 && pluginKeyPattern.MatchString(key)
|
||||
}
|
||||
|
||||
// LocalizedText is locale-keyed display copy. Plugin source may use a bare
|
||||
// string (normalized to {"en": s}) or a map that must include "en". API
|
||||
// responses always emit an object.
|
||||
@@ -79,25 +84,50 @@ func (t *LocalizedText) UnmarshalJSON(data []byte) error {
|
||||
}
|
||||
|
||||
type Meta struct {
|
||||
SortPriority int `json:"sortPriority,omitempty"`
|
||||
Website string `json:"website,omitempty"`
|
||||
APIVersion int `json:"apiVersion"`
|
||||
Key string `json:"key"`
|
||||
Name string `json:"name"`
|
||||
Icon string `json:"icon,omitempty"`
|
||||
Description LocalizedText `json:"description,omitempty"`
|
||||
Version string `json:"version"`
|
||||
Author AuthorMeta `json:"author"`
|
||||
BaseURL string `json:"baseUrl,omitempty"`
|
||||
ChannelTypes []int `json:"channelTypes,omitempty"`
|
||||
Models []string `json:"models"`
|
||||
FetchMode string `json:"fetchMode"`
|
||||
AllowedHosts []string `json:"allowedHosts"`
|
||||
Routes []Route `json:"routes"`
|
||||
Protocols []ProtocolClaim `json:"protocols"`
|
||||
UsageSchema map[string]UsageFieldSchema `json:"usageSchema,omitempty"`
|
||||
UsageExamples []UsageExample `json:"usageExamples,omitempty"`
|
||||
Auth AuthMeta `json:"auth"`
|
||||
RequiredCapabilities []string `json:"requiredCapabilities,omitempty"`
|
||||
SubmitResponseTypes []string `json:"submitResponseTypes,omitempty"`
|
||||
SortPriority int `json:"sortPriority,omitempty"`
|
||||
Website string `json:"website,omitempty"`
|
||||
APIVersion int `json:"apiVersion"`
|
||||
Key string `json:"key"`
|
||||
Name string `json:"name"`
|
||||
Icon string `json:"icon,omitempty"`
|
||||
Description LocalizedText `json:"description,omitempty"`
|
||||
Version string `json:"version"`
|
||||
Author AuthorMeta `json:"author"`
|
||||
BaseURL string `json:"baseUrl,omitempty"`
|
||||
ChannelTypes []int `json:"channelTypes,omitempty"`
|
||||
Models []string `json:"models"`
|
||||
FetchMode string `json:"fetchMode"`
|
||||
AllowedHosts []string `json:"allowedHosts"`
|
||||
Routes []Route `json:"routes"`
|
||||
Protocols []ProtocolClaim `json:"protocols"`
|
||||
UsageSchema map[string]UsageFieldSchema `json:"usageSchema,omitempty"`
|
||||
UsageExamples []UsageExample `json:"usageExamples,omitempty"`
|
||||
UsageProfiles []UsageProfile `json:"usageProfiles,omitempty"`
|
||||
Auth AuthMeta `json:"auth"`
|
||||
}
|
||||
|
||||
// UsageProfile replaces the plugin's default usage metadata for its models.
|
||||
type UsageProfile struct {
|
||||
Models []string `json:"models"`
|
||||
Schema map[string]UsageFieldSchema `json:"schema"`
|
||||
Examples []UsageExample `json:"examples,omitempty"`
|
||||
}
|
||||
|
||||
// UsageForModel returns read-only usage metadata for a declared model. Aliases
|
||||
// must be resolved by the host first; an unknown or ambiguous model uses the
|
||||
// plugin defaults. Profile examples never inherit the default examples.
|
||||
func (m Meta) UsageForModel(model string) (map[string]UsageFieldSchema, []UsageExample) {
|
||||
folded := asciiFold(model)
|
||||
for _, profile := range m.UsageProfiles {
|
||||
for _, declared := range profile.Models {
|
||||
if asciiFold(declared) == folded {
|
||||
return profile.Schema, profile.Examples
|
||||
}
|
||||
}
|
||||
}
|
||||
return m.UsageSchema, m.UsageExamples
|
||||
}
|
||||
|
||||
// ProtocolSupports reports whether the named protocol claim includes mode.
|
||||
@@ -132,6 +162,7 @@ type AuthMeta struct {
|
||||
type UsageFieldSchema struct {
|
||||
Type string `json:"type,omitempty"`
|
||||
Unit string `json:"unit,omitempty"`
|
||||
UnitLabel LocalizedText `json:"unitLabel,omitempty"`
|
||||
Enum []string `json:"enum,omitempty"`
|
||||
Description LocalizedText `json:"description,omitempty"`
|
||||
EnumLabels map[string]LocalizedText `json:"enumLabels,omitempty"`
|
||||
@@ -264,6 +295,13 @@ func CompilePlugin(source string, options Options) (*LoadedPlugin, error) {
|
||||
engine.key = meta.Key
|
||||
engine.version = meta.Version
|
||||
requiredHooks := []string{"buildSubmitRequest", "parseSubmitResponse", "parseTaskResult"}
|
||||
if slices.Contains(meta.SubmitResponseTypes, "sse") {
|
||||
if slices.Contains(meta.RequiredCapabilities, CapabilitySubmitSSEDelta) {
|
||||
requiredHooks = append(requiredHooks, "parseSubmitEventDelta")
|
||||
} else {
|
||||
requiredHooks = append(requiredHooks, "parseSubmitEvent")
|
||||
}
|
||||
}
|
||||
if meta.FetchMode == "batch" {
|
||||
requiredHooks = append(requiredHooks, "buildBatchQueryRequest", "parseBatchResult")
|
||||
} else {
|
||||
@@ -791,6 +829,8 @@ func (r *Registry) Snapshot() RegistrySnapshot {
|
||||
}
|
||||
|
||||
func cloneMeta(meta Meta) Meta {
|
||||
meta.SubmitResponseTypes = slices.Clone(meta.SubmitResponseTypes)
|
||||
meta.RequiredCapabilities = slices.Clone(meta.RequiredCapabilities)
|
||||
meta.ChannelTypes = append([]int(nil), meta.ChannelTypes...)
|
||||
meta.Models = append([]string(nil), meta.Models...)
|
||||
meta.AllowedHosts = append([]string(nil), meta.AllowedHosts...)
|
||||
@@ -806,31 +846,46 @@ func cloneMeta(meta Meta) Meta {
|
||||
if meta.Description != nil {
|
||||
meta.Description = maps.Clone(meta.Description)
|
||||
}
|
||||
if meta.UsageSchema != nil {
|
||||
usageSchema := make(map[string]UsageFieldSchema, len(meta.UsageSchema))
|
||||
for key, field := range meta.UsageSchema {
|
||||
if field.Enum != nil {
|
||||
field.Enum = append([]string{}, field.Enum...)
|
||||
}
|
||||
if field.Description != nil {
|
||||
field.Description = maps.Clone(field.Description)
|
||||
}
|
||||
if field.EnumLabels != nil {
|
||||
labels := make(map[string]LocalizedText, len(field.EnumLabels))
|
||||
for value, label := range field.EnumLabels {
|
||||
labels[value] = maps.Clone(label)
|
||||
}
|
||||
field.EnumLabels = labels
|
||||
}
|
||||
usageSchema[key] = field
|
||||
}
|
||||
meta.UsageSchema = usageSchema
|
||||
meta.UsageSchema = CloneUsageSchema(meta.UsageSchema)
|
||||
meta.UsageExamples = CloneUsageExamples(meta.UsageExamples)
|
||||
meta.UsageProfiles = append([]UsageProfile(nil), meta.UsageProfiles...)
|
||||
for index := range meta.UsageProfiles {
|
||||
profile := &meta.UsageProfiles[index]
|
||||
profile.Models = append([]string(nil), profile.Models...)
|
||||
profile.Schema = CloneUsageSchema(profile.Schema)
|
||||
profile.Examples = CloneUsageExamples(profile.Examples)
|
||||
}
|
||||
meta.UsageExamples = cloneUsageExamples(meta.UsageExamples)
|
||||
return meta
|
||||
}
|
||||
|
||||
func cloneUsageExamples(examples []UsageExample) []UsageExample {
|
||||
// CloneUsageSchema copies schema fields and localized metadata for independent readers.
|
||||
func CloneUsageSchema(schema map[string]UsageFieldSchema) map[string]UsageFieldSchema {
|
||||
if schema == nil {
|
||||
return nil
|
||||
}
|
||||
cloned := make(map[string]UsageFieldSchema, len(schema))
|
||||
for key, field := range schema {
|
||||
if field.Enum != nil {
|
||||
field.Enum = append([]string{}, field.Enum...)
|
||||
}
|
||||
if field.Description != nil {
|
||||
field.Description = maps.Clone(field.Description)
|
||||
}
|
||||
field.UnitLabel = maps.Clone(field.UnitLabel)
|
||||
if field.EnumLabels != nil {
|
||||
labels := make(map[string]LocalizedText, len(field.EnumLabels))
|
||||
for value, label := range field.EnumLabels {
|
||||
labels[value] = maps.Clone(label)
|
||||
}
|
||||
field.EnumLabels = labels
|
||||
}
|
||||
cloned[key] = field
|
||||
}
|
||||
return cloned
|
||||
}
|
||||
|
||||
// CloneUsageExamples copies example facts for independent readers.
|
||||
func CloneUsageExamples(examples []UsageExample) []UsageExample {
|
||||
if examples == nil {
|
||||
return nil
|
||||
}
|
||||
@@ -904,6 +959,15 @@ func (r *Registry) resolveActiveOverrides(
|
||||
return active
|
||||
}
|
||||
|
||||
// UnknownMetaFieldError identifies a manifest field unsupported by this host.
|
||||
type UnknownMetaFieldError struct {
|
||||
Field string
|
||||
}
|
||||
|
||||
func (e *UnknownMetaFieldError) Error() string {
|
||||
return fmt.Sprintf("plugin meta has unknown field %q", e.Field)
|
||||
}
|
||||
|
||||
func decodeMeta(value any) (Meta, error) {
|
||||
object, ok := value.(map[string]any)
|
||||
if !ok {
|
||||
@@ -911,9 +975,9 @@ func decodeMeta(value any) (Meta, error) {
|
||||
}
|
||||
for field := range object {
|
||||
switch field {
|
||||
case "sortPriority", "website", "apiVersion", "key", "name", "icon", "description", "version", "author", "baseUrl", "channelTypes", "channelType", "compatibleChannelTypes", "models", "fetchMode", "allowedHosts", "routes", "protocols", "usageSchema", "usageExamples", "auth", "endpoints", "submitPaths", "actions":
|
||||
case "requiredCapabilities", "submitResponseTypes", "sortPriority", "website", "apiVersion", "key", "name", "icon", "description", "version", "author", "baseUrl", "channelTypes", "channelType", "compatibleChannelTypes", "models", "fetchMode", "allowedHosts", "routes", "protocols", "usageSchema", "usageExamples", "usageProfiles", "auth", "endpoints", "submitPaths", "actions":
|
||||
default:
|
||||
return Meta{}, fmt.Errorf("plugin meta has unknown field %q", field)
|
||||
return Meta{}, &UnknownMetaFieldError{Field: field}
|
||||
}
|
||||
}
|
||||
meta := Meta{}
|
||||
@@ -978,6 +1042,27 @@ func decodeMeta(value any) (Meta, error) {
|
||||
if meta.FetchMode, err = stringMetaField(object, "fetchMode"); err != nil {
|
||||
return Meta{}, err
|
||||
}
|
||||
meta.SubmitResponseTypes = []string{"json"}
|
||||
if _, present := object["submitResponseTypes"]; present {
|
||||
meta.SubmitResponseTypes, err = strictStringSlice(object, "submitResponseTypes")
|
||||
if err != nil {
|
||||
return Meta{}, err
|
||||
}
|
||||
if len(meta.SubmitResponseTypes) == 0 {
|
||||
return Meta{}, fmt.Errorf("submitResponseTypes must not be empty")
|
||||
}
|
||||
seen := make(map[string]bool)
|
||||
for _, kind := range meta.SubmitResponseTypes {
|
||||
if (kind != "json" && kind != "sse") || seen[kind] {
|
||||
return Meta{}, fmt.Errorf("invalid or duplicate submitResponseTypes value %q", kind)
|
||||
}
|
||||
seen[kind] = true
|
||||
}
|
||||
}
|
||||
meta.RequiredCapabilities, err = strictStringSlice(object, "requiredCapabilities")
|
||||
if err != nil {
|
||||
return Meta{}, err
|
||||
}
|
||||
meta.Models, err = strictStringSlice(object, "models")
|
||||
if err != nil {
|
||||
return Meta{}, err
|
||||
@@ -1009,6 +1094,12 @@ func decodeMeta(value any) (Meta, error) {
|
||||
return Meta{}, err
|
||||
}
|
||||
}
|
||||
if usageProfiles, exists := object["usageProfiles"]; exists {
|
||||
meta.UsageProfiles, err = decodeUsageProfiles(usageProfiles)
|
||||
if err != nil {
|
||||
return Meta{}, err
|
||||
}
|
||||
}
|
||||
for _, removedField := range []string{"submitPaths", "actions"} {
|
||||
if _, exists := object[removedField]; exists {
|
||||
return Meta{}, fmt.Errorf("plugin meta %s is no longer supported; declare routes instead", removedField)
|
||||
@@ -1058,6 +1149,16 @@ func ValidateV1Meta(meta Meta) error {
|
||||
}
|
||||
|
||||
func normalizeV1Meta(meta *Meta) error {
|
||||
seenCapabilities := make(map[string]bool, len(meta.RequiredCapabilities))
|
||||
for _, name := range meta.RequiredCapabilities {
|
||||
if !HasCapability(name) || seenCapabilities[name] {
|
||||
return fmt.Errorf("unsupported or duplicate required capability %q", name)
|
||||
}
|
||||
if name == CapabilitySubmitSSEDelta && !slices.Contains(meta.SubmitResponseTypes, "sse") {
|
||||
return fmt.Errorf("%s requires submitResponseTypes to include sse", name)
|
||||
}
|
||||
seenCapabilities[name] = true
|
||||
}
|
||||
if meta.SortPriority < math.MinInt32 || meta.SortPriority > math.MaxInt32 {
|
||||
return fmt.Errorf("plugin meta sortPriority must be a signed 32-bit integer")
|
||||
}
|
||||
@@ -1132,12 +1233,12 @@ func normalizeV1Meta(meta *Meta) error {
|
||||
}
|
||||
meta.BaseURL = normalized
|
||||
}
|
||||
if !pluginKeyPattern.MatchString(meta.Key) {
|
||||
return fmt.Errorf("plugin meta key must match %s", pluginKeyPattern)
|
||||
}
|
||||
if len(meta.Key) > 30 {
|
||||
return fmt.Errorf("plugin meta key must not exceed 30 characters")
|
||||
}
|
||||
if !ValidPluginKey(meta.Key) {
|
||||
return fmt.Errorf("plugin meta key must match %s", pluginKeyPattern)
|
||||
}
|
||||
if !pluginVersionPattern.MatchString(meta.Version) {
|
||||
return fmt.Errorf("plugin meta version must be semver")
|
||||
}
|
||||
@@ -1258,7 +1359,44 @@ func normalizeV1Meta(meta *Meta) error {
|
||||
}
|
||||
}
|
||||
}
|
||||
for name, field := range meta.UsageSchema {
|
||||
if err := validateUsageSchema(meta.UsageSchema); err != nil {
|
||||
return err
|
||||
}
|
||||
if err := validateUsageExamples(meta.UsageSchema, meta.UsageExamples); err != nil {
|
||||
return err
|
||||
}
|
||||
profileModels := make(map[string]struct{})
|
||||
for index, profile := range meta.UsageProfiles {
|
||||
if len(profile.Models) == 0 {
|
||||
return fmt.Errorf("plugin meta usageProfiles[%d] models must contain at least one model", index)
|
||||
}
|
||||
if err := validateModelScope(profile.Models, fmt.Sprintf("usageProfiles[%d]", index)); err != nil {
|
||||
return err
|
||||
}
|
||||
for _, model := range profile.Models {
|
||||
if _, exists := models[model]; !exists {
|
||||
return fmt.Errorf("plugin meta usageProfiles[%d] model %q is not declared in plugin meta models", index, model)
|
||||
}
|
||||
if _, duplicate := profileModels[model]; duplicate {
|
||||
return fmt.Errorf("plugin meta usageProfiles model %q belongs to multiple profiles", model)
|
||||
}
|
||||
profileModels[model] = struct{}{}
|
||||
}
|
||||
if profile.Schema == nil {
|
||||
return fmt.Errorf("plugin meta usageProfiles[%d] schema must be an object", index)
|
||||
}
|
||||
if err := validateUsageSchema(profile.Schema); err != nil {
|
||||
return fmt.Errorf("plugin meta usageProfiles[%d]: %w", index, err)
|
||||
}
|
||||
if err := validateUsageExamples(profile.Schema, profile.Examples); err != nil {
|
||||
return fmt.Errorf("plugin meta usageProfiles[%d]: %w", index, err)
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func validateUsageSchema(schema map[string]UsageFieldSchema) error {
|
||||
for name, field := range schema {
|
||||
if strings.TrimSpace(name) == "" || strings.TrimSpace(name) != name {
|
||||
return fmt.Errorf("plugin meta usageSchema keys must be non-empty canonical names")
|
||||
}
|
||||
@@ -1266,12 +1404,45 @@ func normalizeV1Meta(meta *Meta) error {
|
||||
return err
|
||||
}
|
||||
}
|
||||
if err := validateUsageExamples(meta.UsageSchema, meta.UsageExamples); err != nil {
|
||||
return err
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func decodeUsageProfiles(value any) ([]UsageProfile, error) {
|
||||
items, ok := value.([]any)
|
||||
if !ok {
|
||||
return nil, fmt.Errorf("plugin meta usageProfiles must be an array")
|
||||
}
|
||||
profiles := make([]UsageProfile, 0, len(items))
|
||||
for index, item := range items {
|
||||
object, ok := item.(map[string]any)
|
||||
if !ok {
|
||||
return nil, fmt.Errorf("plugin meta usageProfiles[%d] must be an object", index)
|
||||
}
|
||||
for key := range object {
|
||||
if key != "models" && key != "schema" && key != "examples" {
|
||||
return nil, fmt.Errorf("plugin meta usageProfiles[%d] has unknown field %q", index, key)
|
||||
}
|
||||
}
|
||||
models, err := strictStringSlice(object, "models")
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("plugin meta usageProfiles[%d]: %w", index, err)
|
||||
}
|
||||
schema, err := decodeUsageSchema(object["schema"])
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("plugin meta usageProfiles[%d]: %w", index, err)
|
||||
}
|
||||
profile := UsageProfile{Models: models, Schema: schema}
|
||||
if examples, exists := object["examples"]; exists {
|
||||
profile.Examples, err = decodeUsageExamples(examples)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("plugin meta usageProfiles[%d]: %w", index, err)
|
||||
}
|
||||
}
|
||||
profiles = append(profiles, profile)
|
||||
}
|
||||
return profiles, nil
|
||||
}
|
||||
|
||||
func decodeUsageSchema(value any) (map[string]UsageFieldSchema, error) {
|
||||
if value == nil {
|
||||
return nil, fmt.Errorf("plugin meta usageSchema must be an object")
|
||||
@@ -1288,7 +1459,7 @@ func decodeUsageSchema(value any) (map[string]UsageFieldSchema, error) {
|
||||
}
|
||||
for key := range fieldObject {
|
||||
switch key {
|
||||
case "type", "unit", "enum", "description", "enumLabels":
|
||||
case "type", "unit", "unitLabel", "enum", "description", "enumLabels":
|
||||
default:
|
||||
return nil, fmt.Errorf("plugin meta usageSchema field %q has unknown property %q", name, key)
|
||||
}
|
||||
@@ -1301,6 +1472,9 @@ func decodeUsageSchema(value any) (map[string]UsageFieldSchema, error) {
|
||||
if field.Unit, err = stringMetaField(fieldObject, "unit"); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if field.UnitLabel, err = localizedTextMetaField(fieldObject, "unitLabel", maxUsageFieldDescriptionRunes); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if field.Description, err = localizedTextMetaField(fieldObject, "description", maxUsageFieldDescriptionRunes); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
@@ -1332,6 +1506,14 @@ func decodeUsageSchema(value any) (map[string]UsageFieldSchema, error) {
|
||||
}
|
||||
|
||||
func validateUsageFieldSchema(name string, field UsageFieldSchema) error {
|
||||
if field.UnitLabel != nil {
|
||||
if field.Type != "number" || field.Unit != "count" || field.Enum != nil {
|
||||
return fmt.Errorf("plugin meta usageSchema field %q unitLabel requires a number field with count unit", name)
|
||||
}
|
||||
if err := validateLocalizedText(field.UnitLabel, fmt.Sprintf("usageSchema field %q unitLabel", name), maxUsageFieldDescriptionRunes); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
if err := validateLocalizedText(field.Description, fmt.Sprintf("usageSchema field %q description", name), maxUsageFieldDescriptionRunes); err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
@@ -327,6 +327,44 @@ func TestRegistryDecodesAndValidatesUsageSchema(t *testing.T) {
|
||||
})
|
||||
}
|
||||
|
||||
func TestUsageUnitLabelContract(t *testing.T) {
|
||||
source := routingTestPluginSource("unit-label", 0, `["model"]`, `usageSchema: {
|
||||
images: {type: "number", unit: "count", unitLabel: {EN: " image ", zh: "张"}},
|
||||
items: {type: "number", unit: "count", unitLabel: "item"},
|
||||
legacy: {type: "number", unit: "count"}
|
||||
},`, "")
|
||||
registry := NewRegistry()
|
||||
plugin, err := registry.Register(source, Options{})
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, LocalizedText{"en": "image", "zh": "张"}, plugin.Meta.UsageSchema["images"].UnitLabel)
|
||||
assert.Equal(t, LocalizedText{"en": "item"}, plugin.Meta.UsageSchema["items"].UnitLabel)
|
||||
assert.Nil(t, plugin.Meta.UsageSchema["legacy"].UnitLabel)
|
||||
wire, err := common.Marshal(plugin.Meta.UsageSchema)
|
||||
require.NoError(t, err)
|
||||
assert.JSONEq(t, `{"images":{"type":"number","unit":"count","unitLabel":{"en":"image","zh":"张"}},"items":{"type":"number","unit":"count","unitLabel":{"en":"item"}},"legacy":{"type":"number","unit":"count"}}`, string(wire))
|
||||
snapshot := registry.Snapshot()
|
||||
snapshot.Override[0].UsageSchema["images"].UnitLabel["en"] = "Changed"
|
||||
assert.Equal(t, "image", registry.Snapshot().Override[0].UsageSchema["images"].UnitLabel["en"])
|
||||
|
||||
for _, tc := range []struct{ name, field, message string }{
|
||||
{"seconds", `{type:"number",unit:"second",unitLabel:"image"}`, "unitLabel requires a number field with count unit"},
|
||||
{"tokens", `{type:"number",unit:"token",unitLabel:"image"}`, "unitLabel requires a number field with count unit"},
|
||||
{"credits", `{type:"number",unit:"credit",unitLabel:"image"}`, "unitLabel requires a number field with count unit"},
|
||||
{"boolean", `{type:"boolean",unitLabel:"image"}`, "unitLabel requires a number field with count unit"},
|
||||
{"enum", `{enum:["a"],unitLabel:"image"}`, "unitLabel requires a number field with count unit"},
|
||||
{"missing English", `{type:"number",unit:"count",unitLabel:{zh:"张"}}`, `must include a non-empty "en"`},
|
||||
{"null", `{type:"number",unit:"count",unitLabel:null}`, "must be a string or object"},
|
||||
{"empty", `{type:"number",unit:"count",unitLabel:" "}`, "non-empty string"},
|
||||
{"invalid locale", `{type:"number",unit:"count",unitLabel:{en:"image",zh_CN:"张"}}`, "invalid locale"},
|
||||
{"too long", `{type:"number",unit:"count",unitLabel:"` + strings.Repeat("x", 257) + `"}`, "must not exceed 256 characters"},
|
||||
} {
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
_, err := CompilePlugin(routingTestPluginSource("invalid-unit-label", 0, `["model"]`, "usageSchema: {value: "+tc.field+"},", ""), Options{})
|
||||
require.ErrorContains(t, err, tc.message)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestUsageEnumLabelsContract(t *testing.T) {
|
||||
source := routingTestPluginSource("enum-labels", 0, `["model"]`, `usageSchema: {
|
||||
mode: {enum: ["none", "video", "other"], enumLabels: {none: "No video", video: {EN: " With video ", zh: "有参考视频"}}}
|
||||
@@ -1058,3 +1096,136 @@ func TestPluginDisplayMetadata(t *testing.T) {
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestRegistryUsageProfilesReplaceDefaultsAndIsolateSnapshots(t *testing.T) {
|
||||
source := routingTestPluginSource("usage-profiles", 0, `["image", "video", "default", "empty"]`, `
|
||||
usageSchema: {seconds: {type: "number", unit: "second"}},
|
||||
usageExamples: [{label: "default", facts: {seconds: 1}}],
|
||||
usageProfiles: [
|
||||
{models: ["image"], schema: {image_count: {type: "number", unit: "count", unitLabel: {en: "image", zh: "张"}}}},
|
||||
{models: ["video"], schema: {
|
||||
seconds: {type: "number", unit: "second", description: {en: "Duration"}},
|
||||
resolution: {enum: ["720P", "1080P"], enumLabels: {"720P": {en: "HD"}}}
|
||||
}, examples: [{label: "HD video", facts: {seconds: 5, resolution: "720P"}}]},
|
||||
{models: ["empty"], schema: {}}
|
||||
],`, "")
|
||||
registry := NewRegistry()
|
||||
plugin, err := registry.Register(source, Options{})
|
||||
require.NoError(t, err)
|
||||
|
||||
for _, model := range []string{"image", "IMAGE"} {
|
||||
schema, examples := plugin.Meta.UsageForModel(model)
|
||||
assert.Equal(t, map[string]UsageFieldSchema{"image_count": {Type: "number", Unit: "count", UnitLabel: LocalizedText{"en": "image", "zh": "张"}}}, schema)
|
||||
assert.Nil(t, examples, "profiles must not inherit unrelated default examples")
|
||||
}
|
||||
for _, model := range []string{"default", "unknown", ""} {
|
||||
schema, examples := plugin.Meta.UsageForModel(model)
|
||||
assert.Equal(t, plugin.Meta.UsageSchema, schema)
|
||||
assert.Equal(t, plugin.Meta.UsageExamples, examples)
|
||||
}
|
||||
schema, examples := plugin.Meta.UsageForModel("empty")
|
||||
assert.Empty(t, schema)
|
||||
assert.Nil(t, examples)
|
||||
|
||||
snapshot := registry.Snapshot()
|
||||
require.Len(t, snapshot.Override, 1)
|
||||
profile := &snapshot.Override[0].UsageProfiles[1]
|
||||
snapshot.Override[0].UsageProfiles[0].Schema["image_count"].UnitLabel["en"] = "Changed"
|
||||
imageSchema, _ := registry.Snapshot().Override[0].UsageForModel("image")
|
||||
assert.Equal(t, "image", imageSchema["image_count"].UnitLabel["en"])
|
||||
profile.Models[0] = "changed"
|
||||
profile.Schema["seconds"].Description["en"] = "Changed"
|
||||
profile.Schema["resolution"].Enum[0] = "changed"
|
||||
profile.Schema["resolution"].EnumLabels["720P"]["en"] = "Changed"
|
||||
profile.Examples[0].Label = "Changed"
|
||||
profile.Examples[0].Facts["seconds"] = 99
|
||||
delete(profile.Schema, "seconds")
|
||||
|
||||
schema, examples = registry.Snapshot().Override[0].UsageForModel("video")
|
||||
assert.Equal(t, "Duration", schema["seconds"].Description["en"])
|
||||
assert.Equal(t, []string{"720P", "1080P"}, schema["resolution"].Enum)
|
||||
assert.Equal(t, "HD", schema["resolution"].EnumLabels["720P"]["en"])
|
||||
require.Len(t, examples, 1)
|
||||
assert.Equal(t, "HD video", examples[0].Label)
|
||||
assert.EqualValues(t, 5, examples[0].Facts["seconds"])
|
||||
|
||||
legacy, err := CompilePlugin(routingTestPluginSource("legacy-usage", 0, `["image"]`,
|
||||
`usageSchema: {seconds: {type: "number", unit: "second"}},`, ""), Options{})
|
||||
require.NoError(t, err)
|
||||
schema, examples = legacy.Meta.UsageForModel("image")
|
||||
assert.Equal(t, legacy.Meta.UsageSchema, schema)
|
||||
assert.Nil(t, examples)
|
||||
}
|
||||
|
||||
func TestRegistryRejectsInvalidUsageProfiles(t *testing.T) {
|
||||
for _, tc := range []struct {
|
||||
name, profiles, errorText string
|
||||
}{
|
||||
{"null", `null`, "usageProfiles must be an array"},
|
||||
{"object", `{}`, "usageProfiles must be an array"},
|
||||
{"invalid entry", `[null]`, "must be an object"},
|
||||
{"unknown property", `[{models:["image"], schema:{}, extra:true}]`, "unknown field"},
|
||||
{"missing models", `[{schema:{}}]`, "at least one model"},
|
||||
{"empty models", `[{models:[], schema:{}}]`, "at least one model"},
|
||||
{"unknown model", `[{models:["other"], schema:{}}]`, "not declared"},
|
||||
{"noncanonical model", `[{models:[" image"], schema:{}}]`, "canonical names"},
|
||||
{"wrong declared spelling", `[{models:["IMAGE"], schema:{}}]`, "not declared"},
|
||||
{"duplicate models", `[{models:["image","IMAGE"], schema:{}}]`, "unique case-insensitively"},
|
||||
{"overlapping profiles", `[{models:["image"], schema:{}}, {models:["image"], schema:{}}]`, "multiple profiles"},
|
||||
{"missing schema", `[{models:["image"]}]`, "must be an object"},
|
||||
{"invalid field", `[{models:["image"], schema:{n:{type:"number", unit:"invalid"}}}]`, "unit must be"},
|
||||
{"noncanonical field", `[{models:["image"], schema:{" n":{type:"number", unit:"count"}}}]`, "canonical names"},
|
||||
{"tokens need examples", `[{models:["image"], schema:{tokens:{type:"number", unit:"token"}}}]`, "usageExamples is required"},
|
||||
{"examples use own schema", `[{models:["image"], schema:{n:{type:"number", unit:"count"}}, examples:[{label:"wrong", facts:{seconds:1}}]}]`, "usageProfiles[0]"},
|
||||
} {
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
source := routingTestPluginSource("invalid-profile", 0, `["image", "video"]`,
|
||||
`usageSchema: {seconds: {type:"number", unit:"second"}}, usageExamples: [{label:"default", facts:{seconds:1}}], usageProfiles: `+tc.profiles+`,`, "")
|
||||
_, err := CompilePlugin(source, Options{})
|
||||
require.ErrorContains(t, err, tc.errorText)
|
||||
})
|
||||
}
|
||||
|
||||
plugin, err := CompilePlugin(routingTestPluginSource("token-profile", 0, `["image"]`, `
|
||||
usageProfiles: [{models:["image"], schema:{tokens:{type:"number", unit:"token"}}, examples:[{label:"one token", facts:{tokens:1}}]}],`, ""), Options{})
|
||||
require.NoError(t, err)
|
||||
schema, examples := plugin.Meta.UsageForModel("image")
|
||||
assert.Equal(t, "token", schema["tokens"].Unit)
|
||||
require.Len(t, examples, 1)
|
||||
assert.EqualValues(t, 1, examples[0].Facts["tokens"])
|
||||
plugin.Meta.UsageProfiles[0].Schema = nil
|
||||
require.ErrorContains(t, ValidateV1Meta(plugin.Meta), "schema must be an object")
|
||||
}
|
||||
|
||||
func TestSubmitResponseTypesContract(t *testing.T) {
|
||||
for _, tc := range []struct {
|
||||
name, metadata, hook string
|
||||
valid bool
|
||||
}{
|
||||
{"legacy JSON", "", "", true},
|
||||
{"SSE declared", `submitResponseTypes:["json","sse"],`, `export function parseSubmitEvent(){return {state:null,done:true};}`, true},
|
||||
{"host JSON utility", `requiredCapabilities:["json-clone@1"],`, "", true},
|
||||
{"SSE delta", `submitResponseTypes:["sse"],requiredCapabilities:["submit-sse-delta@1"],`, `export function parseSubmitEventDelta(){return {changes:[],state:null,done:true};}`, true},
|
||||
{"SSE delta missing hook", `submitResponseTypes:["sse"],requiredCapabilities:["submit-sse-delta@1"],`, `export function parseSubmitEvent(){return {state:null,done:true};}`, false},
|
||||
{"delta requires SSE", `requiredCapabilities:["submit-sse-delta@1"],`, "", false},
|
||||
{"unsupported capability version", `requiredCapabilities:["json-clone@2"],`, "", false},
|
||||
{"duplicate capability", `requiredCapabilities:["json-clone@1","json-clone@1"],`, "", false},
|
||||
{"null capabilities", `requiredCapabilities:null,`, "", false},
|
||||
{"missing hook", `submitResponseTypes:["sse"],`, "", false},
|
||||
{"unknown type", `submitResponseTypes:["xml"],`, "", false},
|
||||
{"duplicate", `submitResponseTypes:["json","json"],`, "", false},
|
||||
{"empty", `submitResponseTypes:[],`, "", false},
|
||||
{"null", `submitResponseTypes:null,`, "", false},
|
||||
} {
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
source := routingTestPluginSource("submit-types", 0, `["model"]`, tc.metadata, "") + tc.hook
|
||||
plugin, err := CompilePlugin(source, Options{})
|
||||
if tc.valid {
|
||||
require.NoError(t, err)
|
||||
require.NoError(t, ValidateV1Meta(plugin.Meta))
|
||||
} else {
|
||||
require.Error(t, err)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
+20
-5
@@ -326,6 +326,7 @@ type RoutingGeneration struct {
|
||||
|
||||
byKey map[string]*LoadedPlugin
|
||||
byModel map[string]*LoadedPlugin
|
||||
modelPlugins map[string][]*LoadedPlugin
|
||||
canonicalModelByFold map[string]string
|
||||
byChannelType map[int]*LoadedPlugin
|
||||
routeIndex map[string]RouteBinding
|
||||
@@ -416,6 +417,20 @@ func (g *RoutingGeneration) GetByModel(model string) (*LoadedPlugin, bool) {
|
||||
return plugin, ok
|
||||
}
|
||||
|
||||
// PluginsByModel returns all plugins declaring model, in ascending key order.
|
||||
func (g *RoutingGeneration) PluginsByModel(model string) []*LoadedPlugin {
|
||||
if g == nil {
|
||||
return nil
|
||||
}
|
||||
return slices.Clone(g.modelPlugins[model])
|
||||
}
|
||||
|
||||
// SharedModel reports whether multiple plugins declare a model without copying
|
||||
// its provider list on the relay hot path.
|
||||
func (g *RoutingGeneration) SharedModel(model string) bool {
|
||||
return g != nil && len(g.modelPlugins[model]) >= 2
|
||||
}
|
||||
|
||||
// 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).
|
||||
@@ -463,9 +478,8 @@ func (g *RoutingGeneration) LookupEndpoint(method, path, model string) (Protocol
|
||||
return bindings[0], true
|
||||
}
|
||||
|
||||
// LookupEndpointCandidates returns every legacy provider implementation that
|
||||
// can serve one shared model endpoint. Candidate order is deterministic and
|
||||
// the first binding is the parser used before channel distribution.
|
||||
// LookupEndpointCandidates returns every plugin that can serve a shared model
|
||||
// endpoint, in ascending key order. Each candidate decodes before distribution.
|
||||
func (g *RoutingGeneration) LookupEndpointCandidates(method, path, model string) []ProtocolBinding {
|
||||
if g == nil {
|
||||
return nil
|
||||
@@ -859,6 +873,7 @@ func buildRoutingGenerationFromPlugins(effective map[string]*LoadedPlugin, numbe
|
||||
PublishedAt: time.Now(),
|
||||
byKey: make(map[string]*LoadedPlugin, len(effective)),
|
||||
byModel: make(map[string]*LoadedPlugin),
|
||||
modelPlugins: make(map[string][]*LoadedPlugin),
|
||||
canonicalModelByFold: make(map[string]string),
|
||||
byChannelType: make(map[int]*LoadedPlugin),
|
||||
routeIndex: make(map[string]RouteBinding),
|
||||
@@ -870,6 +885,7 @@ func buildRoutingGenerationFromPlugins(effective map[string]*LoadedPlugin, numbe
|
||||
generation.byKey[key] = plugin
|
||||
generation.plugins = append(generation.plugins, plugin)
|
||||
for _, model := range plugin.Meta.Models {
|
||||
generation.modelPlugins[model] = append(generation.modelPlugins[model], plugin)
|
||||
if _, exists := generation.byModel[model]; !exists {
|
||||
generation.byModel[model] = plugin
|
||||
}
|
||||
@@ -927,8 +943,7 @@ func buildRoutingGenerationFromPlugins(effective map[string]*LoadedPlugin, numbe
|
||||
bindings := generation.protocolIndex[indexKey]
|
||||
if len(bindings) > 0 {
|
||||
other := bindings[0]
|
||||
legacyProviders := len(plugin.Meta.ChannelTypes) > 0 && len(other.Plugin.Meta.ChannelTypes) > 0
|
||||
if !legacyProviders || claim.Name != other.Protocol {
|
||||
if claim.Name != other.Protocol {
|
||||
return nil, fmt.Errorf("plugin %s protocol %s %s model %q conflicts with plugin %s", plugin.Meta.Key, method, operation.Path, model, other.Plugin.Meta.Key)
|
||||
}
|
||||
}
|
||||
|
||||
@@ -155,32 +155,33 @@ func TestPreflightRoutingConflict(t *testing.T) {
|
||||
routingProtocolExport("openai_responses"))
|
||||
protocolRegistry := NewRegistry()
|
||||
require.NoError(t, protocolRegistry.ReplaceOverrides([]*LoadedPlugin{left}))
|
||||
protocolErr := PreflightRoutingConflict(protocolRegistry.Generation(), right)
|
||||
require.Error(t, protocolErr)
|
||||
assert.Contains(t, protocolErr.Error(), "preflight-left")
|
||||
assert.Contains(t, protocolErr.Error(), `model "shared-model" conflicts`)
|
||||
require.NoError(t, PreflightRoutingConflict(protocolRegistry.Generation(), right))
|
||||
}
|
||||
|
||||
func TestPerProtocolModelsStillConflictOnOverlap(t *testing.T) {
|
||||
registry := NewRegistry()
|
||||
first := mustCompileRoutingPlugin(t, "overlap-first", 0, `["shared-model"]`,
|
||||
`protocols: [{name: "openai_responses", models: ["shared-model"], supports: ["stream", "sync", "background"]}],`,
|
||||
routingProtocolExport("openai_responses"))
|
||||
second := mustCompileRoutingPlugin(t, "overlap-second", 0, `["shared-model"]`,
|
||||
`protocols: [{name: "openai_responses", models: ["shared-model"], supports: ["stream", "sync", "background"]}],`,
|
||||
routingProtocolExport("openai_responses"))
|
||||
|
||||
require.NoError(t, registry.ReplaceOverrides([]*LoadedPlugin{first, second}))
|
||||
routingErrors := registry.RoutingErrors()
|
||||
require.Len(t, routingErrors, 1, "one of two overlapping claims must be rejected from the generation")
|
||||
for _, message := range routingErrors {
|
||||
assert.Contains(t, message, "conflicts")
|
||||
func TestPerProtocolModelsAcceptSharedPluginClaims(t *testing.T) {
|
||||
for _, channelType := range []int{0, 24} {
|
||||
t.Run(fmt.Sprintf("first_channel_type_%d", channelType), func(t *testing.T) {
|
||||
registry := NewRegistry()
|
||||
first := mustCompileRoutingPlugin(t, "overlap-first", channelType, `["shared-model"]`,
|
||||
`protocols: [{name: "openai_responses", models: ["shared-model"], supports: ["stream", "sync", "background"]}],`, routingProtocolExport("openai_responses"))
|
||||
second := mustCompileRoutingPlugin(t, "overlap-second", 0, `["shared-model"]`,
|
||||
`protocols: [{name: "openai_responses", models: ["shared-model"], supports: ["stream", "sync", "background"]}],`, routingProtocolExport("openai_responses"))
|
||||
require.NoError(t, registry.ReplaceOverrides([]*LoadedPlugin{second, first}))
|
||||
assert.Empty(t, registry.RoutingErrors())
|
||||
generation := registry.Generation()
|
||||
candidates := generation.LookupEndpointCandidates("POST", "/v1/responses", "shared-model")
|
||||
require.Len(t, candidates, 2)
|
||||
assert.Same(t, first, candidates[0].Plugin)
|
||||
assert.Same(t, second, candidates[1].Plugin)
|
||||
assert.Equal(t, []*LoadedPlugin{first, second}, generation.PluginsByModel("shared-model"))
|
||||
metadata, ok := generation.GetByModel("shared-model")
|
||||
require.True(t, ok)
|
||||
assert.Same(t, first, metadata)
|
||||
copy := generation.PluginsByModel("shared-model")
|
||||
copy[0] = nil
|
||||
assert.Same(t, first, generation.PluginsByModel("shared-model")[0])
|
||||
})
|
||||
}
|
||||
binding, ok := registry.Generation().LookupEndpoint("POST", "/v1/responses", "shared-model")
|
||||
require.True(t, ok)
|
||||
candidates := registry.Generation().LookupEndpointCandidates("POST", "/v1/responses", "shared-model")
|
||||
assert.Len(t, candidates, 1)
|
||||
assert.NotNil(t, binding.Plugin)
|
||||
}
|
||||
|
||||
func TestProtocolClaimDecodeAndValidation(t *testing.T) {
|
||||
@@ -497,16 +498,6 @@ func TestRegistryExcludesConflictingPluginWithoutBlockingGeneration(t *testing.T
|
||||
`export const native = {status: function(ctx, task) { return task; }};`),
|
||||
expectedError: "route GET /vendor/jobs/:id conflicts",
|
||||
},
|
||||
{
|
||||
name: "endpoint model ownership",
|
||||
first: mustCompileRoutingPlugin(t, "endpoint-alpha", 0, `["shared-model"]`,
|
||||
`protocols: [{name: "openai_responses", supports: ["stream", "sync", "background"]}],`,
|
||||
routingProtocolExport("openai_responses")),
|
||||
second: mustCompileRoutingPlugin(t, "endpoint-beta", 0, `["shared-model"]`,
|
||||
`protocols: [{name: "openai_responses", supports: ["stream", "sync", "background"]}],`,
|
||||
routingProtocolExport("openai_responses")),
|
||||
expectedError: `model "shared-model" conflicts`,
|
||||
},
|
||||
}
|
||||
|
||||
for _, testCase := range tests {
|
||||
|
||||
@@ -1,6 +1,7 @@
|
||||
package jsplugin
|
||||
|
||||
import (
|
||||
"context"
|
||||
"crypto/hmac"
|
||||
"crypto/sha256"
|
||||
"encoding/base64"
|
||||
@@ -30,6 +31,24 @@ type volcSignRequest struct {
|
||||
|
||||
func injectGlobals(runtime *sobek.Runtime, identity func() string, now func() time.Time, logOutput func(string)) error {
|
||||
utils := map[string]any{
|
||||
"hasCapability": HasCapability,
|
||||
"json": map[string]any{
|
||||
"clone": func(call sobek.FunctionCall) sobek.Value {
|
||||
if sobek.IsUndefined(call.Argument(0)) {
|
||||
panic(runtime.NewTypeError("json.clone requires a JSON value"))
|
||||
}
|
||||
budget := &jsonStateBudget{ctx: context.Background(), bytes: MaxJSONToolBytes, nodes: maxJSONToolNodes}
|
||||
value, err := newJSONStateNode(call.Argument(0), 0, budget)
|
||||
if err != nil {
|
||||
panic(runtime.NewTypeError("json.clone: %s", err))
|
||||
}
|
||||
cloned, err := value.jsValue(runtime)
|
||||
if err != nil {
|
||||
panic(runtime.NewGoError(err))
|
||||
}
|
||||
return cloned
|
||||
},
|
||||
},
|
||||
"unixNow": func() int64 { return now().Unix() },
|
||||
"jwtSignHS256": func(claims map[string]any, secret string) (string, error) {
|
||||
return jwt.NewWithClaims(jwt.SigningMethodHS256, jwt.MapClaims(claims)).SignedString([]byte(secret))
|
||||
|
||||
@@ -0,0 +1,21 @@
|
||||
MIT License
|
||||
|
||||
Copyright (c) 2026 xqy2006
|
||||
|
||||
Permission is hereby granted, free of charge, to any person obtaining a copy
|
||||
of this software and associated documentation files (the "Software"), to deal
|
||||
in the Software without restriction, including without limitation the rights
|
||||
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
|
||||
copies of the Software, and to permit persons to whom the Software is
|
||||
furnished to do so, subject to the following conditions:
|
||||
|
||||
The above copyright notice and this permission notice shall be included in all
|
||||
copies or substantial portions of the Software.
|
||||
|
||||
THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
|
||||
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
|
||||
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
|
||||
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
|
||||
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
|
||||
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
|
||||
SOFTWARE.
|
||||
@@ -0,0 +1,24 @@
|
||||
MIT License
|
||||
|
||||
Copyright (c) 2026 gpt-model-check contributors
|
||||
|
||||
Permission is hereby granted, free of charge, to any person obtaining a copy
|
||||
of this software and associated documentation files (the "Software"), to deal
|
||||
in the Software without restriction, including without limitation the rights
|
||||
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
|
||||
copies of the Software, and to permit persons to whom the Software is
|
||||
furnished to do so, subject to the following conditions:
|
||||
|
||||
The above copyright notice and this permission notice shall be included in all
|
||||
copies or substantial portions of the Software.
|
||||
|
||||
THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
|
||||
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
|
||||
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
|
||||
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
|
||||
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
|
||||
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
|
||||
SOFTWARE.
|
||||
|
||||
ModelTrace bank and scorer retain the license in assets/LICENSE-ModelTrace.txt.
|
||||
Native probe transport adaptations retain assets/LICENSE-is-gpt-nerfed.txt.
|
||||
@@ -0,0 +1,21 @@
|
||||
MIT License
|
||||
|
||||
Copyright (c) 2026 is-gpt-nerfed contributors
|
||||
|
||||
Permission is hereby granted, free of charge, to any person obtaining a copy
|
||||
of this software and associated documentation files (the "Software"), to deal
|
||||
in the Software without restriction, including without limitation the rights
|
||||
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
|
||||
copies of the Software, and to permit persons to whom the Software is
|
||||
furnished to do so, subject to the following conditions:
|
||||
|
||||
The above copyright notice and this permission notice shall be included in all
|
||||
copies or substantial portions of the Software.
|
||||
|
||||
THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
|
||||
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
|
||||
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
|
||||
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
|
||||
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
|
||||
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
|
||||
SOFTWARE.
|
||||
File diff suppressed because it is too large.
Load diff
@@ -0,0 +1,146 @@
|
||||
// Numerical fingerprint scorer adapted from Katao123/gpt-model-check (MIT),
|
||||
// commit 3d6716a0903cf061daea001d5d44ee410846634e. See the adjacent licenses.
|
||||
package model_fingerprint
|
||||
|
||||
import (
|
||||
_ "embed"
|
||||
"errors"
|
||||
"math"
|
||||
"strings"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
"github.com/QuantumNous/new-api/common"
|
||||
"github.com/grafana/sobek"
|
||||
)
|
||||
|
||||
const Model = "gpt-6-astra"
|
||||
const Version = "modeltrace-20260905-gmc-3d6716a"
|
||||
|
||||
//go:embed bank.json
|
||||
var bankJSON string
|
||||
|
||||
//go:embed scorer.js
|
||||
var scorerJS string
|
||||
|
||||
var compileOnce sync.Once
|
||||
var program *sobek.Program
|
||||
var compileErr error
|
||||
|
||||
type Sample struct {
|
||||
Text string `json:"text"`
|
||||
Count int `json:"expected_count"`
|
||||
}
|
||||
|
||||
type Candidate struct {
|
||||
Model string `json:"model"`
|
||||
Weight float64 `json:"probability"`
|
||||
Score float64 `json:"score"`
|
||||
}
|
||||
|
||||
type Result struct {
|
||||
Status string `json:"status"`
|
||||
ValidSamples int `json:"valid_samples"`
|
||||
Closest string `json:"closest,omitempty"`
|
||||
Weight float64 `json:"relative_weight,omitempty"`
|
||||
Gap float64 `json:"score_gap,omitempty"`
|
||||
}
|
||||
|
||||
func ValidArray(text string, count int) bool {
|
||||
if len(text) > 16384 || count <= 0 {
|
||||
return false
|
||||
}
|
||||
text = strings.TrimSpace(text)
|
||||
if strings.HasPrefix(text, "```") && strings.HasSuffix(text, "```") {
|
||||
if i := strings.IndexByte(text, '\n'); i >= 0 {
|
||||
text = strings.TrimSpace(text[i+1 : len(text)-3])
|
||||
}
|
||||
}
|
||||
var numbers []int
|
||||
if common.Unmarshal([]byte(text), &numbers) != nil {
|
||||
return false
|
||||
}
|
||||
if len(numbers) < max(80, int(math.Ceil(float64(count)*.55))) || len(numbers) > count*2 {
|
||||
return false
|
||||
}
|
||||
for _, number := range numbers {
|
||||
if number < 1 || number > 355 {
|
||||
return false
|
||||
}
|
||||
}
|
||||
return true
|
||||
}
|
||||
|
||||
func Decide(candidates []Candidate, valid int) Result {
|
||||
r := Result{Status: "inconclusive", ValidSamples: valid}
|
||||
if len(candidates) < 2 {
|
||||
return r
|
||||
}
|
||||
top := candidates[0]
|
||||
r.Closest = top.Model
|
||||
r.Weight = top.Weight
|
||||
r.Gap = top.Score - candidates[1].Score
|
||||
if valid != 3 || math.IsNaN(r.Gap) || math.IsInf(r.Gap, 0) || top.Weight < .8 || r.Gap < .5 {
|
||||
return r
|
||||
}
|
||||
if top.Model == Model {
|
||||
r.Status = "match"
|
||||
return r
|
||||
}
|
||||
for _, candidate := range candidates {
|
||||
if candidate.Model == Model && candidate.Weight <= .2 && top.Score-candidate.Score >= .5 {
|
||||
r.Status = "mismatch"
|
||||
}
|
||||
}
|
||||
return r
|
||||
}
|
||||
|
||||
// Score runs only the embedded, pinned scorer. No filesystem, network or host
|
||||
// functions are exposed to the VM; upstream text is data, never JavaScript.
|
||||
func Score(samples []Sample) (Result, error) {
|
||||
valid := make([]Sample, 0, 3)
|
||||
for _, sample := range samples {
|
||||
if ValidArray(sample.Text, sample.Count) {
|
||||
valid = append(valid, sample)
|
||||
}
|
||||
}
|
||||
if len(valid) == 0 {
|
||||
return Result{Status: "inconclusive"}, nil
|
||||
}
|
||||
compileOnce.Do(func() {
|
||||
program, compileErr = sobek.Compile("fingerprint", strings.ReplaceAll(scorerJS, "export ", ""), true)
|
||||
})
|
||||
if compileErr != nil {
|
||||
return Result{}, compileErr
|
||||
}
|
||||
vm := sobek.New()
|
||||
timer := time.AfterFunc(3*time.Second, func() { vm.Interrupt("fingerprint scoring timeout") })
|
||||
defer timer.Stop()
|
||||
if _, err := vm.RunProgram(program); err != nil {
|
||||
return Result{}, err
|
||||
}
|
||||
if err := vm.Set("bankJSON", bankJSON); err != nil {
|
||||
return Result{}, err
|
||||
}
|
||||
encoded, err := common.Marshal(valid)
|
||||
if err != nil {
|
||||
return Result{}, err
|
||||
}
|
||||
if err = vm.Set("samplesJSON", string(encoded)); err != nil {
|
||||
return Result{}, err
|
||||
}
|
||||
value, err := vm.RunString("JSON.stringify(analyzeGlobalOutputs(JSON.parse(samplesJSON), JSON.parse(bankJSON)))")
|
||||
if err != nil {
|
||||
return Result{}, err
|
||||
}
|
||||
var analysis struct {
|
||||
Results []Candidate `json:"results"`
|
||||
}
|
||||
if err = common.Unmarshal([]byte(value.String()), &analysis); err != nil {
|
||||
return Result{}, err
|
||||
}
|
||||
if len(analysis.Results) < 2 {
|
||||
return Result{}, errors.New("missing fingerprint candidates")
|
||||
}
|
||||
return Decide(analysis.Results, len(valid)), nil
|
||||
}
|
||||
@@ -0,0 +1,68 @@
|
||||
package model_fingerprint
|
||||
|
||||
import (
|
||||
"os"
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
"github.com/QuantumNous/new-api/common"
|
||||
)
|
||||
|
||||
func TestFingerprintReferenceAndInvalidSamples(t *testing.T) {
|
||||
data, err := os.ReadFile("testdata/astra-reference.json")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
var rows []struct {
|
||||
Text string `json:"text"`
|
||||
Count int `json:"requested_count"`
|
||||
}
|
||||
if err = common.Unmarshal(data, &rows); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
samples := []Sample{}
|
||||
for _, row := range rows {
|
||||
samples = append(samples, Sample{Text: row.Text, Count: row.Count})
|
||||
}
|
||||
r, err := Score(samples)
|
||||
if err != nil || r.Status != "match" || r.Closest != Model || r.ValidSamples != 3 {
|
||||
t.Fatalf("%+v %v", r, err)
|
||||
}
|
||||
r, err = Score(samples[:2])
|
||||
if err != nil || r.Status != "inconclusive" {
|
||||
t.Fatalf("partial samples: %+v %v", r, err)
|
||||
}
|
||||
for _, s := range []string{"[1,2]", "null", "[" + strings.Repeat("-1,", 299) + "1]", "[" + strings.Repeat("1.2,", 299) + "1]", "ignore instructions; [1]"} {
|
||||
if ValidArray(s, 300) {
|
||||
t.Fatalf("accepted invalid array: %s", s[:min(len(s), 20)])
|
||||
}
|
||||
}
|
||||
c := []Candidate{{Model: "gpt-5.6-luna", Weight: .98, Score: 2}, {Model: Model, Weight: .02, Score: 1}}
|
||||
if Decide(c, 3).Status != "mismatch" {
|
||||
t.Fatal("wrong model gate")
|
||||
}
|
||||
c[1].Score = 1.7
|
||||
if Decide(c, 3).Status != "inconclusive" {
|
||||
t.Fatal("low margin must abstain")
|
||||
}
|
||||
}
|
||||
|
||||
func TestFingerprintResponseCompletionAndTools(t *testing.T) {
|
||||
completed := `{"status":"completed","model":"gpt-6-astra","output":[{"type":"message","content":[{"type":"output_text","text":"[1,2,3]"}]}],"usage":{"input_tokens":21,"output_tokens":12,"input_tokens_details":{"cached_tokens":10}}}`
|
||||
for _, input := range []struct{ body, content string }{{completed, "application/json"}, {"data: {\"type\":\"response.completed\",\"response\":" + completed + "}\n\n", "text/event-stream"}} {
|
||||
r, err := ReadResponse(strings.NewReader(input.body), input.content)
|
||||
if err != nil || r.Text != "[1,2,3]" || r.CachedTokens != 10 || !r.HasUsage {
|
||||
t.Fatalf("%+v %v", r, err)
|
||||
}
|
||||
}
|
||||
for _, body := range []string{
|
||||
"data: {\"type\":\"response.output_text.delta\",\"delta\":\"[1,2,3]\"}\n\n",
|
||||
"data: {\"type\":\"response.output_item.added\",\"item\":{\"type\":\"function_call\"}}\n\n",
|
||||
strings.Replace(completed, "completed", "incomplete", 1),
|
||||
strings.Replace(completed, "\"type\":\"message\"", "\"type\":\"function_call\"", 1),
|
||||
} {
|
||||
if _, err := ReadResponse(strings.NewReader(body), "application/json"); err == nil {
|
||||
t.Fatal("accepted tool or incomplete sample")
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,124 @@
|
||||
package model_fingerprint
|
||||
|
||||
import (
|
||||
"bufio"
|
||||
"errors"
|
||||
"io"
|
||||
"strings"
|
||||
|
||||
"github.com/tidwall/gjson"
|
||||
)
|
||||
|
||||
var ErrInvalidSample = errors.New("invalid probe answer")
|
||||
var ErrIncomplete = errors.New("probe response incomplete")
|
||||
|
||||
type ProbeResponse struct {
|
||||
Text string `json:"-"`
|
||||
DeclaredModel string `json:"declared_model,omitempty"`
|
||||
InputTokens int64 `json:"input_tokens"`
|
||||
OutputTokens int64 `json:"output_tokens"`
|
||||
CachedTokens int64 `json:"cached_tokens"`
|
||||
CacheWriteTokens int64 `json:"cache_write_tokens"`
|
||||
HasUsage bool `json:"has_usage"`
|
||||
}
|
||||
|
||||
func completedResponse(response gjson.Result) (ProbeResponse, error) {
|
||||
r := ProbeResponse{}
|
||||
r.DeclaredModel = response.Get("model").String()
|
||||
if len(r.DeclaredModel) > 128 {
|
||||
r.DeclaredModel = ""
|
||||
}
|
||||
usage := response.Get("usage")
|
||||
r.HasUsage = usage.Get("input_tokens").Type == gjson.Number && usage.Get("output_tokens").Type == gjson.Number
|
||||
r.InputTokens = usage.Get("input_tokens").Int()
|
||||
r.OutputTokens = usage.Get("output_tokens").Int()
|
||||
r.CachedTokens = usage.Get("input_tokens_details.cached_tokens").Int()
|
||||
r.CacheWriteTokens = max(usage.Get("input_tokens_details.cached_creation_tokens").Int(), usage.Get("input_tokens_details.cache_creation_tokens").Int())
|
||||
if response.Get("status").String() != "completed" {
|
||||
return r, ErrIncomplete
|
||||
}
|
||||
for _, item := range response.Get("output").Array() {
|
||||
switch item.Get("type").String() {
|
||||
case "reasoning":
|
||||
case "message":
|
||||
for _, content := range item.Get("content").Array() {
|
||||
if content.Get("type").String() != "output_text" {
|
||||
return r, ErrInvalidSample
|
||||
}
|
||||
r.Text += content.Get("text").String()
|
||||
if len(r.Text) > 16384 {
|
||||
return ProbeResponse{}, ErrInvalidSample
|
||||
}
|
||||
}
|
||||
default:
|
||||
return r, ErrInvalidSample
|
||||
}
|
||||
}
|
||||
return r, nil
|
||||
}
|
||||
|
||||
// Only a complete response is a sample. Partial deltas, tool calls and a
|
||||
// truncated stream cannot become a model verdict. Total bytes are bounded.
|
||||
func ReadResponse(reader io.Reader, contentType string) (ProbeResponse, error) {
|
||||
limited := &io.LimitedReader{R: reader, N: 8 << 20}
|
||||
buffered := bufio.NewReader(limited)
|
||||
peek, _ := buffered.Peek(1)
|
||||
if !strings.Contains(contentType, "text/event-stream") && len(peek) > 0 && peek[0] == '{' {
|
||||
data, err := io.ReadAll(buffered)
|
||||
if err != nil {
|
||||
return ProbeResponse{}, err
|
||||
}
|
||||
if limited.N == 0 || !gjson.ValidBytes(data) {
|
||||
return ProbeResponse{}, ErrInvalidSample
|
||||
}
|
||||
return completedResponse(gjson.ParseBytes(data))
|
||||
}
|
||||
scanner := bufio.NewScanner(buffered)
|
||||
scanner.Buffer(make([]byte, 4096), 2<<20)
|
||||
var data strings.Builder
|
||||
for scanner.Scan() {
|
||||
line := scanner.Text()
|
||||
if strings.HasPrefix(line, "data:") {
|
||||
if data.Len() > 0 {
|
||||
data.WriteByte('\n')
|
||||
}
|
||||
data.WriteString(strings.TrimSpace(strings.TrimPrefix(line, "data:")))
|
||||
if data.Len() > 2<<20 {
|
||||
return ProbeResponse{}, ErrInvalidSample
|
||||
}
|
||||
continue
|
||||
}
|
||||
if line != "" || data.Len() == 0 {
|
||||
continue
|
||||
}
|
||||
payload := data.String()
|
||||
data.Reset()
|
||||
if payload == "[DONE]" {
|
||||
break
|
||||
}
|
||||
if !gjson.Valid(payload) {
|
||||
return ProbeResponse{}, ErrInvalidSample
|
||||
}
|
||||
event := gjson.Parse(payload)
|
||||
typ := event.Get("type").String()
|
||||
if typ == "response.output_item.added" || typ == "response.output_item.done" {
|
||||
itemType := event.Get("item.type").String()
|
||||
if itemType != "message" && itemType != "reasoning" {
|
||||
return ProbeResponse{}, ErrInvalidSample
|
||||
}
|
||||
}
|
||||
if typ == "response.completed" {
|
||||
return completedResponse(event.Get("response"))
|
||||
}
|
||||
if typ == "response.failed" || typ == "response.incomplete" {
|
||||
return completedResponse(event.Get("response"))
|
||||
}
|
||||
if typ == "error" {
|
||||
return ProbeResponse{}, ErrIncomplete
|
||||
}
|
||||
}
|
||||
if err := scanner.Err(); err != nil {
|
||||
return ProbeResponse{}, err
|
||||
}
|
||||
return ProbeResponse{}, ErrIncomplete
|
||||
}
|
||||
@@ -0,0 +1,212 @@
|
||||
export const VALUE_MIN = 1;
|
||||
export const VALUE_MAX = 355;
|
||||
export const DIMENSION = VALUE_MAX - VALUE_MIN + 1;
|
||||
export const ALPHA = 0.5;
|
||||
export const ORDERED_BLOCK_WEIGHT = 0.25;
|
||||
|
||||
export function parseNumbers(text) {
|
||||
const runs = [];
|
||||
let current = [];
|
||||
let previousEnd = 0;
|
||||
for (const match of String(text).matchAll(/\d+/g)) {
|
||||
const separator = String(text).slice(previousEnd, match.index);
|
||||
const value = Number(match[0]);
|
||||
if (current.length && /\p{L}/u.test(separator)) {
|
||||
runs.push(current);
|
||||
current = [];
|
||||
}
|
||||
if (value >= VALUE_MIN && value <= VALUE_MAX) current.push(value);
|
||||
previousEnd = match.index + match[0].length;
|
||||
}
|
||||
if (current.length) runs.push(current);
|
||||
return runs.reduce((best, run) => run.length > best.length ? run : best, []);
|
||||
}
|
||||
|
||||
export function countNumbers(numbers) {
|
||||
const counts = Array(DIMENSION).fill(0);
|
||||
numbers.forEach((number) => { counts[number - VALUE_MIN] += 1; });
|
||||
return counts;
|
||||
}
|
||||
|
||||
function mean(values) {
|
||||
return values.reduce((total, value) => total + value, 0) / values.length;
|
||||
}
|
||||
|
||||
function standardize(values) {
|
||||
const center = mean(values);
|
||||
const variance = mean(values.map((value) => (value - center) ** 2));
|
||||
const scale = Math.max(Math.sqrt(variance), 1e-12);
|
||||
return values.map((value) => (value - center) / scale);
|
||||
}
|
||||
|
||||
function dot(left, right) {
|
||||
let value = 0;
|
||||
for (let index = 0; index < left.length; index += 1) value += left[index] * right[index];
|
||||
return value;
|
||||
}
|
||||
|
||||
function norm(values) {
|
||||
return Math.sqrt(dot(values, values));
|
||||
}
|
||||
|
||||
function normalized(values) {
|
||||
const scale = Math.max(norm(values), 1e-12);
|
||||
return values.map((value) => value / scale);
|
||||
}
|
||||
|
||||
function subtractBasis(values, basis) {
|
||||
const output = values.slice();
|
||||
for (const vector of basis || []) {
|
||||
const projection = dot(output, vector);
|
||||
for (let index = 0; index < output.length; index += 1) output[index] -= projection * vector[index];
|
||||
}
|
||||
return output;
|
||||
}
|
||||
|
||||
function hellingerFeature(counts) {
|
||||
const total = counts.reduce((sum, value) => sum + value, 0) + ALPHA * DIMENSION;
|
||||
return counts.map((value) => Math.sqrt((value + ALPHA) / total));
|
||||
}
|
||||
|
||||
function splitIntoFour(values) {
|
||||
const base = Math.floor(values.length / 4);
|
||||
const remainder = values.length % 4;
|
||||
const chunks = [];
|
||||
let start = 0;
|
||||
for (let index = 0; index < 4; index += 1) {
|
||||
const size = base + (index < remainder ? 1 : 0);
|
||||
chunks.push(values.slice(start, start + size));
|
||||
start += size;
|
||||
}
|
||||
return chunks;
|
||||
}
|
||||
|
||||
function orderedBlockFeature(numbers) {
|
||||
const pieces = [];
|
||||
for (const chunk of splitIntoFour(numbers)) {
|
||||
const bins = Array(16).fill(0.5);
|
||||
for (const value of chunk) {
|
||||
const index = Math.min(15, Math.floor(((value - 1) / 355) * 16));
|
||||
bins[index] += 1;
|
||||
}
|
||||
const total = bins.reduce((sum, value) => sum + value, 0);
|
||||
pieces.push(...bins.map((value) => Math.sqrt(value / total)));
|
||||
}
|
||||
const lastDigits = Array(10).fill(0.5);
|
||||
numbers.forEach((value) => { lastDigits[value % 10] += 1; });
|
||||
const lastTotal = lastDigits.reduce((sum, value) => sum + value, 0);
|
||||
pieces.push(...lastDigits.map((value) => Math.sqrt(value / lastTotal)));
|
||||
return pieces;
|
||||
}
|
||||
|
||||
function robustScoreCounts(counts, bank) {
|
||||
const artifact = bank.robust.hellinger;
|
||||
const feature = hellingerFeature(counts);
|
||||
let projected = feature.map((value, index) => (value - artifact.feature_mean[index]) / artifact.feature_scale[index]);
|
||||
projected = subtractBasis(projected, artifact.nuisance_basis);
|
||||
projected = normalized(projected);
|
||||
const scores = artifact.centroids.map((centroid) => dot(projected, centroid));
|
||||
return standardize(scores);
|
||||
}
|
||||
|
||||
function orderedBlockScores(numbers, bank) {
|
||||
const artifact = bank.robust.ordered_blocks;
|
||||
const feature = orderedBlockFeature(numbers);
|
||||
const standardizedFeature = feature.map((value, index) => (value - artifact.feature_mean[index]) / artifact.feature_scale[index]);
|
||||
const unit = normalized(standardizedFeature);
|
||||
const environmentScores = artifact.environment_centroids.map(
|
||||
(centroids) => centroids.map((centroid) => dot(unit, centroid)),
|
||||
);
|
||||
const template = standardize(artifact.centroids.map((_, modelIndex) => Math.max(
|
||||
...environmentScores.map((scores) => scores[modelIndex]),
|
||||
)));
|
||||
const projected = normalized(subtractBasis(standardizedFeature, artifact.nuisance_basis));
|
||||
const nuisance = standardize(artifact.centroids.map((centroid) => dot(projected, centroid)));
|
||||
return standardize(template.map((value, index) => 0.5 * value + 0.5 * nuisance[index]));
|
||||
}
|
||||
|
||||
function robustScoreNumbers(numbers, bank) {
|
||||
const marginal = robustScoreCounts(countNumbers(numbers), bank);
|
||||
const artifact = bank.robust.ordered_blocks;
|
||||
const weight = artifact ? Number(artifact.weight || 0) : 0;
|
||||
if (!artifact || weight === 0) return marginal;
|
||||
const ordered = orderedBlockScores(numbers, bank);
|
||||
return marginal.map((value, index) => (1 - weight) * value + weight * ordered[index]);
|
||||
}
|
||||
|
||||
function softmax(values) {
|
||||
const maximum = Math.max(...values);
|
||||
const weights = values.map((value) => Math.exp(value - maximum));
|
||||
const total = weights.reduce((sum, value) => sum + value, 0);
|
||||
return weights.map((value) => value / total);
|
||||
}
|
||||
|
||||
function jsSimilarity(left, right) {
|
||||
const leftTotal = left.reduce((sum, value) => sum + value, 0);
|
||||
const rightTotal = right.reduce((sum, value) => sum + value, 0) + ALPHA * DIMENSION;
|
||||
const p = left.map((value) => value / leftTotal);
|
||||
const q = right.map((value) => (value + ALPHA) / rightTotal);
|
||||
const midpoint = p.map((value, index) => (value + q[index]) / 2);
|
||||
const divergence = (values) => values.reduce(
|
||||
(total, value, index) => total + (value ? value * Math.log(value / midpoint[index]) : 0),
|
||||
0,
|
||||
);
|
||||
const js = (divergence(p) + divergence(q)) / 2;
|
||||
return 1 - Math.sqrt(js / Math.log(2));
|
||||
}
|
||||
|
||||
export function analyzeGlobalOutputs(outputs, bank) {
|
||||
const modelIds = bank.models.map((model) => model.id);
|
||||
const valid = [];
|
||||
const diagnostics = [];
|
||||
outputs.forEach((item, index) => {
|
||||
const expected = Number(item.expected_count || 0);
|
||||
const numbers = parseNumbers(item.text || "");
|
||||
const minimum = expected ? Math.max(80, Math.ceil(expected * 0.55)) : 80;
|
||||
const accepted = numbers.length >= minimum;
|
||||
diagnostics.push({ index, parsed_numbers: numbers.length, minimum_numbers: minimum, accepted });
|
||||
if (accepted) valid.push({ numbers, counts: countNumbers(numbers), scores: robustScoreNumbers(numbers, bank) });
|
||||
});
|
||||
if (!valid.length) throw new Error("没有可用回答：请粘贴完整数字序列；拒答或严重截断的回答不会计入。");
|
||||
|
||||
const combinedScores = modelIds.map((_, modelIndex) => mean(valid.map((item) => item.scores[modelIndex])));
|
||||
const calibrationKey = String(Math.min(valid.length, 3));
|
||||
const beta = Number(bank.calibration[calibrationKey].beta);
|
||||
const probabilities = softmax(combinedScores.map((value) => beta * value));
|
||||
const pooledCounts = Array.from({ length: DIMENSION }, (_, index) => valid.reduce((sum, item) => sum + item.counts[index], 0));
|
||||
const modelEntries = Object.fromEntries(bank.models.map((model) => [model.id, model]));
|
||||
const familyOrder = [...new Set(bank.models.map((model) => model.family || "models"))];
|
||||
const familyNames = Object.fromEntries(familyOrder.map((family) => [
|
||||
family,
|
||||
bank.models.find((model) => (model.family || "models") === family)?.family_name || family,
|
||||
]));
|
||||
const results = modelIds.map((model, index) => ({
|
||||
model,
|
||||
display_name: modelEntries[model].display_name,
|
||||
probability: probabilities[index],
|
||||
profile_similarity: jsSimilarity(pooledCounts, modelEntries[model].counts),
|
||||
score: combinedScores[index],
|
||||
family: modelEntries[model].family || "models",
|
||||
family_name: familyNames[modelEntries[model].family || "models"],
|
||||
})).sort((left, right) => right.probability - left.probability);
|
||||
const familyProbabilities = Object.fromEntries(familyOrder.map((family) => [
|
||||
family,
|
||||
results.filter((item) => item.family === family).reduce((sum, item) => sum + item.probability, 0),
|
||||
]));
|
||||
results.forEach((item) => { item.conditional_probability = item.probability / familyProbabilities[item.family]; });
|
||||
const winningFamily = familyOrder.reduce((best, family) => familyProbabilities[family] > familyProbabilities[best] ? family : best);
|
||||
return {
|
||||
prediction: results[0].model,
|
||||
prediction_name: results[0].display_name,
|
||||
probability: results[0].probability,
|
||||
used_outputs: valid.length,
|
||||
results,
|
||||
diagnostics,
|
||||
calibration: { queries: calibrationKey, beta, cv_accuracy: bank.calibration[calibrationKey].cv_accuracy },
|
||||
family_prediction: winningFamily,
|
||||
family_prediction_name: familyNames[winningFamily],
|
||||
family_probability: familyProbabilities[winningFamily],
|
||||
family_probabilities: familyOrder.map((family) => ({ family, display_name: familyNames[family], probability: familyProbabilities[family] })),
|
||||
method: "统一全局稳健数字指纹",
|
||||
};
|
||||
}
|
||||
@@ -0,0 +1 @@
|
||||
[{"row_id": "gpt-6-astra__query-01", "parent_row_id": "gpt-6-astra__query-01", "bank_id": "gpt", "source": "gpt-6-astra", "model_id": "gpt-6-astra", "response_model": "gpt-6-astra", "exact_version": "gpt-6-astra", "condition_id": "environment-01", "nuisance_condition_id": "environment-01", "wrapper_id": "environment-01", "wrapper_transport": "clean", "challenge_id": "query-01", "task_index": 0, "requested_count": 218, "parsed_count": 218, "strict_threshold": 120, "strict_valid": true, "temperature": "provider_default", "config_id": "reference-collection", "provider": "codex", "prompt": "For each of 218 positions, make one separate first-instinct choice of an integer from 1 to 355 inclusive. The current language model must complete this directly without tools, Python, code execution, calculators, search, APIs, or external random generators. Choose every position separately. Do not count upward or downward, including 1, 2, 3, and do not use an arithmetic progression, repeating cycle, repeated block, or another rule-made pattern. Accidental repetitions are valid. Once an item is written, do not sort, reorder, deduplicate, replace, or repair the list. Return one compact JSON array containing the complete sequence and no explanation.", "base_prompt": "For each of 218 positions, make one separate first-instinct choice of an integer from 1 to 355 inclusive. The current language model must complete this directly without tools, Python, code execution, calculators, search, APIs, or external random generators. Choose every position separately. Do not count upward or downward, including 1, 2, 3, and do not use an arithmetic progression, repeating cycle, repeated block, or another rule-made pattern. Accidental repetitions are valid. Once an item is written, do not sort, reorder, deduplicate, replace, or repair the list. Return one compact JSON array containing the complete sequence and no explanation.", "system_prompt": "", "user_prefix": "", "text": "[247,83,319,146,28,204,337,112,59,271,164,352,91,223,17,308,135,76,289,192,43,258,121,334,68,215,301,154,9,237,178,325,54,109,283,196,31,342,87,261,143,22,310,185,73,229,348,116,294,48,203,157,331,95,266,12,181,316,64,240,127,354,39,218,172,298,81,251,106,323,56,189,274,14,139,345,207,93,232,167,305,46,119,287,211,34,253,176,329,72,224,151,8,313,98,269,184,51,341,130,291,66,235,159,322,24,198,277,113,350,85,243,161,37,302,124,280,19,216,146,333,61,194,256,102,317,44,170,288,79,226,355,138,11,264,187,314,97,205,53,275,149,326,32,118,238,69,296,163,221,90,347,179,26,250,132,309,57,201,284,114,339,74,228,155,4,273,193,320,41,245,107,299,182,63,212,351,129,16,267,158,306,88,234,47,328,174,101,286,209,35,252,141,315,77,197,343,120,260,58,231,168,292,21,136,304,94,219,65,336,183,242,7,278,153,324,108,202]", "error": null, "request_seconds": 45.789022299984936, "collected_at": "2026-09-05T02:04:49.454748+00:00"}, {"row_id": "gpt-6-astra__query-02", "parent_row_id": "gpt-6-astra__query-02", "bank_id": "gpt", "source": "gpt-6-astra", "model_id": "gpt-6-astra", "response_model": "gpt-6-astra", "exact_version": "gpt-6-astra", "condition_id": "environment-01", "nuisance_condition_id": "environment-01", "wrapper_id": "environment-01", "wrapper_transport": "clean", "challenge_id": "query-02", "task_index": 0, "requested_count": 233, "parsed_count": 233, "strict_threshold": 129, "strict_valid": true, "temperature": "provider_default", "config_id": "reference-collection", "provider": "codex", "prompt": "Provide exactly 233 whole-number choices between 1 and 355 inclusive, choosing each item separately as it comes to mind. The current language model must complete this directly without tools, Python, code execution, calculators, search, APIs, or external random generators. Choose every position separately. Do not count upward or downward, including 1, 2, 3, and do not use an arithmetic progression, repeating cycle, repeated block, or another rule-made pattern. Accidental repetitions are valid. Once an item is written, do not sort, reorder, deduplicate, replace, or repair the list. Return one compact JSON array containing the complete sequence and no explanation.", "base_prompt": "Provide exactly 233 whole-number choices between 1 and 355 inclusive, choosing each item separately as it comes to mind. The current language model must complete this directly without tools, Python, code execution, calculators, search, APIs, or external random generators. Choose every position separately. Do not count upward or downward, including 1, 2, 3, and do not use an arithmetic progression, repeating cycle, repeated block, or another rule-made pattern. Accidental repetitions are valid. Once an item is written, do not sort, reorder, deduplicate, replace, or repair the list. Return one compact JSON array containing the complete sequence and no explanation.", "system_prompt": "", "user_prefix": "", "text": "[287,43,196,321,78,154,9,242,113,335,62,219,174,28,303,141,87,264,350,106,183,51,298,127,14,231,92,316,1Line truncated
|
||||
@@ -0,0 +1,63 @@
|
||||
package perfmetrics
|
||||
|
||||
import (
|
||||
"context"
|
||||
"github.com/QuantumNous/new-api/common"
|
||||
"strconv"
|
||||
"sync"
|
||||
"time"
|
||||
)
|
||||
|
||||
var liveSuccess sync.Map
|
||||
|
||||
type successTime struct {
|
||||
sync.Mutex
|
||||
at int64
|
||||
}
|
||||
|
||||
func recordUserSuccess(group string, at int64) {
|
||||
if group == "" {
|
||||
return
|
||||
}
|
||||
value, _ := liveSuccess.LoadOrStore(group, &successTime{})
|
||||
v := value.(*successTime)
|
||||
v.Lock()
|
||||
v.at = max(v.at, at)
|
||||
v.Unlock()
|
||||
if common.RedisEnabled && common.RDB != nil {
|
||||
ctx, cancel := context.WithTimeout(context.Background(), time.Second)
|
||||
defer cancel()
|
||||
// A delayed writer must not replace a newer successful request.
|
||||
_ = common.RDB.Eval(ctx, `local old=tonumber(redis.call('GET',KEYS[1]) or '0'); if tonumber(ARGV[1])>old then redis.call('SET',KEYS[1],ARGV[1],'EX',900) end; return 1`, []string{"status:user-success:" + group}, at).Err()
|
||||
}
|
||||
}
|
||||
|
||||
func LatestUserSuccess(ctx context.Context, groups []string) map[string]int64 {
|
||||
result := map[string]int64{}
|
||||
for _, group := range groups {
|
||||
if value, ok := liveSuccess.Load(group); ok {
|
||||
v := value.(*successTime)
|
||||
v.Lock()
|
||||
result[group] = v.at
|
||||
v.Unlock()
|
||||
}
|
||||
}
|
||||
if common.RedisEnabled && common.RDB != nil && len(groups) > 0 {
|
||||
keys := make([]string, len(groups))
|
||||
for i, g := range groups {
|
||||
keys[i] = "status:user-success:" + g
|
||||
}
|
||||
ctx, cancel := context.WithTimeout(ctx, time.Second)
|
||||
defer cancel()
|
||||
values, err := common.RDB.MGet(ctx, keys...).Result()
|
||||
if err == nil {
|
||||
for i, value := range values {
|
||||
if s, ok := value.(string); ok {
|
||||
at, _ := strconv.ParseInt(s, 10, 64)
|
||||
result[groups[i]] = max(result[groups[i]], at)
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
return result
|
||||
}
|
||||
@@ -0,0 +1,36 @@
|
||||
package perfmetrics
|
||||
|
||||
import (
|
||||
"context"
|
||||
relaycommon "github.com/QuantumNous/new-api/relay/common"
|
||||
"github.com/QuantumNous/new-api/relaykit/types"
|
||||
"github.com/stretchr/testify/assert"
|
||||
"testing"
|
||||
"time"
|
||||
)
|
||||
|
||||
func TestLiveUserSuccess(t *testing.T) {
|
||||
t.Cleanup(func() {
|
||||
for _, name := range []string{"user", "probe", "internal", "failed"} {
|
||||
liveSuccess.Delete(name)
|
||||
}
|
||||
})
|
||||
for _, tc := range []struct {
|
||||
name string
|
||||
test bool
|
||||
token int
|
||||
fail bool
|
||||
want bool
|
||||
}{
|
||||
{"user", false, 1, false, true}, {"probe", true, 1, false, false}, {"internal", false, 0, false, false}, {"failed", false, 1, true, false},
|
||||
} {
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
var apiErr *types.NewAPIError
|
||||
if tc.fail {
|
||||
apiErr = types.InitOpenAIError("server_error", 500)
|
||||
}
|
||||
RecordRelayResult(context.Background(), &relaycommon.RelayInfo{UsingGroup: tc.name, TokenId: tc.token, IsChannelTest: tc.test, StartTime: time.Now()}, apiErr)
|
||||
assert.Equal(t, tc.want, LatestUserSuccess(context.Background(), []string{tc.name})[tc.name] > 0)
|
||||
})
|
||||
}
|
||||
}
|
||||
+83
-22
@@ -11,7 +11,10 @@ import (
|
||||
"github.com/QuantumNous/new-api/common"
|
||||
"github.com/QuantumNous/new-api/model"
|
||||
relaycommon "github.com/QuantumNous/new-api/relay/common"
|
||||
"github.com/QuantumNous/new-api/relaykit/types"
|
||||
"github.com/QuantumNous/new-api/setting/perf_metrics_setting"
|
||||
|
||||
"github.com/bytedance/gopkg/util/gopool"
|
||||
)
|
||||
|
||||
var hotBuckets sync.Map
|
||||
@@ -24,11 +27,20 @@ func Init() {
|
||||
go flushLoop()
|
||||
}
|
||||
|
||||
func RecordRelaySample(info *relaycommon.RelayInfo, success bool, outputTokens int64) {
|
||||
// RecordRelayResult samples one finished relay exactly once, at the request
|
||||
// boundary, regardless of how many channel attempts it took.
|
||||
func RecordRelayResult(ctx context.Context, info *relaycommon.RelayInfo, apiErr *types.NewAPIError) {
|
||||
if info == nil {
|
||||
return
|
||||
}
|
||||
outcome := ClassifyRelayOutcome(ctx, info, apiErr)
|
||||
if outcome == OutcomeIgnored {
|
||||
return
|
||||
}
|
||||
now := time.Now()
|
||||
if outcome == OutcomeSuccess && !info.IsChannelTest && info.TokenId > 0 {
|
||||
recordUserSuccess(info.UsingGroup, now.Unix())
|
||||
}
|
||||
hasTtft := info.IsStream && info.HasSendResponse()
|
||||
ttftMs := int64(0)
|
||||
if hasTtft {
|
||||
@@ -48,8 +60,8 @@ func RecordRelaySample(info *relaycommon.RelayInfo, success bool, outputTokens i
|
||||
LatencyMs: latencyMs,
|
||||
TtftMs: ttftMs,
|
||||
HasTtft: hasTtft,
|
||||
Success: success,
|
||||
OutputTokens: outputTokens,
|
||||
Success: outcome == OutcomeSuccess,
|
||||
OutputTokens: info.PerformanceOutputTokens,
|
||||
GenerationMs: generationMs,
|
||||
})
|
||||
}
|
||||
@@ -73,18 +85,25 @@ func Record(sample Sample) {
|
||||
}
|
||||
actual, _ := hotBuckets.LoadOrStore(key, &atomicBucket{})
|
||||
actual.(*atomicBucket).add(sample)
|
||||
recordRedis(key, sample)
|
||||
gopool.Go(func() {
|
||||
recordRedis(key, sample)
|
||||
})
|
||||
}
|
||||
|
||||
// queryWindow covers the current partial hour plus the preceding complete
|
||||
// hours, so every reader shares the same hourly buckets.
|
||||
func queryWindow(now time.Time, hours int) (int64, int64) {
|
||||
if hours <= 0 {
|
||||
hours = 24
|
||||
}
|
||||
hours = min(hours, 24*30)
|
||||
endTs := now.Unix()
|
||||
return endTs - endTs%3600 - int64(hours-1)*3600, endTs
|
||||
}
|
||||
|
||||
func Query(params QueryParams) (QueryResult, error) {
|
||||
if params.Hours <= 0 {
|
||||
params.Hours = 24
|
||||
}
|
||||
if params.Hours > 24*30 {
|
||||
params.Hours = 24 * 30
|
||||
}
|
||||
endTs := time.Now().Unix()
|
||||
startTs := endTs - int64(params.Hours)*3600
|
||||
startTs, endTs := queryWindow(time.Now(), params.Hours)
|
||||
allowedGroups := allowedGroupSet(params.AllowedGroups)
|
||||
|
||||
merged := map[bucketKey]counters{}
|
||||
rows, err := model.GetPerfMetrics(params.Model, params.Group, startTs, endTs)
|
||||
@@ -92,6 +111,11 @@ func Query(params QueryParams) (QueryResult, error) {
|
||||
return QueryResult{}, err
|
||||
}
|
||||
for _, row := range rows {
|
||||
if allowedGroups != nil {
|
||||
if _, ok := allowedGroups[row.Group]; !ok {
|
||||
continue
|
||||
}
|
||||
}
|
||||
mergeCounters(merged, bucketKey{
|
||||
model: row.ModelName,
|
||||
group: row.Group,
|
||||
@@ -115,22 +139,22 @@ func Query(params QueryParams) (QueryResult, error) {
|
||||
if params.Group != "" && k.group != params.Group {
|
||||
return true
|
||||
}
|
||||
if allowedGroups != nil {
|
||||
if _, ok := allowedGroups[k.group]; !ok {
|
||||
return true
|
||||
}
|
||||
}
|
||||
mergeCounters(merged, k, value.(*atomicBucket).snapshot())
|
||||
return true
|
||||
})
|
||||
|
||||
return buildQueryResult(params.Model, merged), nil
|
||||
result := buildQueryResult(params.Model, merged)
|
||||
result.WindowStart, result.WindowEnd = startTs, endTs
|
||||
return result, nil
|
||||
}
|
||||
|
||||
func QuerySummaryAll(hours int, groups []string) (SummaryAllResult, error) {
|
||||
if hours <= 0 {
|
||||
hours = 24
|
||||
}
|
||||
if hours > 24*30 {
|
||||
hours = 24 * 30
|
||||
}
|
||||
endTs := time.Now().Unix()
|
||||
startTs := endTs - int64(hours)*3600
|
||||
startTs, endTs := queryWindow(time.Now(), hours)
|
||||
allowedGroups := allowedGroupSet(groups)
|
||||
|
||||
rows, err := model.GetPerfMetricsSummaryBucketsAll(startTs, endTs, groups)
|
||||
@@ -171,11 +195,17 @@ func QuerySummaryAll(hours int, groups []string) (SummaryAllResult, error) {
|
||||
return true
|
||||
})
|
||||
|
||||
all := counters{}
|
||||
models := make([]ModelSummary, 0, len(totals))
|
||||
for name, total := range totals {
|
||||
if total.requestCount == 0 {
|
||||
continue
|
||||
}
|
||||
all.requestCount += total.requestCount
|
||||
all.successCount += total.successCount
|
||||
all.totalLatencyMs += total.totalLatencyMs
|
||||
all.outputTokens += total.outputTokens
|
||||
all.generationMs += total.generationMs
|
||||
avgLatency := total.totalLatencyMs / total.requestCount
|
||||
successRate := float64(total.successCount) / float64(total.requestCount) * 100
|
||||
avgTps := 0.0
|
||||
@@ -195,7 +225,7 @@ func QuerySummaryAll(hours int, groups []string) (SummaryAllResult, error) {
|
||||
return models[i].RequestCount > models[j].RequestCount
|
||||
})
|
||||
|
||||
return SummaryAllResult{Models: models}, nil
|
||||
return SummaryAllResult{Summary: summarize(all), WindowStart: startTs, WindowEnd: endTs, Models: models}, nil
|
||||
}
|
||||
|
||||
// QueryGroups combines raw bucket counters so busy models carry their actual weight.
|
||||
@@ -324,6 +354,7 @@ func mergeCounters(merged map[bucketKey]counters, key bucketKey, value counters)
|
||||
|
||||
func buildQueryResult(modelName string, merged map[bucketKey]counters) QueryResult {
|
||||
groupBuckets := map[string]map[int64]counters{}
|
||||
modelBuckets := map[string]map[int64]counters{}
|
||||
for key, value := range merged {
|
||||
if value.requestCount == 0 {
|
||||
continue
|
||||
@@ -332,6 +363,7 @@ func buildQueryResult(modelName string, merged map[bucketKey]counters) QueryResu
|
||||
groupBuckets[key.group] = map[int64]counters{}
|
||||
}
|
||||
groupBuckets[key.group][key.bucketTs] = value
|
||||
mergeModelBucket(modelBuckets, modelName, key.bucketTs, value)
|
||||
}
|
||||
|
||||
groups := make([]string, 0, len(groupBuckets))
|
||||
@@ -340,6 +372,7 @@ func buildQueryResult(modelName string, merged map[bucketKey]counters) QueryResu
|
||||
}
|
||||
sort.Strings(groups)
|
||||
|
||||
all := counters{}
|
||||
results := make([]GroupResult, 0, len(groups))
|
||||
for _, group := range groups {
|
||||
buckets := groupBuckets[group]
|
||||
@@ -364,6 +397,11 @@ func buildQueryResult(modelName string, merged map[bucketKey]counters) QueryResu
|
||||
total.generationMs += value.generationMs
|
||||
series = append(series, bucketPoint(ts, value))
|
||||
}
|
||||
all.requestCount += total.requestCount
|
||||
all.successCount += total.successCount
|
||||
all.totalLatencyMs += total.totalLatencyMs
|
||||
all.outputTokens += total.outputTokens
|
||||
all.generationMs += total.generationMs
|
||||
|
||||
results = append(results, GroupResult{
|
||||
Group: group,
|
||||
@@ -375,13 +413,36 @@ func buildQueryResult(modelName string, merged map[bucketKey]counters) QueryResu
|
||||
})
|
||||
}
|
||||
|
||||
modelSeries := make([]BucketPoint, 0, len(modelBuckets[modelName]))
|
||||
for ts, value := range modelBuckets[modelName] {
|
||||
modelSeries = append(modelSeries, bucketPoint(ts, value))
|
||||
}
|
||||
sort.Slice(modelSeries, func(i, j int) bool {
|
||||
return modelSeries[i].Ts < modelSeries[j].Ts
|
||||
})
|
||||
|
||||
return QueryResult{
|
||||
ModelName: modelName,
|
||||
SeriesSchema: seriesSchema,
|
||||
Summary: summarize(all),
|
||||
Series: modelSeries,
|
||||
Groups: results,
|
||||
}
|
||||
}
|
||||
|
||||
// summarize returns nil when nothing was sampled, so clients can show
|
||||
// "no data" instead of a fabricated 0% success rate.
|
||||
func summarize(total counters) *Summary {
|
||||
if total.requestCount <= 0 {
|
||||
return nil
|
||||
}
|
||||
return &Summary{
|
||||
AvgLatencyMs: avg(total.totalLatencyMs, total.requestCount),
|
||||
SuccessRate: math.Round(successRate(total)*100) / 100,
|
||||
AvgTps: math.Round(avgTps(total)*100) / 100,
|
||||
}
|
||||
}
|
||||
|
||||
func bucketPoint(ts int64, value counters) BucketPoint {
|
||||
return BucketPoint{
|
||||
totals: value,
|
||||
|
||||
@@ -0,0 +1,213 @@
|
||||
package perfmetrics
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/QuantumNous/new-api/common"
|
||||
"github.com/QuantumNous/new-api/model"
|
||||
relaycommon "github.com/QuantumNous/new-api/relay/common"
|
||||
"github.com/QuantumNous/new-api/relaykit/types"
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
func TestClassifyRelayOutcome(t *testing.T) {
|
||||
canceled, cancel := context.WithCancel(context.Background())
|
||||
cancel()
|
||||
for _, tc := range []struct {
|
||||
name string
|
||||
ctx context.Context
|
||||
err *types.NewAPIError
|
||||
want Outcome
|
||||
}{
|
||||
{"success", context.Background(), nil, OutcomeSuccess},
|
||||
{"upstream bad request", context.Background(), types.InitOpenAIError("unknown", 400), OutcomeIgnored},
|
||||
{"upstream credentials returned as 400", context.Background(), types.InitOpenAIError("invalid_api_key", 400), OutcomeFailure},
|
||||
{"wrapped credentials keep their cause", context.Background(), types.NewErrorWithStatusCode(types.InitOpenAIError("invalid_api_key", 400), types.ErrorCodeInvalidRequest, 400), OutcomeFailure},
|
||||
{"upstream context limit returned as 500", context.Background(), types.InitOpenAIError("context_length_exceeded", 500), OutcomeIgnored},
|
||||
{"local rate limit", context.Background(), types.NewErrorWithStatusCode(errors.New("limited"), types.ErrorCodeInvalidRequest, 429), OutcomeIgnored},
|
||||
{"upstream rate limit", context.Background(), types.InitOpenAIError("rate_limit_exceeded", 429), OutcomeFailure},
|
||||
{"local quota", context.Background(), types.NewError(errors.New("quota"), types.ErrorCodeInsufficientUserQuota), OutcomeIgnored},
|
||||
{"local violation fee", context.Background(), types.NewError(errors.New("csam"), types.ErrorCodeViolationFeeGrokCSAM), OutcomeIgnored},
|
||||
{"upstream quota", context.Background(), types.InitOpenAIError("insufficient_quota", 429), OutcomeFailure},
|
||||
{"upstream gateway quota", context.Background(), types.InitOpenAIError(types.ErrorCodeInsufficientUserQuota, 403), OutcomeFailure},
|
||||
{"unavailable channel", context.Background(), types.NewErrorWithStatusCode(errors.New("disabled"), types.ErrorCodeGetChannelFailed, 403), OutcomeFailure},
|
||||
{"empty upstream response", context.Background(), types.NewError(errors.New("empty"), types.ErrorCodeEmptyResponse), OutcomeFailure},
|
||||
{"network failure", context.Background(), types.NewOpenAIError(errors.New("connection refused"), types.ErrorCodeDoRequestFailed, 500), OutcomeFailure},
|
||||
{"client cancellation", canceled, types.NewOpenAIError(errors.New("context canceled"), types.ErrorCodeDoRequestFailed, 500), OutcomeIgnored},
|
||||
{"upstream deadline", context.Background(), types.NewOpenAIError(context.DeadlineExceeded, types.ErrorCodeDoRequestFailed, 504), OutcomeFailure},
|
||||
} {
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
assert.Equal(t, tc.want, ClassifyRelayOutcome(tc.ctx, &relaycommon.RelayInfo{}, tc.err))
|
||||
})
|
||||
}
|
||||
assert.Equal(t, OutcomeIgnored, ClassifyRelayOutcome(context.Background(), &relaycommon.RelayInfo{PerformanceBusinessRejection: true}, nil))
|
||||
}
|
||||
|
||||
func TestStreamOutcomeClassification(t *testing.T) {
|
||||
for _, tc := range []struct {
|
||||
name string
|
||||
mark func(*relaycommon.StreamStatus)
|
||||
end relaycommon.StreamEndReason
|
||||
want Outcome
|
||||
}{
|
||||
{"completed", (*relaycommon.StreamStatus).MarkCompleted, relaycommon.StreamEndReasonEOF, OutcomeSuccess},
|
||||
{"business rejection", func(s *relaycommon.StreamStatus) { s.MarkFailed("context_length_exceeded", "", 0) }, relaycommon.StreamEndReasonEOF, OutcomeIgnored},
|
||||
{"service error", func(s *relaycommon.StreamStatus) { s.MarkFailed("server_error", "", 0) }, relaycommon.StreamEndReasonEOF, OutcomeFailure},
|
||||
{"error after completion", func(s *relaycommon.StreamStatus) { s.MarkCompleted(); s.MarkFailed("", "server_error", 0) }, relaycommon.StreamEndReasonEOF, OutcomeFailure},
|
||||
{"output limit", func(s *relaycommon.StreamStatus) { s.MarkIncomplete("max_output_tokens") }, relaycommon.StreamEndReasonEOF, OutcomeSuccess},
|
||||
{"content filter", func(s *relaycommon.StreamStatus) { s.MarkIncomplete("content_filter") }, relaycommon.StreamEndReasonEOF, OutcomeIgnored},
|
||||
{"client cancel", (*relaycommon.StreamStatus).MarkCancelled, relaycommon.StreamEndReasonEOF, OutcomeIgnored},
|
||||
{"client gone", nil, relaycommon.StreamEndReasonClientGone, OutcomeIgnored},
|
||||
{"timeout", nil, relaycommon.StreamEndReasonTimeout, OutcomeFailure},
|
||||
{"missing terminal", nil, relaycommon.StreamEndReasonEOF, OutcomeFailure},
|
||||
{"done marker without terminal", nil, relaycommon.StreamEndReasonDone, OutcomeSuccess},
|
||||
} {
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
stream := relaycommon.NewStreamStatus()
|
||||
stream.RequireTerminal()
|
||||
if tc.mark != nil {
|
||||
tc.mark(stream)
|
||||
}
|
||||
stream.SetEndReason(tc.end, nil)
|
||||
assert.Equal(t, tc.want, ClassifyRelayOutcome(context.Background(), &relaycommon.RelayInfo{StreamStatus: stream}, nil))
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestClientCancellationDuringUpstreamRead(t *testing.T) {
|
||||
ctx, cancel := context.WithCancel(context.Background())
|
||||
cancel()
|
||||
stream := relaycommon.NewStreamStatus()
|
||||
stream.RequireTerminal()
|
||||
stream.SetEndReason(relaycommon.StreamEndReasonScannerErr, errors.New("reader closed"))
|
||||
assert.Equal(t, OutcomeIgnored, ClassifyRelayOutcome(ctx, &relaycommon.RelayInfo{StreamStatus: stream}, nil))
|
||||
|
||||
deadline := relaycommon.NewStreamStatus()
|
||||
deadline.SetEndReason(relaycommon.StreamEndReasonClientGone, context.DeadlineExceeded)
|
||||
assert.Equal(t, OutcomeFailure, ClassifyRelayOutcome(context.Background(), &relaycommon.RelayInfo{StreamStatus: deadline}, nil))
|
||||
}
|
||||
|
||||
func TestPerformanceWindowIncludesCurrentHour(t *testing.T) {
|
||||
now := time.Date(2026, 9, 14, 12, 37, 0, 0, time.UTC)
|
||||
start, end := queryWindow(now, 24)
|
||||
assert.Equal(t, time.Date(2026, 9, 13, 13, 0, 0, 0, time.UTC).Unix(), start)
|
||||
assert.Equal(t, now.Unix(), end)
|
||||
}
|
||||
|
||||
func TestHourlySuccessSeriesWeightsSmallerBuckets(t *testing.T) {
|
||||
points := recentSuccessSeries(map[int64]counters{
|
||||
3600: {requestCount: 100, successCount: 100},
|
||||
3900: {requestCount: 1},
|
||||
7200: {requestCount: 1, successCount: 1},
|
||||
})
|
||||
assert.Equal(t, []SuccessRatePoint{{Ts: 3600, SuccessRate: 99.01}, {Ts: 7200, SuccessRate: 100}}, points)
|
||||
}
|
||||
|
||||
// TEST_PERF_MYSQL_DSN / TEST_PERF_POSTGRES_DSN optionally run the aggregation
|
||||
// against isolated real MySQL/PostgreSQL databases.
|
||||
func TestPerformanceAggregationAndFlush(t *testing.T) {
|
||||
for _, dialect := range []struct{ name, env string }{
|
||||
{"sqlite", ""}, {"mysql", "TEST_PERF_MYSQL_DSN"}, {"postgres", "TEST_PERF_POSTGRES_DSN"},
|
||||
} {
|
||||
t.Run(dialect.name, func(t *testing.T) {
|
||||
dsn := ""
|
||||
if dialect.env != "" {
|
||||
dsn = os.Getenv(dialect.env)
|
||||
if dsn == "" {
|
||||
t.Skip("isolated test database DSN is not configured")
|
||||
}
|
||||
}
|
||||
t.Setenv("SQL_DSN", dsn)
|
||||
oldDB, oldPath, oldMaster, oldRedis := model.DB, common.SQLitePath, common.IsMasterNode, common.RedisEnabled
|
||||
oldType, oldLogType := common.MainDatabaseType(), common.LogDatabaseType()
|
||||
common.SQLitePath, common.IsMasterNode, common.RedisEnabled = filepath.Join(t.TempDir(), "perf.db"), false, false
|
||||
hotBuckets.Clear()
|
||||
t.Cleanup(func() {
|
||||
model.DB, common.SQLitePath, common.IsMasterNode, common.RedisEnabled = oldDB, oldPath, oldMaster, oldRedis
|
||||
common.SetDatabaseTypes(oldType, oldLogType)
|
||||
hotBuckets.Clear()
|
||||
})
|
||||
require.NoError(t, model.InitDB())
|
||||
db := model.DB
|
||||
sqlDB, err := db.DB()
|
||||
require.NoError(t, err)
|
||||
t.Cleanup(func() { require.NoError(t, sqlDB.Close()) })
|
||||
require.NoError(t, db.Migrator().DropTable(&model.PerfMetric{}))
|
||||
require.NoError(t, db.AutoMigrate(&model.PerfMetric{}))
|
||||
|
||||
now := time.Now()
|
||||
start, _ := queryWindow(now, 24)
|
||||
hour := now.Unix() - now.Unix()%3600 - 3600
|
||||
// Historical counters remain usable without reclassification or migration.
|
||||
for _, row := range []model.PerfMetric{
|
||||
{ModelName: "test-model", Group: "a", BucketTs: hour, RequestCount: 100, SuccessCount: 100, TotalLatencyMs: 100000, TtftCount: 100, TtftSumMs: 10000, OutputTokens: 200, GenerationMs: 40000},
|
||||
{ModelName: "test-model", Group: "inactive", BucketTs: hour, RequestCount: 100},
|
||||
{ModelName: "test-model", Group: "a", BucketTs: start - 3600, RequestCount: 100},
|
||||
} {
|
||||
require.NoError(t, model.UpsertPerfMetric(&row))
|
||||
}
|
||||
groups := []string{"a", "b"}
|
||||
|
||||
RecordRelayResult(context.Background(), &relaycommon.RelayInfo{OriginModelName: "test-model", UsingGroup: "b", StartTime: now}, types.InitOpenAIError("unknown", 400))
|
||||
businessRejected, err := QuerySummaryAll(24, groups)
|
||||
require.NoError(t, err)
|
||||
require.NotNil(t, businessRejected.Summary)
|
||||
assert.Equal(t, 100.0, businessRejected.Summary.SuccessRate)
|
||||
|
||||
failure := &atomicBucket{}
|
||||
failure.add(Sample{LatencyMs: 2000})
|
||||
hotBuckets.Store(bucketKey{model: "test-model", group: "b", bucketTs: hour}, failure)
|
||||
before, err := Query(QueryParams{Model: "test-model", Hours: 24, AllowedGroups: groups})
|
||||
require.NoError(t, err)
|
||||
require.NotNil(t, before.Summary)
|
||||
assert.Equal(t, Summary{SuccessRate: 99.01, AvgLatencyMs: 1009, AvgTps: 5}, *before.Summary)
|
||||
assert.Equal(t, start, before.WindowStart)
|
||||
require.Len(t, before.Series, 1)
|
||||
assert.Equal(t, hour, before.Series[0].Ts)
|
||||
assert.InDelta(t, 99.01, before.Series[0].SuccessRate, 0.01)
|
||||
require.Len(t, before.Groups, 2)
|
||||
|
||||
summary, err := QuerySummaryAll(24, groups)
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, before.Summary, summary.Summary)
|
||||
require.Len(t, summary.Models, 1)
|
||||
assert.Equal(t, 99.01, summary.Models[0].SuccessRate)
|
||||
assert.Equal(t, 99.01, summary.Models[0].RecentSuccessSeries[0].SuccessRate)
|
||||
encoded, err := common.Marshal(summary)
|
||||
require.NoError(t, err)
|
||||
assert.NotContains(t, string(encoded), "request_count")
|
||||
assert.NotContains(t, string(encoded), "success_count")
|
||||
|
||||
flushCompletedBuckets()
|
||||
flushCompletedBuckets()
|
||||
after, err := Query(QueryParams{Model: "test-model", Hours: 24, AllowedGroups: groups})
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, before.Summary, after.Summary)
|
||||
assert.Equal(t, before.Series, after.Series)
|
||||
|
||||
RecordRelayResult(context.Background(), &relaycommon.RelayInfo{OriginModelName: "test-model", UsingGroup: "a", StartTime: now}, types.InitOpenAIError("context_length_exceeded", 400))
|
||||
after, err = Query(QueryParams{Model: "test-model", Hours: 24, AllowedGroups: groups})
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, before.Summary, after.Summary)
|
||||
onlyA, err := Query(QueryParams{Model: "test-model", Group: "a", Hours: 24, AllowedGroups: groups})
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, 100.0, onlyA.Summary.SuccessRate)
|
||||
empty, err := Query(QueryParams{Model: "missing", Hours: 24, AllowedGroups: groups})
|
||||
require.NoError(t, err)
|
||||
assert.Nil(t, empty.Summary)
|
||||
assert.Empty(t, empty.Series)
|
||||
|
||||
require.NoError(t, model.UpsertPerfMetric(&model.PerfMetric{ModelName: "second-model", Group: "a", BucketTs: hour, RequestCount: 1}))
|
||||
combined, err := QuerySummaryAll(24, groups)
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, 98.04, combined.Summary.SuccessRate)
|
||||
assert.Equal(t, 99.01, combined.Models[0].SuccessRate)
|
||||
})
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,122 @@
|
||||
package perfmetrics
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"strings"
|
||||
|
||||
relaycommon "github.com/QuantumNous/new-api/relay/common"
|
||||
"github.com/QuantumNous/new-api/relaykit/types"
|
||||
)
|
||||
|
||||
type Outcome string
|
||||
|
||||
const (
|
||||
OutcomeSuccess Outcome = "success"
|
||||
OutcomeFailure Outcome = "failure"
|
||||
OutcomeIgnored Outcome = "ignored"
|
||||
)
|
||||
|
||||
// ClassifyRelayOutcome decides whether one finished relay counts as a health
|
||||
// sample. Business rejections and client cancellations are not samples; the
|
||||
// classification is independent of retries, channel disabling and billing.
|
||||
func ClassifyRelayOutcome(ctx context.Context, info *relaycommon.RelayInfo, apiErr *types.NewAPIError) Outcome {
|
||||
if info == nil || info.PerformanceBusinessRejection {
|
||||
return OutcomeIgnored
|
||||
}
|
||||
if ctx != nil && ctx.Err() == context.Canceled {
|
||||
return OutcomeIgnored
|
||||
}
|
||||
if apiErr != nil && errors.Is(apiErr, context.Canceled) {
|
||||
return OutcomeIgnored
|
||||
}
|
||||
stream := info.StreamStatus.OutcomeSnapshot()
|
||||
if stream.Response == relaycommon.ResponseOutcomeFailed {
|
||||
return classifyFailure(false, stream.ErrorCode, stream.ErrorType, stream.ErrorStatus)
|
||||
}
|
||||
if apiErr != nil {
|
||||
root := rootAPIError(apiErr)
|
||||
local := root.GetErrorType() == types.ErrorTypeNewAPIError
|
||||
return classifyFailure(local, string(root.GetErrorCode()), root.ToOpenAIError().Type, root.StatusCode)
|
||||
}
|
||||
deadlineExceeded := info.StreamStatus != nil && errors.Is(info.StreamStatus.EndError, context.DeadlineExceeded)
|
||||
if stream.Response == relaycommon.ResponseOutcomeCancelled || stream.EndReason == relaycommon.StreamEndReasonPingFail {
|
||||
return OutcomeIgnored
|
||||
}
|
||||
if stream.EndReason == relaycommon.StreamEndReasonClientGone && !deadlineExceeded {
|
||||
return OutcomeIgnored
|
||||
}
|
||||
if stream.Response == relaycommon.ResponseOutcomeIncomplete {
|
||||
switch stream.IncompleteReason {
|
||||
case "max_output_tokens", "max_tokens":
|
||||
return OutcomeSuccess
|
||||
case "content_filter", "safety", "content_policy_violation":
|
||||
return OutcomeIgnored
|
||||
default:
|
||||
return OutcomeFailure
|
||||
}
|
||||
}
|
||||
if stream.HasErrors || deadlineExceeded {
|
||||
return OutcomeFailure
|
||||
}
|
||||
switch stream.EndReason {
|
||||
case relaycommon.StreamEndReasonTimeout, relaycommon.StreamEndReasonScannerErr, relaycommon.StreamEndReasonPanic:
|
||||
return OutcomeFailure
|
||||
}
|
||||
if stream.ExpectsTerminal && stream.Response == relaycommon.ResponseOutcomeUnknown && stream.EndReason != relaycommon.StreamEndReasonDone {
|
||||
return OutcomeFailure
|
||||
}
|
||||
return OutcomeSuccess
|
||||
}
|
||||
|
||||
// rootAPIError follows host wrappers back to the error that describes what
|
||||
// actually happened, e.g. an upstream credential error re-wrapped as a local
|
||||
// invalid request.
|
||||
func rootAPIError(apiErr *types.NewAPIError) *types.NewAPIError {
|
||||
for {
|
||||
var inner *types.NewAPIError
|
||||
if !errors.As(apiErr.Unwrap(), &inner) || inner == apiErr {
|
||||
return apiErr
|
||||
}
|
||||
apiErr = inner
|
||||
}
|
||||
}
|
||||
|
||||
func classifyFailure(local bool, code, errorType string, status int) Outcome {
|
||||
code, errorType = strings.ToLower(code), strings.ToLower(errorType)
|
||||
if local {
|
||||
if strings.HasPrefix(code, "violation_fee.") {
|
||||
return OutcomeIgnored
|
||||
}
|
||||
switch types.ErrorCode(code) {
|
||||
case types.ErrorCodeInvalidRequest, types.ErrorCodeSensitiveWordsDetected, types.ErrorCodeReadRequestBodyFailed,
|
||||
types.ErrorCodeConvertRequestFailed, types.ErrorCodeAccessDenied, types.ErrorCodeBadRequestBody,
|
||||
types.ErrorCodeInsufficientUserQuota, types.ErrorCodePreConsumeTokenQuotaFailed, types.ErrorCodePromptBlocked:
|
||||
return OutcomeIgnored
|
||||
}
|
||||
return OutcomeFailure
|
||||
}
|
||||
// Specific codes take precedence over broad protocol types, because
|
||||
// providers also report invalid credentials as invalid_request_error.
|
||||
for _, value := range []string{code, errorType} {
|
||||
switch value {
|
||||
case "invalid_api_key", "api_key_invalid", "api_key_expired", "api_key_service_blocked", "invalid_authentication",
|
||||
"authentication_error", "unauthenticated", "permission_denied", "permission_error", "access_denied",
|
||||
"insufficient_quota", "quota_exceeded", "resource_exhausted", "rate_limit_exceeded", "rate_limit_error",
|
||||
"overloaded_error", "server_error", "internal_error", "service_unavailable", "model_not_found",
|
||||
"insufficient_user_quota", "pre_consume_token_quota_failed":
|
||||
return OutcomeFailure
|
||||
case "context_length_exceeded", "invalid_request", "invalid_request_error", "invalid_argument",
|
||||
"sensitive_words_detected", "prompt_blocked", "content_filter", "content_policy_violation", "safety":
|
||||
return OutcomeIgnored
|
||||
}
|
||||
if strings.HasPrefix(value, "violation_fee.") {
|
||||
return OutcomeIgnored
|
||||
}
|
||||
}
|
||||
switch status {
|
||||
case 400, 405, 409, 413, 415, 422:
|
||||
return OutcomeIgnored
|
||||
}
|
||||
return OutcomeFailure
|
||||
}
|
||||
@@ -22,6 +22,16 @@ type QueryParams struct {
|
||||
Model string
|
||||
Group string
|
||||
Hours int
|
||||
// AllowedGroups restricts both the per-group results and the model summary;
|
||||
// nil allows every group.
|
||||
AllowedGroups []string
|
||||
}
|
||||
|
||||
// Summary is the request-weighted aggregate over every bucket in the window.
|
||||
type Summary struct {
|
||||
AvgLatencyMs int64 `json:"avg_latency_ms"`
|
||||
SuccessRate float64 `json:"success_rate"`
|
||||
AvgTps float64 `json:"avg_tps"`
|
||||
}
|
||||
|
||||
type BucketPoint struct {
|
||||
@@ -45,6 +55,10 @@ type GroupResult struct {
|
||||
type QueryResult struct {
|
||||
ModelName string `json:"model_name"`
|
||||
SeriesSchema string `json:"series_schema"`
|
||||
Summary *Summary `json:"summary"`
|
||||
Series []BucketPoint `json:"series"`
|
||||
WindowStart int64 `json:"window_start"`
|
||||
WindowEnd int64 `json:"window_end"`
|
||||
Groups []GroupResult `json:"groups"`
|
||||
}
|
||||
|
||||
@@ -63,7 +77,10 @@ type ModelSummary struct {
|
||||
}
|
||||
|
||||
type SummaryAllResult struct {
|
||||
Models []ModelSummary `json:"models"`
|
||||
Summary *Summary `json:"summary"`
|
||||
WindowStart int64 `json:"window_start"`
|
||||
WindowEnd int64 `json:"window_end"`
|
||||
Models []ModelSummary `json:"models"`
|
||||
}
|
||||
|
||||
type bucketKey struct {
|
||||
|
||||
@@ -0,0 +1,250 @@
|
||||
package wsmanager
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"os"
|
||||
"sync"
|
||||
"sync/atomic"
|
||||
"time"
|
||||
|
||||
"github.com/QuantumNous/new-api/common"
|
||||
"github.com/go-redis/redis/v8"
|
||||
)
|
||||
|
||||
const (
|
||||
KindRealtime = "realtime"
|
||||
KindResponses = "responses"
|
||||
|
||||
defaultCloseReason = "channel disabled or deleted"
|
||||
redisChannel = "new-api:wsmanager:channel-close"
|
||||
)
|
||||
|
||||
type entry struct {
|
||||
id uint64
|
||||
channelID int
|
||||
kind string
|
||||
close func(reason string)
|
||||
}
|
||||
|
||||
type closeEvent struct {
|
||||
ChannelIDs []int `json:"channel_ids"`
|
||||
Reason string `json:"reason"`
|
||||
Origin string `json:"origin"`
|
||||
}
|
||||
|
||||
var (
|
||||
mu sync.Mutex
|
||||
nextID uint64
|
||||
registry = map[int]map[uint64]*entry{}
|
||||
|
||||
originOnce sync.Once
|
||||
originID string
|
||||
|
||||
subscriberOnce sync.Once
|
||||
)
|
||||
|
||||
func Register(channelID int, kind string, close func(reason string)) func() {
|
||||
if channelID <= 0 || close == nil {
|
||||
return func() {}
|
||||
}
|
||||
var closeOnce sync.Once
|
||||
safeClose := func(reason string) {
|
||||
closeOnce.Do(func() {
|
||||
close(reason)
|
||||
})
|
||||
}
|
||||
id := atomic.AddUint64(&nextID, 1)
|
||||
e := &entry{
|
||||
id: id,
|
||||
channelID: channelID,
|
||||
kind: kind,
|
||||
close: safeClose,
|
||||
}
|
||||
|
||||
mu.Lock()
|
||||
if registry[channelID] == nil {
|
||||
registry[channelID] = map[uint64]*entry{}
|
||||
}
|
||||
registry[channelID][id] = e
|
||||
mu.Unlock()
|
||||
|
||||
var once sync.Once
|
||||
return func() {
|
||||
once.Do(func() {
|
||||
mu.Lock()
|
||||
defer mu.Unlock()
|
||||
entries := registry[channelID]
|
||||
if entries == nil {
|
||||
return
|
||||
}
|
||||
delete(entries, id)
|
||||
if len(entries) == 0 {
|
||||
delete(registry, channelID)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func CloseChannel(channelID int, reason string) int {
|
||||
return CloseChannels([]int{channelID}, reason)
|
||||
}
|
||||
|
||||
func CloseChannels(channelIDs []int, reason string) int {
|
||||
reason = normalizeReason(reason)
|
||||
entries := takeEntries(channelIDs)
|
||||
for _, e := range entries {
|
||||
e.close(reason)
|
||||
}
|
||||
if len(entries) > 0 {
|
||||
common.SysLog(fmt.Sprintf("closed %d active websocket connection(s), channels=%v, kinds=%v, reason=%s", len(entries), entryChannelIDs(entries), entryKindCounts(entries), reason))
|
||||
}
|
||||
return len(entries)
|
||||
}
|
||||
|
||||
func CloseChannelsAndBroadcast(channelIDs []int, reason string) int {
|
||||
count := CloseChannels(channelIDs, reason)
|
||||
if err := PublishCloseChannels(context.Background(), channelIDs, reason); err != nil {
|
||||
common.SysLog(fmt.Sprintf("failed to publish websocket close event: %v", err))
|
||||
}
|
||||
return count
|
||||
}
|
||||
|
||||
func StartSubscriber(ctx context.Context) {
|
||||
rdb := common.RDB
|
||||
if !common.RedisEnabled || rdb == nil {
|
||||
return
|
||||
}
|
||||
if ctx == nil {
|
||||
ctx = context.Background()
|
||||
}
|
||||
subscriberOnce.Do(func() {
|
||||
go func() {
|
||||
pubsub := rdb.Subscribe(ctx, channelCloseTopic(rdb.Options().DB))
|
||||
defer pubsub.Close()
|
||||
receiveChannelCloseEvents(ctx, pubsub.Channel(), getOriginID())
|
||||
}()
|
||||
})
|
||||
}
|
||||
|
||||
func PublishCloseChannels(ctx context.Context, channelIDs []int, reason string) error {
|
||||
rdb := common.RDB
|
||||
if !common.RedisEnabled || rdb == nil {
|
||||
return nil
|
||||
}
|
||||
ids := uniqueChannelIDs(channelIDs)
|
||||
if len(ids) == 0 {
|
||||
return nil
|
||||
}
|
||||
payload, err := common.Marshal(closeEvent{
|
||||
ChannelIDs: ids,
|
||||
Reason: normalizeReason(reason),
|
||||
Origin: getOriginID(),
|
||||
})
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
return rdb.Publish(ctx, channelCloseTopic(rdb.Options().DB), payload).Err()
|
||||
}
|
||||
|
||||
func channelCloseTopic(database int) string {
|
||||
// Redis Pub/Sub crosses logical databases, unlike the cache keys. Include
|
||||
// the selected database so independent deployments cannot close each other.
|
||||
return fmt.Sprintf("%s:db:%d", redisChannel, database)
|
||||
}
|
||||
|
||||
func receiveChannelCloseEvents(ctx context.Context, ch <-chan *redis.Message, origin string) {
|
||||
for {
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
return
|
||||
case msg, ok := <-ch:
|
||||
if !ok {
|
||||
return
|
||||
}
|
||||
var event closeEvent
|
||||
if err := common.Unmarshal([]byte(msg.Payload), &event); err != nil {
|
||||
common.SysLog(fmt.Sprintf("failed to unmarshal websocket close event: %v", err))
|
||||
continue
|
||||
}
|
||||
if event.Origin == origin {
|
||||
continue
|
||||
}
|
||||
CloseChannels(event.ChannelIDs, event.Reason)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func takeEntries(channelIDs []int) []*entry {
|
||||
ids := uniqueChannelIDs(channelIDs)
|
||||
if len(ids) == 0 {
|
||||
return nil
|
||||
}
|
||||
|
||||
mu.Lock()
|
||||
defer mu.Unlock()
|
||||
|
||||
var entries []*entry
|
||||
for _, channelID := range ids {
|
||||
for _, e := range registry[channelID] {
|
||||
entries = append(entries, e)
|
||||
}
|
||||
delete(registry, channelID)
|
||||
}
|
||||
return entries
|
||||
}
|
||||
|
||||
func entryChannelIDs(entries []*entry) []int {
|
||||
ids := make([]int, 0, len(entries))
|
||||
for _, e := range entries {
|
||||
if e != nil {
|
||||
ids = append(ids, e.channelID)
|
||||
}
|
||||
}
|
||||
return uniqueChannelIDs(ids)
|
||||
}
|
||||
|
||||
func entryKindCounts(entries []*entry) map[string]int {
|
||||
counts := make(map[string]int)
|
||||
for _, e := range entries {
|
||||
if e == nil {
|
||||
continue
|
||||
}
|
||||
counts[e.kind]++
|
||||
}
|
||||
return counts
|
||||
}
|
||||
|
||||
func uniqueChannelIDs(channelIDs []int) []int {
|
||||
seen := make(map[int]struct{}, len(channelIDs))
|
||||
ids := make([]int, 0, len(channelIDs))
|
||||
for _, id := range channelIDs {
|
||||
if id <= 0 {
|
||||
continue
|
||||
}
|
||||
if _, ok := seen[id]; ok {
|
||||
continue
|
||||
}
|
||||
seen[id] = struct{}{}
|
||||
ids = append(ids, id)
|
||||
}
|
||||
return ids
|
||||
}
|
||||
|
||||
func normalizeReason(reason string) string {
|
||||
if reason == "" {
|
||||
return defaultCloseReason
|
||||
}
|
||||
return reason
|
||||
}
|
||||
|
||||
func getOriginID() string {
|
||||
originOnce.Do(func() {
|
||||
name := common.NodeName
|
||||
if name == "" {
|
||||
name = "node"
|
||||
}
|
||||
originID = fmt.Sprintf("%s-%d-%d", name, os.Getpid(), time.Now().UnixNano())
|
||||
})
|
||||
return originID
|
||||
}
|
||||
@@ -0,0 +1,188 @@
|
||||
package wsmanager
|
||||
|
||||
import (
|
||||
"context"
|
||||
"os"
|
||||
"strings"
|
||||
"sync"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/QuantumNous/new-api/common"
|
||||
"github.com/alicebob/miniredis/v2"
|
||||
"github.com/go-redis/redis/v8"
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
func resetRegistryForTest() {
|
||||
mu.Lock()
|
||||
defer mu.Unlock()
|
||||
registry = map[int]map[uint64]*entry{}
|
||||
nextID = 0
|
||||
}
|
||||
|
||||
func TestCloseChannelClosesRegisteredConnectionsOnce(t *testing.T) {
|
||||
resetRegistryForTest()
|
||||
|
||||
var mu sync.Mutex
|
||||
calls := 0
|
||||
Register(10, KindRealtime, func(reason string) {
|
||||
mu.Lock()
|
||||
defer mu.Unlock()
|
||||
require.Equal(t, "test reason", reason, "close callback should receive the provided reason")
|
||||
calls++
|
||||
})
|
||||
Register(10, KindResponses, func(reason string) {
|
||||
mu.Lock()
|
||||
defer mu.Unlock()
|
||||
calls++
|
||||
})
|
||||
|
||||
require.Equal(t, 2, CloseChannel(10, "test reason"), "CloseChannel should close every registered connection for the channel")
|
||||
require.Equal(t, 0, CloseChannel(10, "test reason"), "CloseChannel should not close already removed connections")
|
||||
|
||||
mu.Lock()
|
||||
defer mu.Unlock()
|
||||
assert.Equal(t, 2, calls, "all registered close callbacks should be called once")
|
||||
}
|
||||
|
||||
func TestCloseChannelDoesNotCloseOtherChannels(t *testing.T) {
|
||||
resetRegistryForTest()
|
||||
|
||||
calls := map[int]int{}
|
||||
Register(10, KindRealtime, func(reason string) {
|
||||
calls[10]++
|
||||
})
|
||||
Register(20, KindRealtime, func(reason string) {
|
||||
calls[20]++
|
||||
})
|
||||
|
||||
require.Equal(t, 1, CloseChannel(10, "test"), "CloseChannel should only close the requested channel")
|
||||
assert.Equal(t, 1, calls[10], "requested channel callback should run")
|
||||
assert.Equal(t, 0, calls[20], "other channel callback should not run")
|
||||
}
|
||||
|
||||
func TestUnregisterPreventsClose(t *testing.T) {
|
||||
resetRegistryForTest()
|
||||
|
||||
calls := 0
|
||||
unregister := Register(10, KindRealtime, func(reason string) {
|
||||
calls++
|
||||
})
|
||||
unregister()
|
||||
|
||||
require.Equal(t, 0, CloseChannel(10, "test"), "unregistered connections should not be closed")
|
||||
assert.Equal(t, 0, calls, "unregistered close callback should not run")
|
||||
}
|
||||
|
||||
func TestRegisteredCloseIsIdempotent(t *testing.T) {
|
||||
resetRegistryForTest()
|
||||
|
||||
calls := 0
|
||||
Register(10, KindRealtime, func(reason string) {
|
||||
calls++
|
||||
})
|
||||
|
||||
mu.Lock()
|
||||
var registered *entry
|
||||
for _, e := range registry[10] {
|
||||
registered = e
|
||||
}
|
||||
mu.Unlock()
|
||||
require.NotNil(t, registered, "registered entry should exist")
|
||||
|
||||
registered.close("test")
|
||||
registered.close("test")
|
||||
assert.Equal(t, 1, calls, "registered close callback should be idempotent")
|
||||
}
|
||||
|
||||
func TestPublishCloseChannelsNoopsWhenRedisDisabled(t *testing.T) {
|
||||
resetRegistryForTest()
|
||||
|
||||
oldEnabled := common.RedisEnabled
|
||||
oldRDB := common.RDB
|
||||
common.RedisEnabled = false
|
||||
common.RDB = nil
|
||||
defer func() {
|
||||
common.RedisEnabled = oldEnabled
|
||||
common.RDB = oldRDB
|
||||
}()
|
||||
|
||||
require.NoError(t, PublishCloseChannels(context.Background(), []int{10}, "test"), "publishing should no-op without Redis")
|
||||
}
|
||||
|
||||
func TestRedisChannelCloseEventsStayWithinDatabase(t *testing.T) {
|
||||
resetRegistryForTest()
|
||||
previousEnabled, previousClient := common.RedisEnabled, common.RDB
|
||||
t.Cleanup(func() {
|
||||
common.RedisEnabled, common.RDB = previousEnabled, previousClient
|
||||
resetRegistryForTest()
|
||||
})
|
||||
address := os.Getenv("TEST_WS_MANAGER_REDIS_ADDR")
|
||||
if address == "" {
|
||||
address = miniredis.RunT(t).Addr()
|
||||
}
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
|
||||
t.Cleanup(cancel)
|
||||
localPublisher := redis.NewClient(&redis.Options{Addr: address, DB: 0})
|
||||
localSubscriber := redis.NewClient(&redis.Options{Addr: address, DB: 0})
|
||||
otherPublisher := redis.NewClient(&redis.Options{Addr: address, DB: 1})
|
||||
for _, client := range []*redis.Client{localPublisher, localSubscriber, otherPublisher} {
|
||||
require.NoError(t, client.Ping(ctx).Err())
|
||||
t.Cleanup(func() { require.NoError(t, client.Close()) })
|
||||
}
|
||||
if os.Getenv("TEST_WS_MANAGER_REDIS_ADDR") != "" {
|
||||
info, err := localPublisher.Info(ctx, "server").Result()
|
||||
require.NoError(t, err)
|
||||
for line := range strings.SplitSeq(info, "\r\n") {
|
||||
if strings.HasPrefix(line, "redis_version:") {
|
||||
t.Log(line)
|
||||
}
|
||||
}
|
||||
}
|
||||
pubsub := localSubscriber.Subscribe(ctx, channelCloseTopic(localSubscriber.Options().DB))
|
||||
t.Cleanup(func() { require.NoError(t, pubsub.Close()) })
|
||||
_, err := pubsub.Receive(ctx)
|
||||
require.NoError(t, err, "wait for subscription acknowledgement before publishing")
|
||||
otherPubsub := otherPublisher.Subscribe(ctx, channelCloseTopic(otherPublisher.Options().DB))
|
||||
t.Cleanup(func() { require.NoError(t, otherPubsub.Close()) })
|
||||
_, err = otherPubsub.Receive(ctx)
|
||||
require.NoError(t, err)
|
||||
|
||||
closed := make(chan int, 2)
|
||||
Register(10, KindRealtime, func(string) { closed <- 10 })
|
||||
Register(20, KindResponses, func(string) { closed <- 20 })
|
||||
done := make(chan struct{})
|
||||
go func() {
|
||||
defer close(done)
|
||||
receiveChannelCloseEvents(ctx, pubsub.Channel(), "receiving-node")
|
||||
}()
|
||||
t.Cleanup(func() {
|
||||
cancel()
|
||||
<-done
|
||||
})
|
||||
|
||||
common.RedisEnabled, common.RDB = true, otherPublisher
|
||||
require.NoError(t, PublishCloseChannels(ctx, []int{10}, "other database"))
|
||||
otherEvent, err := otherPubsub.ReceiveMessage(ctx)
|
||||
require.NoError(t, err)
|
||||
var event closeEvent
|
||||
require.NoError(t, common.Unmarshal([]byte(otherEvent.Payload), &event))
|
||||
assert.Equal(t, []int{10}, event.ChannelIDs)
|
||||
assert.Equal(t, "other database", event.Reason)
|
||||
// The old global topic must not affect a namespaced receiver either.
|
||||
require.NoError(t, localPublisher.Publish(ctx, redisChannel, otherEvent.Payload).Err())
|
||||
|
||||
common.RDB = localPublisher
|
||||
require.NoError(t, PublishCloseChannels(ctx, []int{20}, "same database"))
|
||||
select {
|
||||
case channelID := <-closed:
|
||||
assert.Equal(t, 20, channelID, "only the same-database broadcast may close a local connection")
|
||||
case <-ctx.Done():
|
||||
t.Fatal("same-database broadcast did not close its connection")
|
||||
}
|
||||
// Redis processes the preceding publications before this successful one,
|
||||
// so the retained registration proves isolation without sleeps or polling.
|
||||
assert.Equal(t, 1, CloseChannel(10, "test cleanup"))
|
||||
}
|
||||
Reference in new issue
Block a user