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
@@ -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("应取回同一预算实例")
}
}