feat: publish v1.0.0-rc.36-custom.1
This commit is contained in:
commit
3a7e8f3b54
2381 files changed
+551494
No files matched your search
File diff suppressed because it is too large.
Load diff
@@ -0,0 +1,331 @@
|
||||
package billingexpr
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"math"
|
||||
"strings"
|
||||
"sync"
|
||||
|
||||
"github.com/expr-lang/expr"
|
||||
"github.com/expr-lang/expr/ast"
|
||||
"github.com/expr-lang/expr/vm"
|
||||
)
|
||||
|
||||
const maxCacheSize = 256
|
||||
|
||||
// DefaultExprVersion is used when an expression string has no version prefix.
|
||||
const DefaultExprVersion = 1
|
||||
|
||||
const (
|
||||
requestRuleTraceFunction = "_trace"
|
||||
requestRuleTraceIntFunction = "_trace_int"
|
||||
)
|
||||
|
||||
// ParseExprVersion extracts the version tag and body from an expression string.
|
||||
// Format: "v1:tier(...)" → version=1, body="tier(...)".
|
||||
// No prefix defaults to DefaultExprVersion.
|
||||
func ParseExprVersion(exprStr string) (version int, body string) {
|
||||
if strings.HasPrefix(exprStr, "v1:") {
|
||||
return 1, exprStr[3:]
|
||||
}
|
||||
return DefaultExprVersion, exprStr
|
||||
}
|
||||
|
||||
// requestRulePatcher adds trace side effects to existing request multipliers
|
||||
// without changing the stored expression or its numeric result.
|
||||
type requestRulePatcher struct {
|
||||
requestRules []RequestRuleTrace
|
||||
restrictedIdentifier string
|
||||
}
|
||||
|
||||
func (p *requestRulePatcher) Visit(node *ast.Node) {
|
||||
if identifier, ok := (*node).(*ast.IdentifierNode); ok {
|
||||
switch identifier.Value {
|
||||
case requestRuleTraceFunction, requestRuleTraceIntFunction:
|
||||
p.restrictedIdentifier = identifier.Value
|
||||
}
|
||||
return
|
||||
}
|
||||
|
||||
conditional, ok := (*node).(*ast.ConditionalNode)
|
||||
if !ok || !conditional.Ternary || !usesRequestProbe(conditional.Cond) {
|
||||
return
|
||||
}
|
||||
multiplier, ok := requestRuleNumber(conditional.Exp1)
|
||||
fallback, fallbackOK := requestRuleNumber(conditional.Exp2)
|
||||
if !ok || !fallbackOK || fallback != 1 {
|
||||
return
|
||||
}
|
||||
|
||||
ruleIndex := len(p.requestRules)
|
||||
p.requestRules = append(p.requestRules, RequestRuleTrace{
|
||||
Cond: conditional.Cond.String(),
|
||||
Multiplier: multiplier,
|
||||
})
|
||||
|
||||
traceFunction := requestRuleTraceFunction
|
||||
var multiplierNode ast.Node = &ast.FloatNode{Value: multiplier}
|
||||
if _, multiplierIsInt := conditional.Exp1.(*ast.IntegerNode); multiplierIsInt {
|
||||
if _, fallbackIsInt := conditional.Exp2.(*ast.IntegerNode); fallbackIsInt {
|
||||
traceFunction = requestRuleTraceIntFunction
|
||||
multiplierNode = conditional.Exp1
|
||||
}
|
||||
}
|
||||
|
||||
ast.Patch(node, &ast.CallNode{
|
||||
Callee: &ast.IdentifierNode{Value: traceFunction},
|
||||
Arguments: []ast.Node{
|
||||
&ast.IntegerNode{Value: ruleIndex},
|
||||
conditional.Cond,
|
||||
multiplierNode,
|
||||
},
|
||||
})
|
||||
}
|
||||
|
||||
func requestRuleNumber(node ast.Node) (float64, bool) {
|
||||
switch value := node.(type) {
|
||||
case *ast.IntegerNode:
|
||||
return float64(value.Value), true
|
||||
case *ast.FloatNode:
|
||||
return value.Value, true
|
||||
default:
|
||||
return 0, false
|
||||
}
|
||||
}
|
||||
|
||||
func usesRequestProbe(node ast.Node) bool {
|
||||
return ast.Find(node, func(node ast.Node) bool {
|
||||
identifier, ok := node.(*ast.IdentifierNode)
|
||||
if !ok {
|
||||
return false
|
||||
}
|
||||
switch identifier.Value {
|
||||
case "param", "header", "hour", "minute", "weekday", "month", "day":
|
||||
return true
|
||||
default:
|
||||
return false
|
||||
}
|
||||
}) != nil
|
||||
}
|
||||
|
||||
type cachedEntry struct {
|
||||
prog *vm.Program
|
||||
usedVars map[string]bool
|
||||
usedUsageKeys map[string]bool
|
||||
requestRules []RequestRuleTrace
|
||||
version int
|
||||
}
|
||||
|
||||
var (
|
||||
cacheMu sync.RWMutex
|
||||
cache = make(map[string]*cachedEntry, 64)
|
||||
)
|
||||
|
||||
// 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,
|
||||
}
|
||||
|
||||
func getCompileEnv(version int) map[string]any {
|
||||
switch version {
|
||||
default:
|
||||
return compileEnvPrototypeV1
|
||||
}
|
||||
}
|
||||
|
||||
// CompileFromCache compiles an expression string, using a cached program when
|
||||
// available. The cache is keyed by the SHA-256 hex digest of the expression.
|
||||
func CompileFromCache(exprStr string) (*vm.Program, error) {
|
||||
return compileFromCacheByHash(exprStr, ExprHashString(exprStr))
|
||||
}
|
||||
|
||||
// CompileFromCacheByHash is like CompileFromCache but accepts a pre-computed
|
||||
// hash, useful when the caller already has the BillingSnapshot.ExprHash.
|
||||
func CompileFromCacheByHash(exprStr, hash string) (*vm.Program, error) {
|
||||
return compileFromCacheByHash(exprStr, hash)
|
||||
}
|
||||
|
||||
func compileFromCacheByHash(exprStr, hash string) (*vm.Program, error) {
|
||||
entry, err := compileEntryFromCacheByHash(exprStr, hash)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return entry.prog, nil
|
||||
}
|
||||
|
||||
func compileEntryFromCacheByHash(exprStr, hash string) (*cachedEntry, error) {
|
||||
cacheMu.RLock()
|
||||
if entry, ok := cache[hash]; ok {
|
||||
cacheMu.RUnlock()
|
||||
return entry, nil
|
||||
}
|
||||
cacheMu.RUnlock()
|
||||
|
||||
version, body := ParseExprVersion(exprStr)
|
||||
patcher := &requestRulePatcher{}
|
||||
prog, err := expr.Compile(body, expr.Env(getCompileEnv(version)), expr.Patch(patcher), expr.AsFloat64())
|
||||
if patcher.restrictedIdentifier != "" {
|
||||
return nil, fmt.Errorf("expr compile error: identifier %q is reserved for internal use", patcher.restrictedIdentifier)
|
||||
}
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("expr compile error: %w", err)
|
||||
}
|
||||
|
||||
entry := &cachedEntry{
|
||||
prog: prog,
|
||||
usedVars: extractUsedVars(prog),
|
||||
usedUsageKeys: extractUsedUsageKeys(prog),
|
||||
requestRules: patcher.requestRules,
|
||||
version: version,
|
||||
}
|
||||
cacheMu.Lock()
|
||||
if len(cache) >= maxCacheSize {
|
||||
cache = make(map[string]*cachedEntry, 64)
|
||||
}
|
||||
cache[hash] = entry
|
||||
cacheMu.Unlock()
|
||||
|
||||
return entry, nil
|
||||
}
|
||||
|
||||
// ExprVersion returns the version of a cached expression. Returns DefaultExprVersion
|
||||
// if the expression hasn't been compiled yet or is empty.
|
||||
func ExprVersion(exprStr string) int {
|
||||
if exprStr == "" {
|
||||
return DefaultExprVersion
|
||||
}
|
||||
hash := ExprHashString(exprStr)
|
||||
cacheMu.RLock()
|
||||
if entry, ok := cache[hash]; ok {
|
||||
cacheMu.RUnlock()
|
||||
return entry.version
|
||||
}
|
||||
cacheMu.RUnlock()
|
||||
v, _ := ParseExprVersion(exprStr)
|
||||
return v
|
||||
}
|
||||
|
||||
func extractUsedVars(prog *vm.Program) map[string]bool {
|
||||
vars := make(map[string]bool)
|
||||
node := prog.Node()
|
||||
ast.Find(node, func(n ast.Node) bool {
|
||||
if id, ok := n.(*ast.IdentifierNode); ok {
|
||||
switch id.Value {
|
||||
case requestRuleTraceFunction, requestRuleTraceIntFunction:
|
||||
return false
|
||||
}
|
||||
vars[id.Value] = true
|
||||
}
|
||||
return false
|
||||
})
|
||||
return vars
|
||||
}
|
||||
|
||||
func extractUsedUsageKeys(prog *vm.Program) map[string]bool {
|
||||
keys := make(map[string]bool)
|
||||
ast.Find(prog.Node(), func(node ast.Node) bool {
|
||||
call, ok := node.(*ast.CallNode)
|
||||
if !ok || len(call.Arguments) != 1 {
|
||||
return false
|
||||
}
|
||||
callee, ok := call.Callee.(*ast.IdentifierNode)
|
||||
if !ok || callee.Value != "u" {
|
||||
return false
|
||||
}
|
||||
literal, ok := call.Arguments[0].(*ast.StringNode)
|
||||
if !ok {
|
||||
return false
|
||||
}
|
||||
keys[strings.TrimSpace(literal.Value)] = true
|
||||
return false
|
||||
})
|
||||
return keys
|
||||
}
|
||||
|
||||
// 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 {
|
||||
if exprStr == "" {
|
||||
return nil
|
||||
}
|
||||
hash := ExprHashString(exprStr)
|
||||
cacheMu.RLock()
|
||||
if entry, ok := cache[hash]; ok {
|
||||
cacheMu.RUnlock()
|
||||
return entry.usedVars
|
||||
}
|
||||
cacheMu.RUnlock()
|
||||
|
||||
// Compile (and cache) to populate usedVars
|
||||
if _, err := compileFromCacheByHash(exprStr, hash); err != nil {
|
||||
return nil
|
||||
}
|
||||
cacheMu.RLock()
|
||||
entry, ok := cache[hash]
|
||||
cacheMu.RUnlock()
|
||||
if ok {
|
||||
return entry.usedVars
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// UsedUsageKeys returns literal keys referenced by u("...") calls. Calls with
|
||||
// dynamic arguments are intentionally omitted because they cannot be
|
||||
// validated statically.
|
||||
func UsedUsageKeys(exprStr string) map[string]bool {
|
||||
if exprStr == "" {
|
||||
return nil
|
||||
}
|
||||
hash := ExprHashString(exprStr)
|
||||
cacheMu.RLock()
|
||||
if entry, ok := cache[hash]; ok {
|
||||
cacheMu.RUnlock()
|
||||
return entry.usedUsageKeys
|
||||
}
|
||||
cacheMu.RUnlock()
|
||||
|
||||
if _, err := compileFromCacheByHash(exprStr, hash); err != nil {
|
||||
return nil
|
||||
}
|
||||
cacheMu.RLock()
|
||||
entry, ok := cache[hash]
|
||||
cacheMu.RUnlock()
|
||||
if ok {
|
||||
return entry.usedUsageKeys
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// InvalidateCache clears the compiled-expression cache.
|
||||
// Called when billing rules are updated.
|
||||
func InvalidateCache() {
|
||||
cacheMu.Lock()
|
||||
cache = make(map[string]*cachedEntry, 64)
|
||||
cacheMu.Unlock()
|
||||
}
|
||||
@@ -0,0 +1,21 @@
|
||||
package billingexpr_test
|
||||
|
||||
import (
|
||||
"testing"
|
||||
|
||||
"github.com/QuantumNous/new-api/pkg/billingexpr"
|
||||
"github.com/stretchr/testify/assert"
|
||||
)
|
||||
|
||||
func TestUsedUsageKeysExtractsLiteralCallsAndSkipsDynamicArguments(t *testing.T) {
|
||||
expression := `tier("base", u(" seconds ") * 0.4 + (u("clips") > 0 ? u("clips") * 0.1 : 0)) + (u(header("usage-key")) == nil ? 0 : 0)`
|
||||
|
||||
keys := billingexpr.UsedUsageKeys(expression)
|
||||
|
||||
assert.Equal(t, map[string]bool{"seconds": true, "clips": true}, keys)
|
||||
}
|
||||
|
||||
func TestUsedUsageKeysReturnsNilForEmptyOrInvalidExpressions(t *testing.T) {
|
||||
assert.Nil(t, billingexpr.UsedUsageKeys(""))
|
||||
assert.Nil(t, billingexpr.UsedUsageKeys(`tier("broken",`))
|
||||
}
|
||||
@@ -0,0 +1,353 @@
|
||||
# Billing Expression System (billingexpr)
|
||||
|
||||
## Design Philosophy
|
||||
|
||||
**One expression, one truth.** A single expression string completely defines a model's billing logic — pricing, tier conditions, cache/image/audio differentiation, time-based discounts, request-aware multipliers — all in one line. No scattered configuration, no implicit rules, no magic numbers.
|
||||
|
||||
The expression is the billing contract between the administrator and the system. What you write is what gets executed. The system's job is to evaluate it faithfully, not to interpret it.
|
||||
|
||||
### Core Principles
|
||||
|
||||
1. **Expression is self-contained** — The expression string alone determines billing. No external ratio tables, no implicit completion multipliers, no hidden conversion factors. Given the same token counts and request context, the same expression always produces the same cost.
|
||||
|
||||
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.
|
||||
|
||||
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.
|
||||
|
||||
5. **Version-aware** — Expressions carry a version tag (`v1:`, default when omitted). The version controls the compile environment, token normalization, and quota conversion formula, enabling future evolution without breaking existing expressions.
|
||||
|
||||
---
|
||||
|
||||
## Expression Language
|
||||
|
||||
Powered by [expr-lang/expr](https://github.com/expr-lang/expr). Expressions are compiled, cached, and evaluated against a runtime environment.
|
||||
|
||||
### Token Variables
|
||||
|
||||
**输入侧变量:**
|
||||
|
||||
| 变量 | 含义 |
|
||||
|------|------|
|
||||
| `p` | 输入 token 数(**计价用**)。**自动排除**表达式中单独计价的子类别(见下方说明) |
|
||||
| `len` | 输入上下文总长度(**条件判断用**)。不受自动排除影响,始终反映完整输入长度。非 Claude:等于原始 `prompt_tokens`;Claude:等于文本输入 + 缓存读取 + 缓存创建 |
|
||||
| `cr` | 缓存命中(读取)token 数 |
|
||||
| `cc` | 缓存创建 token 数(Claude 5分钟 TTL / 通用) |
|
||||
| `cc1h` | 缓存创建 token 数 — 1小时 TTL(Claude 专用) |
|
||||
| `img` | 图片输入 token 数 |
|
||||
| `ai` | 音频输入 token 数 |
|
||||
|
||||
**输出侧变量:**
|
||||
|
||||
| 变量 | 含义 |
|
||||
|------|------|
|
||||
| `c` | 输出 token 数。**自动排除**表达式中单独计价的子类别(见下方说明) |
|
||||
| `img_o` | 图片输出 token 数 |
|
||||
| `ao` | 音频输出 token 数 |
|
||||
|
||||
#### `p` 和 `c` 的自动排除机制
|
||||
|
||||
`p` 和 `c` 是"兜底变量"——它们代表**所有没有被表达式单独定价的 token**。系统会根据表达式实际使用了哪些变量,自动从 `p` / `c` 中减去对应的子类别 token,避免重复计费。
|
||||
|
||||
**规则:如果表达式使用了某个子类别变量,对应的 token 就从 `p` 或 `c` 中扣除;如果没使用,那些 token 就留在 `p` 或 `c` 里按基础价格计费。**
|
||||
|
||||
> **重要:`len` 不受自动排除影响。** `len` 始终代表完整的输入上下文长度,不管表达式是否单独对缓存/图片/音频定价。因此**阶梯条件应使用 `len` 而非 `p`**,以避免缓存命中导致 `p` 降低而误判档位。
|
||||
|
||||
举例说明(假设上游返回的原始数据:prompt_tokens=1000,其中包含 200 cache read、100 image):
|
||||
|
||||
| 表达式 | `p` 的值 | 说明 |
|
||||
|--------|---------|------|
|
||||
| `p * 3 + c * 15` | 1000 | 没用 `cr`/`img`,所以缓存和图片都包含在 `p` 里,全按 $3 计费 |
|
||||
| `p * 3 + c * 15 + cr * 0.3` | 800 | 用了 `cr`,缓存 200 从 `p` 中扣除,按 $0.3 单独计费;图片仍在 `p` 里按 $3 计费 |
|
||||
| `p * 3 + c * 15 + cr * 0.3 + img * 2` | 700 | 用了 `cr` 和 `img`,都从 `p` 中扣除,各自按自己的价格计费 |
|
||||
|
||||
输出侧同理(假设 completion_tokens=500,其中包含 100 audio output):
|
||||
|
||||
| 表达式 | `c` 的值 | 说明 |
|
||||
|--------|---------|------|
|
||||
| `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 本身就只包含纯文本)不做任何减法。系统根据上游返回格式自动判断,表达式作者无需关心。
|
||||
|
||||
### Built-in Functions
|
||||
|
||||
| Function | Signature | Purpose |
|
||||
|----------|-----------|---------|
|
||||
| `tier` | `tier(name, value) → float64` | Records which pricing tier matched; must wrap the cost expression |
|
||||
| `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 |
|
||||
| `hour` | `hour(tz) → int` | Current hour in timezone (0-23) |
|
||||
| `minute` | `minute(tz) → int` | Current minute (0-59) |
|
||||
| `weekday` | `weekday(tz) → int` | Day of week (0=Sunday, 6=Saturday) |
|
||||
| `month` | `month(tz) → int` | Month (1-12) |
|
||||
| `day` | `day(tz) → int` | Day of month (1-31) |
|
||||
| `max` | `max(a, b) → float64` | Math max |
|
||||
| `min` | `min(a, b) → float64` | Math min |
|
||||
| `abs` | `abs(x) → float64` | Absolute value |
|
||||
| `ceil` | `ceil(x) → float64` | Ceiling |
|
||||
| `floor` | `floor(x) → float64` | Floor |
|
||||
|
||||
### Expression Examples
|
||||
|
||||
```
|
||||
# Simple flat pricing
|
||||
tier("base", p * 2.5 + c * 15 + cr * 0.25)
|
||||
|
||||
# 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)
|
||||
: tier("long_context", p * 6 + c * 22.5 + cr * 0.6 + cc * 7.5 + cc1h * 12)
|
||||
|
||||
# Image model (no separate cache/audio pricing — those tokens stay in p/c)
|
||||
tier("base", p * 2 + c * 8 + img * 2.5)
|
||||
|
||||
# Multimodal with audio
|
||||
tier("base", p * 0.43 + c * 3.06 + img * 0.78 + ai * 3.81 + ao * 15.11)
|
||||
```
|
||||
|
||||
### Request Rules (appended after `|||`)
|
||||
|
||||
Request-conditional multipliers are appended to the expression after a `|||` separator:
|
||||
|
||||
```
|
||||
tier("base", p * 5 + c * 25)|||when(header("anthropic-beta") has "fast-mode") * 6
|
||||
```
|
||||
|
||||
These factors are stored as ordinary multiplication in the final expression (for example, `(tier(...)) * (condition ? 6 : 1)`) and run in the same billing program.
|
||||
|
||||
### Request Rule Tracing
|
||||
|
||||
At compile time, the engine instruments ternary factors with this exact shape:
|
||||
|
||||
```
|
||||
<request-probe condition> ? <numeric literal> : 1
|
||||
```
|
||||
|
||||
The condition must reference at least one request probe (`param`, `header`, `hour`, `minute`, `weekday`, `month`, or `day`). Both branches must be numeric literals and the fallback must equal `1`. Other conditionals, including `(condition ? 2 : 1.5)`, are evaluated normally but are not traced. Integer-only factors use an integer-preserving trace callback, so instrumentation does not change expressions that require an integer operand (for example, `%`). The internal trace callback names are reserved and cannot be used in stored expressions.
|
||||
|
||||
The compiled cache stores the canonical condition and multiplier for every instrumented node. Each run starts with the full detected rule list marked as unmatched; callbacks mark rules that actually evaluate true. Rules skipped by normal expression short-circuiting remain unmatched. This keeps the expression's numeric result unchanged and avoids reparsing it on each request.
|
||||
|
||||
Settlement copies the actual run's traces into the consume log as:
|
||||
|
||||
```json
|
||||
{
|
||||
"request_rules": [
|
||||
{ "cond": "param(\"service_tier\") == \"fast\"", "multiplier": 2, "matched": true }
|
||||
]
|
||||
}
|
||||
```
|
||||
|
||||
The usage-log UI treats `request_rules` as the authoritative rule list and renders directly from it. It parses `cond` only to produce a friendly label and falls back to the canonical condition text when that parser does not recognize the condition. Pricing pages without log context continue to parse the stored expression for display.
|
||||
|
||||
---
|
||||
|
||||
## Task Usage Expressions
|
||||
|
||||
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.
|
||||
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
|
||||
affects evaluation. `token` is the host unit for upstream billing tokens (for
|
||||
example doubao `usage.completion_tokens`). `credit` is the host unit for vendor
|
||||
resource-pack units (for example kling `final_unit_deduction`). Both share the
|
||||
int32 quota bound, not the 3600-second / 128-count limits.
|
||||
|
||||
Task usage billing has a deliberately different conversion rule from token
|
||||
billing:
|
||||
|
||||
```
|
||||
task quota = expression output in USD * QuotaPerUnit * groupRatio
|
||||
token quota = expression output in $/1M tokens / 1,000,000 * QuotaPerUnit * groupRatio
|
||||
```
|
||||
|
||||
In other words, a task expression already returns the request's dollar cost.
|
||||
For example, `u("seconds") * 0.4` means $0.40 per second. Engine semantics do
|
||||
not divide task output by one million.
|
||||
|
||||
The visual editor generates, and public pricing displays recognize, exactly
|
||||
these canonical task shapes. Expressions outside these shapes remain valid in
|
||||
raw mode but fall back to the special-expression display:
|
||||
|
||||
```
|
||||
# Flat unit pricing
|
||||
tier("base", u("seconds") * 0.4)
|
||||
|
||||
# Enum tiers (conditions may combine enum comparisons with &&)
|
||||
u("mode") == "pro"
|
||||
? tier("pro", u("seconds") * 0.8)
|
||||
: tier("std", u("seconds") * 0.4)
|
||||
|
||||
# Optional constant plus multiple numeric usage terms
|
||||
tier("base", 0.1 + u("seconds") * 0.4 + u("clips") * 0.05)
|
||||
|
||||
# Upstream token overlay (doubao Seedance tokens)
|
||||
# The editor takes a $/1M token input and emits the / 1000000 literal.
|
||||
# Engine semantics are unchanged: the expression still returns USD.
|
||||
tier("base", u("tokens") * 9.8 / 1000000)
|
||||
|
||||
# Vendor credit overlay (kling resource-pack units)
|
||||
# The coefficient is the real $/credit price; no /1M scale.
|
||||
tier("base", u("units") * 0.14)
|
||||
```
|
||||
|
||||
The tier body is an optional non-negative constant plus one or more
|
||||
`u("<number field>") * <unit price>` terms. Token-unit fields use the
|
||||
canonical scaled shape `u("<field>") * <dollars per 1M tokens> / 1000000`.
|
||||
Credit, second, and count fields keep the bare `u("<field>") * <unit price>`
|
||||
shape. Tier conditions are equality checks
|
||||
between declared enum fields and values, optionally joined by `&&`, with
|
||||
chained ternaries following the same ordering rules as token tiers. Numeric
|
||||
range tiers are not part of the current canonical shape. Request rules after
|
||||
`|||` remain orthogonal and use the same syntax as token expressions.
|
||||
|
||||
Before saving, the host compiles every expression, rejects literal `u()` keys
|
||||
that the selected task plugin did not declare, and smoke-tests usage vectors.
|
||||
Every numeric field is exercised at 0, 1, and its host-owned unit ceiling
|
||||
(`second` 3600, `count` 128, `token`/`credit` int32 max). Enum values are exercised as a
|
||||
Cartesian product. Smoke vectors are capped at
|
||||
64, reducing oversized enum dimensions to their first and last values. Every
|
||||
evaluated result must be finite and non-negative.
|
||||
|
||||
Submission evaluates the expression with request-derived usage facts and
|
||||
freezes both the expression and those facts in the billing snapshot. On
|
||||
completion, the host overlays completion facts on the frozen submission facts
|
||||
key by key, so measured values replace estimates while facts omitted by the
|
||||
completion hook retain their submission values. The same expression is then
|
||||
evaluated again; a changed fact can therefore produce a settlement delta and a
|
||||
different matched tier. Evaluation failure keeps the pre-consumed charge.
|
||||
|
||||
---
|
||||
|
||||
## Architecture
|
||||
|
||||
### Data Flow
|
||||
|
||||
```
|
||||
Frontend Editor → Storage → Pre-consume → Settlement → Log Display
|
||||
```
|
||||
|
||||
### 1. Frontend Editor
|
||||
|
||||
**File**: `web/src/pages/Setting/Ratio/components/TieredPricingEditor.jsx`
|
||||
|
||||
Two editing modes:
|
||||
- **Visual mode**: Fill in prices per variable, conditions per tier. Generates expression via `generateExprFromVisualConfig()`.
|
||||
- **Raw mode**: Edit the expression string directly. Includes preset templates for common models.
|
||||
|
||||
The editor outputs a billing expression string and an optional request rule expression string. These are combined via `combineBillingExpr(billingExpr, requestRuleExpr)` before storage.
|
||||
|
||||
### 2. Storage
|
||||
|
||||
**File**: `setting/billing_setting/tiered_billing.go`
|
||||
|
||||
Two option maps stored in the `options` DB table:
|
||||
- `ModelBillingMode`: `{ "model-name": "tiered_expr" }` — activates tiered billing for a model
|
||||
- `ModelBillingExpr`: `{ "model-name": "tier(\"base\", p * 2.5 + c * 15)" }` — the expression
|
||||
|
||||
On save, the expression is validated:
|
||||
1. Compiled via `billingexpr.CompileFromCache()` — syntax check
|
||||
2. Smoke-tested with sample token vectors — ensures non-negative results
|
||||
|
||||
### 3. Pre-consume (Quota Estimation)
|
||||
|
||||
**File**: `relay/helper/price.go` → `modelPriceHelperTiered()`
|
||||
|
||||
When a request arrives and the model uses `tiered_expr` billing:
|
||||
1. Loads expression from `billing_setting.GetBillingExpr()`
|
||||
2. Builds `RequestInput` (headers + body) for `param()` / `header()` functions
|
||||
3. Runs expression with estimated tokens: `RunExprWithRequest(expr, {P, C}, requestInput)`
|
||||
4. Converts output to quota: `rawCost / 1,000,000 * QuotaPerUnit`
|
||||
5. Creates `BillingSnapshot` and stores it on `RelayInfo`. Expression and request state stay frozen for settlement. An auto-group retry refreshes group-dependent fields from the selected group before the next upstream attempt. If a free initial group skipped pre-consume and the retry selects a paid group, the billing session is created before that attempt. If an existing session moves to a more expensive group, its reservation is raised to that group's estimate before sending; cheaper groups are refunded only after actual usage is settled.
|
||||
|
||||
### 4. Settlement (Actual Billing)
|
||||
|
||||
**Files**: `service/tiered_settle.go`, `pkg/billingexpr/settle.go`
|
||||
|
||||
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
|
||||
|
||||
2. `TryTieredSettle(relayInfo, params)`:
|
||||
- Uses the captured `BillingSnapshot`, whose group-dependent fields have been refreshed from the final selected group
|
||||
- Re-runs the expression with actual token counts
|
||||
- Converts via `quotaConversion()` (version-dispatched)
|
||||
- Returns actual quota
|
||||
|
||||
### 5. Log Display
|
||||
|
||||
**Files**: `service/log_info_generate.go`, `web/src/helpers/render.jsx`
|
||||
|
||||
Backend: `InjectTieredBillingInfo()` adds `billing_mode`, `expr_b64` (base64 expression), `matched_tier`, and the structured `request_rules` trace list to the log's `other` JSON.
|
||||
|
||||
Frontend: Detects `billing_mode === "tiered_expr"`, decodes `expr_b64`, parses tiers via shared `parseTiersFromExpr()`, and renders request multipliers from `request_rules` when present. Without log traces, it falls back to parsing the stored expression.
|
||||
|
||||
---
|
||||
|
||||
## Key Design Decisions
|
||||
|
||||
### Token Normalization via AST Introspection
|
||||
|
||||
Different upstream APIs report `prompt_tokens` differently:
|
||||
- **OpenAI/GPT**: `prompt_tokens` = total (text + cache + image + audio)
|
||||
- **Claude**: `input_tokens` = text only (cache reported separately)
|
||||
|
||||
The system normalizes `p` to mean "tokens not separately priced" by subtracting sub-categories **only when the expression references them**. This is determined by walking the compiled AST to find `IdentifierNode` references — zero runtime cost after first compilation (cached).
|
||||
|
||||
Example: `p * 2.5 + c * 15 + cr * 0.25`
|
||||
- Expression uses `cr` → cache read tokens subtracted from `p`
|
||||
- Expression doesn't use `img` → image tokens stay in `p`, priced at $2.50
|
||||
|
||||
### `len` — Context Length Variable
|
||||
|
||||
`len` represents the total input context length, designed for **tier condition evaluation** (e.g. `len <= 200000 ? ...`). Unlike `p`, `len` is never reduced by sub-category exclusion.
|
||||
|
||||
**Computation rules:**
|
||||
- **Non-Claude (GPT/OpenAI format)**: `len = prompt_tokens` (the raw total from the upstream response)
|
||||
- **Claude format**: `len = input_tokens + cache_read_tokens + cache_creation_tokens` (since Claude's `input_tokens` is text-only, cache must be added back to reflect full context length)
|
||||
|
||||
This ensures that heavy cache usage doesn't cause the tier condition to incorrectly evaluate to a lower tier. For example, if a request has 300K total context but 250K is cached, `p` with cache subtracted would be only 50K (standard tier), while `len` correctly reports 300K (long-context tier).
|
||||
|
||||
### Quota Conversion
|
||||
|
||||
Expression coefficients are $/1M tokens. Conversion to internal quota:
|
||||
|
||||
```
|
||||
quota = exprOutput / 1,000,000 * QuotaPerUnit * groupRatio
|
||||
```
|
||||
|
||||
This matches the per-call billing pattern: `quota = modelPrice * QuotaPerUnit * groupRatio`.
|
||||
|
||||
### Expression Versioning
|
||||
|
||||
Expressions can carry a version prefix: `v1:tier(...)`. No prefix = v1.
|
||||
|
||||
Version controls:
|
||||
- Compile environment (available variables and functions)
|
||||
- Token normalization logic
|
||||
- Quota conversion formula
|
||||
|
||||
This enables future evolution without breaking existing expressions.
|
||||
|
||||
---
|
||||
|
||||
## File Map
|
||||
|
||||
| Layer | Files |
|
||||
|-------|-------|
|
||||
| Expression engine | `pkg/billingexpr/compile.go`, `run.go`, `settle.go`, `round.go`, `types.go` |
|
||||
| Storage | `setting/billing_setting/tiered_billing.go` |
|
||||
| Pre-consume | `relay/helper/price.go`, `relay/helper/billing_expr_request.go` |
|
||||
| Settlement | `service/tiered_settle.go`, `service/quota.go` |
|
||||
| Log injection | `service/log_info_generate.go` |
|
||||
| Frontend editor | `web/src/pages/Setting/Ratio/components/TieredPricingEditor.jsx` |
|
||||
| Frontend display | `web/src/helpers/render.jsx`, `web/src/helpers/utils.jsx` |
|
||||
| Model detail | `web/src/components/table/model-pricing/modal/components/DynamicPricingBreakdown.jsx` |
|
||||
| Log display | `web/src/hooks/usage-logs/useUsageLogsData.jsx`, `web/src/components/table/usage-logs/UsageLogsColumnDefs.jsx` |
|
||||
@@ -0,0 +1,19 @@
|
||||
package billingexpr
|
||||
|
||||
import "github.com/QuantumNous/new-api/common"
|
||||
|
||||
// QuotaRound converts a float64 quota value to int using half-away-from-zero
|
||||
// rounding with int32 saturation. Every tiered billing path (pre-consume,
|
||||
// settlement, breakdown validation, log fields) MUST use this function to
|
||||
// avoid +-1 discrepancies.
|
||||
//
|
||||
// It delegates to common.QuotaRound so all quota rounding/conversion shares
|
||||
// one saturation + logging policy (see common/quota_math.go).
|
||||
func QuotaRound(f float64) int {
|
||||
return common.QuotaRound(f)
|
||||
}
|
||||
|
||||
// QuotaRoundStrict rejects an unrepresentable pre-consume estimate.
|
||||
func QuotaRoundStrict(f float64) (int, error) {
|
||||
return common.QuotaRoundStrict(f)
|
||||
}
|
||||
@@ -0,0 +1,166 @@
|
||||
package billingexpr
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"math"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"github.com/expr-lang/expr"
|
||||
"github.com/expr-lang/expr/vm"
|
||||
"github.com/tidwall/gjson"
|
||||
)
|
||||
|
||||
// RunExpr compiles (with cache) and executes an expression string.
|
||||
// The environment exposes:
|
||||
// - p, c — prompt / completion tokens (auto-excluding separately-priced sub-categories)
|
||||
// - len — total input context length for tier conditions (never reduced by sub-category exclusion)
|
||||
// - cr, cc, cc1h — cache read / creation / creation-1h tokens
|
||||
// - tier(name, value) — trace callback that records which tier matched
|
||||
// - max, min, abs, ceil, floor — standard math helpers
|
||||
//
|
||||
// Returns the resulting float64 quota (before group ratio) and a TraceResult
|
||||
// with side-channel info captured by tier() during execution.
|
||||
func RunExpr(exprStr string, params TokenParams) (float64, TraceResult, error) {
|
||||
return RunExprWithRequest(exprStr, params, RequestInput{})
|
||||
}
|
||||
|
||||
func RunExprWithRequest(exprStr string, params TokenParams, request RequestInput) (float64, TraceResult, error) {
|
||||
entry, err := compileEntryFromCacheByHash(exprStr, ExprHashString(exprStr))
|
||||
if err != nil {
|
||||
return 0, TraceResult{}, err
|
||||
}
|
||||
return runProgram(entry.prog, entry.requestRules, params, request)
|
||||
}
|
||||
|
||||
// RunExprByHash is like RunExpr but accepts a pre-computed hash for the cache
|
||||
// lookup, avoiding a redundant SHA-256 computation when the caller already
|
||||
// holds BillingSnapshot.ExprHash.
|
||||
func RunExprByHash(exprStr, hash string, params TokenParams) (float64, TraceResult, error) {
|
||||
return RunExprByHashWithRequest(exprStr, hash, params, RequestInput{})
|
||||
}
|
||||
|
||||
func RunExprByHashWithRequest(exprStr, hash string, params TokenParams, request RequestInput) (float64, TraceResult, error) {
|
||||
entry, err := compileEntryFromCacheByHash(exprStr, hash)
|
||||
if err != nil {
|
||||
return 0, TraceResult{}, err
|
||||
}
|
||||
return runProgram(entry.prog, entry.requestRules, params, request)
|
||||
}
|
||||
|
||||
func runProgram(prog *vm.Program, requestRules []RequestRuleTrace, params TokenParams, request RequestInput) (float64, TraceResult, error) {
|
||||
trace := TraceResult{
|
||||
RequestRules: append([]RequestRuleTrace(nil), requestRules...),
|
||||
}
|
||||
headers := normalizeHeaders(request.Headers)
|
||||
|
||||
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,
|
||||
"tier": func(name string, value float64) float64 {
|
||||
trace.MatchedTier = name
|
||||
trace.Cost = value
|
||||
return value
|
||||
},
|
||||
requestRuleTraceFunction: func(ruleIndex int, matched bool, multiplier float64) float64 {
|
||||
if matched && ruleIndex >= 0 && ruleIndex < len(trace.RequestRules) {
|
||||
trace.RequestRules[ruleIndex].Matched = true
|
||||
}
|
||||
if matched {
|
||||
return multiplier
|
||||
}
|
||||
return 1
|
||||
},
|
||||
requestRuleTraceIntFunction: func(ruleIndex int, matched bool, multiplier int) int {
|
||||
if matched && ruleIndex >= 0 && ruleIndex < len(trace.RequestRules) {
|
||||
trace.RequestRules[ruleIndex].Matched = true
|
||||
}
|
||||
if matched {
|
||||
return multiplier
|
||||
}
|
||||
return 1
|
||||
},
|
||||
"header": func(key string) string {
|
||||
return headers[strings.ToLower(strings.TrimSpace(key))]
|
||||
},
|
||||
"param": func(path string) any {
|
||||
path = strings.TrimSpace(path)
|
||||
if path == "" || len(request.Body) == 0 {
|
||||
return nil
|
||||
}
|
||||
result := gjson.GetBytes(request.Body, path)
|
||||
if !result.Exists() {
|
||||
return nil
|
||||
}
|
||||
return result.Value()
|
||||
},
|
||||
"u": func(name string) any {
|
||||
if request.Usage == nil {
|
||||
return nil
|
||||
}
|
||||
return request.Usage[strings.TrimSpace(name)]
|
||||
},
|
||||
"has": func(source any, substr string) bool {
|
||||
if source == nil || substr == "" {
|
||||
return false
|
||||
}
|
||||
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() },
|
||||
"max": math.Max,
|
||||
"min": math.Min,
|
||||
"abs": math.Abs,
|
||||
"ceil": math.Ceil,
|
||||
"floor": math.Floor,
|
||||
}
|
||||
|
||||
out, err := expr.Run(prog, env)
|
||||
if err != nil {
|
||||
return 0, trace, fmt.Errorf("expr run error: %w", err)
|
||||
}
|
||||
f, ok := out.(float64)
|
||||
if !ok {
|
||||
return 0, trace, fmt.Errorf("expr result is %T, want float64", out)
|
||||
}
|
||||
return f, trace, nil
|
||||
}
|
||||
|
||||
func timeInZone(tz string) time.Time {
|
||||
tz = strings.TrimSpace(tz)
|
||||
if tz == "" {
|
||||
return time.Now().UTC()
|
||||
}
|
||||
loc, err := time.LoadLocation(tz)
|
||||
if err != nil {
|
||||
return time.Now().UTC()
|
||||
}
|
||||
return time.Now().In(loc)
|
||||
}
|
||||
|
||||
func normalizeHeaders(headers map[string]string) map[string]string {
|
||||
if len(headers) == 0 {
|
||||
return map[string]string{}
|
||||
}
|
||||
normalized := make(map[string]string, len(headers))
|
||||
for key, value := range headers {
|
||||
k := strings.ToLower(strings.TrimSpace(key))
|
||||
v := strings.TrimSpace(value)
|
||||
if k == "" || v == "" {
|
||||
continue
|
||||
}
|
||||
normalized[k] = v
|
||||
}
|
||||
return normalized
|
||||
}
|
||||
@@ -0,0 +1,42 @@
|
||||
package billingexpr
|
||||
|
||||
import "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
|
||||
// that may use a different conversion formula.
|
||||
func quotaConversion(exprOutput float64, snap *BillingSnapshot) float64 {
|
||||
if snap.TaskUsageBilling {
|
||||
return exprOutput * snap.QuotaPerUnit
|
||||
}
|
||||
switch snap.ExprVersion {
|
||||
default: // v1: coefficients are $/1M tokens prices
|
||||
return exprOutput / 1_000_000 * snap.QuotaPerUnit
|
||||
}
|
||||
}
|
||||
|
||||
// ComputeTieredQuota runs the Expr from a frozen BillingSnapshot against
|
||||
// actual token counts and returns the settlement result.
|
||||
func ComputeTieredQuota(snap *BillingSnapshot, params TokenParams) (TieredResult, error) {
|
||||
return ComputeTieredQuotaWithRequest(snap, params, RequestInput{})
|
||||
}
|
||||
|
||||
func ComputeTieredQuotaWithRequest(snap *BillingSnapshot, params TokenParams, request RequestInput) (TieredResult, error) {
|
||||
cost, trace, err := RunExprByHashWithRequest(snap.ExprString, snap.ExprHash, params, request)
|
||||
if err != nil {
|
||||
return TieredResult{}, err
|
||||
}
|
||||
|
||||
quotaBeforeGroup := quotaConversion(cost, snap)
|
||||
afterGroup, clamp := common.QuotaRoundChecked(quotaBeforeGroup * snap.GroupRatio)
|
||||
crossed := trace.MatchedTier != snap.EstimatedTier
|
||||
|
||||
return TieredResult{
|
||||
ActualQuotaBeforeGroup: quotaBeforeGroup,
|
||||
ActualQuotaAfterGroup: afterGroup,
|
||||
MatchedTier: trace.MatchedTier,
|
||||
RequestRules: trace.RequestRules,
|
||||
CrossedTier: crossed,
|
||||
Clamp: clamp,
|
||||
}, nil
|
||||
}
|
||||
@@ -0,0 +1,52 @@
|
||||
package billingexpr_test
|
||||
|
||||
import (
|
||||
"testing"
|
||||
|
||||
"github.com/QuantumNous/new-api/common"
|
||||
"github.com/QuantumNous/new-api/pkg/billingexpr"
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
// TestComputeTieredQuota_ClampOnOverflow guards the billing-safety invariant
|
||||
// that an oversized tiered settlement clamps to the single-request max instead of
|
||||
// wrapping into a credit, and that the saturation event is surfaced on the
|
||||
// result so callers can record it for admin auditing.
|
||||
func TestComputeTieredQuota_ClampOnOverflow(t *testing.T) {
|
||||
// exprOutput = p * 1e12 = 1e21; quotaBeforeGroup = 1e21 / 1e6 * 5e5 = 5e20,
|
||||
// which far exceeds the supported single-request range and must saturate.
|
||||
exprStr := `tier("base", p * 1000000000000)`
|
||||
snap := &billingexpr.BillingSnapshot{
|
||||
BillingMode: "tiered_expr",
|
||||
ExprString: exprStr,
|
||||
ExprHash: billingexpr.ExprHashString(exprStr),
|
||||
GroupRatio: 1.0,
|
||||
QuotaPerUnit: 500_000,
|
||||
}
|
||||
|
||||
result, err := billingexpr.ComputeTieredQuota(snap, billingexpr.TokenParams{P: 1_000_000_000})
|
||||
require.NoError(t, err)
|
||||
|
||||
assert.Equal(t, common.MaxQuota, result.ActualQuotaAfterGroup, "oversized quota must clamp, never wrap negative")
|
||||
require.NotNil(t, result.Clamp, "clamp event must be surfaced so it can be audited")
|
||||
assert.Equal(t, common.QuotaClampOverflow, result.Clamp.Kind)
|
||||
assert.Equal(t, common.MaxQuota, result.Clamp.Clamped)
|
||||
}
|
||||
|
||||
// TestComputeTieredQuota_NoClampInRange confirms an in-range settlement leaves
|
||||
// Clamp nil, so the audit path is a no-op in the common case.
|
||||
func TestComputeTieredQuota_NoClampInRange(t *testing.T) {
|
||||
exprStr := `tier("base", p * 2 + c * 10)`
|
||||
snap := &billingexpr.BillingSnapshot{
|
||||
BillingMode: "tiered_expr",
|
||||
ExprString: exprStr,
|
||||
ExprHash: billingexpr.ExprHashString(exprStr),
|
||||
GroupRatio: 1.0,
|
||||
QuotaPerUnit: 500_000,
|
||||
}
|
||||
|
||||
result, err := billingexpr.ComputeTieredQuota(snap, billingexpr.TokenParams{P: 1000, C: 500})
|
||||
require.NoError(t, err)
|
||||
assert.Nil(t, result.Clamp, "in-range settlement must not report a clamp")
|
||||
}
|
||||
@@ -0,0 +1,19 @@
|
||||
package billingexpr
|
||||
|
||||
import (
|
||||
"testing"
|
||||
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
func TestTaskUsageExpressionUsesFactsAndTaskQuotaConversion(t *testing.T) {
|
||||
expression := `tier("1080p", u("seconds") * (u("resolution") == "1080p" ? 0.4 : 0.2))`
|
||||
cost, trace, err := RunExprWithRequest(expression, TokenParams{}, RequestInput{Usage: map[string]any{"seconds": 10.0, "resolution": "1080p"}})
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, 4.0, cost)
|
||||
assert.Equal(t, "1080p", trace.MatchedTier)
|
||||
result, err := ComputeTieredQuotaWithRequest(&BillingSnapshot{ExprString: expression, ExprHash: ExprHashString(expression), GroupRatio: 2, QuotaPerUnit: 500000, ExprVersion: 1, TaskUsageBilling: true}, TokenParams{}, RequestInput{Usage: map[string]any{"seconds": 10.0, "resolution": "1080p"}})
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, 4_000_000, result.ActualQuotaAfterGroup)
|
||||
}
|
||||
@@ -0,0 +1,85 @@
|
||||
package billingexpr
|
||||
|
||||
import (
|
||||
"crypto/sha256"
|
||||
"fmt"
|
||||
|
||||
"github.com/QuantumNous/new-api/common"
|
||||
)
|
||||
|
||||
type RequestInput struct {
|
||||
Headers map[string]string
|
||||
Body []byte
|
||||
Usage map[string]any
|
||||
}
|
||||
|
||||
// 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
|
||||
}
|
||||
|
||||
// RequestRuleTrace describes one request-dependent multiplier detected at compile time.
|
||||
type RequestRuleTrace struct {
|
||||
Cond string `json:"cond"`
|
||||
Multiplier float64 `json:"multiplier"`
|
||||
Matched bool `json:"matched"`
|
||||
}
|
||||
|
||||
// TraceResult holds side-channel info captured while an expression runs.
|
||||
type TraceResult struct {
|
||||
MatchedTier string `json:"matched_tier"`
|
||||
RequestRules []RequestRuleTrace `json:"request_rules,omitempty"`
|
||||
Cost float64 `json:"cost"`
|
||||
}
|
||||
|
||||
// BillingSnapshot captures billing state at pre-consume time. Expression and
|
||||
// request fields stay frozen; group-dependent fields are refreshed before an
|
||||
// auto-group retry and settlement. It is fully serializable and contains no
|
||||
// compiled program pointers.
|
||||
type BillingSnapshot struct {
|
||||
BillingMode string `json:"billing_mode"`
|
||||
ModelName string `json:"model_name"`
|
||||
ExprString string `json:"expr_string"`
|
||||
ExprHash string `json:"expr_hash"`
|
||||
GroupRatio float64 `json:"group_ratio"`
|
||||
EstimatedPromptTokens int `json:"estimated_prompt_tokens"`
|
||||
EstimatedCompletionTokens int `json:"estimated_completion_tokens"`
|
||||
EstimatedQuotaBeforeGroup float64 `json:"estimated_quota_before_group"`
|
||||
EstimatedQuotaAfterGroup int `json:"estimated_quota_after_group"`
|
||||
EstimatedTier string `json:"estimated_tier"`
|
||||
QuotaPerUnit float64 `json:"quota_per_unit"`
|
||||
ExprVersion int `json:"expr_version"`
|
||||
TaskUsageBilling bool `json:"task_usage_billing,omitempty"`
|
||||
UsageFacts map[string]any `json:"usage_facts,omitempty"`
|
||||
}
|
||||
|
||||
// TieredResult holds everything needed after running tiered settlement.
|
||||
type TieredResult struct {
|
||||
ActualQuotaBeforeGroup float64 `json:"actual_quota_before_group"`
|
||||
ActualQuotaAfterGroup int `json:"actual_quota_after_group"`
|
||||
MatchedTier string `json:"matched_tier"`
|
||||
RequestRules []RequestRuleTrace `json:"request_rules,omitempty"`
|
||||
CrossedTier bool `json:"crossed_tier"`
|
||||
// Clamp records a single-request saturation event during quota conversion so the
|
||||
// caller can surface it on the consume log for admin auditing. Nil when no
|
||||
// clamping occurred. Not serialized: the marker is attached separately via
|
||||
// the shared quota-saturation audit path.
|
||||
Clamp *common.QuotaClamp `json:"-"`
|
||||
}
|
||||
|
||||
// ExprHashString returns the SHA-256 hex digest of an expression string.
|
||||
func ExprHashString(expr string) string {
|
||||
h := sha256.Sum256([]byte(expr))
|
||||
return fmt.Sprintf("%x", h)
|
||||
}
|
||||
@@ -0,0 +1,53 @@
|
||||
package cachex
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"strconv"
|
||||
"strings"
|
||||
)
|
||||
|
||||
type ValueCodec[V any] interface {
|
||||
Encode(v V) (string, error)
|
||||
Decode(s string) (V, error)
|
||||
}
|
||||
|
||||
type IntCodec struct{}
|
||||
|
||||
func (c IntCodec) Encode(v int) (string, error) {
|
||||
return strconv.Itoa(v), nil
|
||||
}
|
||||
|
||||
func (c IntCodec) Decode(s string) (int, error) {
|
||||
s = strings.TrimSpace(s)
|
||||
if s == "" {
|
||||
return 0, fmt.Errorf("empty int value")
|
||||
}
|
||||
return strconv.Atoi(s)
|
||||
}
|
||||
|
||||
type StringCodec struct{}
|
||||
|
||||
func (c StringCodec) Encode(v string) (string, error) { return v, nil }
|
||||
func (c StringCodec) Decode(s string) (string, error) { return s, nil }
|
||||
|
||||
type JSONCodec[V any] struct{}
|
||||
|
||||
func (c JSONCodec[V]) Encode(v V) (string, error) {
|
||||
b, err := json.Marshal(v)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
return string(b), nil
|
||||
}
|
||||
|
||||
func (c JSONCodec[V]) Decode(s string) (V, error) {
|
||||
var v V
|
||||
if strings.TrimSpace(s) == "" {
|
||||
return v, fmt.Errorf("empty json value")
|
||||
}
|
||||
if err := json.Unmarshal([]byte(s), &v); err != nil {
|
||||
return v, err
|
||||
}
|
||||
return v, nil
|
||||
}
|
||||
@@ -0,0 +1,285 @@
|
||||
package cachex
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"strings"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
"github.com/go-redis/redis/v8"
|
||||
"github.com/samber/hot"
|
||||
)
|
||||
|
||||
const (
|
||||
defaultRedisOpTimeout = 2 * time.Second
|
||||
defaultRedisScanTimeout = 30 * time.Second
|
||||
defaultRedisDelTimeout = 10 * time.Second
|
||||
)
|
||||
|
||||
type HybridCacheConfig[V any] struct {
|
||||
Namespace Namespace
|
||||
|
||||
// Redis is used when RedisEnabled returns true (or RedisEnabled is nil) and Redis is not nil.
|
||||
Redis *redis.Client
|
||||
RedisCodec ValueCodec[V]
|
||||
RedisEnabled func() bool
|
||||
|
||||
// Memory builds a hot cache used when Redis is disabled. Keys stored in memory are fully namespaced.
|
||||
Memory func() *hot.HotCache[string, V]
|
||||
}
|
||||
|
||||
// HybridCache is a small helper that uses Redis when enabled, otherwise falls back to in-memory hot cache.
|
||||
type HybridCache[V any] struct {
|
||||
ns Namespace
|
||||
|
||||
redis *redis.Client
|
||||
redisCodec ValueCodec[V]
|
||||
redisEnabled func() bool
|
||||
|
||||
memOnce sync.Once
|
||||
memInit func() *hot.HotCache[string, V]
|
||||
mem *hot.HotCache[string, V]
|
||||
}
|
||||
|
||||
func NewHybridCache[V any](cfg HybridCacheConfig[V]) *HybridCache[V] {
|
||||
return &HybridCache[V]{
|
||||
ns: cfg.Namespace,
|
||||
redis: cfg.Redis,
|
||||
redisCodec: cfg.RedisCodec,
|
||||
redisEnabled: cfg.RedisEnabled,
|
||||
memInit: cfg.Memory,
|
||||
}
|
||||
}
|
||||
|
||||
func (c *HybridCache[V]) FullKey(key string) string {
|
||||
return c.ns.FullKey(key)
|
||||
}
|
||||
|
||||
func (c *HybridCache[V]) redisOn() bool {
|
||||
if c.redis == nil || c.redisCodec == nil {
|
||||
return false
|
||||
}
|
||||
if c.redisEnabled == nil {
|
||||
return true
|
||||
}
|
||||
return c.redisEnabled()
|
||||
}
|
||||
|
||||
func (c *HybridCache[V]) memCache() *hot.HotCache[string, V] {
|
||||
c.memOnce.Do(func() {
|
||||
if c.memInit == nil {
|
||||
c.mem = hot.NewHotCache[string, V](hot.LRU, 1).Build()
|
||||
return
|
||||
}
|
||||
c.mem = c.memInit()
|
||||
})
|
||||
return c.mem
|
||||
}
|
||||
|
||||
func (c *HybridCache[V]) Get(key string) (value V, found bool, err error) {
|
||||
full := c.ns.FullKey(key)
|
||||
if full == "" {
|
||||
var zero V
|
||||
return zero, false, nil
|
||||
}
|
||||
|
||||
if c.redisOn() {
|
||||
ctx, cancel := context.WithTimeout(context.Background(), defaultRedisOpTimeout)
|
||||
defer cancel()
|
||||
|
||||
raw, e := c.redis.Get(ctx, full).Result()
|
||||
if e == nil {
|
||||
v, decErr := c.redisCodec.Decode(raw)
|
||||
if decErr != nil {
|
||||
var zero V
|
||||
return zero, false, decErr
|
||||
}
|
||||
return v, true, nil
|
||||
}
|
||||
if errors.Is(e, redis.Nil) {
|
||||
var zero V
|
||||
return zero, false, nil
|
||||
}
|
||||
var zero V
|
||||
return zero, false, e
|
||||
}
|
||||
|
||||
return c.memCache().Get(full)
|
||||
}
|
||||
|
||||
func (c *HybridCache[V]) SetWithTTL(key string, v V, ttl time.Duration) error {
|
||||
full := c.ns.FullKey(key)
|
||||
if full == "" {
|
||||
return nil
|
||||
}
|
||||
|
||||
if c.redisOn() {
|
||||
raw, err := c.redisCodec.Encode(v)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
ctx, cancel := context.WithTimeout(context.Background(), defaultRedisOpTimeout)
|
||||
defer cancel()
|
||||
return c.redis.Set(ctx, full, raw, ttl).Err()
|
||||
}
|
||||
|
||||
c.memCache().SetWithTTL(full, v, ttl)
|
||||
return nil
|
||||
}
|
||||
|
||||
// Keys returns keys with valid values. In Redis, it returns all matching keys.
|
||||
func (c *HybridCache[V]) Keys() ([]string, error) {
|
||||
if c.redisOn() {
|
||||
return c.scanKeys(c.ns.MatchPattern())
|
||||
}
|
||||
return c.memCache().Keys(), nil
|
||||
}
|
||||
|
||||
func (c *HybridCache[V]) scanKeys(match string) ([]string, error) {
|
||||
ctx, cancel := context.WithTimeout(context.Background(), defaultRedisScanTimeout)
|
||||
defer cancel()
|
||||
|
||||
var cursor uint64
|
||||
keys := make([]string, 0, 1024)
|
||||
for {
|
||||
k, next, err := c.redis.Scan(ctx, cursor, match, 1000).Result()
|
||||
if err != nil {
|
||||
return keys, err
|
||||
}
|
||||
keys = append(keys, k...)
|
||||
cursor = next
|
||||
if cursor == 0 {
|
||||
break
|
||||
}
|
||||
}
|
||||
return keys, nil
|
||||
}
|
||||
|
||||
func (c *HybridCache[V]) Purge() error {
|
||||
if c.redisOn() {
|
||||
keys, err := c.scanKeys(c.ns.MatchPattern())
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if len(keys) == 0 {
|
||||
return nil
|
||||
}
|
||||
_, err = c.DeleteMany(keys)
|
||||
return err
|
||||
}
|
||||
|
||||
c.memCache().Purge()
|
||||
return nil
|
||||
}
|
||||
|
||||
func (c *HybridCache[V]) DeleteByPrefix(prefix string) (int, error) {
|
||||
fullPrefix := c.ns.FullKey(prefix)
|
||||
if fullPrefix == "" {
|
||||
return 0, nil
|
||||
}
|
||||
if !strings.HasSuffix(fullPrefix, ":") {
|
||||
fullPrefix += ":"
|
||||
}
|
||||
|
||||
if c.redisOn() {
|
||||
match := fullPrefix + "*"
|
||||
keys, err := c.scanKeys(match)
|
||||
if err != nil {
|
||||
return 0, err
|
||||
}
|
||||
if len(keys) == 0 {
|
||||
return 0, nil
|
||||
}
|
||||
|
||||
res, err := c.DeleteMany(keys)
|
||||
if err != nil {
|
||||
return 0, err
|
||||
}
|
||||
deleted := 0
|
||||
for _, ok := range res {
|
||||
if ok {
|
||||
deleted++
|
||||
}
|
||||
}
|
||||
return deleted, nil
|
||||
}
|
||||
|
||||
// In memory, we filter keys and bulk delete.
|
||||
allKeys := c.memCache().Keys()
|
||||
keys := make([]string, 0, 128)
|
||||
for _, k := range allKeys {
|
||||
if strings.HasPrefix(k, fullPrefix) {
|
||||
keys = append(keys, k)
|
||||
}
|
||||
}
|
||||
if len(keys) == 0 {
|
||||
return 0, nil
|
||||
}
|
||||
res, _ := c.DeleteMany(keys)
|
||||
deleted := 0
|
||||
for _, ok := range res {
|
||||
if ok {
|
||||
deleted++
|
||||
}
|
||||
}
|
||||
return deleted, nil
|
||||
}
|
||||
|
||||
// DeleteMany accepts either fully namespaced keys or raw keys and deletes them.
|
||||
// It returns a map keyed by fully namespaced keys.
|
||||
func (c *HybridCache[V]) DeleteMany(keys []string) (map[string]bool, error) {
|
||||
res := make(map[string]bool, len(keys))
|
||||
if len(keys) == 0 {
|
||||
return res, nil
|
||||
}
|
||||
|
||||
fullKeys := make([]string, 0, len(keys))
|
||||
for _, k := range keys {
|
||||
k = c.ns.FullKey(k)
|
||||
if k == "" {
|
||||
continue
|
||||
}
|
||||
fullKeys = append(fullKeys, k)
|
||||
}
|
||||
if len(fullKeys) == 0 {
|
||||
return res, nil
|
||||
}
|
||||
|
||||
if c.redisOn() {
|
||||
ctx, cancel := context.WithTimeout(context.Background(), defaultRedisDelTimeout)
|
||||
defer cancel()
|
||||
|
||||
pipe := c.redis.Pipeline()
|
||||
cmds := make([]*redis.IntCmd, 0, len(fullKeys))
|
||||
for _, k := range fullKeys {
|
||||
// UNLINK is non-blocking vs DEL for large key batches.
|
||||
cmds = append(cmds, pipe.Unlink(ctx, k))
|
||||
}
|
||||
_, err := pipe.Exec(ctx)
|
||||
if err != nil && !errors.Is(err, redis.Nil) {
|
||||
return res, err
|
||||
}
|
||||
for i, cmd := range cmds {
|
||||
deleted := cmd != nil && cmd.Err() == nil && cmd.Val() > 0
|
||||
res[fullKeys[i]] = deleted
|
||||
}
|
||||
return res, nil
|
||||
}
|
||||
|
||||
return c.memCache().DeleteMany(fullKeys), nil
|
||||
}
|
||||
|
||||
func (c *HybridCache[V]) Capacity() (mainCacheCapacity int, missingCacheCapacity int) {
|
||||
if c.redisOn() {
|
||||
return 0, 0
|
||||
}
|
||||
return c.memCache().Capacity()
|
||||
}
|
||||
|
||||
func (c *HybridCache[V]) Algorithm() (mainCacheAlgorithm string, missingCacheAlgorithm string) {
|
||||
if c.redisOn() {
|
||||
return "redis", ""
|
||||
}
|
||||
return c.memCache().Algorithm()
|
||||
}
|
||||
@@ -0,0 +1,38 @@
|
||||
package cachex
|
||||
|
||||
import "strings"
|
||||
|
||||
// Namespace isolates keys between different cache use-cases. (e.g. "channel_affinity:v1").
|
||||
type Namespace string
|
||||
|
||||
func (n Namespace) prefix() string {
|
||||
ns := strings.TrimSpace(string(n))
|
||||
ns = strings.TrimRight(ns, ":")
|
||||
if ns == "" {
|
||||
return ""
|
||||
}
|
||||
return ns + ":"
|
||||
}
|
||||
|
||||
func (n Namespace) FullKey(key string) string {
|
||||
key = strings.TrimSpace(key)
|
||||
if key == "" {
|
||||
return ""
|
||||
}
|
||||
p := n.prefix()
|
||||
if p == "" {
|
||||
return strings.TrimLeft(key, ":")
|
||||
}
|
||||
if strings.HasPrefix(key, p) {
|
||||
return key
|
||||
}
|
||||
return p + strings.TrimLeft(key, ":")
|
||||
}
|
||||
|
||||
func (n Namespace) MatchPattern() string {
|
||||
p := n.prefix()
|
||||
if p == "" {
|
||||
return "*"
|
||||
}
|
||||
return p + "*"
|
||||
}
|
||||
@@ -0,0 +1,219 @@
|
||||
package ionet
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"net/http"
|
||||
"net/url"
|
||||
"strconv"
|
||||
"time"
|
||||
)
|
||||
|
||||
const (
|
||||
DefaultEnterpriseBaseURL = "https://api.io.solutions/enterprise/v1/io-cloud/caas"
|
||||
DefaultBaseURL = "https://api.io.solutions/v1/io-cloud/caas"
|
||||
DefaultTimeout = 30 * time.Second
|
||||
)
|
||||
|
||||
// DefaultHTTPClient is the default HTTP client implementation
|
||||
type DefaultHTTPClient struct {
|
||||
client *http.Client
|
||||
}
|
||||
|
||||
// NewDefaultHTTPClient creates a new default HTTP client
|
||||
func NewDefaultHTTPClient(timeout time.Duration) *DefaultHTTPClient {
|
||||
return &DefaultHTTPClient{
|
||||
client: &http.Client{
|
||||
Timeout: timeout,
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
// Do executes an HTTP request
|
||||
func (c *DefaultHTTPClient) Do(req *HTTPRequest) (*HTTPResponse, error) {
|
||||
httpReq, err := http.NewRequest(req.Method, req.URL, bytes.NewReader(req.Body))
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failed to create HTTP request: %w", err)
|
||||
}
|
||||
|
||||
// Set headers
|
||||
for key, value := range req.Headers {
|
||||
httpReq.Header.Set(key, value)
|
||||
}
|
||||
|
||||
resp, err := c.client.Do(httpReq)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("HTTP request failed: %w", err)
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
|
||||
// Read response body
|
||||
var body bytes.Buffer
|
||||
_, err = body.ReadFrom(resp.Body)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failed to read response body: %w", err)
|
||||
}
|
||||
|
||||
// Convert headers
|
||||
headers := make(map[string]string)
|
||||
for key, values := range resp.Header {
|
||||
if len(values) > 0 {
|
||||
headers[key] = values[0]
|
||||
}
|
||||
}
|
||||
|
||||
return &HTTPResponse{
|
||||
StatusCode: resp.StatusCode,
|
||||
Headers: headers,
|
||||
Body: body.Bytes(),
|
||||
}, nil
|
||||
}
|
||||
|
||||
// NewEnterpriseClient creates a new IO.NET API client targeting the enterprise API base URL.
|
||||
func NewEnterpriseClient(apiKey string) *Client {
|
||||
return NewClientWithConfig(apiKey, DefaultEnterpriseBaseURL, nil)
|
||||
}
|
||||
|
||||
// NewClient creates a new IO.NET API client targeting the public API base URL.
|
||||
func NewClient(apiKey string) *Client {
|
||||
return NewClientWithConfig(apiKey, DefaultBaseURL, nil)
|
||||
}
|
||||
|
||||
// NewClientWithConfig creates a new IO.NET API client with custom configuration
|
||||
func NewClientWithConfig(apiKey, baseURL string, httpClient HTTPClient) *Client {
|
||||
if baseURL == "" {
|
||||
baseURL = DefaultBaseURL
|
||||
}
|
||||
if httpClient == nil {
|
||||
httpClient = NewDefaultHTTPClient(DefaultTimeout)
|
||||
}
|
||||
return &Client{
|
||||
BaseURL: baseURL,
|
||||
APIKey: apiKey,
|
||||
HTTPClient: httpClient,
|
||||
}
|
||||
}
|
||||
|
||||
// makeRequest performs an HTTP request and handles common response processing
|
||||
func (c *Client) makeRequest(method, endpoint string, body interface{}) (*HTTPResponse, error) {
|
||||
var reqBody []byte
|
||||
var err error
|
||||
|
||||
if body != nil {
|
||||
reqBody, err = json.Marshal(body)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failed to marshal request body: %w", err)
|
||||
}
|
||||
}
|
||||
|
||||
headers := map[string]string{
|
||||
"X-API-KEY": c.APIKey,
|
||||
"Content-Type": "application/json",
|
||||
}
|
||||
|
||||
req := &HTTPRequest{
|
||||
Method: method,
|
||||
URL: c.BaseURL + endpoint,
|
||||
Headers: headers,
|
||||
Body: reqBody,
|
||||
}
|
||||
|
||||
resp, err := c.HTTPClient.Do(req)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("request failed: %w", err)
|
||||
}
|
||||
|
||||
// Handle API errors
|
||||
if resp.StatusCode >= 400 {
|
||||
var apiErr APIError
|
||||
if len(resp.Body) > 0 {
|
||||
// Try to parse the actual error format: {"detail": "message"}
|
||||
var errorResp struct {
|
||||
Detail string `json:"detail"`
|
||||
}
|
||||
if err := json.Unmarshal(resp.Body, &errorResp); err == nil && errorResp.Detail != "" {
|
||||
apiErr = APIError{
|
||||
Code: resp.StatusCode,
|
||||
Message: errorResp.Detail,
|
||||
}
|
||||
} else {
|
||||
// Fallback: use raw body as details
|
||||
apiErr = APIError{
|
||||
Code: resp.StatusCode,
|
||||
Message: fmt.Sprintf("API request failed with status %d", resp.StatusCode),
|
||||
Details: string(resp.Body),
|
||||
}
|
||||
}
|
||||
} else {
|
||||
apiErr = APIError{
|
||||
Code: resp.StatusCode,
|
||||
Message: fmt.Sprintf("API request failed with status %d", resp.StatusCode),
|
||||
}
|
||||
}
|
||||
return nil, &apiErr
|
||||
}
|
||||
|
||||
return resp, nil
|
||||
}
|
||||
|
||||
// buildQueryParams builds query parameters for GET requests
|
||||
func buildQueryParams(params map[string]interface{}) string {
|
||||
if len(params) == 0 {
|
||||
return ""
|
||||
}
|
||||
|
||||
values := url.Values{}
|
||||
for key, value := range params {
|
||||
if value == nil {
|
||||
continue
|
||||
}
|
||||
switch v := value.(type) {
|
||||
case string:
|
||||
if v != "" {
|
||||
values.Add(key, v)
|
||||
}
|
||||
case int:
|
||||
if v != 0 {
|
||||
values.Add(key, strconv.Itoa(v))
|
||||
}
|
||||
case int64:
|
||||
if v != 0 {
|
||||
values.Add(key, strconv.FormatInt(v, 10))
|
||||
}
|
||||
case float64:
|
||||
if v != 0 {
|
||||
values.Add(key, strconv.FormatFloat(v, 'f', -1, 64))
|
||||
}
|
||||
case bool:
|
||||
values.Add(key, strconv.FormatBool(v))
|
||||
case time.Time:
|
||||
if !v.IsZero() {
|
||||
values.Add(key, v.Format(time.RFC3339))
|
||||
}
|
||||
case *time.Time:
|
||||
if v != nil && !v.IsZero() {
|
||||
values.Add(key, v.Format(time.RFC3339))
|
||||
}
|
||||
case []int:
|
||||
if len(v) > 0 {
|
||||
if encoded, err := json.Marshal(v); err == nil {
|
||||
values.Add(key, string(encoded))
|
||||
}
|
||||
}
|
||||
case []string:
|
||||
if len(v) > 0 {
|
||||
if encoded, err := json.Marshal(v); err == nil {
|
||||
values.Add(key, string(encoded))
|
||||
}
|
||||
}
|
||||
default:
|
||||
values.Add(key, fmt.Sprint(v))
|
||||
}
|
||||
}
|
||||
|
||||
if len(values) > 0 {
|
||||
return "?" + values.Encode()
|
||||
}
|
||||
return ""
|
||||
}
|
||||
@@ -0,0 +1,302 @@
|
||||
package ionet
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"github.com/samber/lo"
|
||||
)
|
||||
|
||||
// ListContainers retrieves all containers for a specific deployment
|
||||
func (c *Client) ListContainers(deploymentID string) (*ContainerList, error) {
|
||||
if deploymentID == "" {
|
||||
return nil, fmt.Errorf("deployment ID cannot be empty")
|
||||
}
|
||||
|
||||
endpoint := fmt.Sprintf("/deployment/%s/containers", deploymentID)
|
||||
|
||||
resp, err := c.makeRequest("GET", endpoint, nil)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failed to list containers: %w", err)
|
||||
}
|
||||
|
||||
var containerList ContainerList
|
||||
if err := decodeDataWithFlexibleTimes(resp.Body, &containerList); err != nil {
|
||||
return nil, fmt.Errorf("failed to parse containers list: %w", err)
|
||||
}
|
||||
|
||||
return &containerList, nil
|
||||
}
|
||||
|
||||
// GetContainerDetails retrieves detailed information about a specific container
|
||||
func (c *Client) GetContainerDetails(deploymentID, containerID string) (*Container, error) {
|
||||
if deploymentID == "" {
|
||||
return nil, fmt.Errorf("deployment ID cannot be empty")
|
||||
}
|
||||
if containerID == "" {
|
||||
return nil, fmt.Errorf("container ID cannot be empty")
|
||||
}
|
||||
|
||||
endpoint := fmt.Sprintf("/deployment/%s/container/%s", deploymentID, containerID)
|
||||
|
||||
resp, err := c.makeRequest("GET", endpoint, nil)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failed to get container details: %w", err)
|
||||
}
|
||||
|
||||
// API response format not documented, assuming direct format
|
||||
var container Container
|
||||
if err := decodeWithFlexibleTimes(resp.Body, &container); err != nil {
|
||||
return nil, fmt.Errorf("failed to parse container details: %w", err)
|
||||
}
|
||||
|
||||
return &container, nil
|
||||
}
|
||||
|
||||
// GetContainerJobs retrieves containers jobs for a specific container (similar to containers endpoint)
|
||||
func (c *Client) GetContainerJobs(deploymentID, containerID string) (*ContainerList, error) {
|
||||
if deploymentID == "" {
|
||||
return nil, fmt.Errorf("deployment ID cannot be empty")
|
||||
}
|
||||
if containerID == "" {
|
||||
return nil, fmt.Errorf("container ID cannot be empty")
|
||||
}
|
||||
|
||||
endpoint := fmt.Sprintf("/deployment/%s/containers-jobs/%s", deploymentID, containerID)
|
||||
|
||||
resp, err := c.makeRequest("GET", endpoint, nil)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failed to get container jobs: %w", err)
|
||||
}
|
||||
|
||||
var containerList ContainerList
|
||||
if err := decodeDataWithFlexibleTimes(resp.Body, &containerList); err != nil {
|
||||
return nil, fmt.Errorf("failed to parse container jobs: %w", err)
|
||||
}
|
||||
|
||||
return &containerList, nil
|
||||
}
|
||||
|
||||
// buildLogEndpoint constructs the request path for fetching logs
|
||||
func buildLogEndpoint(deploymentID, containerID string, opts *GetLogsOptions) (string, error) {
|
||||
if deploymentID == "" {
|
||||
return "", fmt.Errorf("deployment ID cannot be empty")
|
||||
}
|
||||
if containerID == "" {
|
||||
return "", fmt.Errorf("container ID cannot be empty")
|
||||
}
|
||||
|
||||
params := make(map[string]interface{})
|
||||
|
||||
if opts != nil {
|
||||
if opts.Level != "" {
|
||||
params["level"] = opts.Level
|
||||
}
|
||||
if opts.Stream != "" {
|
||||
params["stream"] = opts.Stream
|
||||
}
|
||||
if opts.Limit > 0 {
|
||||
params["limit"] = opts.Limit
|
||||
}
|
||||
if opts.Cursor != "" {
|
||||
params["cursor"] = opts.Cursor
|
||||
}
|
||||
if opts.Follow {
|
||||
params["follow"] = true
|
||||
}
|
||||
|
||||
if opts.StartTime != nil {
|
||||
params["start_time"] = opts.StartTime
|
||||
}
|
||||
if opts.EndTime != nil {
|
||||
params["end_time"] = opts.EndTime
|
||||
}
|
||||
}
|
||||
|
||||
endpoint := fmt.Sprintf("/deployment/%s/log/%s", deploymentID, containerID)
|
||||
endpoint += buildQueryParams(params)
|
||||
|
||||
return endpoint, nil
|
||||
}
|
||||
|
||||
// GetContainerLogs retrieves logs for containers in a deployment and normalizes them
|
||||
func (c *Client) GetContainerLogs(deploymentID, containerID string, opts *GetLogsOptions) (*ContainerLogs, error) {
|
||||
raw, err := c.GetContainerLogsRaw(deploymentID, containerID, opts)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
logs := &ContainerLogs{
|
||||
ContainerID: containerID,
|
||||
}
|
||||
|
||||
if raw == "" {
|
||||
return logs, nil
|
||||
}
|
||||
|
||||
normalized := strings.ReplaceAll(raw, "\r\n", "\n")
|
||||
lines := strings.Split(normalized, "\n")
|
||||
logs.Logs = lo.FilterMap(lines, func(line string, _ int) (LogEntry, bool) {
|
||||
if strings.TrimSpace(line) == "" {
|
||||
return LogEntry{}, false
|
||||
}
|
||||
return LogEntry{Message: line}, true
|
||||
})
|
||||
|
||||
return logs, nil
|
||||
}
|
||||
|
||||
// GetContainerLogsRaw retrieves the raw text logs for a specific container
|
||||
func (c *Client) GetContainerLogsRaw(deploymentID, containerID string, opts *GetLogsOptions) (string, error) {
|
||||
endpoint, err := buildLogEndpoint(deploymentID, containerID, opts)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
|
||||
resp, err := c.makeRequest("GET", endpoint, nil)
|
||||
if err != nil {
|
||||
return "", fmt.Errorf("failed to get container logs: %w", err)
|
||||
}
|
||||
|
||||
return string(resp.Body), nil
|
||||
}
|
||||
|
||||
// StreamContainerLogs streams real-time logs for a specific container
|
||||
// This method uses a callback function to handle incoming log entries
|
||||
func (c *Client) StreamContainerLogs(deploymentID, containerID string, opts *GetLogsOptions, callback func(*LogEntry) error) error {
|
||||
if deploymentID == "" {
|
||||
return fmt.Errorf("deployment ID cannot be empty")
|
||||
}
|
||||
if containerID == "" {
|
||||
return fmt.Errorf("container ID cannot be empty")
|
||||
}
|
||||
if callback == nil {
|
||||
return fmt.Errorf("callback function cannot be nil")
|
||||
}
|
||||
|
||||
// Set follow to true for streaming
|
||||
if opts == nil {
|
||||
opts = &GetLogsOptions{}
|
||||
}
|
||||
opts.Follow = true
|
||||
|
||||
endpoint, err := buildLogEndpoint(deploymentID, containerID, opts)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
// Note: This is a simplified implementation. In a real scenario, you might want to use
|
||||
// Server-Sent Events (SSE) or WebSocket for streaming logs
|
||||
for {
|
||||
resp, err := c.makeRequest("GET", endpoint, nil)
|
||||
if err != nil {
|
||||
return fmt.Errorf("failed to stream container logs: %w", err)
|
||||
}
|
||||
|
||||
var logs ContainerLogs
|
||||
if err := decodeWithFlexibleTimes(resp.Body, &logs); err != nil {
|
||||
return fmt.Errorf("failed to parse container logs: %w", err)
|
||||
}
|
||||
|
||||
// Call the callback for each log entry
|
||||
for _, logEntry := range logs.Logs {
|
||||
if err := callback(&logEntry); err != nil {
|
||||
return fmt.Errorf("callback error: %w", err)
|
||||
}
|
||||
}
|
||||
|
||||
// If there are no more logs or we have a cursor, continue polling
|
||||
if !logs.HasMore && logs.NextCursor == "" {
|
||||
break
|
||||
}
|
||||
|
||||
// Update cursor for next request
|
||||
if logs.NextCursor != "" {
|
||||
opts.Cursor = logs.NextCursor
|
||||
endpoint, err = buildLogEndpoint(deploymentID, containerID, opts)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
|
||||
// Wait a bit before next poll to avoid overwhelming the API
|
||||
time.Sleep(2 * time.Second)
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
// RestartContainer restarts a specific container (if supported by the API)
|
||||
func (c *Client) RestartContainer(deploymentID, containerID string) error {
|
||||
if deploymentID == "" {
|
||||
return fmt.Errorf("deployment ID cannot be empty")
|
||||
}
|
||||
if containerID == "" {
|
||||
return fmt.Errorf("container ID cannot be empty")
|
||||
}
|
||||
|
||||
endpoint := fmt.Sprintf("/deployment/%s/container/%s/restart", deploymentID, containerID)
|
||||
|
||||
_, err := c.makeRequest("POST", endpoint, nil)
|
||||
if err != nil {
|
||||
return fmt.Errorf("failed to restart container: %w", err)
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
// StopContainer stops a specific container (if supported by the API)
|
||||
func (c *Client) StopContainer(deploymentID, containerID string) error {
|
||||
if deploymentID == "" {
|
||||
return fmt.Errorf("deployment ID cannot be empty")
|
||||
}
|
||||
if containerID == "" {
|
||||
return fmt.Errorf("container ID cannot be empty")
|
||||
}
|
||||
|
||||
endpoint := fmt.Sprintf("/deployment/%s/container/%s/stop", deploymentID, containerID)
|
||||
|
||||
_, err := c.makeRequest("POST", endpoint, nil)
|
||||
if err != nil {
|
||||
return fmt.Errorf("failed to stop container: %w", err)
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
// ExecuteInContainer executes a command in a specific container (if supported by the API)
|
||||
func (c *Client) ExecuteInContainer(deploymentID, containerID string, command []string) (string, error) {
|
||||
if deploymentID == "" {
|
||||
return "", fmt.Errorf("deployment ID cannot be empty")
|
||||
}
|
||||
if containerID == "" {
|
||||
return "", fmt.Errorf("container ID cannot be empty")
|
||||
}
|
||||
if len(command) == 0 {
|
||||
return "", fmt.Errorf("command cannot be empty")
|
||||
}
|
||||
|
||||
reqBody := map[string]interface{}{
|
||||
"command": command,
|
||||
}
|
||||
|
||||
endpoint := fmt.Sprintf("/deployment/%s/container/%s/exec", deploymentID, containerID)
|
||||
|
||||
resp, err := c.makeRequest("POST", endpoint, reqBody)
|
||||
if err != nil {
|
||||
return "", fmt.Errorf("failed to execute command in container: %w", err)
|
||||
}
|
||||
|
||||
var result map[string]interface{}
|
||||
if err := json.Unmarshal(resp.Body, &result); err != nil {
|
||||
return "", fmt.Errorf("failed to parse execution result: %w", err)
|
||||
}
|
||||
|
||||
if output, ok := result["output"].(string); ok {
|
||||
return output, nil
|
||||
}
|
||||
|
||||
return string(resp.Body), nil
|
||||
}
|
||||
@@ -0,0 +1,377 @@
|
||||
package ionet
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"strings"
|
||||
|
||||
"github.com/samber/lo"
|
||||
)
|
||||
|
||||
// DeployContainer deploys a new container with the specified configuration
|
||||
func (c *Client) DeployContainer(req *DeploymentRequest) (*DeploymentResponse, error) {
|
||||
if req == nil {
|
||||
return nil, fmt.Errorf("deployment request cannot be nil")
|
||||
}
|
||||
|
||||
// Validate required fields
|
||||
if req.ResourcePrivateName == "" {
|
||||
return nil, fmt.Errorf("resource_private_name is required")
|
||||
}
|
||||
if len(req.LocationIDs) == 0 {
|
||||
return nil, fmt.Errorf("location_ids is required")
|
||||
}
|
||||
if req.HardwareID <= 0 {
|
||||
return nil, fmt.Errorf("hardware_id is required")
|
||||
}
|
||||
if req.RegistryConfig.ImageURL == "" {
|
||||
return nil, fmt.Errorf("registry_config.image_url is required")
|
||||
}
|
||||
if req.GPUsPerContainer < 1 {
|
||||
return nil, fmt.Errorf("gpus_per_container must be at least 1")
|
||||
}
|
||||
if req.DurationHours < 1 {
|
||||
return nil, fmt.Errorf("duration_hours must be at least 1")
|
||||
}
|
||||
if req.ContainerConfig.ReplicaCount < 1 {
|
||||
return nil, fmt.Errorf("container_config.replica_count must be at least 1")
|
||||
}
|
||||
|
||||
resp, err := c.makeRequest("POST", "/deploy", req)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failed to deploy container: %w", err)
|
||||
}
|
||||
|
||||
// API returns direct format:
|
||||
// {"status": "string", "deployment_id": "..."}
|
||||
var deployResp DeploymentResponse
|
||||
if err := json.Unmarshal(resp.Body, &deployResp); err != nil {
|
||||
return nil, fmt.Errorf("failed to parse deployment response: %w", err)
|
||||
}
|
||||
|
||||
return &deployResp, nil
|
||||
}
|
||||
|
||||
// ListDeployments retrieves a list of deployments with optional filtering
|
||||
func (c *Client) ListDeployments(opts *ListDeploymentsOptions) (*DeploymentList, error) {
|
||||
params := make(map[string]interface{})
|
||||
|
||||
if opts != nil {
|
||||
params["status"] = opts.Status
|
||||
params["location_id"] = opts.LocationID
|
||||
params["page"] = opts.Page
|
||||
params["page_size"] = opts.PageSize
|
||||
params["sort_by"] = opts.SortBy
|
||||
params["sort_order"] = opts.SortOrder
|
||||
}
|
||||
|
||||
endpoint := "/deployments" + buildQueryParams(params)
|
||||
|
||||
resp, err := c.makeRequest("GET", endpoint, nil)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failed to list deployments: %w", err)
|
||||
}
|
||||
|
||||
var deploymentList DeploymentList
|
||||
if err := decodeData(resp.Body, &deploymentList); err != nil {
|
||||
return nil, fmt.Errorf("failed to parse deployments list: %w", err)
|
||||
}
|
||||
|
||||
deploymentList.Deployments = lo.Map(deploymentList.Deployments, func(deployment Deployment, _ int) Deployment {
|
||||
deployment.GPUCount = deployment.HardwareQuantity
|
||||
deployment.Replicas = deployment.HardwareQuantity // Assuming 1:1 mapping for now
|
||||
return deployment
|
||||
})
|
||||
|
||||
return &deploymentList, nil
|
||||
}
|
||||
|
||||
// GetDeployment retrieves detailed information about a specific deployment
|
||||
func (c *Client) GetDeployment(deploymentID string) (*DeploymentDetail, error) {
|
||||
if deploymentID == "" {
|
||||
return nil, fmt.Errorf("deployment ID cannot be empty")
|
||||
}
|
||||
|
||||
endpoint := fmt.Sprintf("/deployment/%s", deploymentID)
|
||||
|
||||
resp, err := c.makeRequest("GET", endpoint, nil)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failed to get deployment details: %w", err)
|
||||
}
|
||||
|
||||
var deploymentDetail DeploymentDetail
|
||||
if err := decodeDataWithFlexibleTimes(resp.Body, &deploymentDetail); err != nil {
|
||||
return nil, fmt.Errorf("failed to parse deployment details: %w", err)
|
||||
}
|
||||
|
||||
return &deploymentDetail, nil
|
||||
}
|
||||
|
||||
// UpdateDeployment updates the configuration of an existing deployment
|
||||
func (c *Client) UpdateDeployment(deploymentID string, req *UpdateDeploymentRequest) (*UpdateDeploymentResponse, error) {
|
||||
if deploymentID == "" {
|
||||
return nil, fmt.Errorf("deployment ID cannot be empty")
|
||||
}
|
||||
if req == nil {
|
||||
return nil, fmt.Errorf("update request cannot be nil")
|
||||
}
|
||||
|
||||
endpoint := fmt.Sprintf("/deployment/%s", deploymentID)
|
||||
|
||||
resp, err := c.makeRequest("PATCH", endpoint, req)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failed to update deployment: %w", err)
|
||||
}
|
||||
|
||||
// API returns direct format:
|
||||
// {"status": "string", "deployment_id": "..."}
|
||||
var updateResp UpdateDeploymentResponse
|
||||
if err := json.Unmarshal(resp.Body, &updateResp); err != nil {
|
||||
return nil, fmt.Errorf("failed to parse update deployment response: %w", err)
|
||||
}
|
||||
|
||||
return &updateResp, nil
|
||||
}
|
||||
|
||||
// ExtendDeployment extends the duration of an existing deployment
|
||||
func (c *Client) ExtendDeployment(deploymentID string, req *ExtendDurationRequest) (*DeploymentDetail, error) {
|
||||
if deploymentID == "" {
|
||||
return nil, fmt.Errorf("deployment ID cannot be empty")
|
||||
}
|
||||
if req == nil {
|
||||
return nil, fmt.Errorf("extend request cannot be nil")
|
||||
}
|
||||
if req.DurationHours < 1 {
|
||||
return nil, fmt.Errorf("duration_hours must be at least 1")
|
||||
}
|
||||
|
||||
endpoint := fmt.Sprintf("/deployment/%s/extend", deploymentID)
|
||||
|
||||
resp, err := c.makeRequest("POST", endpoint, req)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failed to extend deployment: %w", err)
|
||||
}
|
||||
|
||||
var deploymentDetail DeploymentDetail
|
||||
if err := decodeDataWithFlexibleTimes(resp.Body, &deploymentDetail); err != nil {
|
||||
return nil, fmt.Errorf("failed to parse extended deployment details: %w", err)
|
||||
}
|
||||
|
||||
return &deploymentDetail, nil
|
||||
}
|
||||
|
||||
// DeleteDeployment deletes an active deployment
|
||||
func (c *Client) DeleteDeployment(deploymentID string) (*UpdateDeploymentResponse, error) {
|
||||
if deploymentID == "" {
|
||||
return nil, fmt.Errorf("deployment ID cannot be empty")
|
||||
}
|
||||
|
||||
endpoint := fmt.Sprintf("/deployment/%s", deploymentID)
|
||||
|
||||
resp, err := c.makeRequest("DELETE", endpoint, nil)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failed to delete deployment: %w", err)
|
||||
}
|
||||
|
||||
// API returns direct format:
|
||||
// {"status": "string", "deployment_id": "..."}
|
||||
var deleteResp UpdateDeploymentResponse
|
||||
if err := json.Unmarshal(resp.Body, &deleteResp); err != nil {
|
||||
return nil, fmt.Errorf("failed to parse delete deployment response: %w", err)
|
||||
}
|
||||
|
||||
return &deleteResp, nil
|
||||
}
|
||||
|
||||
// GetPriceEstimation calculates the estimated cost for a deployment
|
||||
func (c *Client) GetPriceEstimation(req *PriceEstimationRequest) (*PriceEstimationResponse, error) {
|
||||
if req == nil {
|
||||
return nil, fmt.Errorf("price estimation request cannot be nil")
|
||||
}
|
||||
|
||||
// Validate required fields
|
||||
if len(req.LocationIDs) == 0 {
|
||||
return nil, fmt.Errorf("location_ids is required")
|
||||
}
|
||||
if req.HardwareID == 0 {
|
||||
return nil, fmt.Errorf("hardware_id is required")
|
||||
}
|
||||
if req.ReplicaCount < 1 {
|
||||
return nil, fmt.Errorf("replica_count must be at least 1")
|
||||
}
|
||||
|
||||
currency := strings.TrimSpace(req.Currency)
|
||||
if currency == "" {
|
||||
currency = "usdc"
|
||||
}
|
||||
|
||||
durationType := strings.TrimSpace(req.DurationType)
|
||||
if durationType == "" {
|
||||
durationType = "hour"
|
||||
}
|
||||
durationType = strings.ToLower(durationType)
|
||||
|
||||
apiDurationType := ""
|
||||
|
||||
durationQty := req.DurationQty
|
||||
if durationQty < 1 {
|
||||
durationQty = req.DurationHours
|
||||
}
|
||||
if durationQty < 1 {
|
||||
return nil, fmt.Errorf("duration_qty must be at least 1")
|
||||
}
|
||||
|
||||
hardwareQty := req.HardwareQty
|
||||
if hardwareQty < 1 {
|
||||
hardwareQty = req.GPUsPerContainer
|
||||
}
|
||||
if hardwareQty < 1 {
|
||||
return nil, fmt.Errorf("hardware_qty must be at least 1")
|
||||
}
|
||||
|
||||
durationHoursForRate := req.DurationHours
|
||||
if durationHoursForRate < 1 {
|
||||
durationHoursForRate = durationQty
|
||||
}
|
||||
switch durationType {
|
||||
case "hour", "hours", "hourly":
|
||||
durationHoursForRate = durationQty
|
||||
apiDurationType = "hourly"
|
||||
case "day", "days", "daily":
|
||||
durationHoursForRate = durationQty * 24
|
||||
apiDurationType = "daily"
|
||||
case "week", "weeks", "weekly":
|
||||
durationHoursForRate = durationQty * 24 * 7
|
||||
apiDurationType = "weekly"
|
||||
case "month", "months", "monthly":
|
||||
durationHoursForRate = durationQty * 24 * 30
|
||||
apiDurationType = "monthly"
|
||||
}
|
||||
if durationHoursForRate < 1 {
|
||||
durationHoursForRate = 1
|
||||
}
|
||||
if apiDurationType == "" {
|
||||
apiDurationType = "hourly"
|
||||
}
|
||||
|
||||
params := map[string]interface{}{
|
||||
"location_ids": req.LocationIDs,
|
||||
"hardware_id": req.HardwareID,
|
||||
"hardware_qty": hardwareQty,
|
||||
"gpus_per_container": req.GPUsPerContainer,
|
||||
"duration_type": apiDurationType,
|
||||
"duration_qty": durationQty,
|
||||
"duration_hours": req.DurationHours,
|
||||
"replica_count": req.ReplicaCount,
|
||||
"currency": currency,
|
||||
}
|
||||
|
||||
endpoint := "/price" + buildQueryParams(params)
|
||||
|
||||
resp, err := c.makeRequest("GET", endpoint, nil)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failed to get price estimation: %w", err)
|
||||
}
|
||||
|
||||
// Parse according to the actual API response format from docs:
|
||||
// {
|
||||
// "data": {
|
||||
// "replica_count": 0,
|
||||
// "gpus_per_container": 0,
|
||||
// "available_replica_count": [0],
|
||||
// "discount": 0,
|
||||
// "ionet_fee": 0,
|
||||
// "ionet_fee_percent": 0,
|
||||
// "currency_conversion_fee": 0,
|
||||
// "currency_conversion_fee_percent": 0,
|
||||
// "total_cost_usdc": 0
|
||||
// }
|
||||
// }
|
||||
var pricingData struct {
|
||||
ReplicaCount int `json:"replica_count"`
|
||||
GPUsPerContainer int `json:"gpus_per_container"`
|
||||
AvailableReplicaCount []int `json:"available_replica_count"`
|
||||
Discount float64 `json:"discount"`
|
||||
IonetFee float64 `json:"ionet_fee"`
|
||||
IonetFeePercent float64 `json:"ionet_fee_percent"`
|
||||
CurrencyConversionFee float64 `json:"currency_conversion_fee"`
|
||||
CurrencyConversionFeePercent float64 `json:"currency_conversion_fee_percent"`
|
||||
TotalCostUSDC float64 `json:"total_cost_usdc"`
|
||||
}
|
||||
|
||||
if err := decodeData(resp.Body, &pricingData); err != nil {
|
||||
return nil, fmt.Errorf("failed to parse price estimation response: %w", err)
|
||||
}
|
||||
|
||||
// Convert to our internal format
|
||||
durationHoursFloat := float64(durationHoursForRate)
|
||||
if durationHoursFloat <= 0 {
|
||||
durationHoursFloat = 1
|
||||
}
|
||||
|
||||
priceResp := &PriceEstimationResponse{
|
||||
EstimatedCost: pricingData.TotalCostUSDC,
|
||||
Currency: strings.ToUpper(currency),
|
||||
EstimationValid: true,
|
||||
PriceBreakdown: PriceBreakdown{
|
||||
ComputeCost: pricingData.TotalCostUSDC - pricingData.IonetFee - pricingData.CurrencyConversionFee,
|
||||
TotalCost: pricingData.TotalCostUSDC,
|
||||
HourlyRate: pricingData.TotalCostUSDC / durationHoursFloat,
|
||||
},
|
||||
}
|
||||
|
||||
return priceResp, nil
|
||||
}
|
||||
|
||||
// CheckClusterNameAvailability checks if a cluster name is available
|
||||
func (c *Client) CheckClusterNameAvailability(clusterName string) (bool, error) {
|
||||
if clusterName == "" {
|
||||
return false, fmt.Errorf("cluster name cannot be empty")
|
||||
}
|
||||
|
||||
params := map[string]interface{}{
|
||||
"cluster_name": clusterName,
|
||||
}
|
||||
|
||||
endpoint := "/clusters/check_cluster_name_availability" + buildQueryParams(params)
|
||||
|
||||
resp, err := c.makeRequest("GET", endpoint, nil)
|
||||
if err != nil {
|
||||
return false, fmt.Errorf("failed to check cluster name availability: %w", err)
|
||||
}
|
||||
|
||||
var availabilityResp bool
|
||||
if err := json.Unmarshal(resp.Body, &availabilityResp); err != nil {
|
||||
return false, fmt.Errorf("failed to parse cluster name availability response: %w", err)
|
||||
}
|
||||
|
||||
return availabilityResp, nil
|
||||
}
|
||||
|
||||
// UpdateClusterName updates the name of an existing cluster/deployment
|
||||
func (c *Client) UpdateClusterName(clusterID string, req *UpdateClusterNameRequest) (*UpdateClusterNameResponse, error) {
|
||||
if clusterID == "" {
|
||||
return nil, fmt.Errorf("cluster ID cannot be empty")
|
||||
}
|
||||
if req == nil {
|
||||
return nil, fmt.Errorf("update cluster name request cannot be nil")
|
||||
}
|
||||
if req.Name == "" {
|
||||
return nil, fmt.Errorf("cluster name cannot be empty")
|
||||
}
|
||||
|
||||
endpoint := fmt.Sprintf("/clusters/%s/update-name", clusterID)
|
||||
|
||||
resp, err := c.makeRequest("PUT", endpoint, req)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failed to update cluster name: %w", err)
|
||||
}
|
||||
|
||||
// Parse the response directly without data wrapper based on API docs
|
||||
var updateResp UpdateClusterNameResponse
|
||||
if err := json.Unmarshal(resp.Body, &updateResp); err != nil {
|
||||
return nil, fmt.Errorf("failed to parse update cluster name response: %w", err)
|
||||
}
|
||||
|
||||
return &updateResp, nil
|
||||
}
|
||||
@@ -0,0 +1,202 @@
|
||||
package ionet
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"strings"
|
||||
|
||||
"github.com/samber/lo"
|
||||
)
|
||||
|
||||
// GetAvailableReplicas retrieves available replicas per location for specified hardware
|
||||
func (c *Client) GetAvailableReplicas(hardwareID int, gpuCount int) (*AvailableReplicasResponse, error) {
|
||||
if hardwareID <= 0 {
|
||||
return nil, fmt.Errorf("hardware_id must be greater than 0")
|
||||
}
|
||||
if gpuCount < 1 {
|
||||
return nil, fmt.Errorf("gpu_count must be at least 1")
|
||||
}
|
||||
|
||||
params := map[string]interface{}{
|
||||
"hardware_id": hardwareID,
|
||||
"hardware_qty": gpuCount,
|
||||
}
|
||||
|
||||
endpoint := "/available-replicas" + buildQueryParams(params)
|
||||
|
||||
resp, err := c.makeRequest("GET", endpoint, nil)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failed to get available replicas: %w", err)
|
||||
}
|
||||
|
||||
type availableReplicaPayload struct {
|
||||
ID int `json:"id"`
|
||||
ISO2 string `json:"iso2"`
|
||||
Name string `json:"name"`
|
||||
AvailableReplicas int `json:"available_replicas"`
|
||||
}
|
||||
var payload []availableReplicaPayload
|
||||
|
||||
if err := decodeData(resp.Body, &payload); err != nil {
|
||||
return nil, fmt.Errorf("failed to parse available replicas response: %w", err)
|
||||
}
|
||||
|
||||
replicas := lo.Map(payload, func(item availableReplicaPayload, _ int) AvailableReplica {
|
||||
return AvailableReplica{
|
||||
LocationID: item.ID,
|
||||
LocationName: item.Name,
|
||||
HardwareID: hardwareID,
|
||||
HardwareName: "",
|
||||
AvailableCount: item.AvailableReplicas,
|
||||
MaxGPUs: gpuCount,
|
||||
}
|
||||
})
|
||||
|
||||
return &AvailableReplicasResponse{Replicas: replicas}, nil
|
||||
}
|
||||
|
||||
// GetMaxGPUsPerContainer retrieves the maximum number of GPUs available per hardware type
|
||||
func (c *Client) GetMaxGPUsPerContainer() (*MaxGPUResponse, error) {
|
||||
resp, err := c.makeRequest("GET", "/hardware/max-gpus-per-container", nil)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failed to get max GPUs per container: %w", err)
|
||||
}
|
||||
|
||||
var maxGPUResp MaxGPUResponse
|
||||
if err := decodeData(resp.Body, &maxGPUResp); err != nil {
|
||||
return nil, fmt.Errorf("failed to parse max GPU response: %w", err)
|
||||
}
|
||||
|
||||
return &maxGPUResp, nil
|
||||
}
|
||||
|
||||
// ListHardwareTypes retrieves available hardware types using the max GPUs endpoint
|
||||
func (c *Client) ListHardwareTypes() ([]HardwareType, int, error) {
|
||||
maxGPUResp, err := c.GetMaxGPUsPerContainer()
|
||||
if err != nil {
|
||||
return nil, 0, fmt.Errorf("failed to list hardware types: %w", err)
|
||||
}
|
||||
|
||||
mapped := lo.Map(maxGPUResp.Hardware, func(hw MaxGPUInfo, _ int) HardwareType {
|
||||
name := strings.TrimSpace(hw.HardwareName)
|
||||
if name == "" {
|
||||
name = fmt.Sprintf("Hardware %d", hw.HardwareID)
|
||||
}
|
||||
|
||||
return HardwareType{
|
||||
ID: hw.HardwareID,
|
||||
Name: name,
|
||||
GPUType: "",
|
||||
GPUMemory: 0,
|
||||
MaxGPUs: hw.MaxGPUsPerContainer,
|
||||
CPU: "",
|
||||
Memory: 0,
|
||||
Storage: 0,
|
||||
HourlyRate: 0,
|
||||
Available: hw.Available > 0,
|
||||
BrandName: strings.TrimSpace(hw.BrandName),
|
||||
AvailableCount: hw.Available,
|
||||
}
|
||||
})
|
||||
|
||||
totalAvailable := maxGPUResp.Total
|
||||
if totalAvailable == 0 {
|
||||
totalAvailable = lo.SumBy(maxGPUResp.Hardware, func(hw MaxGPUInfo) int {
|
||||
return hw.Available
|
||||
})
|
||||
}
|
||||
|
||||
return mapped, totalAvailable, nil
|
||||
}
|
||||
|
||||
// ListLocations retrieves available deployment locations (if supported by the API)
|
||||
func (c *Client) ListLocations() (*LocationsResponse, error) {
|
||||
resp, err := c.makeRequest("GET", "/locations", nil)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failed to list locations: %w", err)
|
||||
}
|
||||
|
||||
var locations LocationsResponse
|
||||
if err := decodeData(resp.Body, &locations); err != nil {
|
||||
return nil, fmt.Errorf("failed to parse locations response: %w", err)
|
||||
}
|
||||
|
||||
locations.Locations = lo.Map(locations.Locations, func(location Location, _ int) Location {
|
||||
location.ISO2 = strings.ToUpper(strings.TrimSpace(location.ISO2))
|
||||
return location
|
||||
})
|
||||
|
||||
if locations.Total == 0 {
|
||||
locations.Total = lo.SumBy(locations.Locations, func(location Location) int {
|
||||
return location.Available
|
||||
})
|
||||
}
|
||||
|
||||
return &locations, nil
|
||||
}
|
||||
|
||||
// GetHardwareType retrieves details about a specific hardware type
|
||||
func (c *Client) GetHardwareType(hardwareID int) (*HardwareType, error) {
|
||||
if hardwareID <= 0 {
|
||||
return nil, fmt.Errorf("hardware ID must be greater than 0")
|
||||
}
|
||||
|
||||
endpoint := fmt.Sprintf("/hardware/types/%d", hardwareID)
|
||||
|
||||
resp, err := c.makeRequest("GET", endpoint, nil)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failed to get hardware type: %w", err)
|
||||
}
|
||||
|
||||
// API response format not documented, assuming direct format
|
||||
var hardwareType HardwareType
|
||||
if err := json.Unmarshal(resp.Body, &hardwareType); err != nil {
|
||||
return nil, fmt.Errorf("failed to parse hardware type: %w", err)
|
||||
}
|
||||
|
||||
return &hardwareType, nil
|
||||
}
|
||||
|
||||
// GetLocation retrieves details about a specific location
|
||||
func (c *Client) GetLocation(locationID int) (*Location, error) {
|
||||
if locationID <= 0 {
|
||||
return nil, fmt.Errorf("location ID must be greater than 0")
|
||||
}
|
||||
|
||||
endpoint := fmt.Sprintf("/locations/%d", locationID)
|
||||
|
||||
resp, err := c.makeRequest("GET", endpoint, nil)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failed to get location: %w", err)
|
||||
}
|
||||
|
||||
// API response format not documented, assuming direct format
|
||||
var location Location
|
||||
if err := json.Unmarshal(resp.Body, &location); err != nil {
|
||||
return nil, fmt.Errorf("failed to parse location: %w", err)
|
||||
}
|
||||
|
||||
return &location, nil
|
||||
}
|
||||
|
||||
// GetLocationAvailability retrieves real-time availability for a specific location
|
||||
func (c *Client) GetLocationAvailability(locationID int) (*LocationAvailability, error) {
|
||||
if locationID <= 0 {
|
||||
return nil, fmt.Errorf("location ID must be greater than 0")
|
||||
}
|
||||
|
||||
endpoint := fmt.Sprintf("/locations/%d/availability", locationID)
|
||||
|
||||
resp, err := c.makeRequest("GET", endpoint, nil)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failed to get location availability: %w", err)
|
||||
}
|
||||
|
||||
// API response format not documented, assuming direct format
|
||||
var availability LocationAvailability
|
||||
if err := json.Unmarshal(resp.Body, &availability); err != nil {
|
||||
return nil, fmt.Errorf("failed to parse location availability: %w", err)
|
||||
}
|
||||
|
||||
return &availability, nil
|
||||
}
|
||||
@@ -0,0 +1,96 @@
|
||||
package ionet
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"github.com/samber/lo"
|
||||
)
|
||||
|
||||
// decodeWithFlexibleTimes unmarshals API responses while tolerating timestamp strings
|
||||
// that omit timezone information by normalizing them to RFC3339Nano.
|
||||
func decodeWithFlexibleTimes(data []byte, target interface{}) error {
|
||||
var intermediate interface{}
|
||||
if err := json.Unmarshal(data, &intermediate); err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
normalized := normalizeTimeValues(intermediate)
|
||||
reencoded, err := json.Marshal(normalized)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
return json.Unmarshal(reencoded, target)
|
||||
}
|
||||
|
||||
func decodeData[T any](data []byte, target *T) error {
|
||||
var wrapper struct {
|
||||
Data T `json:"data"`
|
||||
}
|
||||
if err := json.Unmarshal(data, &wrapper); err != nil {
|
||||
return err
|
||||
}
|
||||
*target = wrapper.Data
|
||||
return nil
|
||||
}
|
||||
|
||||
func decodeDataWithFlexibleTimes[T any](data []byte, target *T) error {
|
||||
var wrapper struct {
|
||||
Data T `json:"data"`
|
||||
}
|
||||
if err := decodeWithFlexibleTimes(data, &wrapper); err != nil {
|
||||
return err
|
||||
}
|
||||
*target = wrapper.Data
|
||||
return nil
|
||||
}
|
||||
|
||||
func normalizeTimeValues(value interface{}) interface{} {
|
||||
switch v := value.(type) {
|
||||
case map[string]interface{}:
|
||||
return lo.MapValues(v, func(val interface{}, _ string) interface{} {
|
||||
return normalizeTimeValues(val)
|
||||
})
|
||||
case []interface{}:
|
||||
return lo.Map(v, func(item interface{}, _ int) interface{} {
|
||||
return normalizeTimeValues(item)
|
||||
})
|
||||
case string:
|
||||
if normalized, changed := normalizeTimeString(v); changed {
|
||||
return normalized
|
||||
}
|
||||
return v
|
||||
default:
|
||||
return value
|
||||
}
|
||||
}
|
||||
|
||||
func normalizeTimeString(input string) (string, bool) {
|
||||
trimmed := strings.TrimSpace(input)
|
||||
if trimmed == "" {
|
||||
return input, false
|
||||
}
|
||||
|
||||
if _, err := time.Parse(time.RFC3339Nano, trimmed); err == nil {
|
||||
return trimmed, trimmed != input
|
||||
}
|
||||
if _, err := time.Parse(time.RFC3339, trimmed); err == nil {
|
||||
return trimmed, trimmed != input
|
||||
}
|
||||
|
||||
layouts := []string{
|
||||
"2006-01-02T15:04:05.999999999",
|
||||
"2006-01-02T15:04:05.999999",
|
||||
"2006-01-02T15:04:05",
|
||||
}
|
||||
|
||||
for _, layout := range layouts {
|
||||
if parsed, err := time.Parse(layout, trimmed); err == nil {
|
||||
return parsed.UTC().Format(time.RFC3339Nano), true
|
||||
}
|
||||
}
|
||||
|
||||
return input, false
|
||||
}
|
||||
@@ -0,0 +1,353 @@
|
||||
package ionet
|
||||
|
||||
import (
|
||||
"time"
|
||||
)
|
||||
|
||||
// Client represents the IO.NET API client
|
||||
type Client struct {
|
||||
BaseURL string
|
||||
APIKey string
|
||||
HTTPClient HTTPClient
|
||||
}
|
||||
|
||||
// HTTPClient interface for making HTTP requests
|
||||
type HTTPClient interface {
|
||||
Do(req *HTTPRequest) (*HTTPResponse, error)
|
||||
}
|
||||
|
||||
// HTTPRequest represents an HTTP request
|
||||
type HTTPRequest struct {
|
||||
Method string
|
||||
URL string
|
||||
Headers map[string]string
|
||||
Body []byte
|
||||
}
|
||||
|
||||
// HTTPResponse represents an HTTP response
|
||||
type HTTPResponse struct {
|
||||
StatusCode int
|
||||
Headers map[string]string
|
||||
Body []byte
|
||||
}
|
||||
|
||||
// DeploymentRequest represents a container deployment request
|
||||
type DeploymentRequest struct {
|
||||
ResourcePrivateName string `json:"resource_private_name"`
|
||||
DurationHours int `json:"duration_hours"`
|
||||
GPUsPerContainer int `json:"gpus_per_container"`
|
||||
HardwareID int `json:"hardware_id"`
|
||||
LocationIDs []int `json:"location_ids"`
|
||||
ContainerConfig ContainerConfig `json:"container_config"`
|
||||
RegistryConfig RegistryConfig `json:"registry_config"`
|
||||
}
|
||||
|
||||
// ContainerConfig represents container configuration
|
||||
type ContainerConfig struct {
|
||||
ReplicaCount int `json:"replica_count"`
|
||||
EnvVariables map[string]string `json:"env_variables,omitempty"`
|
||||
SecretEnvVariables map[string]string `json:"secret_env_variables,omitempty"`
|
||||
Entrypoint []string `json:"entrypoint,omitempty"`
|
||||
TrafficPort int `json:"traffic_port,omitempty"`
|
||||
Args []string `json:"args,omitempty"`
|
||||
}
|
||||
|
||||
// RegistryConfig represents registry configuration
|
||||
type RegistryConfig struct {
|
||||
ImageURL string `json:"image_url"`
|
||||
RegistryUsername string `json:"registry_username,omitempty"`
|
||||
RegistrySecret string `json:"registry_secret,omitempty"`
|
||||
}
|
||||
|
||||
// DeploymentResponse represents the response from deployment creation
|
||||
type DeploymentResponse struct {
|
||||
DeploymentID string `json:"deployment_id"`
|
||||
Status string `json:"status"`
|
||||
}
|
||||
|
||||
// DeploymentDetail represents detailed deployment information
|
||||
type DeploymentDetail struct {
|
||||
ID string `json:"id"`
|
||||
Status string `json:"status"`
|
||||
CreatedAt time.Time `json:"created_at"`
|
||||
StartedAt *time.Time `json:"started_at,omitempty"`
|
||||
FinishedAt *time.Time `json:"finished_at,omitempty"`
|
||||
AmountPaid float64 `json:"amount_paid"`
|
||||
CompletedPercent float64 `json:"completed_percent"`
|
||||
TotalGPUs int `json:"total_gpus"`
|
||||
GPUsPerContainer int `json:"gpus_per_container"`
|
||||
TotalContainers int `json:"total_containers"`
|
||||
HardwareName string `json:"hardware_name"`
|
||||
HardwareID int `json:"hardware_id"`
|
||||
Locations []DeploymentLocation `json:"locations"`
|
||||
BrandName string `json:"brand_name"`
|
||||
ComputeMinutesServed int `json:"compute_minutes_served"`
|
||||
ComputeMinutesRemaining int `json:"compute_minutes_remaining"`
|
||||
ContainerConfig DeploymentContainerConfig `json:"container_config"`
|
||||
}
|
||||
|
||||
// DeploymentLocation represents a location in deployment details
|
||||
type DeploymentLocation struct {
|
||||
ID int `json:"id"`
|
||||
ISO2 string `json:"iso2"`
|
||||
Name string `json:"name"`
|
||||
}
|
||||
|
||||
// DeploymentContainerConfig represents container config in deployment details
|
||||
type DeploymentContainerConfig struct {
|
||||
Entrypoint []string `json:"entrypoint"`
|
||||
EnvVariables map[string]interface{} `json:"env_variables"`
|
||||
TrafficPort int `json:"traffic_port"`
|
||||
ImageURL string `json:"image_url"`
|
||||
}
|
||||
|
||||
// Container represents a container within a deployment
|
||||
type Container struct {
|
||||
DeviceID string `json:"device_id"`
|
||||
ContainerID string `json:"container_id"`
|
||||
Hardware string `json:"hardware"`
|
||||
BrandName string `json:"brand_name"`
|
||||
CreatedAt time.Time `json:"created_at"`
|
||||
UptimePercent int `json:"uptime_percent"`
|
||||
GPUsPerContainer int `json:"gpus_per_container"`
|
||||
Status string `json:"status"`
|
||||
ContainerEvents []ContainerEvent `json:"container_events"`
|
||||
PublicURL string `json:"public_url"`
|
||||
}
|
||||
|
||||
// ContainerEvent represents a container event
|
||||
type ContainerEvent struct {
|
||||
Time time.Time `json:"time"`
|
||||
Message string `json:"message"`
|
||||
}
|
||||
|
||||
// ContainerList represents a list of containers
|
||||
type ContainerList struct {
|
||||
Total int `json:"total"`
|
||||
Workers []Container `json:"workers"`
|
||||
}
|
||||
|
||||
// Deployment represents a deployment in the list
|
||||
type Deployment struct {
|
||||
ID string `json:"id"`
|
||||
Status string `json:"status"`
|
||||
Name string `json:"name"`
|
||||
CompletedPercent float64 `json:"completed_percent"`
|
||||
HardwareQuantity int `json:"hardware_quantity"`
|
||||
BrandName string `json:"brand_name"`
|
||||
HardwareName string `json:"hardware_name"`
|
||||
Served string `json:"served"`
|
||||
Remaining string `json:"remaining"`
|
||||
ComputeMinutesServed int `json:"compute_minutes_served"`
|
||||
ComputeMinutesRemaining int `json:"compute_minutes_remaining"`
|
||||
CreatedAt time.Time `json:"created_at"`
|
||||
GPUCount int `json:"-"` // Derived from HardwareQuantity
|
||||
Replicas int `json:"-"` // Derived from HardwareQuantity
|
||||
}
|
||||
|
||||
// DeploymentList represents a list of deployments with pagination
|
||||
type DeploymentList struct {
|
||||
Deployments []Deployment `json:"deployments"`
|
||||
Total int `json:"total"`
|
||||
Statuses []string `json:"statuses"`
|
||||
}
|
||||
|
||||
// AvailableReplica represents replica availability for a location
|
||||
type AvailableReplica struct {
|
||||
LocationID int `json:"location_id"`
|
||||
LocationName string `json:"location_name"`
|
||||
HardwareID int `json:"hardware_id"`
|
||||
HardwareName string `json:"hardware_name"`
|
||||
AvailableCount int `json:"available_count"`
|
||||
MaxGPUs int `json:"max_gpus"`
|
||||
}
|
||||
|
||||
// AvailableReplicasResponse represents the response for available replicas
|
||||
type AvailableReplicasResponse struct {
|
||||
Replicas []AvailableReplica `json:"replicas"`
|
||||
}
|
||||
|
||||
// MaxGPUResponse represents the response for maximum GPUs per container
|
||||
type MaxGPUResponse struct {
|
||||
Hardware []MaxGPUInfo `json:"hardware"`
|
||||
Total int `json:"total"`
|
||||
}
|
||||
|
||||
// MaxGPUInfo represents max GPU information for a hardware type
|
||||
type MaxGPUInfo struct {
|
||||
MaxGPUsPerContainer int `json:"max_gpus_per_container"`
|
||||
Available int `json:"available"`
|
||||
HardwareID int `json:"hardware_id"`
|
||||
HardwareName string `json:"hardware_name"`
|
||||
BrandName string `json:"brand_name"`
|
||||
}
|
||||
|
||||
// PriceEstimationRequest represents a price estimation request
|
||||
type PriceEstimationRequest struct {
|
||||
LocationIDs []int `json:"location_ids"`
|
||||
HardwareID int `json:"hardware_id"`
|
||||
GPUsPerContainer int `json:"gpus_per_container"`
|
||||
DurationHours int `json:"duration_hours"`
|
||||
ReplicaCount int `json:"replica_count"`
|
||||
Currency string `json:"currency"`
|
||||
DurationType string `json:"duration_type"`
|
||||
DurationQty int `json:"duration_qty"`
|
||||
HardwareQty int `json:"hardware_qty"`
|
||||
}
|
||||
|
||||
// PriceEstimationResponse represents the price estimation response
|
||||
type PriceEstimationResponse struct {
|
||||
EstimatedCost float64 `json:"estimated_cost"`
|
||||
Currency string `json:"currency"`
|
||||
PriceBreakdown PriceBreakdown `json:"price_breakdown"`
|
||||
EstimationValid bool `json:"estimation_valid"`
|
||||
}
|
||||
|
||||
// PriceBreakdown represents detailed cost breakdown
|
||||
type PriceBreakdown struct {
|
||||
ComputeCost float64 `json:"compute_cost"`
|
||||
NetworkCost float64 `json:"network_cost,omitempty"`
|
||||
StorageCost float64 `json:"storage_cost,omitempty"`
|
||||
TotalCost float64 `json:"total_cost"`
|
||||
HourlyRate float64 `json:"hourly_rate"`
|
||||
}
|
||||
|
||||
// ContainerLogs represents container log entries
|
||||
type ContainerLogs struct {
|
||||
ContainerID string `json:"container_id"`
|
||||
Logs []LogEntry `json:"logs"`
|
||||
HasMore bool `json:"has_more"`
|
||||
NextCursor string `json:"next_cursor,omitempty"`
|
||||
}
|
||||
|
||||
// LogEntry represents a single log entry
|
||||
type LogEntry struct {
|
||||
Timestamp time.Time `json:"timestamp"`
|
||||
Level string `json:"level,omitempty"`
|
||||
Message string `json:"message"`
|
||||
Source string `json:"source,omitempty"`
|
||||
}
|
||||
|
||||
// UpdateDeploymentRequest represents request to update deployment configuration
|
||||
type UpdateDeploymentRequest struct {
|
||||
EnvVariables map[string]string `json:"env_variables,omitempty"`
|
||||
SecretEnvVariables map[string]string `json:"secret_env_variables,omitempty"`
|
||||
Entrypoint []string `json:"entrypoint,omitempty"`
|
||||
TrafficPort *int `json:"traffic_port,omitempty"`
|
||||
ImageURL string `json:"image_url,omitempty"`
|
||||
RegistryUsername string `json:"registry_username,omitempty"`
|
||||
RegistrySecret string `json:"registry_secret,omitempty"`
|
||||
Args []string `json:"args,omitempty"`
|
||||
Command string `json:"command,omitempty"`
|
||||
}
|
||||
|
||||
// ExtendDurationRequest represents request to extend deployment duration
|
||||
type ExtendDurationRequest struct {
|
||||
DurationHours int `json:"duration_hours"`
|
||||
}
|
||||
|
||||
// UpdateDeploymentResponse represents response from deployment update
|
||||
type UpdateDeploymentResponse struct {
|
||||
Status string `json:"status"`
|
||||
DeploymentID string `json:"deployment_id"`
|
||||
}
|
||||
|
||||
// UpdateClusterNameRequest represents request to update cluster name
|
||||
type UpdateClusterNameRequest struct {
|
||||
Name string `json:"cluster_name"`
|
||||
}
|
||||
|
||||
// UpdateClusterNameResponse represents response from cluster name update
|
||||
type UpdateClusterNameResponse struct {
|
||||
Status string `json:"status"`
|
||||
Message string `json:"message"`
|
||||
}
|
||||
|
||||
// APIError represents an API error response
|
||||
type APIError struct {
|
||||
Code int `json:"code"`
|
||||
Message string `json:"message"`
|
||||
Details string `json:"details,omitempty"`
|
||||
}
|
||||
|
||||
// Error implements the error interface
|
||||
func (e *APIError) Error() string {
|
||||
if e.Details != "" {
|
||||
return e.Message + ": " + e.Details
|
||||
}
|
||||
return e.Message
|
||||
}
|
||||
|
||||
// ListDeploymentsOptions represents options for listing deployments
|
||||
type ListDeploymentsOptions struct {
|
||||
Status string `json:"status,omitempty"` // filter by status
|
||||
LocationID int `json:"location_id,omitempty"` // filter by location
|
||||
Page int `json:"page,omitempty"` // pagination
|
||||
PageSize int `json:"page_size,omitempty"` // pagination
|
||||
SortBy string `json:"sort_by,omitempty"` // sort field
|
||||
SortOrder string `json:"sort_order,omitempty"` // asc/desc
|
||||
}
|
||||
|
||||
// GetLogsOptions represents options for retrieving container logs
|
||||
type GetLogsOptions struct {
|
||||
StartTime *time.Time `json:"start_time,omitempty"`
|
||||
EndTime *time.Time `json:"end_time,omitempty"`
|
||||
Level string `json:"level,omitempty"` // filter by log level
|
||||
Stream string `json:"stream,omitempty"` // filter by stdout/stderr streams
|
||||
Limit int `json:"limit,omitempty"` // max number of log entries
|
||||
Cursor string `json:"cursor,omitempty"` // pagination cursor
|
||||
Follow bool `json:"follow,omitempty"` // stream logs
|
||||
}
|
||||
|
||||
// HardwareType represents a hardware type available for deployment
|
||||
type HardwareType struct {
|
||||
ID int `json:"id"`
|
||||
Name string `json:"name"`
|
||||
Description string `json:"description,omitempty"`
|
||||
GPUType string `json:"gpu_type"`
|
||||
GPUMemory int `json:"gpu_memory"` // in GB
|
||||
MaxGPUs int `json:"max_gpus"`
|
||||
CPU string `json:"cpu,omitempty"`
|
||||
Memory int `json:"memory,omitempty"` // in GB
|
||||
Storage int `json:"storage,omitempty"` // in GB
|
||||
HourlyRate float64 `json:"hourly_rate"`
|
||||
Available bool `json:"available"`
|
||||
BrandName string `json:"brand_name,omitempty"`
|
||||
AvailableCount int `json:"available_count,omitempty"`
|
||||
}
|
||||
|
||||
// Location represents a deployment location
|
||||
type Location struct {
|
||||
ID int `json:"id"`
|
||||
Name string `json:"name"`
|
||||
ISO2 string `json:"iso2,omitempty"`
|
||||
Region string `json:"region,omitempty"`
|
||||
Country string `json:"country,omitempty"`
|
||||
Latitude float64 `json:"latitude,omitempty"`
|
||||
Longitude float64 `json:"longitude,omitempty"`
|
||||
Available int `json:"available,omitempty"`
|
||||
Description string `json:"description,omitempty"`
|
||||
}
|
||||
|
||||
// LocationsResponse represents the list of locations and aggregated metadata.
|
||||
type LocationsResponse struct {
|
||||
Locations []Location `json:"locations"`
|
||||
Total int `json:"total"`
|
||||
}
|
||||
|
||||
// LocationAvailability represents real-time availability for a location
|
||||
type LocationAvailability struct {
|
||||
LocationID int `json:"location_id"`
|
||||
LocationName string `json:"location_name"`
|
||||
Available bool `json:"available"`
|
||||
HardwareAvailability []HardwareAvailability `json:"hardware_availability"`
|
||||
UpdatedAt time.Time `json:"updated_at"`
|
||||
}
|
||||
|
||||
// HardwareAvailability represents availability for specific hardware at a location
|
||||
type HardwareAvailability struct {
|
||||
HardwareID int `json:"hardware_id"`
|
||||
HardwareName string `json:"hardware_name"`
|
||||
AvailableCount int `json:"available_count"`
|
||||
MaxGPUs int `json:"max_gpus"`
|
||||
}
|
||||
@@ -0,0 +1,92 @@
|
||||
package jsplugin
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"io"
|
||||
"os"
|
||||
"strings"
|
||||
)
|
||||
|
||||
// RunCLI implements the `new-api plugin` subcommand: linting a plugin source
|
||||
// and replaying a golden fixture against it. It returns a process exit code.
|
||||
func RunCLI(args []string, stdout, stderr io.Writer) int {
|
||||
if len(args) < 2 {
|
||||
fmt.Fprintln(stderr, "usage: new-api plugin lint <plugin.js> | new-api plugin test <plugin.js> --fixture <fixture.json>")
|
||||
return 2
|
||||
}
|
||||
command, sourcePath := args[0], args[1]
|
||||
source, err := os.ReadFile(sourcePath)
|
||||
if err != nil {
|
||||
fmt.Fprintf(stderr, "read plugin: %v\n", err)
|
||||
return 1
|
||||
}
|
||||
|
||||
switch command {
|
||||
case "lint":
|
||||
if len(args) != 2 {
|
||||
fmt.Fprintln(stderr, "usage: new-api plugin lint <plugin.js>")
|
||||
return 2
|
||||
}
|
||||
plugin, compileErr := NewRegistry().Register(string(source), Options{Key: sourcePath, Version: "lint"})
|
||||
if compileErr != nil {
|
||||
fmt.Fprintf(stderr, "plugin lint failed: %v\n", compileErr)
|
||||
return 1
|
||||
}
|
||||
warnParseTaskResultInProgressFallback(string(source), stderr)
|
||||
fmt.Fprintf(stdout, "plugin %s@%s is valid\n", plugin.Meta.Key, plugin.Meta.Version)
|
||||
return 0
|
||||
case "test":
|
||||
if len(args) != 4 || args[2] != "--fixture" || strings.TrimSpace(args[3]) == "" {
|
||||
fmt.Fprintln(stderr, "usage: new-api plugin test <plugin.js> --fixture <fixture.json>")
|
||||
return 2
|
||||
}
|
||||
fixture, readErr := os.ReadFile(args[3])
|
||||
if readErr != nil {
|
||||
fmt.Fprintf(stderr, "read fixture: %v\n", readErr)
|
||||
return 1
|
||||
}
|
||||
report, replayErr := ReplayFixture(context.Background(), string(source), fixture)
|
||||
if replayErr != nil {
|
||||
fmt.Fprintf(stderr, "plugin fixture failed after %d/%d cases: %v\n", report.Passed, report.Total, replayErr)
|
||||
return 1
|
||||
}
|
||||
fmt.Fprintf(stdout, "plugin fixture passed: %d/%d cases\n", report.Passed, report.Total)
|
||||
return 0
|
||||
default:
|
||||
fmt.Fprintf(stderr, "unknown plugin command %q\n", command)
|
||||
return 2
|
||||
}
|
||||
}
|
||||
|
||||
func warnParseTaskResultInProgressFallback(source string, stderr io.Writer) {
|
||||
body := parseTaskResultFunctionBody(source)
|
||||
if strings.Contains(body, `|| "IN_PROGRESS"`) || strings.Contains(body, `|| 'IN_PROGRESS'`) {
|
||||
fmt.Fprintln(stderr, `warning: parseTaskResult uses || "IN_PROGRESS" fallback; return UNKNOWN for unrecognized statuses`)
|
||||
}
|
||||
}
|
||||
|
||||
func parseTaskResultFunctionBody(source string) string {
|
||||
marker := strings.Index(source, "function parseTaskResult")
|
||||
if marker < 0 {
|
||||
return ""
|
||||
}
|
||||
brace := strings.Index(source[marker:], "{")
|
||||
if brace < 0 {
|
||||
return ""
|
||||
}
|
||||
start := marker + brace
|
||||
depth := 0
|
||||
for i := start; i < len(source); i++ {
|
||||
switch source[i] {
|
||||
case '{':
|
||||
depth++
|
||||
case '}':
|
||||
depth--
|
||||
if depth == 0 {
|
||||
return source[start : i+1]
|
||||
}
|
||||
}
|
||||
}
|
||||
return ""
|
||||
}
|
||||
@@ -0,0 +1,57 @@
|
||||
package jsplugin
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"testing"
|
||||
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
func TestPluginCLI(t *testing.T) {
|
||||
tempDir := t.TempDir()
|
||||
pluginPath := filepath.Join(tempDir, "fixture.js")
|
||||
fixturePath := filepath.Join(tempDir, "fixture.json")
|
||||
require.NoError(t, os.WriteFile(pluginPath, []byte(cliFixturePluginSource), 0o600))
|
||||
require.NoError(t, os.WriteFile(fixturePath, []byte(`{"unixNow":42,"cases":[{"hook":"value","args":[],"expected":42}]}`), 0o600))
|
||||
|
||||
var stdout bytes.Buffer
|
||||
var stderr bytes.Buffer
|
||||
assert.Equal(t, 0, RunCLI([]string{"lint", pluginPath}, &stdout, &stderr))
|
||||
assert.Contains(t, stdout.String(), "plugin cli-fixture@1.0.0 is valid")
|
||||
assert.Empty(t, stderr.String())
|
||||
|
||||
stdout.Reset()
|
||||
assert.Equal(t, 0, RunCLI([]string{"test", pluginPath, "--fixture", fixturePath}, &stdout, &stderr))
|
||||
assert.Contains(t, stdout.String(), "1/1 cases")
|
||||
}
|
||||
|
||||
func TestPluginCLIWarnsOnParseTaskResultInProgressFallback(t *testing.T) {
|
||||
tempDir := t.TempDir()
|
||||
pluginPath := filepath.Join(tempDir, "fallback.js")
|
||||
require.NoError(t, os.WriteFile(pluginPath, []byte(`
|
||||
export const meta = { apiVersion: 1, key: "fallback", name: "Fallback", version: "1.0.0", author: {name: "Test"}, models: ["m"], fetchMode: "per_task" };
|
||||
export function buildSubmitRequest(ctx) { return {url: ctx.baseUrl}; }
|
||||
export function parseSubmitResponse() { return {taskId: "task"}; }
|
||||
export function buildQueryRequest(ctx) { return {url: ctx.baseUrl}; }
|
||||
export function parseTaskResult(ctx, body) { return {status: statuses[body.status] || "IN_PROGRESS"}; }
|
||||
const statuses = { done: "SUCCESS" };
|
||||
`), 0o600))
|
||||
|
||||
var stdout bytes.Buffer
|
||||
var stderr bytes.Buffer
|
||||
assert.Equal(t, 0, RunCLI([]string{"lint", pluginPath}, &stdout, &stderr))
|
||||
assert.Contains(t, stdout.String(), "plugin fallback@1.0.0 is valid")
|
||||
assert.Contains(t, stderr.String(), `|| "IN_PROGRESS"`)
|
||||
}
|
||||
|
||||
const cliFixturePluginSource = `
|
||||
export const meta = { apiVersion: 1, key: "cli-fixture", name: "CLI Fixture", version: "1.0.0", author: {name: "Test"}, channelTypes: [1003], models: ["fixture-model"], fetchMode: "per_task" };
|
||||
export function buildSubmitRequest(ctx) { return {url: ctx.baseUrl}; }
|
||||
export function parseSubmitResponse(ctx, resp) { return {taskId: "task"}; }
|
||||
export function buildQueryRequest(ctx) { return {url: ctx.baseUrl}; }
|
||||
export function parseTaskResult(ctx, body) { return body; }
|
||||
export function value() { return utils.unixNow(); }
|
||||
`
|
||||
@@ -0,0 +1,574 @@
|
||||
package jsplugin
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"fmt"
|
||||
"regexp"
|
||||
"slices"
|
||||
"strings"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
"github.com/QuantumNous/new-api/logger"
|
||||
"github.com/grafana/sobek"
|
||||
"github.com/grafana/sobek/parser"
|
||||
)
|
||||
|
||||
const (
|
||||
DefaultCallTimeout = 5 * time.Second
|
||||
DefaultConcurrency = 8
|
||||
)
|
||||
|
||||
var ErrCallAdmissionTimeout = errors.New("plugin call admission timed out")
|
||||
|
||||
const hookErrorMessageLimit = 512
|
||||
|
||||
// HookError reports a JavaScript exception thrown by a plugin hook. Message
|
||||
// is the sanitized JS error message with engine prefixes stripped; it is safe
|
||||
// to surface to API callers.
|
||||
type HookError struct {
|
||||
Hook string
|
||||
Message string
|
||||
wrapped error
|
||||
}
|
||||
|
||||
func (e *HookError) Error() string {
|
||||
if e == nil || e.wrapped == nil {
|
||||
return "plugin hook failed"
|
||||
}
|
||||
return e.wrapped.Error()
|
||||
}
|
||||
|
||||
func (e *HookError) Unwrap() error {
|
||||
if e == nil {
|
||||
return nil
|
||||
}
|
||||
return e.wrapped
|
||||
}
|
||||
|
||||
func newHookError(hook, rawMessage string, wrapped error) *HookError {
|
||||
var b strings.Builder
|
||||
b.Grow(len(rawMessage))
|
||||
count := 0
|
||||
for _, r := range rawMessage {
|
||||
if count >= hookErrorMessageLimit {
|
||||
break
|
||||
}
|
||||
if r < 0x20 || (r >= 0x80 && r <= 0x9F) {
|
||||
r = ' '
|
||||
}
|
||||
b.WriteRune(r)
|
||||
count++
|
||||
}
|
||||
message := b.String()
|
||||
if message == "" {
|
||||
message = "plugin hook failed"
|
||||
}
|
||||
return &HookError{Hook: hook, Message: message, wrapped: wrapped}
|
||||
}
|
||||
|
||||
func hookErrorFromException(hook string, exc *sobek.Exception, wrapped error) (hookErr *HookError) {
|
||||
// Reading message/toString executes plugin getters, which can throw again
|
||||
// and panic sobek. By this point the caller's recover is already consumed,
|
||||
// so a second panic would crash the process; fall back to a blank message.
|
||||
defer func() {
|
||||
if recover() != nil {
|
||||
hookErr = newHookError(hook, "", wrapped)
|
||||
}
|
||||
}()
|
||||
raw := ""
|
||||
if exc != nil {
|
||||
if val := exc.Value(); val != nil && !sobek.IsUndefined(val) && !sobek.IsNull(val) {
|
||||
gotMessage := false
|
||||
if obj, ok := val.(*sobek.Object); ok {
|
||||
if msg := obj.Get("message"); msg != nil && !sobek.IsUndefined(msg) && !sobek.IsNull(msg) {
|
||||
raw = msg.String()
|
||||
gotMessage = true
|
||||
}
|
||||
}
|
||||
if !gotMessage {
|
||||
if exported, ok := val.Export().(string); ok {
|
||||
raw = exported
|
||||
} else {
|
||||
raw = val.String()
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
return newHookError(hook, raw, wrapped)
|
||||
}
|
||||
|
||||
var forbiddenSyntax = regexp.MustCompile(`(?m)(^|[^A-Za-z0-9_$])(async|await|import)([^A-Za-z0-9_$]|$)`)
|
||||
|
||||
type Options struct {
|
||||
Key string
|
||||
Version string
|
||||
Timeout time.Duration
|
||||
Concurrency int
|
||||
Now func() time.Time
|
||||
Log func(string)
|
||||
}
|
||||
|
||||
type Engine struct {
|
||||
key string
|
||||
version string
|
||||
timeout time.Duration
|
||||
now func() time.Time
|
||||
log func(string)
|
||||
module *sobek.SourceTextModuleRecord
|
||||
pool sync.Pool
|
||||
semaphore chan struct{}
|
||||
}
|
||||
|
||||
type runtimeInstance struct {
|
||||
runtime *sobek.Runtime
|
||||
module sobek.ModuleInstance
|
||||
logContext *runtimeLogContext
|
||||
}
|
||||
|
||||
type runtimeLogContext struct {
|
||||
context context.Context
|
||||
}
|
||||
|
||||
// Compile performs upload-time syntax checks and compiles an ESM plugin once.
|
||||
// All Sobek-specific module and runtime handling is intentionally kept here.
|
||||
func Compile(source string, options Options) (*Engine, error) {
|
||||
if match := forbiddenSyntax.FindString(sourceWithoutCommentsAndStrings(source)); match != "" {
|
||||
return nil, fmt.Errorf("unsupported plugin syntax %q: plugins must be synchronous and cannot import modules", strings.TrimSpace(match))
|
||||
}
|
||||
|
||||
resolve := func(_ any, specifier string) (sobek.ModuleRecord, error) {
|
||||
return nil, fmt.Errorf("plugin imports are disabled: %s", specifier)
|
||||
}
|
||||
// Plugin source is untrusted; without this option a sourceMappingURL
|
||||
// comment makes the parser read arbitrary server files via os.ReadFile.
|
||||
module, err := sobek.ParseModule(options.Key+".js", source, resolve, parser.WithDisableSourceMaps)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("compile plugin: %w", err)
|
||||
}
|
||||
if err = module.Link(); err != nil {
|
||||
return nil, fmt.Errorf("link plugin: %w", err)
|
||||
}
|
||||
|
||||
timeout := options.Timeout
|
||||
if timeout <= 0 {
|
||||
timeout = DefaultCallTimeout
|
||||
}
|
||||
concurrency := options.Concurrency
|
||||
if concurrency <= 0 {
|
||||
concurrency = DefaultConcurrency
|
||||
}
|
||||
now := options.Now
|
||||
if now == nil {
|
||||
now = time.Now
|
||||
}
|
||||
|
||||
engine := &Engine{
|
||||
key: options.Key,
|
||||
version: options.Version,
|
||||
timeout: timeout,
|
||||
now: now,
|
||||
log: options.Log,
|
||||
module: module,
|
||||
semaphore: make(chan struct{}, concurrency),
|
||||
}
|
||||
instance, err := engine.newRuntime(context.Background())
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
instance.logContext.context = nil
|
||||
engine.pool.Put(instance)
|
||||
return engine, nil
|
||||
}
|
||||
|
||||
// Export returns one module export without exposing Sobek values outside the
|
||||
// engine boundary. It is used for declarative exports such as meta.
|
||||
func (e *Engine) Export(ctx context.Context, exportName string) (result any, err error) {
|
||||
select {
|
||||
case e.semaphore <- struct{}{}:
|
||||
defer func() { <-e.semaphore }()
|
||||
case <-ctx.Done():
|
||||
return nil, ctx.Err()
|
||||
}
|
||||
|
||||
instance, err := e.getRuntime(ctx)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
reusable := true
|
||||
defer func() {
|
||||
instance.runtime.ClearInterrupt()
|
||||
instance.logContext.context = nil
|
||||
if reusable {
|
||||
e.pool.Put(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()
|
||||
defer func() {
|
||||
if recovered := recover(); recovered != nil {
|
||||
reusable = false
|
||||
switch value := recovered.(type) {
|
||||
case *sobek.InterruptedError:
|
||||
err = fmt.Errorf("plugin %s@%s export %s interrupted: %v", e.key, e.version, exportName, value.Value())
|
||||
case *sobek.Exception:
|
||||
err = fmt.Errorf("plugin %s@%s export %s failed: %v", e.key, e.version, exportName, value)
|
||||
default:
|
||||
panic(recovered)
|
||||
}
|
||||
}
|
||||
}()
|
||||
value := instance.module.GetBindingValue(exportName)
|
||||
if value == nil || sobek.IsUndefined(value) {
|
||||
return nil, fmt.Errorf("plugin export %q not found", exportName)
|
||||
}
|
||||
return value.Export(), nil
|
||||
}
|
||||
|
||||
// HasExport reports whether a module export exists. Optional contract hooks
|
||||
// should be detected with this method instead of relying on engine errors.
|
||||
func (e *Engine) HasExport(ctx context.Context, exportName string) (bool, error) {
|
||||
select {
|
||||
case e.semaphore <- struct{}{}:
|
||||
defer func() { <-e.semaphore }()
|
||||
case <-ctx.Done():
|
||||
return false, ctx.Err()
|
||||
}
|
||||
instance, err := e.getRuntime(ctx)
|
||||
if err != nil {
|
||||
return false, err
|
||||
}
|
||||
defer func() {
|
||||
instance.logContext.context = nil
|
||||
e.pool.Put(instance)
|
||||
}()
|
||||
value := instance.module.GetBindingValue(exportName)
|
||||
return value != nil && !sobek.IsUndefined(value), nil
|
||||
}
|
||||
|
||||
// HasCallablePath reports whether an exported value, or a nested member below
|
||||
// it, exists and is callable.
|
||||
func (e *Engine) HasCallablePath(ctx context.Context, exportName string, members ...string) (found bool, err error) {
|
||||
select {
|
||||
case e.semaphore <- struct{}{}:
|
||||
defer func() { <-e.semaphore }()
|
||||
case <-ctx.Done():
|
||||
return false, ctx.Err()
|
||||
}
|
||||
instance, err := e.getRuntime(ctx)
|
||||
if err != nil {
|
||||
return false, err
|
||||
}
|
||||
reusable := true
|
||||
defer func() {
|
||||
instance.runtime.ClearInterrupt()
|
||||
instance.logContext.context = nil
|
||||
if reusable {
|
||||
e.pool.Put(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()
|
||||
defer func() {
|
||||
if recovered := recover(); recovered != nil {
|
||||
reusable = false
|
||||
hookName := strings.Join(append([]string{exportName}, members...), ".")
|
||||
switch value := recovered.(type) {
|
||||
case *sobek.InterruptedError:
|
||||
err = fmt.Errorf("plugin %s@%s hook %s inspection interrupted: %v", e.key, e.version, hookName, value.Value())
|
||||
case *sobek.Exception:
|
||||
err = fmt.Errorf("plugin %s@%s hook %s inspection failed: %v", e.key, e.version, hookName, value)
|
||||
default:
|
||||
panic(recovered)
|
||||
}
|
||||
}
|
||||
}()
|
||||
value, _, found := resolveExportPath(instance, exportName, members)
|
||||
if !found {
|
||||
return false, nil
|
||||
}
|
||||
_, callable := sobek.AssertFunction(value)
|
||||
return callable, nil
|
||||
}
|
||||
|
||||
// Call invokes one named module export and returns its JSON-compatible value.
|
||||
func (e *Engine) Call(ctx context.Context, exportName string, args ...any) (result any, err error) {
|
||||
return e.call(ctx, 0, exportName, nil, args...)
|
||||
}
|
||||
|
||||
// CallMember invokes a function stored on an exported object, such as a
|
||||
// renderer in the renderers export.
|
||||
func (e *Engine) CallMember(ctx context.Context, exportName, memberName string, args ...any) (result any, err error) {
|
||||
return e.call(ctx, 0, exportName, []string{memberName}, args...)
|
||||
}
|
||||
|
||||
// CallPath invokes a function nested below an exported object. It is used for
|
||||
// protocol hooks such as protocols.openai_responses.renderEvents.
|
||||
func (e *Engine) CallPath(ctx context.Context, exportName string, members []string, args ...any) (result any, err error) {
|
||||
return e.call(ctx, 0, exportName, members, args...)
|
||||
}
|
||||
|
||||
// CallPathWithAdmissionTimeout gives long-lived observers a separate bound for
|
||||
// waiting on JavaScript capacity. Once admitted, the hook receives the
|
||||
// engine's full execution timeout instead of inheriting time already spent in
|
||||
// the semaphore queue.
|
||||
func (e *Engine) CallPathWithAdmissionTimeout(
|
||||
ctx context.Context,
|
||||
admissionTimeout time.Duration,
|
||||
exportName string,
|
||||
members []string,
|
||||
args ...any,
|
||||
) (result any, err error) {
|
||||
return e.call(ctx, admissionTimeout, exportName, members, args...)
|
||||
}
|
||||
|
||||
func (e *Engine) call(
|
||||
ctx context.Context,
|
||||
admissionTimeout time.Duration,
|
||||
exportName string,
|
||||
members []string,
|
||||
args ...any,
|
||||
) (result any, err error) {
|
||||
if err = e.acquireCallSlot(ctx, admissionTimeout); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer func() { <-e.semaphore }()
|
||||
|
||||
instance, err := e.getRuntime(ctx)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
reusable := true
|
||||
defer func() {
|
||||
instance.runtime.ClearInterrupt()
|
||||
instance.logContext.context = nil
|
||||
if reusable {
|
||||
e.pool.Put(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()
|
||||
defer func() {
|
||||
if recovered := recover(); recovered != nil {
|
||||
reusable = false
|
||||
switch value := recovered.(type) {
|
||||
case *sobek.InterruptedError:
|
||||
err = fmt.Errorf("plugin %s@%s hook %s interrupted: %v", e.key, e.version, hookName, value.Value())
|
||||
case *sobek.Exception:
|
||||
wrapped := fmt.Errorf("plugin %s@%s hook %s failed: %v", e.key, e.version, hookName, value)
|
||||
err = hookErrorFromException(hookName, value, wrapped)
|
||||
default:
|
||||
panic(recovered)
|
||||
}
|
||||
}
|
||||
}()
|
||||
|
||||
value, resolvedHookName, found := resolveExportPath(instance, exportName, members)
|
||||
hookName = resolvedHookName
|
||||
if !found {
|
||||
if len(members) == 0 {
|
||||
return nil, fmt.Errorf("plugin export %q not found", exportName)
|
||||
}
|
||||
return nil, fmt.Errorf("plugin hook %q not found", hookName)
|
||||
}
|
||||
if value == nil || sobek.IsUndefined(value) {
|
||||
return nil, fmt.Errorf("plugin export %q not found", exportName)
|
||||
}
|
||||
callable, ok := sobek.AssertFunction(value)
|
||||
if !ok {
|
||||
return nil, fmt.Errorf("plugin hook %q is not a function", hookName)
|
||||
}
|
||||
|
||||
callArgs := make([]sobek.Value, len(args))
|
||||
for i, arg := range args {
|
||||
callArgs[i] = instance.runtime.ToValue(arg)
|
||||
}
|
||||
|
||||
value, err = callable(sobek.Undefined(), callArgs...)
|
||||
if err != nil {
|
||||
var interrupted *sobek.InterruptedError
|
||||
if errors.As(err, &interrupted) {
|
||||
reusable = false
|
||||
return nil, fmt.Errorf("plugin %s@%s hook %s failed: %w", e.key, e.version, hookName, err)
|
||||
}
|
||||
wrapped := fmt.Errorf("plugin %s@%s hook %s failed: %w", e.key, e.version, hookName, err)
|
||||
var exc *sobek.Exception
|
||||
if errors.As(err, &exc) {
|
||||
return nil, hookErrorFromException(hookName, exc, wrapped)
|
||||
}
|
||||
return nil, wrapped
|
||||
}
|
||||
return value.Export(), nil
|
||||
}
|
||||
|
||||
func (e *Engine) acquireCallSlot(ctx context.Context, admissionTimeout time.Duration) error {
|
||||
if admissionTimeout <= 0 {
|
||||
select {
|
||||
case e.semaphore <- struct{}{}:
|
||||
return nil
|
||||
case <-ctx.Done():
|
||||
return ctx.Err()
|
||||
}
|
||||
}
|
||||
|
||||
timer := time.NewTimer(admissionTimeout)
|
||||
defer timer.Stop()
|
||||
select {
|
||||
case e.semaphore <- struct{}{}:
|
||||
return nil
|
||||
case <-ctx.Done():
|
||||
return ctx.Err()
|
||||
case <-timer.C:
|
||||
if err := ctx.Err(); err != nil {
|
||||
return err
|
||||
}
|
||||
return fmt.Errorf("%w: plugin %s@%s", ErrCallAdmissionTimeout, e.key, e.version)
|
||||
}
|
||||
}
|
||||
|
||||
func resolveExportPath(instance *runtimeInstance, exportName string, members []string) (sobek.Value, string, bool) {
|
||||
value := instance.module.GetBindingValue(exportName)
|
||||
hookName := exportName
|
||||
if value == nil || sobek.IsUndefined(value) || sobek.IsNull(value) {
|
||||
return nil, hookName, false
|
||||
}
|
||||
for _, member := range members {
|
||||
hookName += "." + member
|
||||
object := value.ToObject(instance.runtime)
|
||||
own := slices.Contains(object.GetOwnPropertyNames(), member)
|
||||
if !own {
|
||||
return nil, hookName, false
|
||||
}
|
||||
value = object.Get(member)
|
||||
if value == nil || sobek.IsUndefined(value) || sobek.IsNull(value) {
|
||||
return nil, hookName, false
|
||||
}
|
||||
}
|
||||
return value, hookName, true
|
||||
}
|
||||
|
||||
func (e *Engine) getRuntime(ctx context.Context) (*runtimeInstance, error) {
|
||||
if pooled := e.pool.Get(); pooled != nil {
|
||||
instance := pooled.(*runtimeInstance)
|
||||
instance.logContext.context = ctx
|
||||
return instance, nil
|
||||
}
|
||||
return e.newRuntime(ctx)
|
||||
}
|
||||
|
||||
func (e *Engine) newRuntime(ctx context.Context) (instance *runtimeInstance, err error) {
|
||||
runtime := sobek.New()
|
||||
logContext := &runtimeLogContext{context: ctx}
|
||||
logOutput := e.log
|
||||
if logOutput == nil {
|
||||
logOutput = func(message string) {
|
||||
logger.LogDebug(logContext.context, "task_plugin subsystem=runtime event=console message=%q", message)
|
||||
}
|
||||
}
|
||||
if err := injectGlobals(runtime, func() string {
|
||||
return fmt.Sprintf("[plugin:%s@%s]", e.key, e.version)
|
||||
}, e.now, logOutput); err != nil {
|
||||
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()
|
||||
defer func() {
|
||||
if recovered := recover(); recovered != nil {
|
||||
if interrupted, ok := recovered.(*sobek.InterruptedError); ok {
|
||||
instance = nil
|
||||
err = fmt.Errorf("initialize plugin %s@%s: %v", e.key, e.version, interrupted.Value())
|
||||
return
|
||||
}
|
||||
panic(recovered)
|
||||
}
|
||||
}()
|
||||
promise := runtime.CyclicModuleRecordEvaluate(e.module, func(_ any, specifier string) (sobek.ModuleRecord, error) {
|
||||
return nil, fmt.Errorf("plugin imports are disabled: %s", specifier)
|
||||
})
|
||||
if promise.State() != sobek.PromiseStateFulfilled {
|
||||
return nil, fmt.Errorf("evaluate plugin: %v", promise.Result().Export())
|
||||
}
|
||||
return &runtimeInstance{
|
||||
runtime: runtime,
|
||||
module: runtime.GetModuleInstance(e.module),
|
||||
logContext: logContext,
|
||||
}, nil
|
||||
}
|
||||
|
||||
func sourceWithoutCommentsAndStrings(source string) string {
|
||||
var output strings.Builder
|
||||
output.Grow(len(source))
|
||||
quote := byte(0)
|
||||
escaped := false
|
||||
lineComment := false
|
||||
blockComment := false
|
||||
for i := 0; i < len(source); i++ {
|
||||
current := source[i]
|
||||
next := byte(0)
|
||||
if i+1 < len(source) {
|
||||
next = source[i+1]
|
||||
}
|
||||
if lineComment {
|
||||
if current == '\n' {
|
||||
lineComment = false
|
||||
output.WriteByte('\n')
|
||||
} else {
|
||||
output.WriteByte(' ')
|
||||
}
|
||||
continue
|
||||
}
|
||||
if blockComment {
|
||||
if current == '*' && next == '/' {
|
||||
blockComment = false
|
||||
output.WriteString(" ")
|
||||
i++
|
||||
} else {
|
||||
output.WriteByte(' ')
|
||||
}
|
||||
continue
|
||||
}
|
||||
if quote != 0 {
|
||||
output.WriteByte(' ')
|
||||
if escaped {
|
||||
escaped = false
|
||||
} else if current == '\\' {
|
||||
escaped = true
|
||||
} else if current == quote {
|
||||
quote = 0
|
||||
}
|
||||
continue
|
||||
}
|
||||
if current == '/' && next == '/' {
|
||||
lineComment = true
|
||||
output.WriteString(" ")
|
||||
i++
|
||||
continue
|
||||
}
|
||||
if current == '/' && next == '*' {
|
||||
blockComment = true
|
||||
output.WriteString(" ")
|
||||
i++
|
||||
continue
|
||||
}
|
||||
if current == '\'' || current == '"' || current == '`' {
|
||||
quote = current
|
||||
output.WriteByte(' ')
|
||||
continue
|
||||
}
|
||||
output.WriteByte(current)
|
||||
}
|
||||
return output.String()
|
||||
}
|
||||
@@ -0,0 +1,375 @@
|
||||
package jsplugin
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"errors"
|
||||
"strings"
|
||||
"testing"
|
||||
"time"
|
||||
"unicode/utf8"
|
||||
|
||||
"github.com/QuantumNous/new-api/common"
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
func TestEngineCallsESMExportWithInjectedUtils(t *testing.T) {
|
||||
t.Parallel()
|
||||
logs := make([]string, 0, 1)
|
||||
engine, err := Compile(`
|
||||
export function sign(ctx) {
|
||||
console.log("called", ctx.name);
|
||||
return {
|
||||
now: utils.unixNow(),
|
||||
digest: utils.hmacSHA256(ctx.message, ctx.secret),
|
||||
encoded: utils.base64(ctx.message),
|
||||
};
|
||||
}
|
||||
export const meta = { apiVersion: 1, key: "mock" };
|
||||
`, Options{
|
||||
Key: "mock", Version: "1.0.0",
|
||||
Now: func() time.Time { return time.Unix(1234, 0) },
|
||||
Log: func(message string) { logs = append(logs, message) },
|
||||
})
|
||||
require.NoError(t, err)
|
||||
|
||||
result, err := engine.Call(context.Background(), "sign", map[string]any{
|
||||
"name": "fixture", "message": "hello", "secret": "secret",
|
||||
})
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, map[string]any{
|
||||
"now": int64(1234), "digest": "88aab3ede8d3adf94d26ab90d3bafd4a2083070c3bcce9c014ee04a443847c0b", "encoded": "aGVsbG8=",
|
||||
}, result)
|
||||
assert.Equal(t, []string{"[plugin:mock@1.0.0] called fixture"}, logs)
|
||||
|
||||
meta, err := engine.Export(context.Background(), "meta")
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, map[string]any{"apiVersion": int64(1), "key": "mock"}, meta)
|
||||
}
|
||||
|
||||
func TestEngineConsoleLogUsesDebugLoggerAndRequestContext(t *testing.T) {
|
||||
previousDebug := common.DebugEnabled
|
||||
common.DebugEnabled = false
|
||||
t.Cleanup(func() { common.DebugEnabled = previousDebug })
|
||||
|
||||
var output bytes.Buffer
|
||||
common.LogWriterMu.Lock()
|
||||
previousWriter := gin.DefaultErrorWriter
|
||||
gin.DefaultErrorWriter = &output
|
||||
common.LogWriterMu.Unlock()
|
||||
t.Cleanup(func() {
|
||||
common.LogWriterMu.Lock()
|
||||
gin.DefaultErrorWriter = previousWriter
|
||||
common.LogWriterMu.Unlock()
|
||||
})
|
||||
|
||||
plugin, err := CompilePlugin(`
|
||||
export const meta = {
|
||||
apiVersion: 1,
|
||||
key: "console-debug",
|
||||
name: "Console debug",
|
||||
version: "1.2.3",
|
||||
author: {name: "Test"},
|
||||
models: ["debug-model"],
|
||||
fetchMode: "per_task",
|
||||
};
|
||||
export function run(label) {
|
||||
console.log("checkpoint", label);
|
||||
return true;
|
||||
}
|
||||
export function buildSubmitRequest() { return {url: "https://example.com"}; }
|
||||
export function parseSubmitResponse() { return {taskId: "one"}; }
|
||||
export function buildQueryRequest() { return {url: "https://example.com"}; }
|
||||
export function parseTaskResult() { return {status: "SUCCESS"}; }
|
||||
`, Options{})
|
||||
require.NoError(t, err)
|
||||
engine := plugin.Engine
|
||||
|
||||
disabledContext := context.WithValue(context.Background(), common.RequestIdKey, "plugin-console-disabled")
|
||||
_, err = engine.Call(disabledContext, "run", "disabled")
|
||||
require.NoError(t, err)
|
||||
assert.Empty(t, output.String())
|
||||
|
||||
common.DebugEnabled = true
|
||||
contextA := context.WithValue(context.Background(), common.RequestIdKey, "plugin-console-request-a")
|
||||
_, err = engine.Call(contextA, "run", "context-a")
|
||||
require.NoError(t, err)
|
||||
logA := output.String()
|
||||
output.Reset()
|
||||
|
||||
contextB := context.WithValue(context.Background(), common.RequestIdKey, "plugin-console-request-b")
|
||||
_, err = engine.Call(contextB, "run", "context-b")
|
||||
require.NoError(t, err)
|
||||
logB := output.String()
|
||||
output.Reset()
|
||||
|
||||
_, err = engine.Call(context.Background(), "run", "background")
|
||||
require.NoError(t, err)
|
||||
logBackground := output.String()
|
||||
|
||||
assert.Contains(t, logA, "plugin-console-request-a")
|
||||
assert.NotContains(t, logA, "plugin-console-request-b")
|
||||
assert.Contains(t, logA, "task_plugin subsystem=runtime event=console")
|
||||
assert.Contains(t, logA, "[plugin:console-debug@1.2.3] checkpoint context-a")
|
||||
assert.NotContains(t, logA, "disabled")
|
||||
|
||||
assert.Contains(t, logB, "plugin-console-request-b")
|
||||
assert.NotContains(t, logB, "plugin-console-request-a")
|
||||
assert.Contains(t, logB, "[plugin:console-debug@1.2.3] checkpoint context-b")
|
||||
|
||||
assert.Contains(t, logBackground, "| SYSTEM |")
|
||||
assert.NotContains(t, logBackground, "plugin-console-request-a")
|
||||
assert.NotContains(t, logBackground, "plugin-console-request-b")
|
||||
assert.Contains(t, logBackground, "[plugin:console-debug@1.2.3] checkpoint background")
|
||||
}
|
||||
|
||||
func TestCompileRejectsAsynchronousAndImportedPlugins(t *testing.T) {
|
||||
t.Parallel()
|
||||
for name, source := range map[string]string{
|
||||
"async": `export async function run() {}`,
|
||||
"static import": `import value from "dependency"; export function run() { return value; }`,
|
||||
"dynamic import": `export function run() { return import("dependency"); }`,
|
||||
"top-level await": `const value = await work(); export function run() { return value; }`,
|
||||
} {
|
||||
t.Run(name, func(t *testing.T) {
|
||||
_, err := Compile(source, Options{Key: "invalid"})
|
||||
require.Error(t, err)
|
||||
assert.Contains(t, err.Error(), "unsupported plugin syntax")
|
||||
})
|
||||
}
|
||||
|
||||
_, err := Compile(`export function run() { return "import async await"; }`, Options{Key: "valid"})
|
||||
require.NoError(t, err)
|
||||
}
|
||||
|
||||
func TestCompileIgnoresSourceMapDirectives(t *testing.T) {
|
||||
t.Parallel()
|
||||
// A sourceMappingURL comment must stay inert. Sobek's default loader
|
||||
// os.ReadFiles the referenced server path during Compile and turns any
|
||||
// load failure into a compile error, so an unresolvable path compiling
|
||||
// cleanly proves the loader is disabled.
|
||||
engine, err := Compile("export function run() { return 1; }\n//# sourceMappingURL=/nonexistent/leak-probe.map\n", Options{Key: "sourcemap"})
|
||||
require.NoError(t, err)
|
||||
|
||||
value, err := engine.Call(context.Background(), "run")
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, int64(1), value)
|
||||
}
|
||||
|
||||
func TestEngineInterruptsLongRunningHook(t *testing.T) {
|
||||
t.Parallel()
|
||||
engine, err := Compile(`export function run() { while (true) {} }`, Options{
|
||||
Key: "loop", Version: "1", Timeout: 20 * time.Millisecond,
|
||||
})
|
||||
require.NoError(t, err)
|
||||
|
||||
_, err = engine.Call(context.Background(), "run")
|
||||
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")
|
||||
}
|
||||
|
||||
func TestEngineHookErrorExtractsSanitizedJSMessage(t *testing.T) {
|
||||
t.Parallel()
|
||||
tests := []struct {
|
||||
name string
|
||||
source string
|
||||
wantMessage string
|
||||
wantLen int
|
||||
}{
|
||||
{
|
||||
name: "Error object",
|
||||
source: `export function run() { throw new Error("model is required"); }`,
|
||||
wantMessage: "model is required",
|
||||
},
|
||||
{
|
||||
name: "raw string throw",
|
||||
source: `export function run() { throw "raw string"; }`,
|
||||
wantMessage: "raw string",
|
||||
},
|
||||
{
|
||||
name: "truncates to 512 runes",
|
||||
source: `export function run() { throw new Error("x".repeat(2000)); }`,
|
||||
wantLen: 512,
|
||||
wantMessage: strings.Repeat("x", 512),
|
||||
},
|
||||
{
|
||||
name: "scrubs control characters",
|
||||
source: "export function run() { throw new Error(\"line1\\nline2\\x1b[31mred\"); }",
|
||||
wantMessage: "line1 line2 [31mred",
|
||||
},
|
||||
{
|
||||
name: "throwing message getter falls back without crashing",
|
||||
source: `export function run() { throw {get message() { throw {get message() { return "deep"; }}; }}; }`,
|
||||
wantMessage: "plugin hook failed",
|
||||
},
|
||||
}
|
||||
for _, testCase := range tests {
|
||||
t.Run(testCase.name, func(t *testing.T) {
|
||||
t.Parallel()
|
||||
engine, err := Compile(testCase.source, Options{Key: "diag", Version: "1.0.0"})
|
||||
require.NoError(t, err)
|
||||
|
||||
_, err = engine.Call(context.Background(), "run")
|
||||
require.Error(t, err)
|
||||
|
||||
var hookErr *HookError
|
||||
require.True(t, errors.As(err, &hookErr))
|
||||
assert.Equal(t, "run", hookErr.Hook)
|
||||
assert.Equal(t, testCase.wantMessage, hookErr.Message)
|
||||
if testCase.wantLen > 0 {
|
||||
assert.Equal(t, testCase.wantLen, utf8.RuneCountInString(hookErr.Message))
|
||||
}
|
||||
assert.Contains(t, hookErr.Error(), "plugin diag@1.0.0")
|
||||
assert.NotContains(t, hookErr.Message, "Error:")
|
||||
assert.NotContains(t, hookErr.Message, "plugin diag@")
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestEngineReportsProtocolAdmissionTimeoutSeparately(t *testing.T) {
|
||||
engine, err := Compile(`
|
||||
export const protocols = {
|
||||
responses: {renderEvents: function() { return {events: [], done: false}; }},
|
||||
};
|
||||
`, Options{Key: "admission", Version: "1", Concurrency: 1})
|
||||
require.NoError(t, err)
|
||||
|
||||
engine.semaphore <- struct{}{}
|
||||
_, err = engine.CallPathWithAdmissionTimeout(
|
||||
context.Background(),
|
||||
time.Nanosecond,
|
||||
"protocols",
|
||||
[]string{"responses", "renderEvents"},
|
||||
)
|
||||
<-engine.semaphore
|
||||
|
||||
require.ErrorIs(t, err, ErrCallAdmissionTimeout)
|
||||
result, err := engine.CallPathWithAdmissionTimeout(
|
||||
context.Background(),
|
||||
time.Second,
|
||||
"protocols",
|
||||
[]string{"responses", "renderEvents"},
|
||||
)
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, map[string]any{"events": []any{}, "done": false}, result)
|
||||
}
|
||||
|
||||
func TestEngineExportInterruptsLongRunningGetter(t *testing.T) {
|
||||
t.Parallel()
|
||||
engine, err := Compile(`
|
||||
export const meta = {
|
||||
apiVersion: 1,
|
||||
get name() { while (true) {} },
|
||||
};
|
||||
`, Options{Key: "meta-loop", Version: "1.0.0", Timeout: 20 * time.Millisecond})
|
||||
require.NoError(t, err)
|
||||
|
||||
_, err = engine.Export(context.Background(), "meta")
|
||||
require.ErrorContains(t, err, "export meta interrupted")
|
||||
}
|
||||
|
||||
func TestEngineExportReturnsThrownGetterError(t *testing.T) {
|
||||
t.Parallel()
|
||||
engine, err := Compile(`
|
||||
export const meta = {
|
||||
apiVersion: 1,
|
||||
get name() { throw new Error("getter failed"); },
|
||||
};
|
||||
`, Options{Key: "meta-throw", Version: "1.0.0"})
|
||||
require.NoError(t, err)
|
||||
|
||||
_, err = engine.Export(context.Background(), "meta")
|
||||
require.ErrorContains(t, err, "export meta failed")
|
||||
assert.Contains(t, err.Error(), "getter failed")
|
||||
}
|
||||
|
||||
func TestEngineNestedHooksRequireOwnProperties(t *testing.T) {
|
||||
t.Parallel()
|
||||
engine, err := Compile(`
|
||||
const inheritedRenderers = {
|
||||
inherited: function(value) { return value; },
|
||||
constructor: function(value) { return value; },
|
||||
toString: function(value) { return value; },
|
||||
["__proto__"]: function(value) { return value; },
|
||||
};
|
||||
export const renderers = Object.create(inheritedRenderers);
|
||||
renderers.own = function(value) { return {id: value.id}; };
|
||||
|
||||
const inheritedProtocol = {
|
||||
renderFinal: function(value) { return value; },
|
||||
};
|
||||
export const protocols = {
|
||||
responses: Object.create(inheritedProtocol),
|
||||
};
|
||||
`, Options{Key: "own-hooks", Version: "1.0.0"})
|
||||
require.NoError(t, err)
|
||||
|
||||
found, err := engine.HasCallablePath(context.Background(), "renderers", "own")
|
||||
require.NoError(t, err)
|
||||
assert.True(t, found)
|
||||
result, err := engine.CallMember(context.Background(), "renderers", "own", map[string]any{"id": "task-1"})
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, map[string]any{"id": "task-1"}, result)
|
||||
|
||||
for _, member := range []string{"inherited", "constructor", "toString", "__proto__"} {
|
||||
t.Run(member, func(t *testing.T) {
|
||||
found, err := engine.HasCallablePath(context.Background(), "renderers", member)
|
||||
require.NoError(t, err)
|
||||
assert.False(t, found)
|
||||
|
||||
_, err = engine.CallMember(context.Background(), "renderers", member)
|
||||
require.ErrorContains(t, err, "not found")
|
||||
})
|
||||
}
|
||||
|
||||
found, err = engine.HasCallablePath(context.Background(), "protocols", "responses", "renderFinal")
|
||||
require.NoError(t, err)
|
||||
assert.False(t, found)
|
||||
_, err = engine.CallPath(context.Background(), "protocols", []string{"responses", "renderFinal"})
|
||||
require.ErrorContains(t, err, "not found")
|
||||
}
|
||||
|
||||
func TestCompileInterruptsLongRunningInitialization(t *testing.T) {
|
||||
t.Parallel()
|
||||
_, err := Compile(`while (true) {}; export function run() {}`, Options{
|
||||
Key: "loop", Version: "1", Timeout: 20 * time.Millisecond,
|
||||
})
|
||||
require.Error(t, err)
|
||||
assert.Contains(t, err.Error(), "initialization timed out")
|
||||
}
|
||||
|
||||
func TestValidateRequestURL(t *testing.T) {
|
||||
t.Parallel()
|
||||
tests := []struct {
|
||||
name string
|
||||
requestURL string
|
||||
baseURL string
|
||||
allowedHosts []string
|
||||
wantError string
|
||||
}{
|
||||
{name: "same host", requestURL: "https://api.example.com/v1/task", baseURL: "https://api.example.com/v1"},
|
||||
{name: "default port", requestURL: "https://api.example.com:443/v1/task", baseURL: "https://api.example.com"},
|
||||
{name: "approved host", requestURL: "https://upload.example.com/task", baseURL: "https://api.example.com", allowedHosts: []string{"upload.example.com"}},
|
||||
{name: "approved host with port", requestURL: "http://192.168.1.10:8080/task", baseURL: "http://192.168.1.10:8000", allowedHosts: []string{"192.168.1.10:8080"}},
|
||||
{name: "approved host with explicit default port", requestURL: "https://upload.example.com/task", baseURL: "https://api.example.com", allowedHosts: []string{"upload.example.com:443"}},
|
||||
{name: "approved host port mismatch", requestURL: "http://192.168.1.10:9000/task", baseURL: "http://192.168.1.10:8000", allowedHosts: []string{"192.168.1.10:8080"}, wantError: "not allowed"},
|
||||
{name: "subdomain is not implicit", requestURL: "https://evil.api.example.com/task", baseURL: "https://api.example.com", wantError: "not allowed"},
|
||||
{name: "userinfo trick", requestURL: "https://api.example.com@evil.example/task", baseURL: "https://api.example.com", wantError: "not allowed"},
|
||||
{name: "relative URL", requestURL: "/v1/task", baseURL: "https://api.example.com", wantError: "absolute"},
|
||||
}
|
||||
for _, test := range tests {
|
||||
t.Run(test.name, func(t *testing.T) {
|
||||
err := ValidateRequestURL(test.requestURL, test.baseURL, test.allowedHosts)
|
||||
if test.wantError == "" {
|
||||
require.NoError(t, err)
|
||||
return
|
||||
}
|
||||
require.Error(t, err)
|
||||
assert.True(t, strings.Contains(err.Error(), test.wantError), err.Error())
|
||||
})
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,116 @@
|
||||
package jsplugin
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"reflect"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"github.com/QuantumNous/new-api/common"
|
||||
)
|
||||
|
||||
// Fixture describes deterministic calls into a plugin. Arguments and expected
|
||||
// values stay as JSON so fixtures remain portable across engine implementations.
|
||||
type Fixture struct {
|
||||
UnixNow *int64 `json:"unixNow"`
|
||||
Cases []FixtureCase `json:"cases"`
|
||||
}
|
||||
|
||||
type FixtureCase struct {
|
||||
Name string `json:"name"`
|
||||
Hook string `json:"hook"`
|
||||
Member string `json:"member,omitempty"`
|
||||
Path []string `json:"path,omitempty"`
|
||||
Args []json.RawMessage `json:"args"`
|
||||
Expected json.RawMessage `json:"expected"`
|
||||
ExpectedError string `json:"expectedError,omitempty"`
|
||||
}
|
||||
|
||||
type FixtureReport struct {
|
||||
Total int
|
||||
Passed int
|
||||
}
|
||||
|
||||
// ReplayFixture compiles a plugin and runs every fixture case in declaration
|
||||
// order. unixNow is fixed by the fixture to keep signing and timestamp hooks
|
||||
// reproducible.
|
||||
func ReplayFixture(ctx context.Context, source string, data []byte) (FixtureReport, error) {
|
||||
var fixture Fixture
|
||||
if err := common.Unmarshal(data, &fixture); err != nil {
|
||||
return FixtureReport{}, fmt.Errorf("decode fixture: %w", err)
|
||||
}
|
||||
if len(fixture.Cases) == 0 {
|
||||
return FixtureReport{}, fmt.Errorf("fixture must contain at least one case")
|
||||
}
|
||||
|
||||
options := Options{Key: "fixture", Version: "fixture"}
|
||||
if fixture.UnixNow != nil {
|
||||
fixed := time.Unix(*fixture.UnixNow, 0)
|
||||
options.Now = func() time.Time { return fixed }
|
||||
}
|
||||
plugin, err := NewRegistry().Register(source, options)
|
||||
if err != nil {
|
||||
return FixtureReport{}, fmt.Errorf("compile plugin: %w", err)
|
||||
}
|
||||
|
||||
report := FixtureReport{Total: len(fixture.Cases)}
|
||||
for index, testCase := range fixture.Cases {
|
||||
caseName := strings.TrimSpace(testCase.Name)
|
||||
if caseName == "" {
|
||||
caseName = fmt.Sprintf("case %d", index+1)
|
||||
}
|
||||
if strings.TrimSpace(testCase.Hook) == "" {
|
||||
return report, fmt.Errorf("%s: hook is required", caseName)
|
||||
}
|
||||
|
||||
args := make([]any, len(testCase.Args))
|
||||
for argumentIndex, raw := range testCase.Args {
|
||||
if err = common.Unmarshal(raw, &args[argumentIndex]); err != nil {
|
||||
return report, fmt.Errorf("%s: decode argument %d: %w", caseName, argumentIndex+1, err)
|
||||
}
|
||||
}
|
||||
var result any
|
||||
if testCase.Member != "" && len(testCase.Path) > 0 {
|
||||
return report, fmt.Errorf("%s: member and path are mutually exclusive", caseName)
|
||||
}
|
||||
if len(testCase.Path) > 0 {
|
||||
result, err = plugin.Engine.CallPath(ctx, testCase.Hook, testCase.Path, args...)
|
||||
} else if testCase.Member == "" {
|
||||
result, err = plugin.Engine.Call(ctx, testCase.Hook, args...)
|
||||
} else {
|
||||
result, err = plugin.Engine.CallMember(ctx, testCase.Hook, testCase.Member, args...)
|
||||
}
|
||||
if testCase.ExpectedError != "" {
|
||||
if err == nil || !strings.Contains(err.Error(), testCase.ExpectedError) {
|
||||
return report, fmt.Errorf("%s: expected error containing %q, got %v", caseName, testCase.ExpectedError, err)
|
||||
}
|
||||
report.Passed++
|
||||
continue
|
||||
}
|
||||
if err != nil {
|
||||
return report, fmt.Errorf("%s: %w", caseName, err)
|
||||
}
|
||||
if len(testCase.Expected) == 0 {
|
||||
return report, fmt.Errorf("%s: expected is required when expectedError is empty", caseName)
|
||||
}
|
||||
var expected any
|
||||
if err = common.Unmarshal(testCase.Expected, &expected); err != nil {
|
||||
return report, fmt.Errorf("%s: decode expected value: %w", caseName, err)
|
||||
}
|
||||
actualData, marshalErr := common.Marshal(result)
|
||||
if marshalErr != nil {
|
||||
return report, fmt.Errorf("%s: encode actual value: %w", caseName, marshalErr)
|
||||
}
|
||||
var actual any
|
||||
if err = common.Unmarshal(actualData, &actual); err != nil {
|
||||
return report, fmt.Errorf("%s: normalize actual value: %w", caseName, err)
|
||||
}
|
||||
if !reflect.DeepEqual(expected, actual) {
|
||||
return report, fmt.Errorf("%s: result mismatch: expected %s, got %s", caseName, testCase.Expected, actualData)
|
||||
}
|
||||
report.Passed++
|
||||
}
|
||||
return report, nil
|
||||
}
|
||||
@@ -0,0 +1,49 @@
|
||||
package jsplugin
|
||||
|
||||
import (
|
||||
"context"
|
||||
"testing"
|
||||
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
const fixturePlugin = `
|
||||
export const meta = { apiVersion: 1, key: "fixture", name: "Fixture", version: "1.0.0", author: {name: "Test"}, channelTypes: [1002], models: ["fixture-model"], fetchMode: "per_task", protocols: [{name: "openai_responses", supports: ["sync", "background"]}] };
|
||||
export function buildSubmitRequest(ctx) { return {url: ctx.baseUrl + "/submit"}; }
|
||||
export function parseSubmitResponse(ctx, resp) { return {taskId: resp.body.id}; }
|
||||
export function buildQueryRequest(ctx) { return {url: ctx.baseUrl + "/task"}; }
|
||||
export function parseTaskResult(ctx, body) { return body; }
|
||||
export function stamp(value) { return {value: value, now: utils.unixNow()}; }
|
||||
export function fail() { throw new Error("fixture failure"); }
|
||||
export const native = { compact: function(value) { return {id: value.task_id}; } };
|
||||
export const protocols = { openai_responses: {
|
||||
decodeRequest: function(value) { return {kind: "submit", model: value.model}; },
|
||||
renderFinal: function(ctx, task) { return task; }
|
||||
} };
|
||||
`
|
||||
|
||||
func TestReplayFixture(t *testing.T) {
|
||||
t.Parallel()
|
||||
report, err := ReplayFixture(context.Background(), fixturePlugin, []byte(`{
|
||||
"unixNow": 1700000000,
|
||||
"cases": [
|
||||
{"name":"deterministic time","hook":"stamp","args":["ok"],"expected":{"value":"ok","now":1700000000}},
|
||||
{"name":"member call","hook":"native","member":"compact","args":[{"task_id":"task-1"}],"expected":{"id":"task-1"}},
|
||||
{"name":"nested path call","hook":"protocols","path":["openai_responses","decodeRequest"],"args":[{"model":"video-1"}],"expected":{"kind":"submit","model":"video-1"}},
|
||||
{"name":"expected failure","hook":"fail","args":[],"expectedError":"fixture failure"}
|
||||
]
|
||||
}`))
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, FixtureReport{Total: 4, Passed: 4}, report)
|
||||
}
|
||||
|
||||
func TestReplayFixtureReportsMismatch(t *testing.T) {
|
||||
t.Parallel()
|
||||
report, err := ReplayFixture(context.Background(), fixturePlugin, []byte(`{
|
||||
"cases": [{"name":"wrong output","hook":"stamp","args":["ok"],"expected":{"value":"different"}}]
|
||||
}`))
|
||||
require.Error(t, err)
|
||||
assert.Contains(t, err.Error(), "wrong output: result mismatch")
|
||||
assert.Equal(t, FixtureReport{Total: 1}, report)
|
||||
}
|
||||
@@ -0,0 +1,126 @@
|
||||
package jsplugin
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"encoding/base64"
|
||||
"encoding/xml"
|
||||
"fmt"
|
||||
"io"
|
||||
"strings"
|
||||
"unicode"
|
||||
)
|
||||
|
||||
// MaxIconDataURIBytes bounds an uploaded plugin logo. It is deliberately
|
||||
// generous: the plugin source itself is capped at 1 MiB, so this only keeps a
|
||||
// logo from dwarfing the plugin it decorates.
|
||||
const MaxIconDataURIBytes = 512 * 1024
|
||||
|
||||
var pngSignature = []byte{0x89, 'P', 'N', 'G', '\r', '\n', 0x1a, '\n'}
|
||||
|
||||
// DecodeIconDataURI parses a plugin logo shipped as a sidecar icon.svg or
|
||||
// icon.png file and re-encoded by the uploader as
|
||||
// data:image/png;base64,... or data:image/svg+xml;base64,.... It returns the
|
||||
// media type and the raw image bytes. The admin UI renders logos only through
|
||||
// <img>, which already blocks scripts and external loads; the SVG checks here
|
||||
// are defense in depth so a hostile payload is rejected at upload instead of
|
||||
// being stored.
|
||||
func DecodeIconDataURI(icon string) (string, []byte, error) {
|
||||
if len(icon) > MaxIconDataURIBytes {
|
||||
return "", nil, fmt.Errorf("plugin icon must not exceed %d bytes", MaxIconDataURIBytes)
|
||||
}
|
||||
mediaType, payload, ok := strings.Cut(strings.TrimPrefix(icon, "data:"), ";base64,")
|
||||
if !strings.HasPrefix(icon, "data:") || !ok || (mediaType != "image/png" && mediaType != "image/svg+xml") {
|
||||
return "", nil, fmt.Errorf("plugin icon must be data:image/png;base64,... or data:image/svg+xml;base64,...")
|
||||
}
|
||||
decoded, err := base64.StdEncoding.Strict().DecodeString(payload)
|
||||
if err != nil {
|
||||
return "", nil, fmt.Errorf("plugin icon payload is not valid base64")
|
||||
}
|
||||
if err := ValidateIconImage(mediaType, decoded); err != nil {
|
||||
return "", nil, err
|
||||
}
|
||||
return mediaType, decoded, nil
|
||||
}
|
||||
|
||||
// ValidateIconImage checks raw PNG or SVG bytes for a plugin logo. PNG must
|
||||
// carry the PNG signature. SVG must be well-formed XML rooted at svg without
|
||||
// script or foreignObject elements, event-handler attributes, directives,
|
||||
// javascript: values, or absolute http(s) references.
|
||||
func ValidateIconImage(mediaType string, data []byte) error {
|
||||
if mediaType == "image/png" {
|
||||
if !bytes.HasPrefix(data, pngSignature) {
|
||||
return fmt.Errorf("plugin icon PNG payload is not a PNG image")
|
||||
}
|
||||
return nil
|
||||
}
|
||||
if mediaType != "image/svg+xml" {
|
||||
return fmt.Errorf("plugin icon must be image/png or image/svg+xml")
|
||||
}
|
||||
decoder := xml.NewDecoder(bytes.NewReader(data))
|
||||
rootSeen := false
|
||||
styleDepth := 0
|
||||
for {
|
||||
token, err := decoder.Token()
|
||||
if err == io.EOF {
|
||||
break
|
||||
}
|
||||
if err != nil {
|
||||
return fmt.Errorf("plugin icon SVG is not well-formed XML")
|
||||
}
|
||||
switch element := token.(type) {
|
||||
case xml.Directive:
|
||||
return fmt.Errorf("plugin icon SVG must not contain a DOCTYPE or other directives")
|
||||
case xml.ProcInst:
|
||||
if element.Target != "xml" {
|
||||
return fmt.Errorf("plugin icon SVG must not contain processing instructions")
|
||||
}
|
||||
case xml.StartElement:
|
||||
local := strings.ToLower(element.Name.Local)
|
||||
if !rootSeen {
|
||||
if local != "svg" {
|
||||
return fmt.Errorf("plugin icon SVG root element must be svg")
|
||||
}
|
||||
rootSeen = true
|
||||
}
|
||||
if local == "script" || local == "foreignobject" {
|
||||
return fmt.Errorf("plugin icon SVG must not contain %s elements", element.Name.Local)
|
||||
}
|
||||
if local == "style" {
|
||||
styleDepth++
|
||||
}
|
||||
for _, attr := range element.Attr {
|
||||
if attr.Name.Space == "xmlns" || (attr.Name.Space == "" && attr.Name.Local == "xmlns") {
|
||||
continue
|
||||
}
|
||||
if strings.HasPrefix(strings.ToLower(attr.Name.Local), "on") {
|
||||
return fmt.Errorf("plugin icon SVG must not contain event handler attributes")
|
||||
}
|
||||
if svgValueReferencesExternal(attr.Value) {
|
||||
return fmt.Errorf("plugin icon SVG must not reference scripts or external resources")
|
||||
}
|
||||
}
|
||||
case xml.EndElement:
|
||||
if strings.ToLower(element.Name.Local) == "style" && styleDepth > 0 {
|
||||
styleDepth--
|
||||
}
|
||||
case xml.CharData:
|
||||
if styleDepth > 0 && (svgValueReferencesExternal(string(element)) || strings.Contains(strings.ToLower(string(element)), "@import")) {
|
||||
return fmt.Errorf("plugin icon SVG must not reference scripts or external resources")
|
||||
}
|
||||
}
|
||||
}
|
||||
if !rootSeen {
|
||||
return fmt.Errorf("plugin icon SVG root element must be svg")
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func svgValueReferencesExternal(value string) bool {
|
||||
compact := strings.Map(func(character rune) rune {
|
||||
if unicode.IsSpace(character) || unicode.IsControl(character) {
|
||||
return -1
|
||||
}
|
||||
return unicode.ToLower(character)
|
||||
}, value)
|
||||
return strings.Contains(compact, "javascript:") || strings.Contains(compact, "http://") || strings.Contains(compact, "https://")
|
||||
}
|
||||
@@ -0,0 +1,26 @@
|
||||
package jsplugin
|
||||
|
||||
// asciiFold maps only 'A'–'Z' onto 'a'–'z'. Every other byte is unchanged.
|
||||
// Unicode case folding is intentionally not applied: U+212A (KELVIN SIGN)
|
||||
// must not become 'k' and impersonate an ASCII model name.
|
||||
func asciiFold(s string) string {
|
||||
var buf []byte
|
||||
for i := 0; i < len(s); i++ {
|
||||
c := s[i]
|
||||
if c >= 'A' && c <= 'Z' {
|
||||
if buf == nil {
|
||||
buf = []byte(s)
|
||||
}
|
||||
buf[i] = c + ('a' - 'A')
|
||||
}
|
||||
}
|
||||
if buf == nil {
|
||||
return s
|
||||
}
|
||||
return string(buf)
|
||||
}
|
||||
|
||||
// ASCIIFold is the exported form of asciiFold for consumers outside this package.
|
||||
func ASCIIFold(s string) string {
|
||||
return asciiFold(s)
|
||||
}
|
||||
@@ -0,0 +1,245 @@
|
||||
package jsplugin
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"testing"
|
||||
|
||||
"github.com/QuantumNous/new-api/common"
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
const (
|
||||
responsesDecodeOnly = `export const protocols = {openai_responses: {
|
||||
decodeRequest: function(ctx) { return ctx; }
|
||||
}};`
|
||||
responsesDecodeEvents = `export const protocols = {openai_responses: {
|
||||
decodeRequest: function(ctx) { return ctx; },
|
||||
renderEvents: function() { return {events: [], state: null, done: false}; }
|
||||
}};`
|
||||
responsesDecodeFinal = `export const protocols = {openai_responses: {
|
||||
decodeRequest: function(ctx) { return ctx; },
|
||||
renderFinal: function(ctx, task) { return task; }
|
||||
}};`
|
||||
responsesDecodeBoth = `export const protocols = {openai_responses: {
|
||||
decodeRequest: function(ctx) { return ctx; },
|
||||
renderEvents: function() { return {events: [], state: null, done: false}; },
|
||||
renderFinal: function(ctx, task) { return task; }
|
||||
}};`
|
||||
videoProtocolExport = `export const protocols = {openai_video: {
|
||||
decodeRequest: function(ctx) { return ctx; },
|
||||
render: function(ctx, task) { return task; }
|
||||
}};
|
||||
export function listArtifacts() { return []; }
|
||||
export function buildContentRequest() { return {}; }`
|
||||
)
|
||||
|
||||
func TestProtocolSupportsLoadErrors(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
protocols string
|
||||
exports string
|
||||
err string
|
||||
}{
|
||||
{
|
||||
name: "bare string",
|
||||
protocols: `["openai_responses"]`,
|
||||
exports: responsesDecodeBoth,
|
||||
err: `plugin acme protocol "openai_responses" must declare supports; replace the bare string with {name: "openai_responses", supports: [...]} choosing from "stream", "sync", "background"`,
|
||||
},
|
||||
{
|
||||
name: "object without supports",
|
||||
protocols: `[{name: "openai_responses"}]`,
|
||||
exports: responsesDecodeBoth,
|
||||
err: `plugin acme protocol "openai_responses" must declare supports; add supports: [...] choosing from "stream", "sync", "background"`,
|
||||
},
|
||||
{
|
||||
name: "supports sync but only renderEvents",
|
||||
protocols: `[{name: "openai_responses", supports: ["sync"]}]`,
|
||||
exports: responsesDecodeEvents,
|
||||
err: `plugin acme protocol "openai_responses" supports "sync" but does not export protocols.openai_responses.renderFinal; implement it or declare supports: ["stream"]`,
|
||||
},
|
||||
{
|
||||
name: "supports sync with only decodeRequest",
|
||||
protocols: `[{name: "openai_responses", supports: ["sync"]}]`,
|
||||
exports: responsesDecodeOnly,
|
||||
err: `plugin acme protocol "openai_responses" supports "sync" but does not export protocols.openai_responses.renderFinal; implement it`,
|
||||
},
|
||||
{
|
||||
name: "supports stream but also exports renderFinal",
|
||||
protocols: `[{name: "openai_responses", supports: ["stream"]}]`,
|
||||
exports: responsesDecodeBoth,
|
||||
err: `plugin acme protocol "openai_responses" exports protocols.openai_responses.renderFinal but no supported mode uses it; add "sync" or "background" to supports or remove the hook`,
|
||||
},
|
||||
{
|
||||
name: "supports sync and background but also exports renderEvents",
|
||||
protocols: `[{name: "openai_responses", supports: ["sync", "background"]}]`,
|
||||
exports: responsesDecodeBoth,
|
||||
err: `plugin acme protocol "openai_responses" exports protocols.openai_responses.renderEvents but no supported mode uses it; add "stream" to supports or remove the hook`,
|
||||
},
|
||||
{
|
||||
name: "empty supports",
|
||||
protocols: `[{name: "openai_responses", supports: []}]`,
|
||||
exports: responsesDecodeBoth,
|
||||
err: `plugin acme protocol "openai_responses" supports must contain at least one of "stream", "sync", "background"`,
|
||||
},
|
||||
{
|
||||
name: "duplicate supports",
|
||||
protocols: `[{name: "openai_responses", supports: ["stream", "stream"]}]`,
|
||||
exports: responsesDecodeBoth,
|
||||
err: `plugin acme protocol "openai_responses" supports must be unique`,
|
||||
},
|
||||
{
|
||||
name: "retrieve is not a mode",
|
||||
protocols: `[{name: "openai_responses", supports: ["retrieve"]}]`,
|
||||
exports: responsesDecodeBoth,
|
||||
err: `plugin acme protocol "openai_responses" has no mode "retrieve"; retrieval of a created response is always available and is never declared`,
|
||||
},
|
||||
{
|
||||
name: "openai_video forbids supports",
|
||||
protocols: `[{name: "openai_video", supports: ["stream"]}]`,
|
||||
exports: videoProtocolExport,
|
||||
err: `plugin acme protocol "openai_video" does not define modes; supports is not allowed`,
|
||||
},
|
||||
{
|
||||
name: "unknown protocol forbids supports",
|
||||
protocols: `[{name: "openai_custom", supports: ["stream"]}]`,
|
||||
err: `plugin acme protocol "openai_custom" does not define modes; supports is not allowed`,
|
||||
},
|
||||
}
|
||||
for _, testCase := range tests {
|
||||
t.Run(testCase.name, func(t *testing.T) {
|
||||
_, err := compileProtocolPlugin(t, "acme", `["model"]`, testCase.protocols, testCase.exports)
|
||||
require.ErrorContains(t, err, testCase.err)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestProtocolSupportsHappyPaths(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
models string
|
||||
protocols string
|
||||
exports string
|
||||
wantProtocols []ProtocolClaim
|
||||
}{
|
||||
{
|
||||
name: "stream only with renderEvents",
|
||||
models: `["model"]`,
|
||||
protocols: `[{name: "openai_responses", supports: ["stream"]}]`,
|
||||
exports: responsesDecodeEvents,
|
||||
wantProtocols: []ProtocolClaim{
|
||||
{Name: "openai_responses", Supports: []string{"stream"}, objectForm: true},
|
||||
},
|
||||
},
|
||||
{
|
||||
name: "sync and background with renderFinal",
|
||||
models: `["model"]`,
|
||||
protocols: `[{name: "openai_responses", supports: ["sync", "background"]}]`,
|
||||
exports: responsesDecodeFinal,
|
||||
wantProtocols: []ProtocolClaim{
|
||||
{Name: "openai_responses", Supports: []string{"sync", "background"}, objectForm: true},
|
||||
},
|
||||
},
|
||||
{
|
||||
name: "all modes normalize to table order",
|
||||
models: `["model"]`,
|
||||
protocols: `[{name: "openai_responses", supports: ["background", "stream", "sync"]}]`,
|
||||
exports: responsesDecodeBoth,
|
||||
wantProtocols: []ProtocolClaim{
|
||||
{Name: "openai_responses", Supports: []string{"stream", "sync", "background"}, objectForm: true},
|
||||
},
|
||||
},
|
||||
{
|
||||
name: "openai_video object without supports",
|
||||
models: `["gpt-5.5", "gpt-5.6"]`,
|
||||
protocols: `[{name: "openai_video", models: ["gpt-5.5"]}]`,
|
||||
exports: videoProtocolExport,
|
||||
wantProtocols: []ProtocolClaim{
|
||||
{Name: "openai_video", Models: []string{"gpt-5.5"}, objectForm: true},
|
||||
},
|
||||
},
|
||||
{
|
||||
name: "bare openai_video",
|
||||
models: `["model"]`,
|
||||
protocols: `["openai_video"]`,
|
||||
exports: videoProtocolExport,
|
||||
wantProtocols: []ProtocolClaim{
|
||||
{Name: "openai_video"},
|
||||
},
|
||||
},
|
||||
}
|
||||
for _, testCase := range tests {
|
||||
t.Run(testCase.name, func(t *testing.T) {
|
||||
plugin, err := compileProtocolPlugin(t, "acme", testCase.models, testCase.protocols, testCase.exports)
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, testCase.wantProtocols, plugin.Meta.Protocols)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestCloneMetaDeepCopiesSupports(t *testing.T) {
|
||||
registry := NewRegistry()
|
||||
_, err := registry.Register(protocolPluginSource(
|
||||
"acme",
|
||||
`["model"]`,
|
||||
`[{name: "openai_responses", supports: ["stream", "sync", "background"]}]`,
|
||||
responsesDecodeBoth,
|
||||
), Options{})
|
||||
require.NoError(t, err)
|
||||
|
||||
snapshot := registry.Snapshot()
|
||||
require.Len(t, snapshot.Override, 1)
|
||||
require.Len(t, snapshot.Override[0].Protocols, 1)
|
||||
snapshot.Override[0].Protocols[0].Supports[0] = "mutated"
|
||||
|
||||
plugin, ok := registry.Get("acme")
|
||||
require.True(t, ok)
|
||||
assert.Equal(t, []string{"stream", "sync", "background"}, plugin.Meta.Protocols[0].Supports)
|
||||
}
|
||||
|
||||
func TestMetaProtocolSupports(t *testing.T) {
|
||||
streamOnly, err := compileProtocolPlugin(t, "acme", `["model"]`,
|
||||
`[{name: "openai_responses", supports: ["stream"]}]`, responsesDecodeEvents)
|
||||
require.NoError(t, err)
|
||||
assert.True(t, streamOnly.Meta.ProtocolSupports("openai_responses", "stream"))
|
||||
assert.False(t, streamOnly.Meta.ProtocolSupports("openai_responses", "sync"))
|
||||
assert.False(t, streamOnly.Meta.ProtocolSupports("openai_responses", "background"))
|
||||
assert.False(t, streamOnly.Meta.ProtocolSupports("openai_responses", "retrieve"))
|
||||
assert.False(t, streamOnly.Meta.ProtocolSupports("openai_video", "stream"))
|
||||
assert.False(t, streamOnly.Meta.ProtocolSupports("missing", "stream"))
|
||||
|
||||
video, err := compileProtocolPlugin(t, "acme-video", `["model"]`, `["openai_video"]`, videoProtocolExport)
|
||||
require.NoError(t, err)
|
||||
assert.False(t, video.Meta.ProtocolSupports("openai_video", "stream"))
|
||||
}
|
||||
|
||||
func TestProtocolClaimMarshalEmitsSupportsInTableOrder(t *testing.T) {
|
||||
plugin, err := compileProtocolPlugin(t, "acme", `["model"]`,
|
||||
`[{name: "openai_responses", supports: ["background", "sync", "stream"]}]`, responsesDecodeBoth)
|
||||
require.NoError(t, err)
|
||||
encoded, err := common.Marshal(plugin.Meta.Protocols[0])
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, `{"name":"openai_responses","supports":["stream","sync","background"]}`, string(encoded))
|
||||
}
|
||||
|
||||
func compileProtocolPlugin(t *testing.T, key, models, protocols, exports string) (*LoadedPlugin, error) {
|
||||
t.Helper()
|
||||
return CompilePlugin(protocolPluginSource(key, models, protocols, exports), Options{})
|
||||
}
|
||||
|
||||
func protocolPluginSource(key, models, protocols, exports string) string {
|
||||
return fmt.Sprintf(`
|
||||
export const meta = {
|
||||
apiVersion: 1, key: %q, name: %q, version: "1.0.0",
|
||||
author: {name: "Test"},
|
||||
models: %s, fetchMode: "per_task",
|
||||
protocols: %s,
|
||||
};
|
||||
export function buildSubmitRequest() { return {}; }
|
||||
export function parseSubmitResponse() { return {}; }
|
||||
export function buildQueryRequest() { return {}; }
|
||||
export function parseTaskResult() { return {}; }
|
||||
%s
|
||||
`, key, key, models, protocols, exports)
|
||||
}
|
||||
File diff suppressed because it is too large.
Load diff
@@ -0,0 +1,105 @@
|
||||
package jsplugin
|
||||
|
||||
import (
|
||||
"testing"
|
||||
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
func TestRegistrySetDisabledFactoryKeysHidesFactoryPlugin(t *testing.T) {
|
||||
registry := NewRegistry()
|
||||
_, err := registry.RegisterFactory(routingTestPluginSource(
|
||||
"factory-off",
|
||||
50,
|
||||
`["factory-off-model"]`,
|
||||
`routes: [
|
||||
{method: "POST", path: "/vendor/factory-off", type: "submit", action: "generate", decode: "decodeVideo", render: "videoCreated"}
|
||||
],
|
||||
protocols: [{name: "openai_responses", supports: ["stream", "sync", "background"]}],`,
|
||||
`export const native = {
|
||||
decodeVideo: function(ctx) { return {kind: "submit", model: "factory-off-model", requestBody: ctx.body.value}; },
|
||||
videoCreated: function(ctx, task) { return task; }
|
||||
};
|
||||
`+routingProtocolExport("openai_responses"),
|
||||
), Options{})
|
||||
require.NoError(t, err)
|
||||
|
||||
_, ok := registry.Get("factory-off")
|
||||
require.True(t, ok)
|
||||
generation := registry.Generation()
|
||||
before := generation.Number
|
||||
_, ok = generation.LookupDeclaredRoute("POST", "/vendor/factory-off")
|
||||
require.True(t, ok)
|
||||
_, ok = generation.LookupEndpoint("POST", "/v1/responses", "factory-off-model")
|
||||
require.True(t, ok)
|
||||
|
||||
registry.SetDisabledFactoryKeys([]string{"factory-off"})
|
||||
_, ok = registry.Get("factory-off")
|
||||
assert.False(t, ok)
|
||||
generation = registry.Generation()
|
||||
assert.Greater(t, generation.Number, before)
|
||||
_, ok = generation.LookupDeclaredRoute("POST", "/vendor/factory-off")
|
||||
assert.False(t, ok)
|
||||
_, ok = generation.LookupEndpoint("POST", "/v1/responses", "factory-off-model")
|
||||
assert.False(t, ok)
|
||||
for _, plugin := range generation.Plugins() {
|
||||
assert.NotEqual(t, "factory-off", plugin.Meta.Key)
|
||||
}
|
||||
for _, route := range generation.Routes() {
|
||||
assert.NotEqual(t, "factory-off", route.Plugin.Meta.Key)
|
||||
}
|
||||
|
||||
snapshot := registry.Snapshot()
|
||||
require.Len(t, snapshot.Factory, 1)
|
||||
assert.Equal(t, "factory-off", snapshot.Factory[0].Key)
|
||||
assert.Equal(t, []string{"factory-off"}, snapshot.DisabledFactory)
|
||||
|
||||
registry.SetDisabledFactoryKeys(nil)
|
||||
plugin, ok := registry.Get("factory-off")
|
||||
require.True(t, ok)
|
||||
assert.Equal(t, "factory-off", plugin.Meta.Key)
|
||||
assert.Empty(t, registry.Snapshot().DisabledFactory)
|
||||
_, ok = registry.Generation().LookupDeclaredRoute("POST", "/vendor/factory-off")
|
||||
assert.True(t, ok)
|
||||
}
|
||||
|
||||
func TestRegistrySetDisabledFactoryKeysLeavesEnabledOverrideServing(t *testing.T) {
|
||||
registry := NewRegistry()
|
||||
require.NoError(t, registerTestPlugin(registry, "1.0.0-factory", true))
|
||||
require.NoError(t, registerTestPlugin(registry, "1.0.0-override", false))
|
||||
|
||||
registry.SetDisabledFactoryKeys([]string{"test"})
|
||||
plugin, ok := registry.Get("test")
|
||||
require.True(t, ok)
|
||||
assert.Equal(t, "1.0.0-override", plugin.Meta.Version)
|
||||
|
||||
snapshot := registry.Snapshot()
|
||||
require.Len(t, snapshot.Factory, 1)
|
||||
assert.Equal(t, "1.0.0-factory", snapshot.Factory[0].Version)
|
||||
assert.Equal(t, []string{"test"}, snapshot.DisabledFactory)
|
||||
}
|
||||
|
||||
func TestRegistrySetDisabledFactoryKeysEmptySetDoesNotBumpGeneration(t *testing.T) {
|
||||
registry := NewRegistry()
|
||||
require.NoError(t, registerTestPlugin(registry, "1.0.0-factory", true))
|
||||
before := registry.Generation().Number
|
||||
|
||||
registry.SetDisabledFactoryKeys(nil)
|
||||
assert.Equal(t, before, registry.Generation().Number)
|
||||
registry.SetDisabledFactoryKeys([]string{})
|
||||
assert.Equal(t, before, registry.Generation().Number)
|
||||
registry.SetDisabledFactoryKeys([]string{"", " "})
|
||||
assert.Equal(t, before, registry.Generation().Number)
|
||||
}
|
||||
|
||||
func TestRegistrySnapshotDisabledFactoryIsSorted(t *testing.T) {
|
||||
registry := NewRegistry()
|
||||
require.NoError(t, registerTestPlugin(registry, "1.0.0", true))
|
||||
_, err := registry.RegisterFactory(routingTestPluginSource("other", 0, `["other-model"]`, "", ""), Options{})
|
||||
require.NoError(t, err)
|
||||
|
||||
registry.SetDisabledFactoryKeys([]string{"other", "test"})
|
||||
assert.Equal(t, []string{"other", "test"}, registry.Snapshot().DisabledFactory)
|
||||
require.Len(t, registry.Snapshot().Factory, 2)
|
||||
}
|
||||
@@ -0,0 +1,141 @@
|
||||
package jsplugin
|
||||
|
||||
import (
|
||||
"testing"
|
||||
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
func TestRegistrySetEnabledHidesFactoryAndOverride(t *testing.T) {
|
||||
registry := NewRegistry()
|
||||
_, err := registry.RegisterFactory(routingTestPluginSource(
|
||||
"master-factory",
|
||||
50,
|
||||
`["master-factory-model"]`,
|
||||
`routes: [
|
||||
{method: "POST", path: "/vendor/master-factory", type: "submit", action: "generate", decode: "decodeVideo", render: "videoCreated"}
|
||||
],
|
||||
protocols: [{name: "openai_responses", supports: ["stream", "sync", "background"]}],`,
|
||||
`export const native = {
|
||||
decodeVideo: function(ctx) { return {kind: "submit", model: "master-factory-model", requestBody: ctx.body.value}; },
|
||||
videoCreated: function(ctx, task) { return task; }
|
||||
};
|
||||
`+routingProtocolExport("openai_responses"),
|
||||
), Options{})
|
||||
require.NoError(t, err)
|
||||
_, err = registry.Register(routingTestPluginSource(
|
||||
"master-override",
|
||||
51,
|
||||
`["master-override-model"]`,
|
||||
`routes: [
|
||||
{method: "POST", path: "/vendor/master-override", type: "submit", action: "generate", decode: "decodeVideo", render: "videoCreated"}
|
||||
],
|
||||
protocols: [{name: "openai_responses", supports: ["stream", "sync", "background"]}],`,
|
||||
`export const native = {
|
||||
decodeVideo: function(ctx) { return {kind: "submit", model: "master-override-model", requestBody: ctx.body.value}; },
|
||||
videoCreated: function(ctx, task) { return task; }
|
||||
};
|
||||
`+routingProtocolExport("openai_responses"),
|
||||
), Options{})
|
||||
require.NoError(t, err)
|
||||
|
||||
_, ok := registry.Get("master-factory")
|
||||
require.True(t, ok)
|
||||
_, ok = registry.Get("master-override")
|
||||
require.True(t, ok)
|
||||
before := registry.Generation().Number
|
||||
|
||||
registry.SetEnabled(false)
|
||||
_, ok = registry.Get("master-factory")
|
||||
assert.False(t, ok)
|
||||
_, ok = registry.Get("master-override")
|
||||
assert.False(t, ok)
|
||||
generation := registry.Generation()
|
||||
assert.Greater(t, generation.Number, before)
|
||||
assert.Empty(t, generation.Plugins())
|
||||
assert.Empty(t, generation.Routes())
|
||||
_, ok = generation.LookupDeclaredRoute("POST", "/vendor/master-factory")
|
||||
assert.False(t, ok)
|
||||
_, ok = generation.LookupDeclaredRoute("POST", "/vendor/master-override")
|
||||
assert.False(t, ok)
|
||||
_, ok = generation.LookupEndpoint("POST", "/v1/responses", "master-factory-model")
|
||||
assert.False(t, ok)
|
||||
_, ok = generation.LookupEndpoint("POST", "/v1/responses", "master-override-model")
|
||||
assert.False(t, ok)
|
||||
|
||||
snapshot := registry.Snapshot()
|
||||
require.Len(t, snapshot.Factory, 1)
|
||||
assert.Equal(t, "master-factory", snapshot.Factory[0].Key)
|
||||
require.Len(t, snapshot.Override, 1)
|
||||
assert.Equal(t, "master-override", snapshot.Override[0].Key)
|
||||
}
|
||||
|
||||
func TestRegistryMutationsWhileDisabledDoNotResurrectEndpoints(t *testing.T) {
|
||||
registry := NewRegistry()
|
||||
require.NoError(t, registerTestPlugin(registry, "1.0.0-factory", true))
|
||||
registry.SetEnabled(false)
|
||||
_, ok := registry.Get("test")
|
||||
require.False(t, ok)
|
||||
|
||||
require.NoError(t, registerTestPlugin(registry, "1.0.0-override", false))
|
||||
registry.SetDisabledFactoryKeys([]string{"other"})
|
||||
_, ok = registry.Get("test")
|
||||
assert.False(t, ok)
|
||||
assert.Empty(t, registry.Generation().Plugins())
|
||||
|
||||
override := registry.OverridePlugins()
|
||||
require.Contains(t, override, "test")
|
||||
require.NoError(t, registry.ReplaceOverrides([]*LoadedPlugin{override["test"]}))
|
||||
_, ok = registry.Get("test")
|
||||
assert.False(t, ok)
|
||||
assert.Empty(t, registry.Generation().Plugins())
|
||||
}
|
||||
|
||||
func TestRegistrySetEnabledTrueRestoresFactoryAndOverride(t *testing.T) {
|
||||
registry := NewRegistry()
|
||||
require.NoError(t, registerTestPlugin(registry, "1.0.0-factory", true))
|
||||
registry.SetEnabled(false)
|
||||
require.NoError(t, registerTestPlugin(registry, "1.0.0-override", false))
|
||||
_, ok := registry.Get("test")
|
||||
require.False(t, ok)
|
||||
|
||||
registry.SetEnabled(true)
|
||||
plugin, ok := registry.Get("test")
|
||||
require.True(t, ok)
|
||||
assert.Equal(t, "1.0.0-override", plugin.Meta.Version)
|
||||
snapshot := registry.Snapshot()
|
||||
require.Len(t, snapshot.Factory, 1)
|
||||
require.Len(t, snapshot.Override, 1)
|
||||
}
|
||||
|
||||
func TestRegistrySetEnabledNoOpDoesNotBumpGeneration(t *testing.T) {
|
||||
registry := NewRegistry()
|
||||
require.NoError(t, registerTestPlugin(registry, "1.0.0", true))
|
||||
before := registry.Generation().Number
|
||||
|
||||
registry.SetEnabled(true)
|
||||
assert.Equal(t, before, registry.Generation().Number)
|
||||
}
|
||||
|
||||
func TestRegistryMasterSwitchIsOrthogonalToLayerFlags(t *testing.T) {
|
||||
registry := NewRegistry()
|
||||
require.NoError(t, registerTestPlugin(registry, "1.0.0-factory", true))
|
||||
|
||||
registry.SetDisabledFactoryKeys([]string{"test"})
|
||||
_, ok := registry.Get("test")
|
||||
assert.False(t, ok)
|
||||
|
||||
registry.SetEnabled(false)
|
||||
_, ok = registry.Get("test")
|
||||
assert.False(t, ok)
|
||||
|
||||
registry.SetEnabled(true)
|
||||
_, ok = registry.Get("test")
|
||||
assert.False(t, ok)
|
||||
|
||||
registry.SetDisabledFactoryKeys(nil)
|
||||
plugin, ok := registry.Get("test")
|
||||
require.True(t, ok)
|
||||
assert.Equal(t, "1.0.0-factory", plugin.Meta.Version)
|
||||
}
|
||||
File diff suppressed because it is too large.
Load diff
@@ -0,0 +1,43 @@
|
||||
package jsplugin
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"net"
|
||||
"net/url"
|
||||
"strings"
|
||||
)
|
||||
|
||||
// ValidateRequestURL prevents plugins from directing a channel credential to
|
||||
// hosts other than the configured base URL or an administrator-approved host.
|
||||
func ValidateRequestURL(requestURL, baseURL string, allowedHosts []string) error {
|
||||
request, err := url.Parse(requestURL)
|
||||
if err != nil || request.Scheme == "" || request.Host == "" {
|
||||
return fmt.Errorf("plugin request URL must be absolute")
|
||||
}
|
||||
base, err := url.Parse(baseURL)
|
||||
if err != nil || base.Host == "" {
|
||||
return fmt.Errorf("channel base URL is invalid")
|
||||
}
|
||||
requestHost := canonicalHost(request)
|
||||
if requestHost == canonicalHost(base) {
|
||||
return nil
|
||||
}
|
||||
for _, allowed := range allowedHosts {
|
||||
// Parse with the request scheme so "host:443" matches an https request
|
||||
// the same way an explicit default port on the request URL does.
|
||||
allowedURL, parseErr := url.Parse(request.Scheme + "://" + strings.TrimSpace(allowed))
|
||||
if parseErr == nil && requestHost == canonicalHost(allowedURL) {
|
||||
return nil
|
||||
}
|
||||
}
|
||||
return fmt.Errorf("plugin request host %q is not allowed", request.Host)
|
||||
}
|
||||
|
||||
func canonicalHost(value *url.URL) string {
|
||||
host := strings.ToLower(value.Hostname())
|
||||
port := value.Port()
|
||||
if port == "" || port == "80" && value.Scheme == "http" || port == "443" && value.Scheme == "https" {
|
||||
return host
|
||||
}
|
||||
return net.JoinHostPort(host, port)
|
||||
}
|
||||
@@ -0,0 +1,963 @@
|
||||
package jsplugin
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"maps"
|
||||
"net/http"
|
||||
"regexp"
|
||||
"slices"
|
||||
"sort"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"github.com/QuantumNous/new-api/constant"
|
||||
)
|
||||
|
||||
type RouteType string
|
||||
|
||||
const (
|
||||
RouteTypeSubmit RouteType = "submit"
|
||||
RouteTypeQuery RouteType = "query"
|
||||
RouteTypeDynamic RouteType = "dynamic"
|
||||
)
|
||||
|
||||
type Route struct {
|
||||
Method string `json:"method"`
|
||||
Path string `json:"path"`
|
||||
Type RouteType `json:"type"`
|
||||
Action string `json:"action,omitempty"`
|
||||
Decode string `json:"decode,omitempty"`
|
||||
Render string `json:"render,omitempty"`
|
||||
TaskIDParam string `json:"taskIdParam,omitempty"`
|
||||
// Models restricts this route to the listed models. The host matches the
|
||||
// canonical top-level "model" body field before any JS hook runs; empty
|
||||
// means unrestricted. Must be a subset of meta.models.
|
||||
Models []string `json:"models,omitempty"`
|
||||
}
|
||||
|
||||
// ProtocolClaim is one entry of meta.protocols. Models narrows the protocol's
|
||||
// endpoint bindings to a subset of meta.models; empty binds every model.
|
||||
// Supports names the request forms a mode-bearing protocol accepts; decode
|
||||
// and normalize rewrite it into host-table order.
|
||||
type ProtocolClaim struct {
|
||||
Name string `json:"name"`
|
||||
Models []string `json:"models,omitempty"`
|
||||
Supports []string `json:"supports,omitempty"`
|
||||
objectForm bool
|
||||
}
|
||||
|
||||
// ProtocolMode is one client request form a host protocol operation accepts
|
||||
// and the plugin hook that implements it.
|
||||
type ProtocolMode struct {
|
||||
Name string
|
||||
Hook string
|
||||
}
|
||||
|
||||
type BodyKind string
|
||||
|
||||
const (
|
||||
BodyNone BodyKind = "none"
|
||||
BodyJSON BodyKind = "json"
|
||||
BodyForm BodyKind = "form"
|
||||
BodyMultipart BodyKind = "multipart"
|
||||
)
|
||||
|
||||
type HostProtocolOperation struct {
|
||||
Name string
|
||||
Methods []string
|
||||
Path string
|
||||
BodyKinds []BodyKind
|
||||
ModelField string
|
||||
RequiredProtocolMembers []string
|
||||
Modes []ProtocolMode
|
||||
RequiredDriverHooks []string
|
||||
}
|
||||
|
||||
type HostProtocolDefinition struct {
|
||||
Name string
|
||||
Operations []HostProtocolOperation
|
||||
}
|
||||
|
||||
var hostProtocols = []HostProtocolDefinition{
|
||||
{Name: "openai_responses", Operations: []HostProtocolOperation{
|
||||
{Name: "create", Methods: []string{http.MethodPost}, Path: "/v1/responses", BodyKinds: []BodyKind{BodyJSON}, ModelField: "model", RequiredProtocolMembers: []string{"decodeRequest"}, Modes: []ProtocolMode{{Name: "stream", Hook: "renderEvents"}, {Name: "sync", Hook: "renderFinal"}, {Name: "background", Hook: "renderFinal"}}},
|
||||
{Name: "retrieve", Methods: []string{http.MethodGet}, Path: "/v1/responses/:response_id", BodyKinds: []BodyKind{BodyNone}},
|
||||
}},
|
||||
{Name: "openai_video", Operations: []HostProtocolOperation{
|
||||
{Name: "create", Methods: []string{http.MethodPost}, Path: "/v1/videos", BodyKinds: []BodyKind{BodyJSON, BodyMultipart}, ModelField: "model", RequiredProtocolMembers: []string{"decodeRequest"}},
|
||||
{Name: "retrieve", Methods: []string{http.MethodGet}, Path: "/v1/videos/:task_id", BodyKinds: []BodyKind{BodyNone}, RequiredProtocolMembers: []string{"render"}},
|
||||
{Name: "content", Methods: []string{http.MethodGet, http.MethodHead}, Path: "/v1/videos/:task_id/content", BodyKinds: []BodyKind{BodyNone}, RequiredDriverHooks: []string{"listArtifacts", "buildContentRequest"}},
|
||||
}},
|
||||
}
|
||||
|
||||
func HostProtocol(name string) (HostProtocolDefinition, bool) {
|
||||
for _, definition := range hostProtocols {
|
||||
if definition.Name == name {
|
||||
return definition, true
|
||||
}
|
||||
}
|
||||
return HostProtocolDefinition{}, false
|
||||
}
|
||||
|
||||
func HostProtocols() []HostProtocolDefinition {
|
||||
definitions := make([]HostProtocolDefinition, len(hostProtocols))
|
||||
for index, definition := range hostProtocols {
|
||||
definitions[index] = definition
|
||||
definitions[index].Operations = append([]HostProtocolOperation(nil), definition.Operations...)
|
||||
for operationIndex := range definitions[index].Operations {
|
||||
operation := &definitions[index].Operations[operationIndex]
|
||||
operation.Methods = append([]string(nil), operation.Methods...)
|
||||
operation.BodyKinds = append([]BodyKind(nil), operation.BodyKinds...)
|
||||
operation.RequiredProtocolMembers = append([]string(nil), operation.RequiredProtocolMembers...)
|
||||
operation.Modes = append([]ProtocolMode(nil), operation.Modes...)
|
||||
operation.RequiredDriverHooks = append([]string(nil), operation.RequiredDriverHooks...)
|
||||
}
|
||||
}
|
||||
return definitions
|
||||
}
|
||||
|
||||
// DefinedModes returns each distinct mode on the protocol in host-table order.
|
||||
func (d HostProtocolDefinition) DefinedModes() []ProtocolMode {
|
||||
seen := make(map[string]struct{})
|
||||
modes := make([]ProtocolMode, 0)
|
||||
for _, operation := range d.Operations {
|
||||
for _, mode := range operation.Modes {
|
||||
if _, exists := seen[mode.Name]; exists {
|
||||
continue
|
||||
}
|
||||
seen[mode.Name] = struct{}{}
|
||||
modes = append(modes, mode)
|
||||
}
|
||||
}
|
||||
return modes
|
||||
}
|
||||
|
||||
func orderProtocolSupports(protocol string, supports []string) []string {
|
||||
if len(supports) == 0 {
|
||||
return supports
|
||||
}
|
||||
definition, ok := HostProtocol(protocol)
|
||||
if !ok {
|
||||
return supports
|
||||
}
|
||||
rank := make(map[string]int)
|
||||
for index, mode := range definition.DefinedModes() {
|
||||
rank[mode.Name] = index
|
||||
}
|
||||
ordered := append([]string(nil), supports...)
|
||||
slices.SortStableFunc(ordered, func(left, right string) int {
|
||||
leftRank, leftKnown := rank[left]
|
||||
rightRank, rightKnown := rank[right]
|
||||
switch {
|
||||
case leftKnown && rightKnown:
|
||||
if leftRank < rightRank {
|
||||
return -1
|
||||
}
|
||||
if leftRank > rightRank {
|
||||
return 1
|
||||
}
|
||||
return 0
|
||||
case leftKnown:
|
||||
return -1
|
||||
case rightKnown:
|
||||
return 1
|
||||
default:
|
||||
return 0
|
||||
}
|
||||
})
|
||||
return ordered
|
||||
}
|
||||
|
||||
func LookupHostProtocolOperation(method, path string) (string, HostProtocolOperation, bool) {
|
||||
method = strings.ToUpper(strings.TrimSpace(method))
|
||||
for _, definition := range hostProtocols {
|
||||
for _, operation := range definition.Operations {
|
||||
if operation.Path != path || operation.ModelField == "" {
|
||||
continue
|
||||
}
|
||||
if slices.Contains(operation.Methods, method) {
|
||||
return definition.Name, operation, true
|
||||
}
|
||||
}
|
||||
}
|
||||
return "", HostProtocolOperation{}, false
|
||||
}
|
||||
|
||||
type RouteBinding struct {
|
||||
Plugin *LoadedPlugin
|
||||
Route Route
|
||||
}
|
||||
|
||||
type ProtocolBinding struct {
|
||||
Plugin *LoadedPlugin
|
||||
Protocol string
|
||||
Operation HostProtocolOperation
|
||||
Model string
|
||||
}
|
||||
|
||||
const (
|
||||
ContextKeyPinnedPlugin = "task_plugin_pinned_plugin"
|
||||
ContextKeyPinnedRoute = "task_plugin_pinned_route"
|
||||
ContextKeyPinnedEndpoint = "task_plugin_pinned_endpoint"
|
||||
ContextKeyRouteRequest = "task_plugin_route_request"
|
||||
ContextKeyProtocolRequest = "task_plugin_protocol_request"
|
||||
)
|
||||
|
||||
type PinnedPlugin struct {
|
||||
Generation *RoutingGeneration
|
||||
Plugin *LoadedPlugin
|
||||
}
|
||||
|
||||
type PinnedRoute struct {
|
||||
Generation *RoutingGeneration
|
||||
Plugin *LoadedPlugin
|
||||
Route Route
|
||||
}
|
||||
|
||||
// PinnedEndpoint carries the exact generation and endpoint candidates selected
|
||||
// before distribution. Plugin initially names the deterministic request parser;
|
||||
// distribution may rebind it to another candidate from the same generation
|
||||
// when multiple legacy providers expose the same model.
|
||||
type PinnedEndpoint struct {
|
||||
Generation *RoutingGeneration
|
||||
Plugin *LoadedPlugin
|
||||
Protocol string
|
||||
Operation HostProtocolOperation
|
||||
Model string
|
||||
MappedModel string
|
||||
Candidates []ProtocolBinding
|
||||
}
|
||||
|
||||
// RouteRequestContext is the canonical request view exposed to declarative
|
||||
// routing hooks. RequestBody contains decoded JSON or multipart text fields;
|
||||
// raw binary and multipart file bytes remain host-owned.
|
||||
type RouteRequestContext struct {
|
||||
Path string `json:"path"`
|
||||
Method string `json:"method"`
|
||||
Params map[string]string `json:"params"`
|
||||
Query map[string][]string `json:"query"`
|
||||
Body any `json:"body"`
|
||||
Files []map[string]any `json:"-"`
|
||||
RequestBody any `json:"-"`
|
||||
}
|
||||
|
||||
func (r RouteRequestContext) JSValue() map[string]any {
|
||||
params := make(map[string]string, len(r.Params))
|
||||
maps.Copy(params, r.Params)
|
||||
query := make(map[string][]string, len(r.Query))
|
||||
for key, values := range r.Query {
|
||||
query[key] = append([]string(nil), values...)
|
||||
}
|
||||
return map[string]any{
|
||||
"path": r.Path,
|
||||
"method": r.Method,
|
||||
"params": params,
|
||||
"query": query,
|
||||
"body": clonePluginRequestValue(r.Body),
|
||||
}
|
||||
}
|
||||
|
||||
func clonePluginRequestValue(value any) any {
|
||||
switch typed := value.(type) {
|
||||
case map[string]any:
|
||||
cloned := make(map[string]any, len(typed))
|
||||
for key, item := range typed {
|
||||
cloned[key] = clonePluginRequestValue(item)
|
||||
}
|
||||
return cloned
|
||||
case []any:
|
||||
cloned := make([]any, len(typed))
|
||||
for index, item := range typed {
|
||||
cloned[index] = clonePluginRequestValue(item)
|
||||
}
|
||||
return cloned
|
||||
case []string:
|
||||
return append([]string(nil), typed...)
|
||||
case map[string][]string:
|
||||
cloned := make(map[string][]string, len(typed))
|
||||
for key, values := range typed {
|
||||
cloned[key] = append([]string(nil), values...)
|
||||
}
|
||||
return cloned
|
||||
case []map[string]any:
|
||||
cloned := make([]map[string]any, len(typed))
|
||||
for index, item := range typed {
|
||||
cloned[index] = clonePluginRequestValue(item).(map[string]any)
|
||||
}
|
||||
return cloned
|
||||
default:
|
||||
return value
|
||||
}
|
||||
}
|
||||
|
||||
type ProtocolRequestContext struct {
|
||||
RouteRequestContext
|
||||
Protocol string `json:"protocol"`
|
||||
Operation string `json:"operation"`
|
||||
Model string `json:"model"`
|
||||
// UpstreamModel is the declared machine identity when Model is a
|
||||
// channel-mapping alias; empty otherwise. Decode hooks that key rate
|
||||
// tables or request shaping by model must use it over Model.
|
||||
UpstreamModel string `json:"upstreamModel,omitempty"`
|
||||
Stream bool `json:"stream"`
|
||||
}
|
||||
|
||||
func (p ProtocolRequestContext) JSValue() map[string]any {
|
||||
value := p.RouteRequestContext.JSValue()
|
||||
value["protocol"] = p.Protocol
|
||||
value["operation"] = p.Operation
|
||||
value["model"] = p.Model
|
||||
if p.UpstreamModel != "" {
|
||||
value["upstreamModel"] = p.UpstreamModel
|
||||
}
|
||||
value["stream"] = p.Stream
|
||||
return value
|
||||
}
|
||||
|
||||
// SupportsHostProtocol reports whether the current host release has a
|
||||
// concrete wire/state machine for an otherwise valid manifest endpoint.
|
||||
func SupportsHostProtocol(protocol string) bool { _, ok := HostProtocol(protocol); return ok }
|
||||
|
||||
// RoutingGeneration is an immutable, request-pinnable view of all effective
|
||||
// plugins and their deterministic routing indexes.
|
||||
type RoutingGeneration struct {
|
||||
Number uint64
|
||||
PublishedAt time.Time
|
||||
|
||||
byKey map[string]*LoadedPlugin
|
||||
byModel map[string]*LoadedPlugin
|
||||
canonicalModelByFold map[string]string
|
||||
byChannelType map[int]*LoadedPlugin
|
||||
routeIndex map[string]RouteBinding
|
||||
protocolIndex map[string][]ProtocolBinding
|
||||
plugins []*LoadedPlugin
|
||||
routes []RouteBinding
|
||||
runtime http.Handler
|
||||
retainCurrent map[string]struct{}
|
||||
}
|
||||
|
||||
var (
|
||||
routeMethodPattern = regexp.MustCompile(`^(GET|POST|PUT|PATCH|DELETE)$`)
|
||||
pathNamePattern = regexp.MustCompile(`^[A-Za-z_][A-Za-z0-9_]*$`)
|
||||
staticSegment = regexp.MustCompile(`^[A-Za-z0-9._~-]+$`)
|
||||
memberNamePattern = regexp.MustCompile(`^[A-Za-z_$][A-Za-z0-9_$]*$`)
|
||||
)
|
||||
|
||||
var reservedRouteNamespaces = []string{
|
||||
"/api",
|
||||
"/assets",
|
||||
"/setup",
|
||||
"/v1/tasks",
|
||||
"/console",
|
||||
"/login",
|
||||
"/forbidden",
|
||||
"/sign-in",
|
||||
"/sign-up",
|
||||
"/forgot-password",
|
||||
"/oauth",
|
||||
"/otp",
|
||||
"/register",
|
||||
"/reset",
|
||||
"/privacy-policy",
|
||||
"/user-agreement",
|
||||
"/about",
|
||||
"/pricing",
|
||||
"/rankings",
|
||||
"/user",
|
||||
"/401",
|
||||
"/403",
|
||||
"/404",
|
||||
"/500",
|
||||
"/503",
|
||||
"/chat2link",
|
||||
"/system-settings",
|
||||
"/channels",
|
||||
"/chat",
|
||||
"/dashboard",
|
||||
"/errors",
|
||||
"/keys",
|
||||
"/models",
|
||||
"/playground",
|
||||
"/profile",
|
||||
"/redemption-codes",
|
||||
"/subscriptions",
|
||||
"/system-info",
|
||||
"/task-plugins",
|
||||
"/usage-logs",
|
||||
"/users",
|
||||
"/wallet",
|
||||
}
|
||||
|
||||
func (g *RoutingGeneration) Get(key string) (*LoadedPlugin, bool) {
|
||||
if g == nil {
|
||||
return nil, false
|
||||
}
|
||||
plugin, ok := g.byKey[key]
|
||||
return plugin, ok
|
||||
}
|
||||
|
||||
func (g *RoutingGeneration) GetByChannelType(channelType int) (*LoadedPlugin, bool) {
|
||||
if g == nil || channelType == 0 || channelType == constant.ChannelTypeTaskPlugin {
|
||||
return nil, false
|
||||
}
|
||||
plugin, ok := g.byChannelType[channelType]
|
||||
return plugin, ok
|
||||
}
|
||||
|
||||
// GetByModel returns the deterministic effective plugin metadata used for a
|
||||
// model-level host concern such as billing. When multiple providers expose the
|
||||
// same model name, the first plugin in generation order owns that shared
|
||||
// metadata view.
|
||||
func (g *RoutingGeneration) GetByModel(model string) (*LoadedPlugin, bool) {
|
||||
if g == nil || model == "" {
|
||||
return nil, false
|
||||
}
|
||||
plugin, ok := g.byModel[model]
|
||||
return plugin, ok
|
||||
}
|
||||
|
||||
// CanonicalModel returns the declared spelling for model. An exact byModel
|
||||
// hit wins and returns the input unchanged; otherwise the ASCII-folded
|
||||
// index is consulted. Miss and nil-receiver return ("", false).
|
||||
func (g *RoutingGeneration) CanonicalModel(model string) (string, bool) {
|
||||
if g == nil || model == "" {
|
||||
return "", false
|
||||
}
|
||||
if _, ok := g.byModel[model]; ok {
|
||||
return model, true
|
||||
}
|
||||
declared, ok := g.canonicalModelByFold[asciiFold(model)]
|
||||
return declared, ok
|
||||
}
|
||||
|
||||
// LookupDeclaredRoute resolves a manifest path declaration. It does not match
|
||||
// an incoming concrete URL; runtime matching is delegated to Gin.
|
||||
func (g *RoutingGeneration) LookupDeclaredRoute(method, path string) (RouteBinding, bool) {
|
||||
if g == nil {
|
||||
return RouteBinding{}, false
|
||||
}
|
||||
normalizedMethod, err := normalizeRouteMethod(method)
|
||||
if err != nil {
|
||||
return RouteBinding{}, false
|
||||
}
|
||||
shape, err := routePathShape(path)
|
||||
if err != nil {
|
||||
return RouteBinding{}, false
|
||||
}
|
||||
binding, ok := g.routeIndex[normalizedMethod+" "+shape]
|
||||
return binding, ok
|
||||
}
|
||||
|
||||
func (g *RoutingGeneration) LookupEndpoint(method, path, model string) (ProtocolBinding, bool) {
|
||||
if g == nil {
|
||||
return ProtocolBinding{}, false
|
||||
}
|
||||
normalizedMethod, err := normalizeRouteMethod(method)
|
||||
if err != nil {
|
||||
return ProtocolBinding{}, false
|
||||
}
|
||||
bindings := g.protocolIndex[endpointIndexKey(normalizedMethod, path, model)]
|
||||
if len(bindings) == 0 {
|
||||
return ProtocolBinding{}, false
|
||||
}
|
||||
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.
|
||||
func (g *RoutingGeneration) LookupEndpointCandidates(method, path, model string) []ProtocolBinding {
|
||||
if g == nil {
|
||||
return nil
|
||||
}
|
||||
normalizedMethod, err := normalizeRouteMethod(method)
|
||||
if err != nil {
|
||||
return nil
|
||||
}
|
||||
bindings := g.protocolIndex[endpointIndexKey(normalizedMethod, path, model)]
|
||||
return append([]ProtocolBinding(nil), bindings...)
|
||||
}
|
||||
|
||||
func (g *RoutingGeneration) Plugins() []*LoadedPlugin {
|
||||
if g == nil {
|
||||
return nil
|
||||
}
|
||||
return append([]*LoadedPlugin(nil), g.plugins...)
|
||||
}
|
||||
|
||||
func (g *RoutingGeneration) Routes() []RouteBinding {
|
||||
if g == nil {
|
||||
return nil
|
||||
}
|
||||
return append([]RouteBinding(nil), g.routes...)
|
||||
}
|
||||
|
||||
// RuntimeHandler is the inner router built for this exact generation. It is
|
||||
// published in the same atomic pointer as the routing indexes.
|
||||
func (g *RoutingGeneration) RuntimeHandler() http.Handler {
|
||||
if g == nil {
|
||||
return nil
|
||||
}
|
||||
return g.runtime
|
||||
}
|
||||
|
||||
func (g *RoutingGeneration) RetainsIncumbent(key string) bool {
|
||||
if g == nil {
|
||||
return false
|
||||
}
|
||||
_, retained := g.retainCurrent[key]
|
||||
return retained
|
||||
}
|
||||
|
||||
// RebuildWithPlugins creates a generation from the supplied exact plugin
|
||||
// pointers. It is used when a rejected hot update must retain the incumbent
|
||||
// runtime object for that key.
|
||||
func (g *RoutingGeneration) RebuildWithPlugins(plugins []*LoadedPlugin) (*RoutingGeneration, error) {
|
||||
if g == nil {
|
||||
return nil, fmt.Errorf("cannot rebuild a nil routing generation")
|
||||
}
|
||||
byKey := make(map[string]*LoadedPlugin, len(plugins))
|
||||
for _, plugin := range plugins {
|
||||
if plugin == nil {
|
||||
return nil, fmt.Errorf("cannot rebuild routing generation with a nil plugin")
|
||||
}
|
||||
if _, exists := g.byKey[plugin.Meta.Key]; !exists {
|
||||
return nil, fmt.Errorf("plugin %q is not present in routing generation %d", plugin.Meta.Key, g.Number)
|
||||
}
|
||||
if _, duplicate := byKey[plugin.Meta.Key]; duplicate {
|
||||
return nil, fmt.Errorf("plugin %q appears more than once in routing generation rebuild", plugin.Meta.Key)
|
||||
}
|
||||
byKey[plugin.Meta.Key] = plugin
|
||||
}
|
||||
rebuilt, err := buildRoutingGeneration(byKey, nil, g.Number)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
rebuilt.PublishedAt = g.PublishedAt
|
||||
rebuilt.retainCurrent = cloneStringSet(g.retainCurrent)
|
||||
return rebuilt, nil
|
||||
}
|
||||
|
||||
// WithRuntime returns a shallow immutable copy carrying the prepared inner
|
||||
// handler. Callers use it before publication; published generations must not
|
||||
// be mutated.
|
||||
func (g *RoutingGeneration) WithRuntime(handler http.Handler) *RoutingGeneration {
|
||||
if g == nil {
|
||||
return nil
|
||||
}
|
||||
prepared := *g
|
||||
prepared.runtime = handler
|
||||
return &prepared
|
||||
}
|
||||
|
||||
func cloneStringSet(source map[string]struct{}) map[string]struct{} {
|
||||
if len(source) == 0 {
|
||||
return nil
|
||||
}
|
||||
clone := make(map[string]struct{}, len(source))
|
||||
for key := range source {
|
||||
clone[key] = struct{}{}
|
||||
}
|
||||
return clone
|
||||
}
|
||||
|
||||
func normalizeRouteMethod(method string) (string, error) {
|
||||
method = strings.ToUpper(strings.TrimSpace(method))
|
||||
if !routeMethodPattern.MatchString(method) {
|
||||
return "", fmt.Errorf("plugin route method %q is not supported", method)
|
||||
}
|
||||
return method, nil
|
||||
}
|
||||
|
||||
// NormalizeRoutePath validates the canonical path syntax used by plugin route
|
||||
// declarations. A trailing slash is allowed and remains significant.
|
||||
func NormalizeRoutePath(routePath string) (string, error) {
|
||||
if routePath == "" || routePath[0] != '/' {
|
||||
return "", fmt.Errorf("plugin route path must start with /")
|
||||
}
|
||||
if routePath == "/" {
|
||||
return "", fmt.Errorf("plugin route path / is reserved")
|
||||
}
|
||||
if strings.ContainsAny(routePath, "?#%") {
|
||||
return "", fmt.Errorf("plugin route path %q must not contain a query, fragment, or percent-encoding", routePath)
|
||||
}
|
||||
if strings.Contains(routePath, "//") {
|
||||
return "", fmt.Errorf("plugin route path %q must not contain empty segments", routePath)
|
||||
}
|
||||
|
||||
segments := strings.Split(strings.TrimPrefix(routePath, "/"), "/")
|
||||
seenNames := make(map[string]struct{})
|
||||
for index, segment := range segments {
|
||||
if segment == "" && index == len(segments)-1 {
|
||||
continue
|
||||
}
|
||||
if segment == "." || segment == ".." {
|
||||
return "", fmt.Errorf("plugin route path %q must not contain dot segments", routePath)
|
||||
}
|
||||
if after, ok := strings.CutPrefix(segment, ":"); ok {
|
||||
name := after
|
||||
if !pathNamePattern.MatchString(name) {
|
||||
return "", fmt.Errorf("plugin route path %q has invalid parameter %q", routePath, segment)
|
||||
}
|
||||
if _, exists := seenNames[name]; exists {
|
||||
return "", fmt.Errorf("plugin route path %q repeats parameter %q", routePath, name)
|
||||
}
|
||||
seenNames[name] = struct{}{}
|
||||
continue
|
||||
}
|
||||
if after, ok := strings.CutPrefix(segment, "*"); ok {
|
||||
name := after
|
||||
if index != len(segments)-1 || !pathNamePattern.MatchString(name) {
|
||||
return "", fmt.Errorf("plugin route path %q has an invalid catch-all segment", routePath)
|
||||
}
|
||||
if _, exists := seenNames[name]; exists {
|
||||
return "", fmt.Errorf("plugin route path %q repeats parameter %q", routePath, name)
|
||||
}
|
||||
seenNames[name] = struct{}{}
|
||||
continue
|
||||
}
|
||||
if !staticSegment.MatchString(segment) {
|
||||
return "", fmt.Errorf("plugin route path %q has invalid segment %q", routePath, segment)
|
||||
}
|
||||
}
|
||||
return routePath, nil
|
||||
}
|
||||
|
||||
func routePathShape(routePath string) (string, error) {
|
||||
normalized, err := NormalizeRoutePath(routePath)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
segments := strings.Split(strings.TrimPrefix(normalized, "/"), "/")
|
||||
for index, segment := range segments {
|
||||
if strings.HasPrefix(segment, ":") {
|
||||
segments[index] = ":"
|
||||
} else if strings.HasPrefix(segment, "*") {
|
||||
segments[index] = "*"
|
||||
}
|
||||
}
|
||||
return "/" + strings.Join(segments, "/"), nil
|
||||
}
|
||||
|
||||
func endpointIndexKey(method, path, model string) string {
|
||||
return method + "\x00" + path + "\x00" + model
|
||||
}
|
||||
|
||||
func intersectingReservedNamespace(routePath string) (string, bool) {
|
||||
for _, namespace := range reservedRouteNamespaces {
|
||||
if routePatternIntersectsNamespace(routePath, namespace) {
|
||||
return namespace, true
|
||||
}
|
||||
}
|
||||
return "", false
|
||||
}
|
||||
|
||||
func routePatternIntersectsNamespace(routePath, namespace string) bool {
|
||||
routeSegments := strings.Split(strings.Trim(strings.TrimPrefix(routePath, "/"), "/"), "/")
|
||||
namespaceSegments := strings.Split(strings.TrimPrefix(namespace, "/"), "/")
|
||||
for index, namespaceSegment := range namespaceSegments {
|
||||
if index >= len(routeSegments) {
|
||||
return false
|
||||
}
|
||||
routeSegment := routeSegments[index]
|
||||
if strings.HasPrefix(routeSegment, "*") {
|
||||
return true
|
||||
}
|
||||
if !strings.HasPrefix(routeSegment, ":") && routeSegment != namespaceSegment {
|
||||
return false
|
||||
}
|
||||
}
|
||||
return true
|
||||
}
|
||||
|
||||
func validateRoute(route *Route) error {
|
||||
if route.Method != strings.ToUpper(strings.TrimSpace(route.Method)) {
|
||||
return fmt.Errorf("plugin route method %q must use canonical uppercase spelling", route.Method)
|
||||
}
|
||||
method, err := normalizeRouteMethod(route.Method)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
route.Method = method
|
||||
route.Path, err = NormalizeRoutePath(route.Path)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if namespace, reserved := intersectingReservedNamespace(route.Path); reserved {
|
||||
return fmt.Errorf("plugin route path %q intersects reserved namespace %s", route.Path, namespace)
|
||||
}
|
||||
switch route.Type {
|
||||
case RouteTypeSubmit:
|
||||
if route.Decode == "" || route.Render == "" || route.TaskIDParam != "" {
|
||||
return fmt.Errorf("submit route %s %s must declare decode and render and must not declare taskIdParam", route.Method, route.Path)
|
||||
}
|
||||
case RouteTypeQuery:
|
||||
if route.Decode != "" || strings.TrimSpace(route.Render) == "" {
|
||||
return fmt.Errorf("query route %s %s must declare render and must not declare decode", route.Method, route.Path)
|
||||
}
|
||||
if route.Action != "" {
|
||||
return fmt.Errorf("query route %s %s must not declare action", route.Method, route.Path)
|
||||
}
|
||||
if route.TaskIDParam == "" {
|
||||
route.TaskIDParam = "task_id"
|
||||
}
|
||||
if !pathNamePattern.MatchString(route.TaskIDParam) {
|
||||
return fmt.Errorf("query route %s %s has invalid taskIdParam %q", route.Method, route.Path, route.TaskIDParam)
|
||||
}
|
||||
if !pathHasParameter(route.Path, route.TaskIDParam) {
|
||||
return fmt.Errorf("query route %s %s must contain :%s", route.Method, route.Path, route.TaskIDParam)
|
||||
}
|
||||
case RouteTypeDynamic:
|
||||
if route.Decode == "" || route.Render == "" || route.TaskIDParam != "" {
|
||||
return fmt.Errorf("dynamic route %s %s must declare decode and render and must not declare taskIdParam", route.Method, route.Path)
|
||||
}
|
||||
default:
|
||||
return fmt.Errorf("plugin route %s %s has unsupported type %q", route.Method, route.Path, route.Type)
|
||||
}
|
||||
if route.Decode != "" && !memberNamePattern.MatchString(route.Decode) {
|
||||
return fmt.Errorf("plugin route %s %s has invalid decode %q", route.Method, route.Path, route.Decode)
|
||||
}
|
||||
if route.Render != "" && !memberNamePattern.MatchString(route.Render) {
|
||||
return fmt.Errorf("plugin route %s %s has invalid render %q", route.Method, route.Path, route.Render)
|
||||
}
|
||||
if strings.TrimSpace(route.Action) != route.Action {
|
||||
return fmt.Errorf("plugin route %s %s action must not have surrounding whitespace", route.Method, route.Path)
|
||||
}
|
||||
if len(route.Models) > 0 {
|
||||
if route.Type == RouteTypeQuery {
|
||||
return fmt.Errorf("query route %s %s must not declare models", route.Method, route.Path)
|
||||
}
|
||||
if err := validateModelScope(route.Models, fmt.Sprintf("route %s %s", route.Method, route.Path)); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// validateModelScope enforces the shared rules for route.models and
|
||||
// per-protocol models entries: non-empty canonical names, no duplicates.
|
||||
func validateModelScope(models []string, subject string) error {
|
||||
seen := make(map[string]struct{}, len(models))
|
||||
for _, model := range models {
|
||||
if strings.TrimSpace(model) == "" || strings.TrimSpace(model) != model {
|
||||
return fmt.Errorf("plugin %s models must contain non-empty canonical names", subject)
|
||||
}
|
||||
folded := asciiFold(model)
|
||||
if _, duplicate := seen[folded]; duplicate {
|
||||
return fmt.Errorf("plugin %s models must be unique case-insensitively", subject)
|
||||
}
|
||||
seen[folded] = struct{}{}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func pathHasParameter(routePath, name string) bool {
|
||||
return slices.Contains(strings.Split(routePath, "/"), ":"+name)
|
||||
}
|
||||
|
||||
func ResolveRouteAction(route Route, resolvedAction string) string {
|
||||
if strings.TrimSpace(resolvedAction) != "" {
|
||||
return resolvedAction
|
||||
}
|
||||
return route.Action
|
||||
}
|
||||
|
||||
func buildRoutingGeneration(factory, override map[string]*LoadedPlugin, number uint64) (*RoutingGeneration, error) {
|
||||
effective := effectivePlugins(factory, override)
|
||||
return buildRoutingGenerationFromPlugins(effective, number)
|
||||
}
|
||||
|
||||
func buildRoutingGenerationAdmitting(
|
||||
factory, override map[string]*LoadedPlugin,
|
||||
number uint64,
|
||||
current *RoutingGeneration,
|
||||
retainCurrent map[string]struct{},
|
||||
) (*RoutingGeneration, map[string]string, error) {
|
||||
candidates := effectivePlugins(factory, override)
|
||||
accepted := make(map[string]*LoadedPlugin, len(candidates))
|
||||
currentByKey := make(map[string]*LoadedPlugin)
|
||||
if current != nil {
|
||||
for _, plugin := range current.plugins {
|
||||
currentByKey[plugin.Meta.Key] = plugin
|
||||
}
|
||||
}
|
||||
generation, err := buildRoutingGenerationFromPlugins(accepted, number)
|
||||
if err != nil {
|
||||
return nil, nil, fmt.Errorf("initialize routing generation: %w", err)
|
||||
}
|
||||
|
||||
unchangedKeys := make([]string, 0)
|
||||
changedKeys := make([]string, 0)
|
||||
newKeys := make([]string, 0)
|
||||
for key, candidate := range candidates {
|
||||
incumbent, exists := currentByKey[key]
|
||||
switch {
|
||||
case exists && candidate == incumbent:
|
||||
unchangedKeys = append(unchangedKeys, key)
|
||||
case exists:
|
||||
changedKeys = append(changedKeys, key)
|
||||
default:
|
||||
newKeys = append(newKeys, key)
|
||||
}
|
||||
}
|
||||
sort.Strings(unchangedKeys)
|
||||
sort.Strings(changedKeys)
|
||||
sort.Strings(newKeys)
|
||||
|
||||
routingErrors := make(map[string]string)
|
||||
orderedKeys := append(unchangedKeys, changedKeys...)
|
||||
orderedKeys = append(orderedKeys, newKeys...)
|
||||
rejectedKeys := make([]string, 0)
|
||||
for _, key := range orderedKeys {
|
||||
candidate := candidates[key]
|
||||
accepted[key] = candidate
|
||||
trial, trialErr := buildRoutingGenerationFromPlugins(accepted, number)
|
||||
if trialErr == nil {
|
||||
generation = trial
|
||||
continue
|
||||
}
|
||||
delete(accepted, key)
|
||||
routingErrors[key] = fmt.Sprintf("plugin %s rejected from routing generation: %v", key, trialErr)
|
||||
rejectedKeys = append(rejectedKeys, key)
|
||||
}
|
||||
|
||||
for _, key := range rejectedKeys {
|
||||
candidate := candidates[key]
|
||||
incumbent, hasIncumbent := currentByKey[key]
|
||||
_, mayRetain := retainCurrent[key]
|
||||
if !hasIncumbent || !mayRetain || incumbent == candidate {
|
||||
continue
|
||||
}
|
||||
accepted[key] = incumbent
|
||||
fallback, fallbackErr := buildRoutingGenerationFromPlugins(accepted, number)
|
||||
if fallbackErr == nil {
|
||||
generation = fallback
|
||||
continue
|
||||
}
|
||||
delete(accepted, key)
|
||||
}
|
||||
generation.retainCurrent = cloneStringSet(retainCurrent)
|
||||
return generation, routingErrors, nil
|
||||
}
|
||||
|
||||
func effectivePlugins(factory, override map[string]*LoadedPlugin) map[string]*LoadedPlugin {
|
||||
effective := make(map[string]*LoadedPlugin, len(factory)+len(override))
|
||||
maps.Copy(effective, factory)
|
||||
maps.Copy(effective, override)
|
||||
return effective
|
||||
}
|
||||
|
||||
func buildRoutingGenerationFromPlugins(effective map[string]*LoadedPlugin, number uint64) (*RoutingGeneration, error) {
|
||||
keys := make([]string, 0, len(effective))
|
||||
for key := range effective {
|
||||
keys = append(keys, key)
|
||||
}
|
||||
sort.Strings(keys)
|
||||
|
||||
generation := &RoutingGeneration{
|
||||
Number: number,
|
||||
PublishedAt: time.Now(),
|
||||
byKey: make(map[string]*LoadedPlugin, len(effective)),
|
||||
byModel: make(map[string]*LoadedPlugin),
|
||||
canonicalModelByFold: make(map[string]string),
|
||||
byChannelType: make(map[int]*LoadedPlugin),
|
||||
routeIndex: make(map[string]RouteBinding),
|
||||
protocolIndex: make(map[string][]ProtocolBinding),
|
||||
plugins: make([]*LoadedPlugin, 0, len(effective)),
|
||||
}
|
||||
for _, key := range keys {
|
||||
plugin := effective[key]
|
||||
generation.byKey[key] = plugin
|
||||
generation.plugins = append(generation.plugins, plugin)
|
||||
for _, model := range plugin.Meta.Models {
|
||||
if _, exists := generation.byModel[model]; !exists {
|
||||
generation.byModel[model] = plugin
|
||||
}
|
||||
folded := asciiFold(model)
|
||||
if existing, exists := generation.canonicalModelByFold[folded]; exists {
|
||||
if existing != model {
|
||||
otherKey := plugin.Meta.Key
|
||||
if other, ok := generation.byModel[existing]; ok {
|
||||
otherKey = other.Meta.Key
|
||||
}
|
||||
return nil, fmt.Errorf("plugin %s model %q conflicts with plugin %s model %q", plugin.Meta.Key, model, otherKey, existing)
|
||||
}
|
||||
continue
|
||||
}
|
||||
generation.canonicalModelByFold[folded] = model
|
||||
}
|
||||
|
||||
for _, channelType := range plugin.Meta.ChannelTypes {
|
||||
if channelType == 0 || channelType == constant.ChannelTypeTaskPlugin {
|
||||
continue
|
||||
}
|
||||
if other, exists := generation.byChannelType[channelType]; exists {
|
||||
return nil, fmt.Errorf("plugin %s channelType %d conflicts with plugin %s", plugin.Meta.Key, channelType, other.Meta.Key)
|
||||
}
|
||||
generation.byChannelType[channelType] = plugin
|
||||
}
|
||||
|
||||
for _, route := range plugin.Meta.Routes {
|
||||
shape, err := routePathShape(route.Path)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
indexKey := route.Method + " " + shape
|
||||
if other, exists := generation.routeIndex[indexKey]; exists {
|
||||
return nil, fmt.Errorf("plugin %s route %s %s conflicts with plugin %s route %s", plugin.Meta.Key, route.Method, route.Path, other.Plugin.Meta.Key, other.Route.Path)
|
||||
}
|
||||
binding := RouteBinding{Plugin: plugin, Route: route}
|
||||
generation.routeIndex[indexKey] = binding
|
||||
generation.routes = append(generation.routes, binding)
|
||||
}
|
||||
|
||||
for _, claim := range plugin.Meta.Protocols {
|
||||
definition, _ := HostProtocol(claim.Name)
|
||||
boundModels := plugin.Meta.Models
|
||||
if len(claim.Models) > 0 {
|
||||
boundModels = claim.Models
|
||||
}
|
||||
for _, operation := range definition.Operations {
|
||||
if operation.ModelField == "" {
|
||||
continue
|
||||
}
|
||||
for _, method := range operation.Methods {
|
||||
for _, model := range boundModels {
|
||||
indexKey := endpointIndexKey(method, operation.Path, model)
|
||||
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 {
|
||||
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)
|
||||
}
|
||||
}
|
||||
generation.protocolIndex[indexKey] = append(bindings, ProtocolBinding{Plugin: plugin, Protocol: claim.Name, Operation: operation, Model: model})
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
return generation, nil
|
||||
}
|
||||
|
||||
// PreflightRoutingConflict reports whether admitting candidate into the
|
||||
// current generation would collide on a channel type, native route, or
|
||||
// protocol-model binding. A same-key entry is replaced first so re-uploading
|
||||
// a plugin (or overriding a factory built-in) does not self-conflict.
|
||||
func PreflightRoutingConflict(current *RoutingGeneration, candidate *LoadedPlugin) error {
|
||||
if candidate == nil {
|
||||
return fmt.Errorf("cannot preflight a nil plugin")
|
||||
}
|
||||
effective := make(map[string]*LoadedPlugin)
|
||||
number := uint64(0)
|
||||
if current != nil {
|
||||
for _, plugin := range current.Plugins() {
|
||||
effective[plugin.Meta.Key] = plugin
|
||||
}
|
||||
number = current.Number
|
||||
}
|
||||
effective[candidate.Meta.Key] = candidate
|
||||
_, err := buildRoutingGenerationFromPlugins(effective, number)
|
||||
return err
|
||||
}
|
||||
File diff suppressed because it is too large.
Load diff
@@ -0,0 +1,139 @@
|
||||
package jsplugin
|
||||
|
||||
import (
|
||||
"crypto/hmac"
|
||||
"crypto/sha256"
|
||||
"encoding/base64"
|
||||
"encoding/hex"
|
||||
"fmt"
|
||||
"net/url"
|
||||
"sort"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"github.com/golang-jwt/jwt/v5"
|
||||
"github.com/google/uuid"
|
||||
"github.com/grafana/sobek"
|
||||
)
|
||||
|
||||
type volcSignRequest struct {
|
||||
Method string `json:"method"`
|
||||
URL string `json:"url"`
|
||||
Headers map[string]string `json:"headers"`
|
||||
Body string `json:"body"`
|
||||
AccessKey string `json:"accessKey"`
|
||||
SecretKey string `json:"secretKey"`
|
||||
Region string `json:"region"`
|
||||
Service string `json:"service"`
|
||||
Timestamp int64 `json:"timestamp"`
|
||||
}
|
||||
|
||||
func injectGlobals(runtime *sobek.Runtime, identity func() string, now func() time.Time, logOutput func(string)) error {
|
||||
utils := map[string]any{
|
||||
"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))
|
||||
},
|
||||
"hmacSHA256": func(message, secret string) string {
|
||||
mac := hmac.New(sha256.New, []byte(secret))
|
||||
_, _ = mac.Write([]byte(message))
|
||||
return hex.EncodeToString(mac.Sum(nil))
|
||||
},
|
||||
"base64": func(value string) string { return base64.StdEncoding.EncodeToString([]byte(value)) },
|
||||
"base64URL": func(value string) string {
|
||||
return base64.RawURLEncoding.EncodeToString([]byte(value))
|
||||
},
|
||||
"base64URLDecode": func(value string) (string, error) {
|
||||
decoded, err := base64.RawURLEncoding.DecodeString(value)
|
||||
return string(decoded), err
|
||||
},
|
||||
"uuid": func() string { return uuid.NewString() },
|
||||
"volcSignV4": func(request volcSignRequest) (map[string]string, error) {
|
||||
return signVolcV4(request, now)
|
||||
},
|
||||
}
|
||||
if err := runtime.Set("utils", utils); err != nil {
|
||||
return err
|
||||
}
|
||||
console := runtime.NewObject()
|
||||
if err := console.Set("log", func(call sobek.FunctionCall) sobek.Value {
|
||||
parts := make([]string, len(call.Arguments))
|
||||
for i, argument := range call.Arguments {
|
||||
parts[i] = argument.String()
|
||||
}
|
||||
if logOutput != nil {
|
||||
logOutput(identity() + " " + strings.Join(parts, " "))
|
||||
}
|
||||
return sobek.Undefined()
|
||||
}); err != nil {
|
||||
return err
|
||||
}
|
||||
return runtime.Set("console", console)
|
||||
}
|
||||
|
||||
func signVolcV4(request volcSignRequest, now func() time.Time) (map[string]string, error) {
|
||||
parsedURL, err := url.Parse(request.URL)
|
||||
if err != nil || parsedURL.Host == "" {
|
||||
return nil, fmt.Errorf("invalid Volcengine signing URL")
|
||||
}
|
||||
region := request.Region
|
||||
if region == "" {
|
||||
region = "cn-north-1"
|
||||
}
|
||||
service := request.Service
|
||||
if service == "" {
|
||||
service = "cv"
|
||||
}
|
||||
timestamp := now().UTC()
|
||||
if request.Timestamp != 0 {
|
||||
timestamp = time.Unix(request.Timestamp, 0).UTC()
|
||||
}
|
||||
xDate := timestamp.Format("20060102T150405Z")
|
||||
shortDate := timestamp.Format("20060102")
|
||||
bodyHash := sha256.Sum256([]byte(request.Body))
|
||||
requestPath := parsedURL.EscapedPath()
|
||||
if requestPath == "" {
|
||||
requestPath = "/"
|
||||
}
|
||||
|
||||
headers := make(map[string]string, len(request.Headers)+3)
|
||||
for name, value := range request.Headers {
|
||||
headers[strings.ToLower(name)] = strings.TrimSpace(value)
|
||||
}
|
||||
headers["host"] = parsedURL.Host
|
||||
headers["x-date"] = xDate
|
||||
headers["x-content-sha256"] = hex.EncodeToString(bodyHash[:])
|
||||
keys := make([]string, 0, len(headers))
|
||||
for name := range headers {
|
||||
keys = append(keys, name)
|
||||
}
|
||||
sort.Strings(keys)
|
||||
var canonicalHeaders strings.Builder
|
||||
for _, name := range keys {
|
||||
canonicalHeaders.WriteString(name)
|
||||
canonicalHeaders.WriteByte(':')
|
||||
canonicalHeaders.WriteString(headers[name])
|
||||
canonicalHeaders.WriteByte('\n')
|
||||
}
|
||||
signedHeaders := strings.Join(keys, ";")
|
||||
canonicalRequest := strings.Join([]string{
|
||||
strings.ToUpper(request.Method), requestPath, parsedURL.Query().Encode(),
|
||||
canonicalHeaders.String(), signedHeaders, hex.EncodeToString(bodyHash[:]),
|
||||
}, "\n")
|
||||
canonicalHash := sha256.Sum256([]byte(canonicalRequest))
|
||||
scope := fmt.Sprintf("%s/%s/%s/request", shortDate, region, service)
|
||||
stringToSign := fmt.Sprintf("HMAC-SHA256\n%s\n%s\n%s", xDate, scope, hex.EncodeToString(canonicalHash[:]))
|
||||
sign := func(key []byte, value string) []byte {
|
||||
mac := hmac.New(sha256.New, key)
|
||||
_, _ = mac.Write([]byte(value))
|
||||
return mac.Sum(nil)
|
||||
}
|
||||
signingKey := sign(sign(sign([]byte(request.SecretKey), shortDate), region), service)
|
||||
signingKey = sign(signingKey, "request")
|
||||
signature := hex.EncodeToString(sign(signingKey, stringToSign))
|
||||
return map[string]string{
|
||||
"Authorization": fmt.Sprintf("HMAC-SHA256 Credential=%s/%s, SignedHeaders=%s, Signature=%s", request.AccessKey, scope, signedHeaders, signature),
|
||||
"X-Date": xDate,
|
||||
"X-Content-Sha256": hex.EncodeToString(bodyHash[:]),
|
||||
}, nil
|
||||
}
|
||||
@@ -0,0 +1,98 @@
|
||||
package perfmetrics
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"strconv"
|
||||
"time"
|
||||
|
||||
"github.com/QuantumNous/new-api/common"
|
||||
"github.com/QuantumNous/new-api/model"
|
||||
"github.com/QuantumNous/new-api/setting/perf_metrics_setting"
|
||||
)
|
||||
|
||||
func flushLoop() {
|
||||
for {
|
||||
interval := perf_metrics_setting.GetFlushIntervalMinutes()
|
||||
time.Sleep(time.Duration(interval) * time.Minute)
|
||||
setting := perf_metrics_setting.GetSetting()
|
||||
if !setting.Enabled {
|
||||
continue
|
||||
}
|
||||
flushCompletedBuckets()
|
||||
cleanupExpiredMetrics(setting.RetentionDays)
|
||||
}
|
||||
}
|
||||
|
||||
func flushCompletedBuckets() {
|
||||
currentBucket := bucketStart(time.Now().Unix())
|
||||
hotBuckets.Range(func(key, value any) bool {
|
||||
k := key.(bucketKey)
|
||||
if k.bucketTs >= currentBucket {
|
||||
return true
|
||||
}
|
||||
|
||||
bucket := value.(*atomicBucket)
|
||||
drained := bucket.drain()
|
||||
if drained.requestCount == 0 {
|
||||
deleteOldEmptyBucket(k, key)
|
||||
return true
|
||||
}
|
||||
|
||||
err := model.UpsertPerfMetric(&model.PerfMetric{
|
||||
ModelName: k.model,
|
||||
Group: k.group,
|
||||
BucketTs: k.bucketTs,
|
||||
RequestCount: drained.requestCount,
|
||||
SuccessCount: drained.successCount,
|
||||
TotalLatencyMs: drained.totalLatencyMs,
|
||||
TtftSumMs: drained.ttftSumMs,
|
||||
TtftCount: drained.ttftCount,
|
||||
OutputTokens: drained.outputTokens,
|
||||
GenerationMs: drained.generationMs,
|
||||
})
|
||||
if err != nil {
|
||||
bucket.addCounters(drained)
|
||||
common.SysError(fmt.Sprintf("failed to flush perf metric bucket model=%s group=%s bucket=%d: %s", k.model, k.group, k.bucketTs, err.Error()))
|
||||
return true
|
||||
}
|
||||
|
||||
deleteOldEmptyBucket(k, key)
|
||||
return true
|
||||
})
|
||||
}
|
||||
|
||||
func deleteOldEmptyBucket(k bucketKey, rawKey any) {
|
||||
if k.bucketTs < bucketStart(time.Now().Add(-24*time.Hour).Unix()) {
|
||||
hotBuckets.Delete(rawKey)
|
||||
}
|
||||
}
|
||||
|
||||
func cleanupExpiredMetrics(retentionDays int) {
|
||||
if retentionDays <= 0 {
|
||||
return
|
||||
}
|
||||
cutoff := time.Now().Add(-time.Duration(retentionDays) * 24 * time.Hour).Unix()
|
||||
if err := model.DeletePerfMetricsBefore(cutoff); err != nil {
|
||||
common.SysError("failed to cleanup expired perf metrics: " + err.Error())
|
||||
}
|
||||
}
|
||||
|
||||
func redisCounters(values map[string]string) counters {
|
||||
return counters{
|
||||
requestCount: parseRedisInt(values["req"]),
|
||||
successCount: parseRedisInt(values["ok"]),
|
||||
totalLatencyMs: parseRedisInt(values["lat"]),
|
||||
ttftSumMs: parseRedisInt(values["ttft"]),
|
||||
ttftCount: parseRedisInt(values["ttft_n"]),
|
||||
outputTokens: parseRedisInt(values["out"]),
|
||||
generationMs: parseRedisInt(values["gen_ms"]),
|
||||
}
|
||||
}
|
||||
|
||||
func parseRedisInt(value string) int64 {
|
||||
if value == "" {
|
||||
return 0
|
||||
}
|
||||
parsed, _ := strconv.ParseInt(value, 10, 64)
|
||||
return parsed
|
||||
}
|
||||
@@ -0,0 +1,18 @@
|
||||
package perfmetrics
|
||||
|
||||
import (
|
||||
"github.com/stretchr/testify/require"
|
||||
"testing"
|
||||
)
|
||||
|
||||
func TestGroupBucketsWeightRequestsAndTtftSamples(t *testing.T) {
|
||||
merged := map[bucketKey]counters{}
|
||||
key := bucketKey{group: "test", bucketTs: 100}
|
||||
mergeCounters(merged, key, bucketPoint(100, counters{requestCount: 1, successCount: 0}).totals)
|
||||
mergeCounters(merged, key, bucketPoint(100, counters{requestCount: 9, successCount: 9, ttftSumMs: 9000, ttftCount: 3, outputTokens: 100, generationMs: 2000}).totals)
|
||||
result := buildQueryResult("", merged)
|
||||
require.Len(t, result.Groups, 1)
|
||||
require.Equal(t, 90.0, result.Groups[0].SuccessRate)
|
||||
require.Equal(t, int64(3000), result.Groups[0].AvgTtftMs)
|
||||
require.Equal(t, 50.0, result.Groups[0].AvgTps)
|
||||
}
|
||||
@@ -0,0 +1,465 @@
|
||||
package perfmetrics
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"math"
|
||||
"sort"
|
||||
"sync"
|
||||
"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/setting/perf_metrics_setting"
|
||||
)
|
||||
|
||||
var hotBuckets sync.Map
|
||||
|
||||
// seriesSchema is a stable client cache/schema marker. Do not change it when
|
||||
// hiding fields or making response-only privacy hardening changes.
|
||||
const seriesSchema = "dbcd0a3c01b55203"
|
||||
|
||||
func Init() {
|
||||
go flushLoop()
|
||||
}
|
||||
|
||||
func RecordRelaySample(info *relaycommon.RelayInfo, success bool, outputTokens int64) {
|
||||
if info == nil {
|
||||
return
|
||||
}
|
||||
now := time.Now()
|
||||
hasTtft := info.IsStream && info.HasSendResponse()
|
||||
ttftMs := int64(0)
|
||||
if hasTtft {
|
||||
ttftMs = info.FirstResponseTime.Sub(info.StartTime).Milliseconds()
|
||||
}
|
||||
latencyMs := now.Sub(info.StartTime).Milliseconds()
|
||||
generationMs := latencyMs
|
||||
if hasTtft {
|
||||
generationMs = now.Sub(info.FirstResponseTime).Milliseconds()
|
||||
}
|
||||
if generationMs <= 0 {
|
||||
generationMs = latencyMs
|
||||
}
|
||||
Record(Sample{
|
||||
Model: info.OriginModelName,
|
||||
Group: info.UsingGroup,
|
||||
LatencyMs: latencyMs,
|
||||
TtftMs: ttftMs,
|
||||
HasTtft: hasTtft,
|
||||
Success: success,
|
||||
OutputTokens: outputTokens,
|
||||
GenerationMs: generationMs,
|
||||
})
|
||||
}
|
||||
|
||||
func Record(sample Sample) {
|
||||
setting := perf_metrics_setting.GetSetting()
|
||||
if !setting.Enabled || sample.Model == "" {
|
||||
return
|
||||
}
|
||||
if sample.Group == "" {
|
||||
sample.Group = "default"
|
||||
}
|
||||
if sample.LatencyMs < 0 {
|
||||
sample.LatencyMs = 0
|
||||
}
|
||||
|
||||
key := bucketKey{
|
||||
model: sample.Model,
|
||||
group: sample.Group,
|
||||
bucketTs: bucketStart(time.Now().Unix()),
|
||||
}
|
||||
actual, _ := hotBuckets.LoadOrStore(key, &atomicBucket{})
|
||||
actual.(*atomicBucket).add(sample)
|
||||
recordRedis(key, sample)
|
||||
}
|
||||
|
||||
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
|
||||
|
||||
merged := map[bucketKey]counters{}
|
||||
rows, err := model.GetPerfMetrics(params.Model, params.Group, startTs, endTs)
|
||||
if err != nil {
|
||||
return QueryResult{}, err
|
||||
}
|
||||
for _, row := range rows {
|
||||
mergeCounters(merged, bucketKey{
|
||||
model: row.ModelName,
|
||||
group: row.Group,
|
||||
bucketTs: row.BucketTs,
|
||||
}, counters{
|
||||
requestCount: row.RequestCount,
|
||||
successCount: row.SuccessCount,
|
||||
totalLatencyMs: row.TotalLatencyMs,
|
||||
ttftSumMs: row.TtftSumMs,
|
||||
ttftCount: row.TtftCount,
|
||||
outputTokens: row.OutputTokens,
|
||||
generationMs: row.GenerationMs,
|
||||
})
|
||||
}
|
||||
|
||||
hotBuckets.Range(func(key, value any) bool {
|
||||
k := key.(bucketKey)
|
||||
if k.model != params.Model || k.bucketTs < startTs || k.bucketTs > endTs {
|
||||
return true
|
||||
}
|
||||
if params.Group != "" && k.group != params.Group {
|
||||
return true
|
||||
}
|
||||
mergeCounters(merged, k, value.(*atomicBucket).snapshot())
|
||||
return true
|
||||
})
|
||||
|
||||
return buildQueryResult(params.Model, merged), 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
|
||||
allowedGroups := allowedGroupSet(groups)
|
||||
|
||||
rows, err := model.GetPerfMetricsSummaryBucketsAll(startTs, endTs, groups)
|
||||
if err != nil {
|
||||
return SummaryAllResult{}, err
|
||||
}
|
||||
|
||||
totals := map[string]counters{}
|
||||
modelBuckets := map[string]map[int64]counters{}
|
||||
for _, row := range rows {
|
||||
value := counters{
|
||||
requestCount: row.RequestCount,
|
||||
successCount: row.SuccessCount,
|
||||
totalLatencyMs: row.TotalLatencyMs,
|
||||
outputTokens: row.OutputTokens,
|
||||
generationMs: row.GenerationMs,
|
||||
}
|
||||
mergeModelTotals(totals, row.ModelName, value)
|
||||
mergeModelBucket(modelBuckets, row.ModelName, row.BucketTs, value)
|
||||
}
|
||||
|
||||
hotBuckets.Range(func(key, value any) bool {
|
||||
k := key.(bucketKey)
|
||||
if k.bucketTs < startTs || k.bucketTs > endTs {
|
||||
return true
|
||||
}
|
||||
if allowedGroups != nil {
|
||||
if _, ok := allowedGroups[k.group]; !ok {
|
||||
return true
|
||||
}
|
||||
}
|
||||
snap := value.(*atomicBucket).snapshot()
|
||||
if snap.requestCount == 0 {
|
||||
return true
|
||||
}
|
||||
mergeModelTotals(totals, k.model, snap)
|
||||
mergeModelBucket(modelBuckets, k.model, k.bucketTs, snap)
|
||||
return true
|
||||
})
|
||||
|
||||
models := make([]ModelSummary, 0, len(totals))
|
||||
for name, total := range totals {
|
||||
if total.requestCount == 0 {
|
||||
continue
|
||||
}
|
||||
avgLatency := total.totalLatencyMs / total.requestCount
|
||||
successRate := float64(total.successCount) / float64(total.requestCount) * 100
|
||||
avgTps := 0.0
|
||||
if total.generationMs > 0 {
|
||||
avgTps = float64(total.outputTokens) / (float64(total.generationMs) / 1000.0)
|
||||
}
|
||||
models = append(models, ModelSummary{
|
||||
ModelName: name,
|
||||
AvgLatencyMs: avgLatency,
|
||||
SuccessRate: math.Round(successRate*100) / 100,
|
||||
AvgTps: math.Round(avgTps*100) / 100,
|
||||
RecentSuccessSeries: recentSuccessSeries(modelBuckets[name]),
|
||||
RequestCount: total.requestCount,
|
||||
})
|
||||
}
|
||||
sort.Slice(models, func(i, j int) bool {
|
||||
return models[i].RequestCount > models[j].RequestCount
|
||||
})
|
||||
|
||||
return SummaryAllResult{Models: models}, nil
|
||||
}
|
||||
|
||||
// QueryGroups combines raw bucket counters so busy models carry their actual weight.
|
||||
func QueryGroups(models []string, hours int) (QueryResult, error) {
|
||||
merged := map[bucketKey]counters{}
|
||||
seen := map[string]bool{}
|
||||
for _, name := range models {
|
||||
if seen[name] {
|
||||
continue
|
||||
}
|
||||
seen[name] = true
|
||||
result, err := Query(QueryParams{Model: name, Hours: hours})
|
||||
if err != nil {
|
||||
return QueryResult{}, err
|
||||
}
|
||||
for _, group := range result.Groups {
|
||||
for _, point := range group.Series {
|
||||
mergeCounters(merged, bucketKey{group: group.Group, bucketTs: point.Ts}, point.totals)
|
||||
}
|
||||
}
|
||||
}
|
||||
return buildQueryResult("", merged), nil
|
||||
}
|
||||
|
||||
func mergeModelTotals(totals map[string]counters, modelName string, value counters) {
|
||||
if value.requestCount == 0 {
|
||||
return
|
||||
}
|
||||
current := totals[modelName]
|
||||
current.requestCount += value.requestCount
|
||||
current.successCount += value.successCount
|
||||
current.totalLatencyMs += value.totalLatencyMs
|
||||
current.ttftSumMs += value.ttftSumMs
|
||||
current.ttftCount += value.ttftCount
|
||||
current.outputTokens += value.outputTokens
|
||||
current.generationMs += value.generationMs
|
||||
totals[modelName] = current
|
||||
}
|
||||
|
||||
func mergeModelBucket(modelBuckets map[string]map[int64]counters, modelName string, bucketTs int64, value counters) {
|
||||
if value.requestCount == 0 {
|
||||
return
|
||||
}
|
||||
if _, ok := modelBuckets[modelName]; !ok {
|
||||
modelBuckets[modelName] = map[int64]counters{}
|
||||
}
|
||||
current := modelBuckets[modelName][bucketTs]
|
||||
current.requestCount += value.requestCount
|
||||
current.successCount += value.successCount
|
||||
current.totalLatencyMs += value.totalLatencyMs
|
||||
current.ttftSumMs += value.ttftSumMs
|
||||
current.ttftCount += value.ttftCount
|
||||
current.outputTokens += value.outputTokens
|
||||
current.generationMs += value.generationMs
|
||||
modelBuckets[modelName][bucketTs] = current
|
||||
}
|
||||
|
||||
func recentSuccessSeries(buckets map[int64]counters) []SuccessRatePoint {
|
||||
if len(buckets) == 0 {
|
||||
return nil
|
||||
}
|
||||
hourly := map[int64]counters{}
|
||||
for ts, value := range buckets {
|
||||
hourTs := ts - ts%3600
|
||||
merged := hourly[hourTs]
|
||||
merged.requestCount += value.requestCount
|
||||
merged.successCount += value.successCount
|
||||
hourly[hourTs] = merged
|
||||
}
|
||||
timestamps := make([]int64, 0, len(hourly))
|
||||
for hourTs, value := range hourly {
|
||||
if value.requestCount == 0 {
|
||||
continue
|
||||
}
|
||||
timestamps = append(timestamps, hourTs)
|
||||
}
|
||||
if len(timestamps) == 0 {
|
||||
return nil
|
||||
}
|
||||
sort.Slice(timestamps, func(i, j int) bool {
|
||||
return timestamps[i] < timestamps[j]
|
||||
})
|
||||
points := make([]SuccessRatePoint, 0, len(timestamps))
|
||||
for _, hourTs := range timestamps {
|
||||
points = append(points, SuccessRatePoint{
|
||||
Ts: hourTs,
|
||||
SuccessRate: math.Round(successRate(hourly[hourTs])*100) / 100,
|
||||
})
|
||||
}
|
||||
return points
|
||||
}
|
||||
|
||||
func allowedGroupSet(groups []string) map[string]struct{} {
|
||||
if groups == nil {
|
||||
return nil
|
||||
}
|
||||
allowed := make(map[string]struct{}, len(groups))
|
||||
for _, group := range groups {
|
||||
allowed[group] = struct{}{}
|
||||
}
|
||||
return allowed
|
||||
}
|
||||
|
||||
func bucketStart(ts int64) int64 {
|
||||
bucketSeconds := perf_metrics_setting.GetBucketSeconds()
|
||||
if bucketSeconds <= 0 {
|
||||
bucketSeconds = 3600
|
||||
}
|
||||
return ts - (ts % bucketSeconds)
|
||||
}
|
||||
|
||||
func mergeCounters(merged map[bucketKey]counters, key bucketKey, value counters) {
|
||||
if value.requestCount == 0 {
|
||||
return
|
||||
}
|
||||
current := merged[key]
|
||||
current.requestCount += value.requestCount
|
||||
current.successCount += value.successCount
|
||||
current.totalLatencyMs += value.totalLatencyMs
|
||||
current.ttftSumMs += value.ttftSumMs
|
||||
current.ttftCount += value.ttftCount
|
||||
current.outputTokens += value.outputTokens
|
||||
current.generationMs += value.generationMs
|
||||
merged[key] = current
|
||||
}
|
||||
|
||||
func buildQueryResult(modelName string, merged map[bucketKey]counters) QueryResult {
|
||||
groupBuckets := map[string]map[int64]counters{}
|
||||
for key, value := range merged {
|
||||
if value.requestCount == 0 {
|
||||
continue
|
||||
}
|
||||
if _, ok := groupBuckets[key.group]; !ok {
|
||||
groupBuckets[key.group] = map[int64]counters{}
|
||||
}
|
||||
groupBuckets[key.group][key.bucketTs] = value
|
||||
}
|
||||
|
||||
groups := make([]string, 0, len(groupBuckets))
|
||||
for group := range groupBuckets {
|
||||
groups = append(groups, group)
|
||||
}
|
||||
sort.Strings(groups)
|
||||
|
||||
results := make([]GroupResult, 0, len(groups))
|
||||
for _, group := range groups {
|
||||
buckets := groupBuckets[group]
|
||||
timestamps := make([]int64, 0, len(buckets))
|
||||
for ts := range buckets {
|
||||
timestamps = append(timestamps, ts)
|
||||
}
|
||||
sort.Slice(timestamps, func(i, j int) bool {
|
||||
return timestamps[i] < timestamps[j]
|
||||
})
|
||||
|
||||
total := counters{}
|
||||
series := make([]BucketPoint, 0, len(timestamps))
|
||||
for _, ts := range timestamps {
|
||||
value := buckets[ts]
|
||||
total.requestCount += value.requestCount
|
||||
total.successCount += value.successCount
|
||||
total.totalLatencyMs += value.totalLatencyMs
|
||||
total.ttftSumMs += value.ttftSumMs
|
||||
total.ttftCount += value.ttftCount
|
||||
total.outputTokens += value.outputTokens
|
||||
total.generationMs += value.generationMs
|
||||
series = append(series, bucketPoint(ts, value))
|
||||
}
|
||||
|
||||
results = append(results, GroupResult{
|
||||
Group: group,
|
||||
AvgTtftMs: avg(total.ttftSumMs, total.ttftCount),
|
||||
AvgLatencyMs: avg(total.totalLatencyMs, total.requestCount),
|
||||
SuccessRate: successRate(total),
|
||||
AvgTps: avgTps(total),
|
||||
Series: series,
|
||||
})
|
||||
}
|
||||
|
||||
return QueryResult{
|
||||
ModelName: modelName,
|
||||
SeriesSchema: seriesSchema,
|
||||
Groups: results,
|
||||
}
|
||||
}
|
||||
|
||||
func bucketPoint(ts int64, value counters) BucketPoint {
|
||||
return BucketPoint{
|
||||
totals: value,
|
||||
Ts: ts,
|
||||
AvgTtftMs: avg(value.ttftSumMs, value.ttftCount),
|
||||
AvgLatencyMs: avg(value.totalLatencyMs, value.requestCount),
|
||||
SuccessRate: successRate(value),
|
||||
AvgTps: avgTps(value),
|
||||
}
|
||||
}
|
||||
|
||||
func avg(sum int64, count int64) int64 {
|
||||
if count <= 0 {
|
||||
return 0
|
||||
}
|
||||
return sum / count
|
||||
}
|
||||
|
||||
func successRate(value counters) float64 {
|
||||
if value.requestCount <= 0 {
|
||||
return 0
|
||||
}
|
||||
return float64(value.successCount) / float64(value.requestCount) * 100
|
||||
}
|
||||
|
||||
func avgTps(value counters) float64 {
|
||||
if value.outputTokens <= 0 || value.generationMs <= 0 {
|
||||
return 0
|
||||
}
|
||||
return float64(value.outputTokens) / (float64(value.generationMs) / 1000)
|
||||
}
|
||||
|
||||
func recordRedis(key bucketKey, sample Sample) {
|
||||
if !common.RedisEnabled || common.RDB == nil {
|
||||
return
|
||||
}
|
||||
ctx, cancel := context.WithTimeout(context.Background(), time.Second)
|
||||
defer cancel()
|
||||
|
||||
redisKey := redisBucketKey(key)
|
||||
pipe := common.RDB.TxPipeline()
|
||||
pipe.HIncrBy(ctx, redisKey, "req", 1)
|
||||
if sample.Success {
|
||||
pipe.HIncrBy(ctx, redisKey, "ok", 1)
|
||||
}
|
||||
if sample.LatencyMs > 0 {
|
||||
pipe.HIncrBy(ctx, redisKey, "lat", sample.LatencyMs)
|
||||
}
|
||||
if sample.HasTtft && sample.TtftMs >= 0 {
|
||||
pipe.HIncrBy(ctx, redisKey, "ttft", sample.TtftMs)
|
||||
pipe.HIncrBy(ctx, redisKey, "ttft_n", 1)
|
||||
}
|
||||
if sample.OutputTokens > 0 && sample.GenerationMs > 0 {
|
||||
pipe.HIncrBy(ctx, redisKey, "out", sample.OutputTokens)
|
||||
pipe.HIncrBy(ctx, redisKey, "gen_ms", sample.GenerationMs)
|
||||
}
|
||||
pipe.Expire(ctx, redisKey, time.Hour)
|
||||
_, _ = pipe.Exec(ctx)
|
||||
}
|
||||
|
||||
func mergeRedisActiveBuckets(merged map[bucketKey]counters, params QueryParams, startTs int64, endTs int64) {
|
||||
if !common.RedisEnabled || common.RDB == nil || params.Model == "" || params.Group == "" {
|
||||
return
|
||||
}
|
||||
active := bucketStart(time.Now().Unix())
|
||||
if active < startTs || active > endTs {
|
||||
return
|
||||
}
|
||||
key := bucketKey{model: params.Model, group: params.Group, bucketTs: active}
|
||||
ctx, cancel := context.WithTimeout(context.Background(), time.Second)
|
||||
defer cancel()
|
||||
values, err := common.RDB.HGetAll(ctx, redisBucketKey(key)).Result()
|
||||
if err != nil || len(values) == 0 {
|
||||
return
|
||||
}
|
||||
mergeCounters(merged, key, redisCounters(values))
|
||||
}
|
||||
|
||||
func redisBucketKey(key bucketKey) string {
|
||||
return fmt.Sprintf("perf:%s:%s:%d", key.model, key.group, key.bucketTs)
|
||||
}
|
||||
@@ -0,0 +1,159 @@
|
||||
package perfmetrics
|
||||
|
||||
import "sync/atomic"
|
||||
|
||||
type Store interface {
|
||||
Record(sample Sample)
|
||||
Query(params QueryParams) (QueryResult, error)
|
||||
}
|
||||
|
||||
type Sample struct {
|
||||
Model string
|
||||
Group string
|
||||
LatencyMs int64
|
||||
TtftMs int64
|
||||
HasTtft bool
|
||||
Success bool
|
||||
OutputTokens int64
|
||||
GenerationMs int64
|
||||
}
|
||||
|
||||
type QueryParams struct {
|
||||
Model string
|
||||
Group string
|
||||
Hours int
|
||||
}
|
||||
|
||||
type BucketPoint struct {
|
||||
totals counters
|
||||
Ts int64 `json:"ts"`
|
||||
AvgTtftMs int64 `json:"avg_ttft_ms"`
|
||||
AvgLatencyMs int64 `json:"avg_latency_ms"`
|
||||
SuccessRate float64 `json:"success_rate"`
|
||||
AvgTps float64 `json:"avg_tps"`
|
||||
}
|
||||
|
||||
type GroupResult struct {
|
||||
Group string `json:"group"`
|
||||
AvgTtftMs int64 `json:"avg_ttft_ms"`
|
||||
AvgLatencyMs int64 `json:"avg_latency_ms"`
|
||||
SuccessRate float64 `json:"success_rate"`
|
||||
AvgTps float64 `json:"avg_tps"`
|
||||
Series []BucketPoint `json:"series"`
|
||||
}
|
||||
|
||||
type QueryResult struct {
|
||||
ModelName string `json:"model_name"`
|
||||
SeriesSchema string `json:"series_schema"`
|
||||
Groups []GroupResult `json:"groups"`
|
||||
}
|
||||
|
||||
type SuccessRatePoint struct {
|
||||
Ts int64 `json:"ts"`
|
||||
SuccessRate float64 `json:"success_rate"`
|
||||
}
|
||||
|
||||
type ModelSummary struct {
|
||||
ModelName string `json:"model_name"`
|
||||
AvgLatencyMs int64 `json:"avg_latency_ms"`
|
||||
SuccessRate float64 `json:"success_rate"`
|
||||
AvgTps float64 `json:"avg_tps"`
|
||||
RecentSuccessSeries []SuccessRatePoint `json:"recent_success_series,omitempty"`
|
||||
RequestCount int64 `json:"-"`
|
||||
}
|
||||
|
||||
type SummaryAllResult struct {
|
||||
Models []ModelSummary `json:"models"`
|
||||
}
|
||||
|
||||
type bucketKey struct {
|
||||
model string
|
||||
group string
|
||||
bucketTs int64
|
||||
}
|
||||
|
||||
type counters struct {
|
||||
requestCount int64
|
||||
successCount int64
|
||||
totalLatencyMs int64
|
||||
ttftSumMs int64
|
||||
ttftCount int64
|
||||
outputTokens int64
|
||||
generationMs int64
|
||||
}
|
||||
|
||||
type atomicBucket struct {
|
||||
requestCount atomic.Int64
|
||||
successCount atomic.Int64
|
||||
totalLatencyMs atomic.Int64
|
||||
ttftSumMs atomic.Int64
|
||||
ttftCount atomic.Int64
|
||||
outputTokens atomic.Int64
|
||||
generationMs atomic.Int64
|
||||
}
|
||||
|
||||
func (b *atomicBucket) add(sample Sample) {
|
||||
b.requestCount.Add(1)
|
||||
if sample.Success {
|
||||
b.successCount.Add(1)
|
||||
}
|
||||
if sample.LatencyMs > 0 {
|
||||
b.totalLatencyMs.Add(sample.LatencyMs)
|
||||
}
|
||||
if sample.HasTtft && sample.TtftMs >= 0 {
|
||||
b.ttftSumMs.Add(sample.TtftMs)
|
||||
b.ttftCount.Add(1)
|
||||
}
|
||||
if sample.OutputTokens > 0 && sample.GenerationMs > 0 {
|
||||
b.outputTokens.Add(sample.OutputTokens)
|
||||
b.generationMs.Add(sample.GenerationMs)
|
||||
}
|
||||
}
|
||||
|
||||
func (b *atomicBucket) snapshot() counters {
|
||||
return counters{
|
||||
requestCount: b.requestCount.Load(),
|
||||
successCount: b.successCount.Load(),
|
||||
totalLatencyMs: b.totalLatencyMs.Load(),
|
||||
ttftSumMs: b.ttftSumMs.Load(),
|
||||
ttftCount: b.ttftCount.Load(),
|
||||
outputTokens: b.outputTokens.Load(),
|
||||
generationMs: b.generationMs.Load(),
|
||||
}
|
||||
}
|
||||
|
||||
func (b *atomicBucket) drain() counters {
|
||||
return counters{
|
||||
requestCount: b.requestCount.Swap(0),
|
||||
successCount: b.successCount.Swap(0),
|
||||
totalLatencyMs: b.totalLatencyMs.Swap(0),
|
||||
ttftSumMs: b.ttftSumMs.Swap(0),
|
||||
ttftCount: b.ttftCount.Swap(0),
|
||||
outputTokens: b.outputTokens.Swap(0),
|
||||
generationMs: b.generationMs.Swap(0),
|
||||
}
|
||||
}
|
||||
|
||||
func (b *atomicBucket) addCounters(c counters) {
|
||||
if c.requestCount != 0 {
|
||||
b.requestCount.Add(c.requestCount)
|
||||
}
|
||||
if c.successCount != 0 {
|
||||
b.successCount.Add(c.successCount)
|
||||
}
|
||||
if c.totalLatencyMs != 0 {
|
||||
b.totalLatencyMs.Add(c.totalLatencyMs)
|
||||
}
|
||||
if c.ttftSumMs != 0 {
|
||||
b.ttftSumMs.Add(c.ttftSumMs)
|
||||
}
|
||||
if c.ttftCount != 0 {
|
||||
b.ttftCount.Add(c.ttftCount)
|
||||
}
|
||||
if c.outputTokens != 0 {
|
||||
b.outputTokens.Add(c.outputTokens)
|
||||
}
|
||||
if c.generationMs != 0 {
|
||||
b.generationMs.Add(c.generationMs)
|
||||
}
|
||||
}
|
||||
Reference in new issue
Block a user