Merge pull request #14518 from SivanCola/feature/root-task-token-budget-14025

Share token budgets through Agent delegation / 委派 Agent 共享 token 预算

.github/workflows/ci.yml

@@ -294,6 +294,39 @@ jobs:
           if-no-files-found: ignore
           retention-days: 7
 
+      - name: test (Windows root token 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(RootToken|EffectRootToken)' ./internal/runtime/agent ./internal/assembly/boot | tee "$RUNNER_TEMP/root-token-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/runtime/agent", "reasonix/internal/assembly/boot"] - $packages | length) == 0 and
+            (["TestRootTokenBudgetIncludesChildProviderUsage", "TestRootTokenBudgetIncludesChildProviderUsage/configured",
+              "TestRootTokenBudgetIncludesChildProviderUsage/context", "TestRootTokenBudgetIncludesChildProviderUsage/unconfigured",
+              "TestRootTokenBudgetIncludesChildProviderUsage/explicit_zero", "TestRootTokenBudgetSurvivesNestedAndConcurrentChildren",
+              "TestRootTokenBudgetSurvivesNestedAndConcurrentChildren/nested", "TestRootTokenBudgetSurvivesNestedAndConcurrentChildren/siblings",
+              "TestRootTokenBudgetSurvivesNestedAndConcurrentChildren/child_override", "TestRootTokenBudgetContinuationAndReset",
+              "TestRootTokenBudgetContinuationAndReset/continuation", "TestRootTokenBudgetContinuationAndReset/new_task",
+              "TestRootTokenBudgetContinuationAndReset/reset", "TestRootTokenScopeAtomicAccountingAndJobCancellation",
+              "TestEffectRootTokenBudgetReachesDelegatedJobs", "TestEffectRootTokenBudgetReachesDelegatedJobs/foreground",
+              "TestEffectRootTokenBudgetReachesDelegatedJobs/foreground_capability", "TestEffectRootTokenBudgetReachesDelegatedJobs/background",
+              "TestEffectRootTokenBudgetReachesDelegatedJobs/background_capability", "TestEffectRootTokenBudgetReachesDelegatedJobs/background_fleet"] - $tests | length) == 0
+          ' "$RUNNER_TEMP/root-token-budget.jsonl"
+
+      - name: upload Windows root token budget inventory
+        if: always() && env.RUN_STEPS == 'true' && runner.os == 'Windows' && github.event_name == 'pull_request'
+        uses: actions/upload-artifact@cf430e030ddbb5b0abf93d22962f4752f3646cd9 # v7
+        with:
+          name: root-token-budget-windows-${{ github.run_attempt }}
+          path: ${{ runner.temp }}/root-token-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

docs/ROOT_TOKEN_BUDGET_SECURITY.md

@@ -0,0 +1,86 @@
+---
+owner: @SivanCola
+backup: @esengine
+status: active
+reviewed: 2026-10-26
+---
+
+# Agent-root token budget / Agent 根任务 token 预算
+
+An explicit `task_token_budget` now counts the main-loop reported prompt and
+completion tokens of an Agent-root delegation tree. Cached input still counts.
+The default threshold remains zero. This is a cooperative landing gate after
+tool rounds, followed by the existing finalization round.
+
+显式 `task_token_budget` 现在累计一个 Agent 根委派树的主循环已报告输入及输出
+token，缓存命中输入同样计入。默认阈值仍为零。这是在工具轮次结束后触发的协作式
+收尾门禁，仍保留既有的最后收尾轮次。
+
+## Identity and lifetime / 身份与寿命
+
+The root Agent owns a process-local atomic counter in its existing task ledger.
+Derived contexts share that counter; a child cannot enlarge an inherited root
+limit through its local override. A root continuation keeps the counter. A new
+root task, root session replacement or explicit root reset starts another counter.
+
+根 Agent 在既有任务账本内持有进程内原子计数器。派生 context 共享计数器；子任务
+不能通过局部 override 放宽继承的根阈值。继续同一根任务会保留计数器；新根任务、
+替换根会话或显式重置根预算则使用另一个计数器。
+
+Task and fleet background jobs project only the token scope through
+`agent.InheritRootTokenBudget`. Job Manager cancellation and deadlines remain
+authoritative. A parent turn ending does not erase the handed-off scope. An old
+job still charges its old counter after the parent starts a different task.
+
+task 和 fleet 后台 job 通过 `agent.InheritRootTokenBudget` 仅传递 token scope，
+取消与截止时间仍由 Job Manager 掌握。父 turn 结束不会抹掉已交出的 scope。父任务
+换成另一项工作后，旧 job 仍记入旧计数器，不会扣减新任务的预算。
+
+## Authority and limits / 权限与边界
+
+The context key is private. The internal Go projection helper copies an existing
+scope; it does not construct a new allowance or copy parent cancellation. Known
+usage is added once per observed Agent round. Negative usage cannot credit the
+counter and unsigned overflow saturates instead of reopening a budget.
+
+context key 是私有类型。内部 Go 投影函数复制既有 scope，不创建新额度，也不复制
+父任务的取消权。每个被观察到的 Agent 轮次只记账一次。负用量不能返还额度，
+无符号整数溢出会饱和，不能重新打开预算。
+
+This is not request admission, token reservation, a hard spending ceiling or
+process termination. Already-running requests and each Agent's finalization
+round can exceed the threshold. Queued work is not rejected before its first
+request by this gate. A missing usage report is unknown, not proof of free work.
+
+该机制不是请求准入、token 预留、硬费用上限或进程终止机制。正在执行的请求以及
+各 Agent 的收尾轮次可能越过阈值。该门禁不会在排队工作发出首个请求前拒绝它。
+缺失用量报告代表未知，不能据此认为工作免费。
+
+Auxiliary provider calls outside `observeRunBudget`, remote processes, durable
+restart accounting and a shared planner/executor Coordinator root are not
+covered. Existing cost and wall-time gates and per-Agent budget telemetry retain
+their meanings. No currency amounts are mixed; unpriced models still count tokens.
+
+`observeRunBudget` 之外的辅助模型调用、远程进程、持久化重启记账，以及跨 planner
+与 executor 的统一 Coordinator 根预算均不在本次范围内。既有费用、时间门禁及
+每 Agent 预算遥测保持原含义；不混加不同币种金额，无定价模型仍累计 token。
+
+## Compatibility and validation / 兼容与验证
+
+There is no persisted format, wire, tool schema, permission or collaboration-mode
+change. Existing positive token limits become shared across this Agent subtree;
+zero still disables the root token gate. Provider-visible system/tool prefixes
+receive no budget metadata. Pauses use the existing host-owned token pause.
+
+不修改持久化格式、wire、工具 schema、权限或协作模式。既有正 token 阈值现在由
+这个 Agent 子树共享，零仍关闭根 token 门禁。系统与工具缓存前缀不写入预算元数据，
+暂停继续使用既有 host-owned token pause。
+
+Regressions exercise actual parent/child provider calls, nested and concurrent
+children, continuations and replacement-task isolation. Real boot covers direct
+and capability-routed foreground/background tasks plus a background fleet, with
+actual file reads and a parent cancellation before delayed background work.
+
+回归覆盖实际父子 provider 调用、嵌套与并发子任务、继续运行和替换任务隔离。
+真实 boot 覆盖直接及 capability 路由的前后台 task 和后台 fleet，验证实际文件
+读取，并让延迟后台工作在父 turn 取消后继续执行。

internal/assembly/boot/effect_root_token_budget_test.go

@@ -0,0 +1,150 @@
+package boot
+
+import (
+	"context"
+	"fmt"
+	"strings"
+	"sync/atomic"
+	"testing"
+	"time"
+
+	"reasonix/internal/contract/event"
+	"reasonix/internal/contract/provider"
+	"reasonix/internal/runtime/agent"
+	"reasonix/internal/session/control"
+)
+
+type rootTokenEffectProvider struct {
+	calls atomic.Int32
+	first func() *provider.ToolCall
+	wait  <-chan struct{}
+}
+
+func (*rootTokenEffectProvider) Name() string { return "root-token-effect" }
+
+func (p *rootTokenEffectProvider) Stream(ctx context.Context, req provider.Request) (<-chan provider.Chunk, error) {
+	if p.wait != nil {
+		select {
+		case <-p.wait:
+		case <-ctx.Done():
+			return nil, ctx.Err()
+		}
+	}
+	round := p.calls.Add(1)
+	ch := make(chan provider.Chunk, 3)
+	switch {
+	case round == 1 && p.first != nil:
+		ch <- provider.Chunk{Type: provider.ChunkToolCall, ToolCall: p.first()}
+	case p.first == nil && len(req.Tools) > 0:
+		ch <- provider.Chunk{Type: provider.ChunkToolCall, ToolCall: browserCall(fmt.Sprintf("read-%d", round), "read_file", map[string]any{"path": "fixture.txt"})}
+	default:
+		ch <- provider.Chunk{Type: provider.ChunkText, Text: "Read finished."}
+	}
+	ch <- provider.Chunk{Type: provider.ChunkUsage, Usage: &provider.Usage{PromptTokens: 1000, CacheHitTokens: 900, CacheMissTokens: 100, CompletionTokens: 100, RequestCount: 1}}
+	ch <- provider.Chunk{Type: provider.ChunkDone}
+	close(ch)
+	return ch, nil
+}
+
+func TestEffectRootTokenBudgetReachesDelegatedJobs(t *testing.T) {
+	for _, mode := range []string{"foreground", "foreground_capability", "background", "background_capability", "background_fleet"} {
+		t.Run(mode, func(t *testing.T) {
+			home, root := statedBootHome(t), t.TempDir()
+			background := mode == "background" || mode == "background_capability" || mode == "background_fleet"
+			unblock := make(chan struct{})
+			children := []*rootTokenEffectProvider{{}, {}}
+			if background {
+				for _, child := range children {
+					child.wait = unblock
+				}
+			}
+			count := 1
+			parent := &rootTokenEffectProvider{first: func() *provider.ToolCall {
+				args := map[string]any{"prompt": "read the fixture", "model": "child-a", "profile": "explore", "tools": []string{"read_file"}, "max_steps": 8, "run_in_background": background}
+				if mode == "background_fleet" {
+					count = 2
+					tasks := []map[string]any{}
+					for _, model := range []string{"child-a", "child-b"} {
+						tasks = append(tasks, map[string]any{"prompt": "read the fixture", "model": model, "read_only": true, "tools": []string{"read_file"}, "max_steps": 8})
+					}
+					return browserCall("delegate", "use_capability", map[string]any{"action": "call", "capability_id": "tool:fleet", "arguments": map[string]any{"tasks": tasks, "run_in_background": true}})
+				}
+				if mode == "foreground_capability" || mode == "background_capability" {
+					return browserCall("delegate", "use_capability", map[string]any{"action": "call", "capability_id": "tool:task", "arguments": args})
+				}
+				return browserCall("delegate", "task", args)
+			}}
+			kind := uniqueKind("root-tokens")
+			provider.Register(kind, func(cfg provider.Config) (provider.Provider, error) {
+				switch cfg.Model {
+				case "child-a":
+					return children[0], nil
+				case "child-b":
+					return children[1], nil
+				default:
+					return parent, nil
+				}
+			})
+			writeFile(t, root, "fixture.txt", "fixture content")
+			writeFile(t, root, "reasonix.toml", `default_model = "parent"
+[agent]
+task_token_budget = 1500
+[codegraph]
+enabled = false
+[[providers]]
+name = "parent"
+kind = "`+kind+`"
+model = "parent"
+[[providers]]
+name = "child-a"
+kind = "`+kind+`"
+model = "child-a"
+[[providers]]
+name = "child-b"
+kind = "`+kind+`"
+model = "child-b"
+`)
+			approveWorkspaceAt(t, home, root)
+			var reads atomic.Int32
+			sink := event.FuncSink(func(e event.Event) {
+				if e.Kind == event.ToolResult && e.Tool.Name == "read_file" && e.Tool.Err == "" && strings.Contains(e.Tool.Output, "fixture content") {
+					reads.Add(1)
+				}
+			})
+			ctrl, err := Build(t.Context(), Options{Home: home, WorkspaceRoot: root, Sink: sink})
+			if err != nil {
+				t.Fatal(err)
+			}
+			t.Cleanup(ctrl.Close)
+			ctrl.SetToolApprovalMode(control.ToolApprovalYolo)
+			runCtx, cancel := context.WithCancel(t.Context())
+			defer cancel()
+			runErr := ctrl.Run(runCtx, "delegate a fixture read")
+			cancel()
+			close(unblock)
+			if !background {
+				pause, ok := agent.InspectRunPause(runErr)
+				if !ok || pause.Kind != "task_budget" || pause.Key != "token" || !pause.HostOwned {
+					t.Errorf("root ended with %v (%+v), want token pause", runErr, pause)
+				}
+			} else if runErr != nil {
+				t.Errorf("background submission parent: %v", runErr)
+			}
+			done := make(chan struct{})
+			ctrl.JobManager().AfterUnfinished(func() { close(done) })
+			select {
+			case <-done:
+			case <-time.After(10 * time.Second):
+				t.Fatal("managed child did not finish")
+			}
+			for i, child := range children[:count] {
+				if child.calls.Load() != 2 {
+					t.Errorf("child%d requests = %d, want one read and one existing finalization round", i, child.calls.Load())
+				}
+			}
+			if reads.Load() != int32(count) || parent.calls.Load() != 2 {
+				t.Errorf("real read results=%d parent requests=%d, want %d/2", reads.Load(), parent.calls.Load(), count)
+			}
+		})
+	}
+}

internal/runtime/agent/agent.go

@@ -898,6 +898,7 @@ func (a *Agent) Run(ctx context.Context, input string) (runErr error) {
 	}
 
 	_, state := a.beginRunTurn(ctx, input)
+	ctx = a.bindRootTokenBudget(ctx, state)
 	if state.snapshot != nil {
 		defer state.snapshot.cancel()
 	}

internal/runtime/agent/root_token_budget.go

@@ -0,0 +1,68 @@
+package agent
+
+import (
+	"context"
+	"fmt"
+	"sync/atomic"
+
+	"reasonix/internal/contract/provider"
+)
+
+type rootTokenBudget struct {
+	used  *atomic.Uint64
+	limit uint64
+}
+
+type rootTokenBudgetKey struct{}
+
+// InheritRootTokenBudget carries only the Agent-root token scope into a job.
+// The job keeps its own cancellation and deadline authority.
+func InheritRootTokenBudget(ctx, parent context.Context) context.Context {
+	if scope, ok := parent.Value(rootTokenBudgetKey{}).(*rootTokenBudget); ok {
+		return context.WithValue(ctx, rootTokenBudgetKey{}, scope)
+	}
+	return ctx
+}
+
+func (a *Agent) bindRootTokenBudget(ctx context.Context, state *turnRuntime) context.Context {
+	scope, inherited := ctx.Value(rootTokenBudgetKey{}).(*rootTokenBudget)
+	if !inherited {
+		previous := a.task.budget.rootTokens
+		if previous == nil {
+			previous = &rootTokenBudget{used: new(atomic.Uint64)}
+			previous.observe(&provider.Usage{PromptTokens: a.task.budget.promptTokens, CompletionTokens: a.task.budget.outputTokens})
+		}
+		scope = &rootTokenBudget{used: previous.used, limit: uint64(max(0, a.taskBudgetLimit(ctx).Tokens))}
+		a.task.budget.rootTokens = scope
+		ctx = context.WithValue(ctx, rootTokenBudgetKey{}, scope)
+	}
+	state.budget.rootTokens = scope
+	return ctx
+}
+
+func (b *rootTokenBudget) observe(usage *provider.Usage) {
+	if b == nil || usage == nil {
+		return
+	}
+	delta := uint64(max(0, usage.PromptTokens)) + uint64(max(0, usage.CompletionTokens))
+	for {
+		old := b.used.Load()
+		next := old + delta
+		if next < old {
+			next = ^uint64(0)
+		}
+		if b.used.CompareAndSwap(old, next) {
+			return
+		}
+	}
+}
+
+func (a *Agent) exceededTaskBudget(ctx context.Context, state *turnRuntime) (string, string) {
+	if axis, detail := a.task.budget.exceeded(a.taskBudgetLimit(ctx)); axis != "" {
+		return axis, detail
+	}
+	if root := state.budget.rootTokens; root != nil && root.limit > 0 && root.used.Load() >= root.limit {
+		return "token", fmt.Sprintf("root task reached its %d token budget across delegated runs", root.limit)
+	}
+	return "", ""
+}

internal/runtime/agent/root_token_budget_test.go

@@ -0,0 +1,119 @@
+package agent
+
+import (
+	"context"
+	"encoding/json"
+	"testing"
+
+	"reasonix/internal/contract/event"
+	"reasonix/internal/contract/tool"
+	"reasonix/internal/state/sessionstore"
+)
+
+type tokenBudgetChildTool struct {
+	readProbe
+	child  *Agent
+	err    error
+	budget *TaskBudget
+}
+
+func (t *tokenBudgetChildTool) Execute(ctx context.Context, _ json.RawMessage) (string, error) {
+	if t.budget != nil {
+		ctx = WithTaskBudget(ctx, *t.budget)
+	}
+	t.err = t.child.Run(ctx, "read a few files")
+	return "child finished", nil
+}
+
+func TestRootTokenBudgetSurvivesNestedAndConcurrentChildren(t *testing.T) {
+	for _, mode := range []string{"nested", "siblings", "child_override"} {
+		t.Run(mode, func(t *testing.T) {
+			leafProvider := &spendingProvider{max: 10}
+			leafTools := tool.NewRegistry()
+			leafTools.Add(readProbe{})
+			leaf := New(leafProvider, leafTools, sessionstore.NewSession("leaf"), Options{}, event.Discard)
+			call := &tokenBudgetChildTool{child: leaf}
+			if mode == "child_override" {
+				call.budget = &TaskBudget{}
+			}
+			rootTools := tool.NewRegistry()
+			rootTools.Add(call)
+			var siblings *tokenBudgetSiblings
+			switch mode {
+			case "nested":
+				middle := New(&spendingProvider{max: 1}, rootTools, sessionstore.NewSession("middle"), Options{}, event.Discard)
+				rootTools = tool.NewRegistry()
+				rootTools.Add(&tokenBudgetChildTool{child: middle})
+			case "siblings":
+				otherTools := tool.NewRegistry()
+				otherTools.Add(readProbe{})
+				other := New(&spendingProvider{max: 10}, otherTools, sessionstore.NewSession("sibling"), Options{}, event.Discard)
+				rootTools = tool.NewRegistry()
+				siblings = &tokenBudgetSiblings{children: []*Agent{leaf, other}}
+				rootTools.Add(siblings)
+			}
+			root := New(&spendingProvider{max: 1}, rootTools, sessionstore.NewSession("root"), Options{TaskBudget: TaskBudget{Tokens: 1500}}, event.Discard)
+			pause, ok := InspectRunPause(root.Run(t.Context(), "delegate work"))
+			if !ok || pause.Kind != "task_budget" || pause.Key != "token" || !pause.HostOwned {
+				t.Fatalf("root pause = %+v (%v), want token budget", pause, ok)
+			}
+			if leafProvider.rounds.Load() != 2 {
+				t.Fatalf("leaf requests = %d, want one tool round and one existing finalization round", leafProvider.rounds.Load())
+			}
+			if siblings != nil {
+				for i, err := range siblings.errors {
+					pause, ok := InspectRunPause(err)
+					if !ok || pause.Key != "token" || siblings.children[i].svc.prov.(*spendingProvider).rounds.Load() != 2 {
+						t.Errorf("sibling%d did not execute and land on the shared token gate: %v", i, err)
+					}
+				}
+			} else if pause, ok := InspectRunPause(call.err); !ok || pause.Key != "token" {
+				t.Errorf("leaf did not pause on tokens: %v", call.err)
+			}
+		})
+	}
+}
+
+func TestRootTokenBudgetIncludesChildProviderUsage(t *testing.T) {
+	for _, test := range []struct {
+		name       string
+		configured int
+		override   *TaskBudget
+		wantChild  int32
+		wantPause  bool
+	}{
+		{name: "configured", configured: 1500, wantChild: 2, wantPause: true},
+		{name: "context", override: &TaskBudget{Tokens: 3000}, wantChild: 3, wantPause: true},
+		{name: "unconfigured", wantChild: 11},
+		{name: "explicit_zero", configured: 1500, override: &TaskBudget{}, wantChild: 11},
+	} {
+		t.Run(test.name, func(t *testing.T) {
+			childProvider := &spendingProvider{max: 10}
+			childTools := tool.NewRegistry()
+			childTools.Add(readProbe{})
+			child := New(childProvider, childTools, sessionstore.NewSession("child"), Options{}, event.Discard)
+			call := &tokenBudgetChildTool{child: child}
+			parentTools := tool.NewRegistry()
+			parentTools.Add(call)
+			parent := New(&spendingProvider{max: 1}, parentTools, sessionstore.NewSession("parent"), Options{TaskBudget: TaskBudget{Tokens: test.configured}}, event.Discard)
+			ctx := t.Context()
+			if test.override != nil {
+				ctx = WithTaskBudget(ctx, *test.override)
+			}
+			err := parent.Run(ctx, "delegate a read")
+			if got := childProvider.rounds.Load(); got != test.wantChild {
+				t.Errorf("child provider requests = %d, want %d including its existing finalization round", got, test.wantChild)
+			}
+			for role, runErr := range map[string]error{"parent": err, "child": call.err} {
+				info, paused := InspectRunPause(runErr)
+				if test.wantPause {
+					if !paused || info.Kind != "task_budget" || info.Key != "token" || !info.HostOwned {
+						t.Errorf("%s pause = %+v (%v), want host-owned token pause", role, info, runErr)
+					}
+				} else if runErr != nil {
+					t.Errorf("%s run: %v", role, runErr)
+				}
+			}
+		})
+	}
+}

internal/runtime/agent/root_token_scope_test.go

@@ -0,0 +1,105 @@
+package agent
+
+import (
+	"context"
+	"encoding/json"
+	"sync"
+	"sync/atomic"
+	"testing"
+
+	"reasonix/internal/contract/event"
+	"reasonix/internal/contract/provider"
+	"reasonix/internal/contract/tool"
+	"reasonix/internal/state/sessionstore"
+)
+
+type tokenBudgetSiblings struct {
+	readProbe
+	children []*Agent
+	errors   []error
+}
+
+func (t *tokenBudgetSiblings) Execute(ctx context.Context, _ json.RawMessage) (string, error) {
+	var group sync.WaitGroup
+	t.errors = make([]error, len(t.children))
+	for i, child := range t.children {
+		group.Go(func() { t.errors[i] = child.Run(ctx, "read a few files") })
+	}
+	group.Wait()
+	return "children ended", nil
+}
+
+func TestRootTokenBudgetContinuationAndReset(t *testing.T) {
+	for _, mode := range []string{"continuation", "new_task", "reset"} {
+		t.Run(mode, func(t *testing.T) {
+			child := New(&spendingProvider{}, tool.NewRegistry(), sessionstore.NewSession("child"), Options{}, event.Discard)
+			tools := tool.NewRegistry()
+			tools.Add(&tokenBudgetChildTool{child: child})
+			prov := &spendingProvider{max: 1}
+			root := New(prov, tools, sessionstore.NewSession("root"), Options{TaskBudget: TaskBudget{Tokens: 4000}}, event.Discard)
+			if err := root.Run(t.Context(), "first work"); err != nil {
+				t.Fatal(err)
+			}
+			old := root.task.budget.rootTokens
+			if old.used.Load() != 3300 {
+				t.Fatalf("root tokens = %d, want parent2200+child1100", old.used.Load())
+			}
+			if mode != "new_task" {
+				root.pending.preserveEvidence = true
+			}
+			if mode == "reset" {
+				root.ResetTaskBudget()
+			}
+			prov.rounds.Store(0)
+			err := root.Run(t.Context(), "second work")
+			if mode == "continuation" {
+				pause, ok := InspectRunPause(err)
+				if !ok || pause.Kind != "task_budget" || pause.Key != "token" {
+					t.Fatalf("continuation = %v (%+v), want root token pause", err, pause)
+				}
+			} else {
+				if err != nil {
+					t.Fatal(err)
+				}
+				old.observe(&provider.Usage{PromptTokens: 5000})
+				if root.task.budget.rootTokens.used.Load() != 3300 {
+					t.Fatal("an old task context charged the replacement task")
+				}
+			}
+		})
+	}
+}
+
+func TestRootTokenScopeAtomicAccountingAndJobCancellation(t *testing.T) {
+	root := &rootTokenBudget{used: new(atomic.Uint64), limit: 500}
+	var group sync.WaitGroup
+	for range 100 {
+		group.Go(func() { root.observe(&provider.Usage{PromptTokens: 7, CompletionTokens: 3}) })
+	}
+	group.Wait()
+	if root.used.Load() != 1000 {
+		t.Fatal("concurrent usage was lost or counted twice")
+	}
+	root.observe(nil)
+	root.observe(&provider.Usage{PromptTokens: -3, CompletionTokens: -4})
+	if root.used.Load() != 1000 {
+		t.Fatal("missing or negative usage credited the root")
+	}
+	root.used.Store(^uint64(0) - 1)
+	root.observe(&provider.Usage{PromptTokens: 10})
+	if root.used.Load() != ^uint64(0) {
+		t.Fatal("usage overflow reopened the budget")
+	}
+	parent, cancelParent := context.WithCancel(context.WithValue(t.Context(), rootTokenBudgetKey{}, root))
+	job, cancelJob := context.WithCancel(t.Context())
+	defer cancelJob()
+	projected := InheritRootTokenBudget(job, parent)
+	cancelParent()
+	if projected.Err() != nil || projected.Value(rootTokenBudgetKey{}) != root {
+		t.Fatal("job inherited parent cancellation or lost the budget")
+	}
+	cancelJob()
+	if projected.Err() != context.Canceled {
+		t.Fatal("budget projection lost job cancellation")
+	}
+}

internal/runtime/agent/run_budget.go

@@ -50,7 +50,8 @@ type runBudget struct {
 	unpricedTurns bool
 	// limit is configuration, not accumulation: it survives the reset that
 	// starts a new task.
-	limit TaskBudget
+	limit      TaskBudget
+	rootTokens *rootTokenBudget
 }
 
 // observe folds one round's provider usage into the turn's running total.
@@ -135,6 +136,7 @@ func (a *Agent) observeRunBudget(state *turnRuntime, usage *provider.Usage) {
 	if state == nil {
 		return
 	}
+	state.budget.rootTokens.observe(usage)
 	state.budget.observe(usage, a.svc.pricing)
 	if a.task.budget.started.IsZero() {
 		a.task.budget.started = state.budget.started

internal/runtime/agent/run_loop.go

@@ -755,7 +755,7 @@ func (a *Agent) handleToolRound(ctx context.Context, state *turnRuntime, step in
 
 	// Spend is checked before rounds: it is the axis a runaway is actually
 	// reported in, so on the turns both would catch it should be the one named.
-	if axis, detail := a.task.budget.exceeded(a.taskBudgetLimit(ctx)); axis != "" {
+	if axis, detail := a.exceededTaskBudget(ctx, state); axis != "" {
 		a.armFinalizationRound(state, landCause{kind: "task_budget", axis: axis, detail: detail})
 		return true, nil
 	}

internal/runtime/delegation/fleet.go

@@ -191,6 +191,7 @@ func (f *FleetTool) Execute(ctx context.Context, args json.RawMessage) (result s
 			}
 		}
 		job := jm.StartForSession(jobs.SessionFromContext(ctx), "fleet", label, func(jobCtx context.Context, _ io.Writer) (string, error) {
+			jobCtx = agent.InheritRootTokenBudget(jobCtx, ctx)
 			// Execute returns as soon as the job is registered, so the job owns
 			// the handed-off merger until every child preview and terminal has
 			// flushed. Closing it in Execute would strand child cards at running.

internal/runtime/delegation/task.go

@@ -653,6 +653,7 @@ func (t *TaskTool) runProfileSpec(ctx context.Context, spec ProfileExecSpec, exi
 		// Capture acquire request by value for the job goroutine.
 		slotReq := acquireReq
 		job := jm.StartForSession(jobs.SessionFromContext(ctx), "task", label, func(jobCtx context.Context, _ io.Writer) (result string, err error) {
+			jobCtx = agent.InheritRootTokenBudget(jobCtx, ctx)
 			jobCtx = agent.WithParentSession(jobCtx, parentSession)
 			jobCtx = evidence.WithLedger(jobCtx, backgroundEvidence)
 			defer func() { jobs.PublishEvidence(jobCtx, backgroundEvidence.Summary()) }()

release-notes/studio/unreleased/14025-root-token-budget.md

@@ -0,0 +1,2 @@
+- 修复：显式 `task_token_budget` 共享累计 Agent task/fleet 子树已报告的 token，包括后台 job 和缓存输入。默认不限，保留软收尾；不是硬预留或金额上限。关联 #14025。
+- Fixed: explicit `task_token_budget` shares reported tokens across Agent task/fleet descendants. Defaults stay unlimited; existing soft landing remains. Refs #14025.