feat: publish v1.0.0-rc.36-custom.1

This commit is contained in:
Codex committed 2026-09-11 05:07:48 +08:00
commit 3a7e8f3b54
2381 files changed
+551494

No files matched your search

File diff suppressed because it is too large. Load diff
+331
View File
@@ -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()
}
+21
View File
@@ -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",`))
}
+353
View File
@@ -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` |
+19
View File
@@ -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)
}
+166
View File
@@ -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
}
+42
View File
@@ -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
}
+52
View File
@@ -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")
}
+19
View File
@@ -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)
}
+85
View File
@@ -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)
}
+53
View File
@@ -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
}
+285
View File
@@ -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()
}
+38
View File
@@ -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 + "*"
}
+219
View File
@@ -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 ""
}
+302
View File
@@ -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
}
+377
View File
@@ -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
}
+202
View File
@@ -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
}
+96
View File
@@ -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
}
+353
View File
@@ -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"`
}
+92
View File
@@ -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 ""
}
+57
View File
@@ -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(); }
`
+574
View File
@@ -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()
}
+375
View File
@@ -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())
})
}
}
+116
View File
@@ -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
}
+49
View File
@@ -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)
}
+126
View File
@@ -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://")
}
+26
View File
@@ -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)
}
+245
View File
@@ -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
+43
View File
@@ -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)
}
+963
View File
@@ -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
+139
View File
@@ -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
}
+98
View File
@@ -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
}
+18
View File
@@ -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)
}
+465
View File
@@ -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)
}
+159
View File
@@ -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)
}
}