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>
2026-06-26 10:00:51 +08:00
parent 2f78fc565e
commit faa1871760
16 changed files with 442 additions and 17 deletions
@@ -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())