diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 83f40fdf75..6c9445a2ed 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -452,6 +452,41 @@ jobs: if-no-files-found: ignore retention-days: 7 + - name: test (Windows bounded root budget) + if: env.RUN_STEPS == 'true' && runner.os == 'Windows' && github.event_name == 'pull_request' + timeout-minutes: 5 + shell: bash + run: | + set -o pipefail + go test -count=1 -timeout=3m -json -run '^Test(AuxiliaryUsageObserver|BoundedUsageObserver|RootBounded|EffectRootBounded)' ./internal/contract/provider ./internal/model/boundedllm ./internal/runtime/agent ./internal/assembly/boot | tee "$RUNNER_TEMP/bounded-root-budget.jsonl" + jq -se ' + [.[] | select(.Test == null and .Action == "pass") | .Package] as $packages | + [.[] | select(.Test != null and .Action == "pass") | .Test] as $tests | + all(.[]; .Action != "skip" and .Action != "fail") and + (["reasonix/internal/contract/provider", "reasonix/internal/model/boundedllm", + "reasonix/internal/runtime/agent", "reasonix/internal/assembly/boot"] - $packages | length) == 0 and + (["TestAuxiliaryUsageObserverReplacesAndCopies", + "TestBoundedUsageObserverOwnsOneTerminal", "TestBoundedUsageObserverOwnsOneTerminal/reported", + "TestBoundedUsageObserverOwnsOneTerminal/billed_failure", "TestBoundedUsageObserverOwnsOneTerminal/without_sink", + "TestBoundedUsageObserverOwnsOneTerminal/missing", "TestBoundedUsageObserverOwnsOneTerminal/input_refused", + "TestRootBoundedReviewerBudget", "TestRootBoundedReviewerBudget/reported", + "TestRootBoundedReviewerBudget/billed_failure", "TestRootBoundedReviewerBudget/missing_usage", + "TestRootBoundedReviewerBudget/disabled", "TestRootBoundedReviewerBudget/no_emission", + "TestRootBoundedReviewerBudget/projected_job_context", "TestRootBoundedNestedBudgetChargesOnce", + "TestRootBoundedLateProjectionStaysWithOldTask", "TestEffectRootBoundedAdvisorBudget", + "TestEffectRootBoundedAdvisorBudget/reported", "TestEffectRootBoundedAdvisorBudget/disabled", + "TestEffectRootBoundedAdvisorBudget/missing_usage"] - $tests | length) == 0 + ' "$RUNNER_TEMP/bounded-root-budget.jsonl" + + - name: upload Windows bounded root budget inventory + if: always() && env.RUN_STEPS == 'true' && runner.os == 'Windows' && github.event_name == 'pull_request' + uses: actions/upload-artifact@cf430e030ddbb5b0abf93d22962f4752f3646cd9 # v7 + with: + name: bounded-root-budget-windows-${{ github.run_attempt }} + path: ${{ runner.temp }}/bounded-root-budget.jsonl + if-no-files-found: ignore + retention-days: 7 + # The general Windows PR smoke list intentionally omits internal/contract/tool. # Keep the session-temp portability contract covered without widening the # leg to every tool test: two real PowerShell launches must share the diff --git a/docs/BOUNDED_ROOT_TOKEN_BUDGET.md b/docs/BOUNDED_ROOT_TOKEN_BUDGET.md new file mode 100644 index 0000000000..bd30570f0b --- /dev/null +++ b/docs/BOUNDED_ROOT_TOKEN_BUDGET.md @@ -0,0 +1,105 @@ +--- +owner: @SivanCola +backup: @esengine +status: active +reviewed: 2026-10-26 +--- + +# Bounded calls in root token budgets / 根 token 预算中的有界调用 + +An active Agent root now receives usage from bounded reviewer/advisor calls +through one optional internal Go context observer. Bounded calls already emit +their own usage; they previously bypassed the root token counter. + +活跃 Agent 根任务现在通过单个可选的内部 Go context observer 接收有界 reviewer/ +advisor 调用的用量。此前这些调用已有自己的用量事件,却绕过了根 token 计数。 + +## Ownership / 所有权 + +Agent root binding installs the observer, and the budget-only job projection +copies it with the opaque scope. A child replaces the inherited observer rather +than stacking callbacks, so one logical bounded terminal charges the root once. + +Agent 根绑定安装 observer;仅投射预算的 job context 将其与不透明 scope 一起携带。 +子任务替换继承 observer,不叠加回调,因此一个逻辑有界调用终态只计一次。 + +The callback captures the original counter, not the Agent's mutable current-task +field. Late projected calls remain charged to the old task after a replacement. +Job cancellation and deadlines remain owned by the target job context. + +回调捕获原 counter,不读取 Agent 可变的当前 task 字段。替换任务后的延后调用 +仍计入旧任务。job 的取消与 deadline 继续由目标 job context 掌握。 + +The provider contract adds two internal Go helpers: one replaces the observer, +and one reports a copy of finalized usage. An observer cannot mutate the record +that the existing usage sink receives. Callbacks are trusted host code and own +their synchronization; this is not a callback panic-isolation boundary. + +The helpers require a non-nil context, as usual in Go. The bounded owner retains +its existing context normalization before it calls them. + +provider 契约新增两个内部 Go helper:一个替换 observer,一个上报终态用量副本。 +observer 不能改写既有 usage sink 收到的记录。回调属于可信 host 代码,自行负责 +同步;本机制不是回调 panic 隔离边界。 + +helper 要求非 nil context;有界 owner 调用前保留既有 context 规范化。 + +## Accounting and limits / 计账与阈值 + +The bounded owner reports its finalized usage once, including usage received +before a stream error. Input refusal without a request does not charge it. +Accounting does not depend on whether the optional usage sink/source is set. + +有界调用 owner 仅上报一次终态用量,包括流错误前收到的用量。请求前的输入拒绝 +不会扣计。计账不依赖可选 usage sink/source 是否启用。 + +Positive prompt/completion counts follow the existing saturating root rule, +including positive estimates. Missing counts remain unknown, not inferred from +request counts. Normal-round local telemetry and cost/currency rules stay as +they were; this does not add a monetary ledger. + +正 prompt/completion 计数沿用既有根饱和规则,包括正估算值。缺失计数仍未知, +不能由请求次数推导。正常轮次的局部遥测与费用/币种规则保持原样,不新增费用账本。 + +The existing cooperative check runs after a tool round and keeps finalization. +Explicit limits can land earlier after an advisor call; zero still means +unbounded. There is no request reservation, provider cancellation or process +termination in this change. + +既有协作式检查仍在工具轮后执行,并保留收尾轮。显式阈值可能在 advisor 调用后 +更早收尾;零值仍不限。本改动不做请求预留、供应商取消或进程终止。 + +## Caller and compatibility boundary / 调用者与兼容边界 + +Advisor, recovery reviewer, Goal evaluator, best-of judge, commit-message and +prompt-refinement callers share the bounded owner. Participation depends on an +active context observer, not their label. Calls before/after Agent Run or +outside its context remain outside that root scope. + +advisor、恢复 reviewer、Goal evaluator、best-of judge、提交信息和提示改写调用者 +共用有界 owner。是否计入取决于活跃 context observer,而不是调用者标签。 +Agent Run 前后或其 context 之外的调用仍不属于该根 scope。 + +The two helpers add no RPC, config, persisted or provider-wire field. Standalone +reviewers retain their two-message, no-tool requests and existing time/output +bounds. Main system messages, tool schemas and prefix construction are unchanged. + +两个 helper 不新增 RPC、配置、持久化或 provider-wire 字段。独立 reviewer 保留 +两条 message、无工具请求及既有时间/输出边界。主 system message、工具 schema +与前缀构造保持不变。 + +Real Agent and boot HTTP/SSE regressions cover actual advisor answers, nested +scope replacement, missing usage, disabled limits and serialized prefix guards. +Late projection is a constructed control, not a production background incident; +mock request equality is not a live-provider cache-rate measurement. + +真实 Agent 与 boot HTTP/SSE 回归覆盖实际 advisor 答案、嵌套 scope 替换、缺失 +用量、禁用阈值与序列化前缀守卫。延后投射属于构造控制,不是生产后台事故;mock +请求相等不是供应商命中率实测。 + +This is not complete auxiliary-provider admission, a shared Coordinator root, +durable accounting, root currency aggregation or the full isolation plan. +See [root token budgets](ROOT_TOKEN_BUDGET_SECURITY.md). + +本机制不等于完整辅助 provider 准入、共享 Coordinator 根预算、持久计账、根币种 +汇总或完整隔离方案。作者分析不等于独立 review 或原生平台执行。 diff --git a/internal/assembly/boot/effect_root_bounded_budget_test.go b/internal/assembly/boot/effect_root_bounded_budget_test.go new file mode 100644 index 0000000000..896025708c --- /dev/null +++ b/internal/assembly/boot/effect_root_bounded_budget_test.go @@ -0,0 +1,142 @@ +package boot + +import ( + "encoding/json" + "fmt" + "io" + "net/http" + "net/http/httptest" + "strings" + "sync" + "sync/atomic" + "testing" + + "reasonix/internal/contract/event" + "reasonix/internal/runtime/agent" + "reasonix/internal/session/control" + "reasonix/internal/state/sessionstore" +) + +func TestEffectRootBoundedAdvisorBudget(t *testing.T) { + for _, tc := range []struct { + name string + limit int + missing bool + wantAdvice int32 + }{ + {name: "reported", limit: 1500, wantAdvice: 1}, + {name: "disabled", wantAdvice: 4}, + {name: "missing_usage", limit: 1500, missing: true, wantAdvice: 2}, + } { + t.Run(tc.name, func(t *testing.T) { + home, root := statedBootHome(t), t.TempDir() + var normal, advice, results, usageEvents atomic.Int32 + var mu sync.Mutex + var systems, schemas []string + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + var req struct { + Model string `json:"model"` + Tools json.RawMessage `json:"tools"` + Messages []struct { + Role string `json:"role"` + Content json.RawMessage `json:"content"` + } `json:"messages"` + } + if err := json.NewDecoder(r.Body).Decode(&req); err != nil { + t.Error(err) + w.WriteHeader(http.StatusBadRequest) + return + } + w.Header().Set("Content-Type", "text/event-stream") + advisor := req.Model == "advisor-model" + if advisor { + advice.Add(1) + if len(req.Messages) != 2 || req.Messages[0].Role != "system" || req.Messages[1].Role != "user" || (len(req.Tools) > 0 && string(req.Tools) != "null") { + t.Error("advisor lost its bounded isolated request shape") + } + _, _ = io.WriteString(w, "data: {\"choices\":[{\"delta\":{\"content\":\"ADVICE: inspect the lock owner\"},\"finish_reason\":\"stop\"}]}\n\n") + } else { + mu.Lock() + for _, m := range req.Messages { + if m.Role == "system" { + systems = append(systems, string(m.Content)) + break + } + } + schemas = append(schemas, string(req.Tools)) + mu.Unlock() + if n := normal.Add(1); n <= 4 { + _, _ = fmt.Fprintf(w, "data: {\"choices\":[{\"delta\":{\"tool_calls\":[{\"index\":0,\"id\":\"advise-%d\",\"type\":\"function\",\"function\":{\"name\":\"advise\",\"arguments\":\"{\\\"question\\\":\\\"Which lock owns this operation?\\\"}\"}}]},\"finish_reason\":\"tool_calls\"}]}\n\n", n) + } else { + _, _ = io.WriteString(w, "data: {\"choices\":[{\"delta\":{\"content\":\"Finished\"},\"finish_reason\":\"stop\"}]}\n\n") + } + } + if !advisor || !tc.missing { + _, _ = io.WriteString(w, "data: {\"choices\":[],\"usage\":{\"prompt_tokens\":1000,\"completion_tokens\":100,\"total_tokens\":1100}}\n\n") + } + _, _ = io.WriteString(w, "data: [DONE]\n\n") + })) + t.Cleanup(srv.Close) + writeFile(t, root, "reasonix.toml", fmt.Sprintf(`default_model = "worker" +[agent] +advisor_model = "advisor" +task_token_budget = %d +[codegraph] +enabled = false +[[providers]] +name = "worker" +kind = "openai" +model = "worker-model" +base_url = %q +api_key_env = "BOUNDED_ROOT_EFFECT_KEY" +[[providers]] +name = "advisor" +kind = "openai" +model = "advisor-model" +base_url = %q +api_key_env = "BOUNDED_ROOT_EFFECT_KEY" +`, tc.limit, srv.URL, srv.URL)) + t.Setenv("BOUNDED_ROOT_EFFECT_KEY", "test") + approveWorkspaceAt(t, home, root) + sink := event.FuncSink(func(e event.Event) { + if e.Kind == event.ToolResult && e.Tool.Name == "advise" && e.Tool.Err == "" && strings.Contains(e.Tool.Output, "ADVICE: inspect") { + results.Add(1) + } + if e.Kind == event.Usage && e.UsageSource == event.UsageSourceAdvisor && e.Usage != nil { + usageEvents.Add(1) + if e.Usage.RequestCount != 1 || e.Usage.Estimated != tc.missing { + t.Errorf("advisor provenance=%+v", e.Usage) + } + } + }) + ctrl, err := Build(t.Context(), Options{Home: home, WorkspaceRoot: root, Sink: sink}) + if err != nil { + t.Fatal(err) + } + t.Cleanup(ctrl.Close) + ctrl.SetSessionPath(sessionstore.NewSessionPath(ctrl.SessionDir(), ctrl.Label())) + ctrl.SetToolApprovalMode(control.ToolApprovalYolo) + err = ctrl.Run(t.Context(), "review lock ownership") + if tc.limit == 0 { + if err != nil { + t.Fatal(err) + } + } else if pause, ok := agent.InspectRunPause(err); !ok || pause.Key != "token" { + t.Fatalf("pause=%+v err=%v", pause, err) + } + if advice.Load() != tc.wantAdvice || results.Load() != tc.wantAdvice || usageEvents.Load() != tc.wantAdvice || normal.Load() != tc.wantAdvice+1 { + t.Fatalf("normal/advisor/result/usage=%d/%d/%d/%d, want%d/%d/%d/%d", normal.Load(), advice.Load(), results.Load(), usageEvents.Load(), tc.wantAdvice+1, tc.wantAdvice, tc.wantAdvice, tc.wantAdvice) + } + mu.Lock() + defer mu.Unlock() + if len(systems) != int(normal.Load()) || len(schemas) != len(systems) { + t.Fatal("missing main request prefix") + } + for i := 1; i < len(systems); i++ { + if systems[i] != systems[0] || schemas[i] != schemas[0] { + t.Fatal("root accounting changed stable system/schema prefix") + } + } + }) + } +} diff --git a/internal/contract/provider/auxiliary_usage.go b/internal/contract/provider/auxiliary_usage.go new file mode 100644 index 0000000000..9609bb9589 --- /dev/null +++ b/internal/contract/provider/auxiliary_usage.go @@ -0,0 +1,26 @@ +package provider + +import "context" + +type auxiliaryUsageObserverKey struct{} + +// WithAuxiliaryUsageObserver binds one trusted host accounting observer. It +// replaces an inherited observer so nested scopes do not charge a root twice. +// A nil observer disables observation. Cancellation and deadlines are unchanged. +func WithAuxiliaryUsageObserver(ctx context.Context, observe func(*Usage)) context.Context { + return context.WithValue(ctx, auxiliaryUsageObserverKey{}, observe) +} + +// ObserveAuxiliaryUsage reports one logical call's finalized usage without +// exposing its caller-owned record to mutation. Observers may run concurrently +// and own their synchronization. No observer or unknown usage is a no-op. +func ObserveAuxiliaryUsage(ctx context.Context, usage *Usage) { + if usage == nil { + return + } + observe, _ := ctx.Value(auxiliaryUsageObserverKey{}).(func(*Usage)) + if observe != nil { + copy := *usage + observe(©) + } +} diff --git a/internal/contract/provider/auxiliary_usage_test.go b/internal/contract/provider/auxiliary_usage_test.go new file mode 100644 index 0000000000..2d227adc35 --- /dev/null +++ b/internal/contract/provider/auxiliary_usage_test.go @@ -0,0 +1,30 @@ +package provider + +import ( + "context" + "errors" + "testing" +) + +func TestAuxiliaryUsageObserverReplacesAndCopies(t *testing.T) { + ctx, cancel := context.WithCancel(t.Context()) + defer cancel() + outer, inner := 0, 0 + ctx = WithAuxiliaryUsageObserver(ctx, func(*Usage) { outer++ }) + ctx = WithAuxiliaryUsageObserver(ctx, func(u *Usage) { inner++; u.PromptTokens = 999; u.Estimated = true }) + cancel() + if !errors.Is(ctx.Err(), context.Canceled) { + t.Fatal("observer changed cancellation") + } + u := Usage{PromptTokens: 10, CompletionTokens: 2, RequestCount: 3} + ObserveAuxiliaryUsage(ctx, &u) + if outer != 0 || inner != 1 || u.PromptTokens != 10 || u.Estimated { + t.Fatalf("observer leaked or duplicated:outer%d inner%d usage%+v", outer, inner, u) + } + ObserveAuxiliaryUsage(ctx, nil) + ObserveAuxiliaryUsage(context.Background(), &u) + ObserveAuxiliaryUsage(WithAuxiliaryUsageObserver(ctx, nil), &u) + if inner != 1 { + t.Fatal("unknown usage or disabled observer caused a charge") + } +} diff --git a/internal/model/boundedllm/bounded.go b/internal/model/boundedllm/bounded.go index 9fa4b8528b..da8407e624 100644 --- a/internal/model/boundedllm/bounded.go +++ b/internal/model/boundedllm/bounded.go @@ -117,6 +117,7 @@ func Call(ctx context.Context, cfg Config, system, evidence string) (string, err var usage *provider.Usage defer func() { usage = provider.UsageWithRequestAttemptCount(callCtx, usage) + provider.ObserveAuxiliaryUsage(callCtx, usage) if usage != nil && cfg.UsageSource != "" && cfg.Sink != nil { cfg.Sink.Emit(event.Event{ Kind: event.Usage, diff --git a/internal/model/boundedllm/usage_observer_test.go b/internal/model/boundedllm/usage_observer_test.go new file mode 100644 index 0000000000..9467115e50 --- /dev/null +++ b/internal/model/boundedllm/usage_observer_test.go @@ -0,0 +1,72 @@ +package boundedllm + +import ( + "errors" + "testing" + + "reasonix/internal/contract/event" + "reasonix/internal/contract/provider" +) + +func TestBoundedUsageObserverOwnsOneTerminal(t *testing.T) { + for _, mode := range []string{"reported", "billed_failure", "without_sink", "missing", "input_refused"} { + t.Run(mode, func(t *testing.T) { + observed, emitted := 0, 0 + original := provider.Usage{PromptTokens: 1000, CompletionTokens: 100, TotalTokens: 1100, RequestCount: 2} + failure := errors.New("review failure") + chunks := chunkProvider{{Type: provider.ChunkText, Text: "reviewed"}} + if mode != "missing" { + chunks = append(chunks, provider.Chunk{Type: provider.ChunkUsage, Usage: &original}) + } + if mode == "billed_failure" { + chunks = append(chunks, provider.Chunk{Type: provider.ChunkError, Err: failure}) + } + ctx := provider.WithAuxiliaryUsageObserver(t.Context(), func(u *provider.Usage) { + observed++ + if u.PromptTokens != 1000 || u.RequestCount != 2 { + t.Errorf("observer usage=%+v", u) + } + u.PromptTokens = 9999 + }) + cfg := Config{Provider: chunks, UsageSource: event.UsageSourceAdvisor, Sink: event.FuncSink(func(e event.Event) { + emitted++ + if e.Usage == nil || e.Usage.PromptTokens != 1000 { + t.Errorf("observer corrupted emitted usage=%+v", e.Usage) + } + })} + if mode == "without_sink" { + cfg.Sink = nil + cfg.UsageSource = "" + } + if mode == "input_refused" { + cfg.MaxSystemBytes = 1 + } + text, err := Call(ctx, cfg, "policy", "evidence") + if mode == "billed_failure" { + if !errors.Is(err, failure) { + t.Fatalf("error cause lost:%v", err) + } + } else if mode == "input_refused" { + if err == nil { + t.Fatal("oversized input was admitted") + } + } else if err != nil || text != "reviewed" { + t.Fatalf("answer=%q err=%v", text, err) + } + want := 1 + if mode == "missing" || mode == "input_refused" { + want = 0 + } + if observed != want || original.PromptTokens != 1000 { + t.Fatalf("observed=%d want%d original=%+v", observed, want, original) + } + wantEmit := want + if mode == "without_sink" { + wantEmit = 0 + } + if emitted != wantEmit { + t.Fatalf("usage events=%d,want%d", emitted, wantEmit) + } + }) + } +} diff --git a/internal/runtime/agent/root_bounded_budget_test.go b/internal/runtime/agent/root_bounded_budget_test.go new file mode 100644 index 0000000000..627c328544 --- /dev/null +++ b/internal/runtime/agent/root_bounded_budget_test.go @@ -0,0 +1,165 @@ +package agent + +import ( + "context" + "encoding/json" + "errors" + "testing" + + "reasonix/internal/contract/event" + "reasonix/internal/contract/provider" + "reasonix/internal/contract/tool" + "reasonix/internal/model/boundedllm" + "reasonix/internal/state/sessionstore" +) + +type boundedBudgetProvider struct { + calls int + missing, fail bool +} + +func (*boundedBudgetProvider) Name() string { return "bounded-budget" } + +func (p *boundedBudgetProvider) Stream(context.Context, provider.Request) (<-chan provider.Chunk, error) { + p.calls++ + ch := make(chan provider.Chunk, 3) + ch <- provider.Chunk{Type: provider.ChunkText, Text: "reviewed"} + if !p.missing { + ch <- provider.Chunk{Type: provider.ChunkUsage, Usage: &provider.Usage{PromptTokens: 1000, CompletionTokens: 100, TotalTokens: 1100}} + } + if p.fail { + ch <- provider.Chunk{Type: provider.ChunkError, Err: errors.New("billed review failure")} + } + close(ch) + return ch, nil +} + +type boundedBudgetTool struct { + readProbe + prov *boundedBudgetProvider + noEmission, projected bool + capture context.Context + failures int +} + +func (p *boundedBudgetTool) Execute(ctx context.Context, _ json.RawMessage) (string, error) { + p.capture = InheritRootTokenBudget(context.Background(), ctx) + if p.projected { + ctx = p.capture + } + cfg := boundedllm.Config{Provider: p.prov, Sink: event.Discard, UsageSource: event.UsageSourceAdvisor} + if p.noEmission { + cfg.Sink = nil + cfg.UsageSource = "" + } + answer, err := boundedllm.Call(ctx, cfg, "review policy", "bounded evidence") + if err != nil { + p.failures++ + return "review failed", nil + } + return answer, nil +} + +func TestRootBoundedReviewerBudget(t *testing.T) { + for _, tc := range []struct { + name string + missing, fail, disabled, noEmission, projected bool + wantCalls int32 + }{ + {name: "reported", wantCalls: 2}, + {name: "billed_failure", fail: true, wantCalls: 2}, + {name: "missing_usage", missing: true, wantCalls: 3}, + {name: "disabled", disabled: true, wantCalls: 5}, + {name: "no_emission", noEmission: true, wantCalls: 2}, + {name: "projected_job_context", projected: true, wantCalls: 2}, + } { + t.Run(tc.name, func(t *testing.T) { + aux := &boundedBudgetProvider{missing: tc.missing, fail: tc.fail} + probe := &boundedBudgetTool{prov: aux, noEmission: tc.noEmission, projected: tc.projected} + reg := tool.NewRegistry() + reg.Add(probe) + main := &spendingProvider{max: 4} + root := New(main, reg, sessionstore.NewSession("root"), Options{TaskBudget: TaskBudget{Tokens: 1500}}, event.Discard) + ctx := t.Context() + if tc.disabled { + ctx = WithTaskBudget(ctx, TaskBudget{}) + } + err := root.Run(ctx, "obtain a review") + if tc.disabled { + if err != nil { + t.Fatal(err) + } + } else if pause, ok := InspectRunPause(err); !ok || pause.Key != "token" { + t.Fatalf("pause=%+v err=%v", pause, err) + } + if main.rounds.Load() != tc.wantCalls || aux.calls != int(tc.wantCalls-1) { + t.Fatalf("normal/reviewer calls=%d/%d, want%d/%d", main.rounds.Load(), aux.calls, tc.wantCalls, tc.wantCalls-1) + } + want := uint64(main.rounds.Load()) * 1100 + if !tc.missing { + want += uint64(aux.calls) * 1100 + } + if root.task.budget.rootTokens.used.Load() != want { + t.Fatalf("root tokens=%d,want%d", root.task.budget.rootTokens.used.Load(), want) + } + if root.task.budget.promptTokens+root.task.budget.outputTokens != int(main.rounds.Load())*1100 { + t.Fatal("bounded usage changed local normal telemetry") + } + if tc.fail && probe.failures != aux.calls { + t.Fatal("fixture did not retain each billed failure") + } + }) + } +} + +func TestRootBoundedNestedBudgetChargesOnce(t *testing.T) { + aux := &boundedBudgetProvider{} + childTools := tool.NewRegistry() + childTools.Add(&boundedBudgetTool{prov: aux}) + childProv := &spendingProvider{max: 4} + child := New(childProv, childTools, sessionstore.NewSession("child"), Options{}, event.Discard) + parentTools := tool.NewRegistry() + parentTools.Add(&tokenBudgetChildTool{child: child}) + parentProv := &spendingProvider{max: 1} + root := New(parentProv, parentTools, sessionstore.NewSession("root"), Options{TaskBudget: TaskBudget{Tokens: 2500}}, event.Discard) + err := root.Run(t.Context(), "delegate the review") + if pause, ok := InspectRunPause(err); !ok || pause.Key != "token" { + t.Fatalf("pause=%+v err=%v", pause, err) + } + if parentProv.rounds.Load() != 2 || childProv.rounds.Load() != 2 || aux.calls != 1 { + t.Fatalf("parent/child/aux=%d/%d/%d,want2/2/1", parentProv.rounds.Load(), childProv.rounds.Load(), aux.calls) + } + if got := root.task.budget.rootTokens.used.Load(); got != 5500 { + t.Fatalf("root tokens=%d,want5500 without duplicate observers", got) + } +} + +func TestRootBoundedLateProjectionStaysWithOldTask(t *testing.T) { + aux := &boundedBudgetProvider{} + probe := &boundedBudgetTool{prov: aux} + reg := tool.NewRegistry() + reg.Add(probe) + root := New(&spendingProvider{max: 1}, reg, sessionstore.NewSession("root"), Options{}, event.Discard) + ctx, cancel := context.WithCancel(t.Context()) + t.Cleanup(cancel) + if err := root.Run(ctx, "capture budget"); err != nil { + t.Fatal(err) + } + old := root.task.budget.rootTokens + before := old.used.Load() + cancel() + if err := root.Run(t.Context(), "new task"); err != nil { + t.Fatal(err) + } + current := root.task.budget.rootTokens + currentBefore := current.used.Load() + if old == current { + t.Fatal("new task kept old scope") + } + if _, err := boundedllm.Call(probe.capture, boundedllm.Config{Provider: aux}, "review policy", "late evidence"); err != nil { + t.Fatal(err) + } + if old.used.Load() != before+1100 || current.used.Load() != currentBefore { + t.Fatalf("late usage went to wrong scope:old%d/%d,new%d/%d", old.used.Load(), before, current.used.Load(), currentBefore) + } +} diff --git a/internal/runtime/agent/root_token_budget.go b/internal/runtime/agent/root_token_budget.go index 2f85f3dc14..82d3722619 100644 --- a/internal/runtime/agent/root_token_budget.go +++ b/internal/runtime/agent/root_token_budget.go @@ -20,7 +20,8 @@ type rootTokenBudgetKey struct{} // Token and wall bounds share its lifetime; job cancellation remains separate. func InheritRootTokenBudget(ctx, parent context.Context) context.Context { if scope, ok := parent.Value(rootTokenBudgetKey{}).(*rootTokenBudget); ok { - return context.WithValue(ctx, rootTokenBudgetKey{}, scope) + ctx = context.WithValue(ctx, rootTokenBudgetKey{}, scope) + return provider.WithAuxiliaryUsageObserver(ctx, scope.observe) } return ctx } @@ -47,7 +48,7 @@ func (a *Agent) bindRootTokenBudget(ctx context.Context, state *turnRuntime) con ctx = context.WithValue(ctx, rootTokenBudgetKey{}, scope) } state.budget.rootTokens = scope - return ctx + return provider.WithAuxiliaryUsageObserver(ctx, scope.observe) } func (b *rootTokenBudget) observe(usage *provider.Usage) { diff --git a/release-notes/studio/unreleased/14025-bounded-root-token-usage.md b/release-notes/studio/unreleased/14025-bounded-root-token-usage.md new file mode 100644 index 0000000000..77223e1068 --- /dev/null +++ b/release-notes/studio/unreleased/14025-bounded-root-token-usage.md @@ -0,0 +1,2 @@ +- 修复:活跃 Agent 根 context 内的有界 reviewer/advisor 用量现在计入既有 token 软预算,保留缺失用量、零阈值、局部遥测和独立请求边界。关联 #14025。 +- Fixed: bounded reviewer/advisor usage in an active Agent root now counts toward its token soft budget. Missing usage, zero limits and isolated requests retain their behavior. Refs #14025.