feat(harness): 成本/Token 预算护栏 —— 单任务硬上限 + 单用户日预算(恒温器最后一环)

token 用量估算计量(CJK≈1/字、ASCII≈1/4字,无需分词器,护栏够用)。

单任务硬上限(dispatcher):Budget 挂 ctx 沿图透传,各 LLM 节点(对话/ReAct/compose/
报告)入口计输入 token、出口计输出,触顶即中止整图——防失控成本(死循环/超大报告)。
报告路径优雅降级:触顶跳过剩余章节出部分稿,不整体失败。预算来源 Meta.token_budget
或 env TASK_TOKEN_BUDGET(默认 20 万)。

单用户日预算(gateway):dispatcher 收尾经 NATS 回写 UsageEvent → 网关按用户按天累计
Redis(48h 过期自滚动)→ 提交前门控 USER_DAILY_TOKEN_BUDGET(0=不限,超额 402)。
/billing 升级为真实用量:当日已用 / 日预算 / 余额。

契约 UsageEvent + MetaTokenBudget + SubjectUsage;bus Publish/SubscribeUsage;
orchestrator SetUsageSink + 预算触顶 failed(不计熔断)。harness budget 5 单测,三模块全绿。
live:单任务 budget=30 → failed(已用约689);用户日 budget=200 → /billing remaining=0 → 402。

至此 harness 由「测温计」完成向「恒温器」的演进(评测闭环/纠偏/忠实度/脱敏/输入护栏/预算六项)。

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
This commit is contained in:
Blizzard
2026-06-26 10:00:51 +08:00
parent 2f78fc565e
commit faa1871760
16 changed files with 442 additions and 17 deletions
+5 -3
View File
@@ -107,8 +107,10 @@ Harness = 围绕 LLM 的可靠性 / 安全 / 质量治理层。4 个组件均为
| 输出脱敏 RedactSecrets | 26 行 + 39 测试 | ⭐⭐⭐ | 4 正则(sk-/AKIA/JWT/Bearer),流式逐片 |
| 输入护栏 Guardrail | guardrail.go + 测试 | ⭐⭐½ | 注入正则(中英) + 体积限制 + 敏感词黑名单(默认空) |
**定性**一个真·生产级熔断器 + 三个"能跑但偏启发式"的护栏/评测。当前更像**「测温计」(检测+记录)**
而非**「恒温器」(检测→决策→纠偏闭环)**——多数信号只打日志、不门控、不纠偏。
**定性(更新)**起初是**「测温计」(检测+记录)**——多数信号只打日志。经评测闭环、低分自动纠偏、
忠实度评测、输出脱敏增强、输入护栏升级、成本预算护栏六项,现已成**「恒温器」(检测→决策→门控/纠偏闭环)**
熔断挡后端雪崩、输入护栏两层挡注入/越狱、评测分级+低分重生成、输出脱敏挡密钥/PII、预算护栏挡失控成本。
下一步重心转 P0/P1 生产硬骨头(优雅停机、HA、容量实测、本地模型)。
**已知短板 / 优化清单**(按性价比,逐项推进):
@@ -123,7 +125,7 @@ Harness = 围绕 LLM 的可靠性 / 安全 / 质量治理层。4 个组件均为
- [x] **P3 坏输出自动纠偏** ✅:poor(<0.5) 触发评语驱动的重生成,重评后**仅采纳更优者(不退步)**,
采纳的修订版落会话历史 + 评测终值带 `corrected` 标记落库。`maxRefineRounds=1``canRefine`(模型就绪且熔断未开)门控;
单 goroutine 串 评测→纠偏→落历史避竞态。单测覆盖 采纳/不退步/非低分不触发,live 验证好答案不误触发。**至此 harness 由「测温计」迈入「恒温器」。**
- [ ] **P3 成本/Token 预算护栏**:单任务/单用户 token 上限与告警
- [x] **P3 成本/Token 预算护栏** ✅:用量估算计量(CJK/ASCII 启发式,无需分词器)。**单任务硬上限**:Budget 挂 ctx 沿图透传,各 LLM 节点(对话/ReAct/compose/报告)入口计输入、出口计输出,触顶即中止(报告优雅降级跳过剩余章节);防失控成本。**单用户日预算**:dispatcher 收尾经 NATS 回写 UsageEvent → 网关按用户按天累计 Redis → 提交前门控(USER_DAILY_TOKEN_BUDGET,超额 402)`/billing` 出当日用量/预算/余额。live:单任务 budget=30 → failed(已用约689);用户日 budget=200 → 402
---
@@ -55,6 +55,7 @@ func main() {
log.Fatalf("[dispatcher] build eino graph: %v", err)
}
orch.SetGuardian(guardian) // 输入护栏 Tier2
orch.SetUsageSink(sub) // 成本护栏:token 用量回写网关累计/计费
// 健康心跳:dispatcher 无 HTTP/工具端点,挂一个 NATS 应答让管理端「服务状态」探到它在线。
startedAt := time.Now()
@@ -63,6 +63,19 @@ func (o *Orchestrator) runComposeConversation(ctx context.Context, taskID string
Upstream: append([]string{}, b.agentOut...), // 前序协作 agent 产出 → 接力
}
msgs, _ := buildMessages(ctx, rc)
// 成本护栏:计入输入 token;触顶则中止整图。
if bud := harness.BudgetFrom(ctx); bud != nil {
for _, m := range msgs {
bud.AddPrompt(m.Content)
}
if bud.Exceeded() {
if b.fatalErr == nil {
b.fatalErr = errBudget
}
tr.emit(node, "system", "error", "token 预算", "已达单任务预算上限,中止", 0)
return
}
}
t0 := time.Now()
// ChatModel 的 start/end 由 composeTracercallbacks)落轨迹,这里不再手写 emit(归一)。
@@ -74,7 +87,7 @@ func (o *Orchestrator) runComposeConversation(ctx context.Context, taskID string
defer sr.Close()
chunks := 0
var produced strings.Builder // 本节点产出(供下游 agent 接力)
var produced strings.Builder // 本节点产出(供下游 agent 接力)
red := harness.NewStreamRedactor() // 输出护栏:跨分片脱敏,杜绝密钥被切断而漏检
emit := func(safe string) {
if safe == "" {
@@ -99,6 +112,9 @@ func (o *Orchestrator) runComposeConversation(ctx context.Context, taskID string
emit(red.Push(chunk.Content))
}
emit(red.Flush()) // 吐出暂留尾部
if bud := harness.BudgetFrom(ctx); bud != nil {
bud.AddComplete(produced.String()) // 成本护栏:计入输出 token
}
o.recordAgentOutput(b, produced.String())
tr.info(node, "system", "compose 图", fmt.Sprintf("%d 段输出 / %d 字(Eino compose 运行时)", chunks, len([]rune(produced.String()))))
}
@@ -273,6 +273,19 @@ func (o *Orchestrator) runAgent(ctx context.Context, taskID string, b *board, sy
Upstream: append([]string{}, b.agentOut...), // 前序协作 agent 产出 → 接力
}
msgs, _ := buildMessages(ctx, rc)
// 成本护栏:计入本节点输入 token;若已触顶则中止整图(防失控成本)。
if bud := harness.BudgetFrom(ctx); bud != nil {
for _, m := range msgs {
bud.AddPrompt(m.Content)
}
if bud.Exceeded() {
if b.fatalErr == nil {
b.fatalErr = errBudget
}
tr.emit(node, "system", "error", "token 预算", "已达单任务预算上限,中止", 0)
return
}
}
tr.emit(node, "model", "start", "模型流式推理", "", 0)
t0 := time.Now()
n := 0
@@ -304,6 +317,9 @@ func (o *Orchestrator) runAgent(ctx context.Context, taskID string, b *board, sy
return
}
emit(red.Flush()) // 吐出暂留的尾部(最后一段疑似密钥的判定)
if bud := harness.BudgetFrom(ctx); bud != nil {
bud.AddComplete(produced.String()) // 成本护栏:计入本节点输出 token
}
if red.Hits() > 0 {
tr.info(node, "system", "输出护栏", fmt.Sprintf("已脱敏 %d 处疑似密钥/PII", red.Hits()))
}
@@ -8,6 +8,8 @@ import (
"fmt"
"log"
"log/slog"
"os"
"strconv"
"strings"
"sync"
"time"
@@ -51,6 +53,14 @@ type EvalSink interface {
PublishEval(ev *contract.EvalEvent) error
}
// UsageSink 回写任务 token 用量供网关累计/计费(由 NATS bus 实现;可为 nil → 不回写)。
type UsageSink interface {
PublishUsage(ev *contract.UsageEvent) error
}
// errBudget 是单任务 token 预算触顶时的哨兵错误:任务以 failed 收尾并附明确原因(防失控成本)。
var errBudget = errors.New("token 预算超限,已中止")
// errRejected 是审批节点拒绝(或超时)时图执行返回的哨兵错误:它是合法终态而非故障,
// Handle 据此判 rejected 并优雅收尾(不计熔断失败)。
var errRejected = errors.New("approval rejected")
@@ -79,16 +89,17 @@ const approvalTimeout = 5 * time.Minute
// Orchestrator 把每个 DSL 任务动态编译为 Eino 图并执行(记忆召回 → 工具节点 → 注入 → 流式)。
type Orchestrator struct {
pool LLM
breaker *harness.CircuitBreaker
eval *harness.Evaluator
sink TokenSink
tools ToolCaller
exec ExecSink
status StatusSink // 任务生命周期状态回写(可为 nil)
approval ApprovalWaiter // HITL 审批等待(可为 nil → 审批节点自动放行)
evalSink EvalSink // 评测结果回写落库(可为 nil → 仅打日志)
guard *harness.Classifier // 输入护栏 Tier2:对网关标记的灰区任务做 LLM 裁决(可为 nil → 不做)
pool LLM
breaker *harness.CircuitBreaker
eval *harness.Evaluator
sink TokenSink
tools ToolCaller
exec ExecSink
status StatusSink // 任务生命周期状态回写(可为 nil)
approval ApprovalWaiter // HITL 审批等待(可为 nil → 审批节点自动放行)
evalSink EvalSink // 评测结果回写落库(可为 nil → 仅打日志)
guard *harness.Classifier // 输入护栏 Tier2:对网关标记的灰区任务做 LLM 裁决(可为 nil → 不做)
usageSink UsageSink // token 用量回写(可为 nil → 不回写)
turnMu sync.Mutex // 保护 turns(攒批计数,多任务 goroutine 共享)
turns map[string]int // sessionID → 累计轮次,用于每 N 轮触发 consolidate
@@ -104,6 +115,52 @@ func NewOrchestrator(pool LLM, breaker *harness.CircuitBreaker, eval *harness.Ev
// SetGuardian 注入输入护栏 Tier2 的 LLM 分类器(可选;不注入则灰区任务直接放行执行)。
func (o *Orchestrator) SetGuardian(c *harness.Classifier) { o.guard = c }
// SetUsageSink 注入 token 用量回写出口(可选;不注入则不上报用量)。
func (o *Orchestrator) SetUsageSink(s UsageSink) { o.usageSink = s }
// taskBudget 取本任务的 token 预算上限:优先 Meta(网关按用户/套餐下发),否则 env TASK_TOKEN_BUDGET(默认 20 万)。
func (o *Orchestrator) taskBudget(t *contract.Task) int {
switch n := t.Meta[contract.MetaTokenBudget].(type) {
case float64:
if n > 0 {
return int(n)
}
case int:
if n > 0 {
return n
}
}
return envInt("TASK_TOKEN_BUDGET", 200000)
}
// emitUsage 任务收尾回写本轮 token 用量(用量为 0 或无出口则跳过)。
func (o *Orchestrator) emitUsage(t *contract.Task, b *harness.Budget) {
if o.usageSink == nil {
return
}
p, c, total := b.Snapshot()
if total == 0 {
return
}
uid, _ := t.Meta[contract.MetaUserID].(string)
if err := o.usageSink.PublishUsage(&contract.UsageEvent{
TaskID: t.ID, UserID: uid, PromptTok: p, CompTok: c, TotalTok: total,
Exceeded: b.Exceeded(), TS: time.Now().UnixMilli(),
}); err != nil {
log.Printf("[usage] 回写用量失败 task=%s: %v", t.ID, err)
}
}
// envInt 读正整数环境变量,缺省回退 def。
func envInt(key string, def int) int {
if v := os.Getenv(key); v != "" {
if n, err := strconv.Atoi(v); err == nil && n > 0 {
return n
}
}
return def
}
// setStatus 回写一次任务状态流转(status 为 nil 时静默跳过)。
func (o *Orchestrator) setStatus(taskID, status, detail string) {
if o.status == nil {
@@ -180,6 +237,11 @@ func (o *Orchestrator) Handle(ctx context.Context, t *contract.Task) error {
tctx, cancel := context.WithTimeout(ctx, taskExecTimeout)
defer cancel()
// 成本护栏:单任务 token 预算挂到 ctx,沿图执行各 LLM 节点计量+封顶;收尾回写用量(计费/日预算)。
budget := harness.NewBudget(o.taskBudget(t))
tctx = harness.WithBudget(tctx, budget)
defer o.emitUsage(t, budget)
// 报告生成走专用多步编排(规划→分章并行检索撰写→汇聚→渲染 Word),而非通用对话图。
if intent, _ := t.Meta[contract.MetaIntent].(string); intent == contract.IntentReport {
err := o.handleReport(tctx, t, tr)
@@ -203,6 +265,19 @@ func (o *Orchestrator) Handle(ctx context.Context, t *contract.Task) error {
o.setStatus(t.ID, contract.TaskRejected, truncate(answer, 120))
return nil
}
if errors.Is(err, errBudget) {
// token 预算触顶:策略性中止,非后端故障。收尾流 + 置 failed(附明确原因),不计熔断、不重投。
_, _, total := budget.Snapshot()
slog.WarnContext(ctx, "task aborted: token budget exceeded", "task_id", t.ID, "tokens", total)
if answer != "" {
_ = o.sink.PublishToken(t.ID, []byte(answer))
}
_ = o.sink.PublishToken(t.ID, []byte("\n\n⚠️ 已达单任务 token 预算上限,自动中止。"))
_ = o.sink.CompleteStream(t.ID)
o.breaker.Report(true)
o.setStatus(t.ID, contract.TaskFailed, fmt.Sprintf("token 预算超限(已用约 %d", total))
return nil
}
if err != nil {
span.RecordError(err)
span.SetStatus(codes.Error, err.Error())
@@ -202,6 +202,19 @@ func (o *Orchestrator) runReactAgent(ctx context.Context, taskID string, b *boar
Upstream: append([]string{}, b.agentOut...), // 前序协作 agent 产出 → 接力
}
msgs, _ := buildMessages(ctx, rc)
// 成本护栏:计入输入 token;触顶则中止整图。
if bud := harness.BudgetFrom(ctx); bud != nil {
for _, m := range msgs {
bud.AddPrompt(m.Content)
}
if bud.Exceeded() {
if b.fatalErr == nil {
b.fatalErr = errBudget
}
tr.emit(node, "system", "error", "token 预算", "已达单任务预算上限,中止", 0)
return
}
}
tr.emit(node, "model", "start", "ReAct 智能体(自主调工具)", fmt.Sprintf("%d 个工具可用", len(tools)), 0)
t0 := time.Now()
@@ -213,7 +226,7 @@ func (o *Orchestrator) runReactAgent(ctx context.Context, taskID string, b *boar
defer sr.Close()
chunks := 0
var produced strings.Builder // 本节点产出(供下游 agent 接力)
var produced strings.Builder // 本节点产出(供下游 agent 接力)
red := harness.NewStreamRedactor() // 输出护栏:跨分片脱敏,杜绝密钥被切断而漏检
emit := func(safe string) {
if safe == "" {
@@ -238,6 +251,9 @@ func (o *Orchestrator) runReactAgent(ctx context.Context, taskID string, b *boar
emit(red.Push(chunk.Content))
}
emit(red.Flush()) // 吐出暂留尾部
if bud := harness.BudgetFrom(ctx); bud != nil {
bud.AddComplete(produced.String()) // 成本护栏:计入输出 token
}
o.recordAgentOutput(b, produced.String())
tr.emit(node, "model", "end", "ReAct 智能体",
fmt.Sprintf("%d 段输出 / %d 字", chunks, len([]rune(produced.String()))), time.Since(t0).Milliseconds())
@@ -10,6 +10,7 @@ import (
"time"
"github.com/sundynix/sundynix-dispatcher/internal/dsl"
"github.com/sundynix/sundynix-dispatcher/internal/harness"
"github.com/sundynix/sundynix-dispatcher/internal/llm"
"github.com/sundynix/sundynix-shared/contract"
)
@@ -204,6 +205,14 @@ func (o *Orchestrator) writeSection(ctx context.Context, topic, kb, heading stri
ub.WriteString("\n")
}
ub.WriteString("请就「本章标题」撰写 200–400 字正文。只输出正文,不要重复标题、不要再列提纲。")
// 成本护栏:计入输入;若已触顶则跳过本章(报告优雅降级出部分稿,不整体失败)。
if bud := harness.BudgetFrom(ctx); bud != nil {
bud.AddPrompt(sys + ub.String())
if bud.Exceeded() {
tr.info(node, "system", "token 预算", "已达预算上限,跳过本章")
return "(已达 token 预算上限,本章自动跳过。)"
}
}
cctx, cancel := llmCtx(ctx)
defer cancel()
txt, err := o.pool.Chat(cctx, []llm.ChatMessage{{Role: "system", Content: sys}, {Role: "user", Content: ub.String()}})
@@ -211,6 +220,9 @@ func (o *Orchestrator) writeSection(ctx context.Context, topic, kb, heading stri
log.Printf("[report] 撰写「%s」失败: %v", heading, err)
return "(本章撰写失败:" + err.Error() + ""
}
if bud := harness.BudgetFrom(ctx); bud != nil {
bud.AddComplete(txt) // 成本护栏:计入输出
}
return strings.TrimSpace(txt)
}
@@ -0,0 +1,82 @@
package harness
import (
"context"
"sync"
"unicode"
)
// EstimateTokens 粗估一段文本的 token 数(无需依赖模型/分词器,作预算护栏足够):
// CJK 字符约 1 token/字;其余(拉丁/数字/空白/标点)约 1 token/4 字符。偏保守(宁高勿低)。
func EstimateTokens(s string) int {
cjk, other := 0, 0
for _, r := range s {
if unicode.Is(unicode.Han, r) || unicode.Is(unicode.Hiragana, r) ||
unicode.Is(unicode.Katakana, r) || unicode.Is(unicode.Hangul, r) {
cjk++
} else {
other++
}
}
return cjk + (other+3)/4 // 向上取整
}
// Budget 是单任务的 token 预算计量与封顶(并发安全:流式回调与节点循环可能并发计量)。
// max ≤ 0 表示不限额(仅计量、不中止)。
type Budget struct {
mu sync.Mutex
max int
prompt int
comp int
}
// NewBudget 建一个上限为 max 的任务预算(max≤0 → 不限额,纯计量)。
func NewBudget(max int) *Budget { return &Budget{max: max} }
// AddPrompt / AddComplete 累计输入/输出 token(按文本估算)。
func (b *Budget) AddPrompt(text string) { b.add(EstimateTokens(text), 0) }
func (b *Budget) AddComplete(text string) { b.add(0, EstimateTokens(text)) }
func (b *Budget) add(p, c int) {
if b == nil {
return
}
b.mu.Lock()
b.prompt += p
b.comp += c
b.mu.Unlock()
}
// Exceeded 报告是否已触顶(max≤0 恒为 false)。
func (b *Budget) Exceeded() bool {
if b == nil {
return false
}
b.mu.Lock()
defer b.mu.Unlock()
return b.max > 0 && b.prompt+b.comp > b.max
}
// Snapshot 返回当前用量快照(prompt / completion / total)。
func (b *Budget) Snapshot() (prompt, comp, total int) {
if b == nil {
return 0, 0, 0
}
b.mu.Lock()
defer b.mu.Unlock()
return b.prompt, b.comp, b.prompt + b.comp
}
// budgetKey 是 context 里承载任务预算的私有键(避免跨包碰撞)。
type budgetKey struct{}
// WithBudget 把任务预算挂到 context,沿图执行透传(免改各节点函数签名)。
func WithBudget(ctx context.Context, b *Budget) context.Context {
return context.WithValue(ctx, budgetKey{}, b)
}
// BudgetFrom 取出 context 里的任务预算(无则返回 nil,调用方按"不限额"处理)。
func BudgetFrom(ctx context.Context) *Budget {
b, _ := ctx.Value(budgetKey{}).(*Budget)
return b
}
@@ -0,0 +1,67 @@
package harness
import (
"context"
"testing"
)
func TestEstimateTokens(t *testing.T) {
if got := EstimateTokens(""); got != 0 {
t.Errorf("空串应 0, got %d", got)
}
// 8 个 ASCII → (8+3)/4 = 2
if got := EstimateTokens("abcdefgh"); got != 2 {
t.Errorf("ASCII 估算 want 2, got %d", got)
}
// 4 个中文 → 4
if got := EstimateTokens("你好世界"); got != 4 {
t.Errorf("中文估算 want 4, got %d", got)
}
}
func TestBudget_ExceedAndSnapshot(t *testing.T) {
b := NewBudget(10)
b.AddPrompt("你好世界") // 4
b.AddComplete("你好") // 2 → total 6
if b.Exceeded() {
t.Fatal("6/10 不应触顶")
}
b.AddComplete("一二三四五") // +5 → 11
if !b.Exceeded() {
t.Fatal("11/10 应触顶")
}
p, c, total := b.Snapshot()
if p != 4 || c != 7 || total != 11 {
t.Errorf("快照 want 4/7/11, got %d/%d/%d", p, c, total)
}
}
func TestBudget_Unlimited(t *testing.T) {
b := NewBudget(0) // 不限额
b.AddComplete("非常非常非常长的一段中文内容反复堆叠占用大量预算额度")
if b.Exceeded() {
t.Error("max≤0 不应触顶")
}
}
func TestBudget_NilSafe(t *testing.T) {
var b *Budget
b.AddPrompt("x") // 不应 panic
if b.Exceeded() {
t.Error("nil 预算视为不限额")
}
if _, _, total := b.Snapshot(); total != 0 {
t.Error("nil 快照应为 0")
}
}
func TestBudget_Context(t *testing.T) {
if BudgetFrom(context.Background()) != nil {
t.Error("无预算的 ctx 应返回 nil")
}
b := NewBudget(100)
ctx := WithBudget(context.Background(), b)
if BudgetFrom(ctx) != b {
t.Error("应取回同一预算实例")
}
}
@@ -86,6 +86,11 @@ func (s *Subscriber) PublishEval(ev *contract.EvalEvent) error {
return s.inner.PublishEval(ev)
}
// PublishUsage 让 Subscriber 满足 eino.UsageSink,把任务 token 用量回写给网关累计/计费。
func (s *Subscriber) PublishUsage(ev *contract.UsageEvent) error {
return s.inner.PublishUsage(ev)
}
// WaitApproval 让 Subscriber 满足 eino.ApprovalWaiter,阻塞等待审批节点的人工决定。
func (s *Subscriber) WaitApproval(ctx context.Context, taskID string, timeout time.Duration) (*contract.ApprovalDecision, error) {
return s.inner.WaitApproval(ctx, taskID, timeout)
+14
View File
@@ -6,6 +6,7 @@ import (
"encoding/json"
"log"
"os"
"time"
"github.com/sundynix/sundynix-gateway/internal/blob"
"github.com/sundynix/sundynix-gateway/internal/nats"
@@ -78,6 +79,19 @@ func main() {
log.Printf("[gateway] subscribe eval: %v", err)
}
// 成本护栏:订阅 dispatcher 回写的任务 token 用量,按用户按天累计到 Redis(供提交前日预算门控 / 计费)。
if _, err := bus.SubscribeUsage(func(ev *contract.UsageEvent) {
if ev.UserID == "" || ev.TotalTok <= 0 {
return
}
day := time.UnixMilli(ev.TS).Format("20060102")
if _, err := cache.AddUsage(context.Background(), ev.UserID, day, ev.TotalTok); err != nil {
log.Printf("[gateway] 累计用量 user=%s 失败: %v", ev.UserID, err)
}
}); err != nil {
log.Printf("[gateway] subscribe usage: %v", err)
}
r := router.New(db, cache, bus, blobStore)
addr := envOr("GATEWAY_ADDR", ":8080")
log.Printf("[gateway] listening on %s", addr)
@@ -7,6 +7,8 @@ import (
"io"
"log"
"net/http"
"os"
"strconv"
"time"
"github.com/gin-contrib/sse"
@@ -42,6 +44,17 @@ func (h *Handler) SubmitTask(c *gin.Context) {
c.JSON(http.StatusUnprocessableEntity, gin.H{"error": err.Error()})
return
}
// 成本护栏:单用户当日 token 日预算门控(USER_DAILY_TOKEN_BUDGET0=不限)。超额则拒绝新任务。
if budget := userDailyTokenBudget(); budget > 0 {
uid := userID(c)
used := h.cache.GetUsage(c.Request.Context(), uid, time.Now().Format("20060102"))
if used >= int64(budget) {
c.JSON(http.StatusPaymentRequired, gin.H{
"error": "已达当日 token 预算上限", "used": used, "budget": budget,
})
return
}
}
// 附上用户标识(召回偏好记忆)与会话标识(召回短期多轮历史)。
// 真实场景由鉴权/会话中间件注入;此处用请求头,缺省匿名/默认会话。
task.Meta[contract.MetaUserID] = userID(c)
@@ -362,12 +375,36 @@ func sessionID(c *gin.Context) string {
return "default"
}
// userDailyTokenBudget 读单用户当日 token 预算(env USER_DAILY_TOKEN_BUDGET,缺省/非法=0 即不限)。
func userDailyTokenBudget() int {
if v := os.Getenv("USER_DAILY_TOKEN_BUDGET"); v != "" {
if n, err := strconv.Atoi(v); err == nil && n > 0 {
return n
}
}
return 0
}
func (h *Handler) Billing(c *gin.Context) {
// TODO: 商业化与计费模块;暂以已提交任务计数演示真实读库。
n, err := h.db.CountTasks(c.Request.Context())
if err != nil {
c.JSON(http.StatusBadGateway, gin.H{"error": err.Error()})
return
}
c.JSON(http.StatusOK, gin.H{"status": "ok", "tasks_submitted": n, "persisted": h.db.Enabled()})
// 成本护栏:当日 token 用量与预算(按用户隔离)。
uid := userID(c)
used := h.cache.GetUsage(c.Request.Context(), uid, time.Now().Format("20060102"))
budget := userDailyTokenBudget()
remaining := -1 // -1 表示不限额
if budget > 0 {
if r := budget - int(used); r > 0 {
remaining = r
} else {
remaining = 0
}
}
c.JSON(http.StatusOK, gin.H{
"status": "ok", "tasks_submitted": n, "persisted": h.db.Enabled(),
"token_used_today": used, "daily_budget": budget, "remaining": remaining,
})
}
@@ -74,6 +74,11 @@ func (b *Bus) SubscribeEval(onEvent func(*contract.EvalEvent)) (func() error, er
return b.inner.SubscribeEval(onEvent)
}
// SubscribeUsage 订阅 dispatcher 回写的任务 token 用量(累计到用户日预算 / 计费)。
func (b *Bus) SubscribeUsage(onEvent func(*contract.UsageEvent)) (func() error, error) {
return b.inner.SubscribeUsage(onEvent)
}
// ServeConfig 让网关作为配置控制面,响应某 kind 的配置请求。
func (b *Bus) ServeConfig(kind string, provide func() *contract.ModelConfig) (func() error, error) {
return b.inner.ServeConfig(kind, provide)
+35
View File
@@ -50,6 +50,41 @@ func (r *Redis) Allow(ctx context.Context, key string, limit int64, window time.
// ---- Token 流持久化(Redis Stream:可回放的追加日志,根治 SSE 连晚/重连丢 token----
// ---- Token 用量日预算(成本护栏:按用户按天累计 token,供提交前门控)----
// usageDailyKey 是某用户某天的累计 token keyday 形如 20260626,便于到期自然滚动)。
func usageDailyKey(userID, day string) string { return "sundynix:usage:" + userID + ":" + day }
// AddUsage 给某用户当天累计 token,并返回累计后总量(首次设 48h 过期自动清理)。降级返回 (0,nil)。
func (r *Redis) AddUsage(ctx context.Context, userID, day string, tokens int) (int64, error) {
if r.rdb == nil || userID == "" {
return 0, nil
}
rk := usageDailyKey(userID, day)
n, err := r.rdb.IncrBy(ctx, rk, int64(tokens)).Result()
if err != nil {
return 0, err
}
if n == int64(tokens) { // 首次累计:设过期
_ = r.rdb.Expire(ctx, rk, 48*time.Hour).Err()
}
return n, nil
}
// GetUsage 取某用户当天已用 token(不存在/降级返回 0)。
func (r *Redis) GetUsage(ctx context.Context, userID, day string) int64 {
if r.rdb == nil || userID == "" {
return 0
}
n, err := r.rdb.Get(ctx, usageDailyKey(userID, day)).Int64()
if err != nil {
return 0
}
return n
}
// ---- Token 流持久化(Redis Stream:可回放的追加日志,根治 SSE 连晚/重连丢 token----
const streamTTL = 10 * time.Minute
func streamKey(taskID string) string { return "sundynix:stream:" + taskID }
+23
View File
@@ -363,6 +363,29 @@ func (b *Bus) SubscribeEval(onEvent func(*contract.EvalEvent)) (unsub func() err
return sub.Unsubscribe, nil
}
// PublishUsage 广播一次任务 token 用量(dispatcher 收尾调用)。
func (b *Bus) PublishUsage(ev *contract.UsageEvent) error {
data, err := json.Marshal(ev)
if err != nil {
return err
}
return b.nc.Publish(contract.SubjectUsage, data)
}
// SubscribeUsage 订阅 token 用量(网关调用,累计到用户日预算 / 计费)。
func (b *Bus) SubscribeUsage(onEvent func(*contract.UsageEvent)) (unsub func() error, err error) {
sub, err := b.nc.Subscribe(contract.SubjectUsage, func(m *nats.Msg) {
var ev contract.UsageEvent
if json.Unmarshal(m.Data, &ev) == nil {
onEvent(&ev)
}
})
if err != nil {
return nil, fmt.Errorf("subscribe usage: %w", err)
}
return sub.Unsubscribe, nil
}
// ---- 人工审批(HITLcore NATS pub-sub----
// PublishApproval 广播一次人工审批决定(网关在收到 UI 的批准/拒绝后调用)。
+19
View File
@@ -40,8 +40,24 @@ const (
// 自动化评测结果回写:dispatcher 评完经此广播,网关订阅落 PG 并供 UI 查询。core NATS pub-sub。
SubjectEval = "sundynix.eval.task"
// Token 用量回写:dispatcher 任务收尾经此广播本轮 token 用量,网关订阅累加到用户日预算并供计费。core NATS pub-sub。
SubjectUsage = "sundynix.usage.task"
)
// UsageEvent 是一次任务的 token 用量(dispatcher 收尾经 SubjectUsage 回流给网关累计/计费)。
// dispatcher 不持有计价,只报原始 token + 模型名;真钱成本由网关据 Pricing 折算。
type UsageEvent struct {
TaskID string `json:"task_id"`
UserID string `json:"user_id,omitempty"`
Model string `json:"model,omitempty"`
PromptTok int `json:"prompt_tok"` // 输入 token(估算)
CompTok int `json:"comp_tok"` // 输出 token(估算)
TotalTok int `json:"total_tok"` // 合计
Exceeded bool `json:"exceeded"` // 是否触顶单任务预算被中止
TS int64 `json:"ts"` // unix 毫秒
}
// 评测质量分级(据综合分 + 忠实度阈值,闭环门控/告警用)。
const (
EvalOK = "ok" // 综合 ≥ 0.75 且无忠实度风险
@@ -106,6 +122,9 @@ const (
// MetaSafetyCheck 是输入护栏「灰区升级」标志:网关 Tier1(归一化+正则)判为疑似但不确定时置 true,
// Dispatcher 执行前据此调 LLM jailbreak 分类器(Tier2)裁决。明确干净/明确恶意的输入不带此标志,不付 LLM 成本。
MetaSafetyCheck = "safety_check"
// MetaTokenBudget 是单任务 token 预算上限(数字):网关按用户/套餐下发,Dispatcher 据此封顶单任务用量,
// 触顶即中止(防失控成本)。缺省用 Dispatcher 的 env TASK_TOKEN_BUDGET。
MetaTokenBudget = "token_budget"
// 配置控制面按 kind 寻址:sundynix.config.<kind>.get / .updated。
// Gateway 持有配置,消费方(Dispatcher/mcp-go)经 NATS 取用/订阅变更。