fix: include completed child usage in lead budget checks (#6218)

* fix: include completed child usage in lead budget checks

* test: require completed usage middleware in lead budget regressions

* docs: fit middleware guidance within the size budget

Condense the subagent, loop, and token budget guard descriptions while retaining their limits, lifecycle rules, and completed child usage ordering.

---------

Co-authored-by: Willem Jiang <willem.jiang@gmail.com>

README.md

@@ -1714,6 +1714,8 @@ Use `/compact` in the Web UI composer to summarize older context for the current
 
 The chat header also shows a context-window gauge when the selected model has a positive `context_window` configured. It estimates the latest materialized checkpoint's message tokens and keeps the previous same-thread percentage visible while data refetches, independently of the cumulative token-usage setting.
 
+When `token_budget.enabled` and `token_usage.enabled` are both enabled, each lead-agent run budget includes completed subagent usage, including the final batch. Hidden goal continuations share the budget; a new user run starts fresh.
+
 ### Sub-Agents
 
 When a sub-agent ends with `return_direct=True` tools, including tools contributed by extension middleware, their outputs are returned in tool-call order. A failed tool marks the task as failed while preserving the batch outputs.

backend/packages/harness/deerflow/agents/lead_agent/agent.py

@@ -47,7 +47,7 @@ from deerflow.agents.middlewares.summarization_middleware import DeerFlowSummari
 from deerflow.agents.middlewares.terminal_response_middleware import TerminalResponseMiddleware
 from deerflow.agents.middlewares.title_middleware import TitleMiddleware
 from deerflow.agents.middlewares.todo_middleware import TodoMiddleware
-from deerflow.agents.middlewares.token_usage_middleware import TokenUsageMiddleware
+from deerflow.agents.middlewares.token_usage_middleware import CompletedSubagentUsageMiddleware, TokenUsageMiddleware
 from deerflow.agents.middlewares.tool_declarations import layer_one_outcome, narrow_declared_tools, verify_declared_tool_view
 from deerflow.agents.middlewares.tool_error_handling_middleware import build_lead_runtime_middlewares
 from deerflow.agents.middlewares.view_image_middleware import ViewImageMiddleware
@@ -746,6 +746,12 @@ def build_middlewares(
 
         middlewares.append(TokenBudgetMiddleware.from_config(token_budget_config))
 
+    # https://docs.langchain.com/oss/python/langchain/middleware/custom#execution-order
+    # Backfill only completed child usage before budget enforcement. Keep
+    # current-step attribution after the guards that can change tool calls.
+    if token_budget_config.enabled and resolved_app_config.token_usage.enabled:
+        middlewares.append(CompletedSubagentUsageMiddleware())
+
     # Inject custom middlewares before ClarificationMiddleware
     if custom_middlewares:
         middlewares.extend(custom_middlewares)

backend/packages/harness/deerflow/agents/middlewares/AGENTS.md

@@ -107,15 +107,14 @@ Before changing a later authorization phase, read the [authorization RFC](../../
 28. **DeferredToolPromotionAuditMiddleware** - Observes final `tool_search` `Command`s; keep it outer of `SkillToolPolicyMiddleware` so denied names are excluded. It atomically claims new names per lead run or subagent execution to dedupe parallel searches, derives subagent attribution from the server-installed recorder, returns the original `Command`, and omits private payloads
 29. **DeferredToolFilterMiddleware** - *(optional, if `tool_search.enabled`)* Hides deferred (MCP) tool schemas from the bound model until `tool_search` or `McpRoutingMiddleware` promotes them (reads per-thread promotions from `ThreadState.promoted`, hash-scoped)
 30. **SystemMessageCoalescingMiddleware** - Merges every SystemMessage into a single leading SystemMessage per request; provider-agnostic fix for strict backends (vLLM/SGLang/Qwen/Anthropic) that reject non-leading system messages. Touches the per-request payload only (checkpoint state unchanged); on midnight crossings only the latest `dynamic_context_reminder` SystemMessage survives. The subagent builder places its date-only context middleware immediately before this coalescer, so the built-in subagent prompt and hidden date reminder still reach providers as one leading system block
-31. **SubagentLimitMiddleware** - *(optional, if `subagent_enabled`)* Truncates excess ordinary `task` tool calls to enforce both the per-response concurrency limit (`max_concurrent_subagents`, resolved against startup `subagent_runtime.max_running` and the 1-64 safety range before construction) and the per-run total delegation cap (`max_total_subagents` runtime override or `subagents.max_total_per_run`, default 6, clamped to 1-50). The total cap counts current-run entries in the durable delegation ledger (entries are tagged with `run_id` when captured), so repeated planning checkpoints in one run cannot keep launching legal-sized batches indefinitely, while later user turns in the same thread get a fresh run budget. Explicit durable `batch_task` calls are a separate mode with persisted total/live/running limits and are not rewritten into ordinary ledger entries. If the ordinary cap is exhausted, the middleware strips remaining `task` calls, forces `finish_reason="stop"`, and appends a visible limit note so the run can synthesize existing results instead of ending with an empty tool-call response.
-   Delegations use `(run_id, tool_call_id)` for capture and retention. Later-run replies cannot close older orphans; missing run IDs match conservatively.
-32. **LoopDetectionMiddleware** - *(optional, if `loop_detection.enabled`)* Detects repeated tool-call loops; hard-stop clears structured, raw, and content-block tool calls before forcing a final text answer; stamps `loop_capped` via `consume_stop_reason` (#3875 Phase 2), symmetric to `TokenBudgetMiddleware`; persists warned-state transitions (first per call hash or per tool-frequency burst) and hard stops as `middleware:loop_detection`, attributed with `is_subagent` and the optional `agent_id`, without tool arguments, message content, tool results, or argument-derived hashes. Ordinary task subagents get dedicated recorder keys through a parent-loop proxy; never pass `RunJournal` into their isolated loop. Durable batch subagents have no parent run journal and do not persist these transitions
-   State is run-scoped: new user runs get fresh budgets; same-run goal
-   continuations share history. Keep sibling warnings isolated and lifecycle
-   hooks topology-stable. Before changing this guard, read
+31. **SubagentLimitMiddleware** - `subagent_enabled`: caps ordinary `task` calls per response (`max_concurrent_subagents`, resolved before construction against startup `subagent_runtime.max_running` and 1-64) and per run (`max_total_subagents` override or `subagents.max_total_per_run`, default 6, clamped 1-50). Count current-run durable ledger entries, tagged with `run_id` on capture: repeated planning shares the cap; later user runs get a fresh budget. Durable `batch_task` has separate persisted total/live/running limits and no ordinary ledger entries. When exhausted, strip remaining `task` calls, force `finish_reason="stop"`, and append a visible limit note so the agent synthesizes existing results.
+   Capture/retain delegations by `(run_id, tool_call_id)`. Later-run replies cannot close older orphans; missing run IDs match conservatively.
+32. **LoopDetectionMiddleware** - `loop_detection.enabled`: detects repeated tool-call loops. Hard stops clear structured/raw/content-block calls, force a final text answer, and stamp `loop_capped` via `consume_stop_reason` (#3875 Phase 2, matching `TokenBudgetMiddleware`). Persist first warned-state transitions per call hash/tool-frequency burst and hard stops as `middleware:loop_detection`, with `is_subagent`/optional `agent_id`; omit arguments, message content, results and argument-derived hashes. Ordinary task subagents use dedicated recorder keys via a parent-loop proxy; never pass `RunJournal` into their isolated loop. Durable batch subagents have no parent journal and persist no transitions.
+   New user runs reset budgets/history; same-run goal continuations share it.
+   Isolate sibling warnings and keep lifecycle hooks topology-stable. Before changing this guard, read
    [Loop detection lifecycle](../../../../../docs/LOOP_DETECTION.md) for
    fallback identity, cleanup/LRU/reset, severity ordering, and test invariants.
-33. **TokenBudgetMiddleware** - `token_budget.enabled`: shares run-ID budgets across continuations; missing/invalid IDs clear invocation state.
+33. **TokenBudgetMiddleware** - `token_budget.enabled`: run-ID budgets survive continuations; missing/invalid IDs reset invocation state. With token usage tracking, place CompletedSubagentUsageMiddleware after the guard: reversed `after_model` hooks backfill terminal child usage before enforcement. Keep TokenUsageMiddleware earlier so current-step attribution reflects calls retained by all guards.
 34. **Custom middlewares** - *(optional)* Any `custom_middlewares` passed to `build_middlewares` are injected here, before config-declared extensions and the terminal-response/safety/clarification tail
 35. **Configured extension middlewares** - `extensions.middlewares` in `config.yaml` or `extensions_config.json` optionally accepts `module.path:ClassName` strings or `{class, kwargs}` objects. `deerflow.reflection.resolve_class` loads `AgentMiddleware` classes; import, class, and constructor errors fail agent creation. `kwargs` must be JSON-compatible; YAML dates/timestamps become ISO strings. Order: built-ins/custom and loop/token guards → extensions → terminal-response/safety/clarification tail. Subagents share the list before their safety tail; separate lead/subagent lists are unsupported. Trusted operator config only: paths instantiate arbitrary code. Gateway skill/MCP toggles preserve it in raw JSON; adding an API write path requires explicit trust-boundary review.
 36. **TerminalResponseMiddleware** - When a provider returns an empty terminal `AIMessage` after tool execution, injects a hidden recovery prompt and retries the model once; a second empty response is replaced in checkpoint state by a visible error fallback marked for the run worker, so the run finishes as an error instead of a silent success

backend/packages/harness/deerflow/agents/middlewares/token_usage_middleware.py

@@ -277,6 +277,53 @@ def _build_attribution(message: AIMessage, todos: list[Todo]) -> dict[str, Any]:
     }
 
 
+def _attribute_completed_subagent_usage(messages: list[Any]) -> dict[int, AIMessage | ToolMessage]:
+    """Return copied dispatch/terminal message updates without current-step attribution."""
+    # Annotate subagent token usage onto the AIMessage that dispatched it.
+    # Terminal usage travels with the current run's ToolMessage, so provider
+    # tool-call IDs reused by another run cannot overwrite attribution state.
+    # Walk backward through consecutive ToolMessages before the new AIMessage
+    # so that multiple concurrent task tool calls all get their subagent tokens
+    # written back to the same dispatch message (merging into one update).
+    state_updates: dict[int, AIMessage | ToolMessage] = {}
+    if len(messages) >= 2:
+        idx = len(messages) - 2
+        while idx >= 0:
+            tool_msg = messages[idx]
+            if not isinstance(tool_msg, ToolMessage) or not tool_msg.tool_call_id:
+                break
+
+            subagent_usage = _subagent_usage_from_tool_message(tool_msg)
+            if subagent_usage:
+                # Search backward from the ToolMessage to find the AIMessage
+                # that dispatched it.  A single model response can dispatch
+                # multiple task tool calls, so we can't assume a fixed offset.
+                dispatch_idx = idx - 1
+                while dispatch_idx >= 0:
+                    candidate = messages[dispatch_idx]
+                    if isinstance(candidate, AIMessage) and _has_tool_call(candidate, tool_msg.tool_call_id):
+                        # Accumulate into an existing update for the same
+                        # AIMessage (multiple task calls in one response),
+                        # or merge fresh from the original message.
+                        existing_update = state_updates.get(dispatch_idx)
+                        prev = existing_update.usage_metadata if isinstance(existing_update, AIMessage) else (getattr(candidate, "usage_metadata", None) or {})
+                        merged = {
+                            **prev,
+                            "input_tokens": prev.get("input_tokens", 0) + subagent_usage["input_tokens"],
+                            "output_tokens": prev.get("output_tokens", 0) + subagent_usage["output_tokens"],
+                            "total_tokens": prev.get("total_tokens", 0) + subagent_usage["total_tokens"],
+                        }
+                        state_updates[dispatch_idx] = candidate.model_copy(update={"usage_metadata": merged})
+                        tool_metadata = dict(tool_msg.additional_kwargs)
+                        tool_metadata[SUBAGENT_TOKEN_USAGE_ATTRIBUTED_KEY] = True
+                        state_updates[idx] = tool_msg.model_copy(update={"additional_kwargs": tool_metadata})
+                        break
+                    dispatch_idx -= 1
+            idx -= 1
+
+    return state_updates
+
+
 class TokenUsageMiddleware(AgentMiddleware):
     """Logs token usage from model responses and annotates the AI step."""
 
@@ -285,47 +332,7 @@ class TokenUsageMiddleware(AgentMiddleware):
         if not messages:
             return None
 
-        # Annotate subagent token usage onto the AIMessage that dispatched it.
-        # Terminal usage travels with the current run's ToolMessage, so provider
-        # tool-call IDs reused by another run cannot overwrite attribution state.
-        # Walk backward through consecutive ToolMessages before the new AIMessage
-        # so that multiple concurrent task tool calls all get their subagent tokens
-        # written back to the same dispatch message (merging into one update).
-        state_updates: dict[int, AIMessage | ToolMessage] = {}
-        if len(messages) >= 2:
-            idx = len(messages) - 2
-            while idx >= 0:
-                tool_msg = messages[idx]
-                if not isinstance(tool_msg, ToolMessage) or not tool_msg.tool_call_id:
-                    break
-
-                subagent_usage = _subagent_usage_from_tool_message(tool_msg)
-                if subagent_usage:
-                    # Search backward from the ToolMessage to find the AIMessage
-                    # that dispatched it.  A single model response can dispatch
-                    # multiple task tool calls, so we can't assume a fixed offset.
-                    dispatch_idx = idx - 1
-                    while dispatch_idx >= 0:
-                        candidate = messages[dispatch_idx]
-                        if isinstance(candidate, AIMessage) and _has_tool_call(candidate, tool_msg.tool_call_id):
-                            # Accumulate into an existing update for the same
-                            # AIMessage (multiple task calls in one response),
-                            # or merge fresh from the original message.
-                            existing_update = state_updates.get(dispatch_idx)
-                            prev = existing_update.usage_metadata if isinstance(existing_update, AIMessage) else (getattr(candidate, "usage_metadata", None) or {})
-                            merged = {
-                                **prev,
-                                "input_tokens": prev.get("input_tokens", 0) + subagent_usage["input_tokens"],
-                                "output_tokens": prev.get("output_tokens", 0) + subagent_usage["output_tokens"],
-                                "total_tokens": prev.get("total_tokens", 0) + subagent_usage["total_tokens"],
-                            }
-                            state_updates[dispatch_idx] = candidate.model_copy(update={"usage_metadata": merged})
-                            tool_metadata = dict(tool_msg.additional_kwargs)
-                            tool_metadata[SUBAGENT_TOKEN_USAGE_ATTRIBUTED_KEY] = True
-                            state_updates[idx] = tool_msg.model_copy(update={"additional_kwargs": tool_metadata})
-                            break
-                        dispatch_idx -= 1
-                idx -= 1
+        state_updates = _attribute_completed_subagent_usage(messages)
 
         last = messages[-1]
         if not isinstance(last, AIMessage):
@@ -370,3 +377,19 @@ class TokenUsageMiddleware(AgentMiddleware):
     @override
     async def aafter_model(self, state: AgentState, runtime: Runtime) -> dict | None:
         return self._apply(state)
+
+
+class CompletedSubagentUsageMiddleware(AgentMiddleware):
+    """Backfill completed subagent usage before the lead budget checks it."""
+
+    def _apply(self, state: AgentState) -> dict | None:
+        state_updates = _attribute_completed_subagent_usage(state.get("messages", []))
+        return {"messages": [state_updates[idx] for idx in sorted(state_updates)]} if state_updates else None
+
+    @override
+    def after_model(self, state: AgentState, runtime: Runtime) -> dict | None:
+        return self._apply(state)
+
+    @override
+    async def aafter_model(self, state: AgentState, runtime: Runtime) -> dict | None:
+        return self._apply(state)

backend/tests/test_lead_token_budget.py

@@ -0,0 +1,220 @@
+"""Regression tests for subagent usage in the assembled lead budget chain."""
+
+from typing import Annotated
+
+import pytest
+from langchain.agents import create_agent
+from langchain_core.language_models.fake_chat_models import FakeMessagesListChatModel
+from langchain_core.messages import AIMessage, HumanMessage
+from langchain_core.tools import InjectedToolCallId, tool
+from langgraph.checkpoint.memory import InMemorySaver
+
+from deerflow.agents.lead_agent.agent import build_middlewares
+from deerflow.agents.middlewares.loop_detection_middleware import LoopDetectionMiddleware
+from deerflow.agents.middlewares.subagent_limit_middleware import SubagentLimitMiddleware
+from deerflow.agents.middlewares.token_budget_middleware import TokenBudgetMiddleware
+from deerflow.agents.middlewares.token_usage_middleware import TOKEN_USAGE_ATTRIBUTION_KEY, CompletedSubagentUsageMiddleware, TokenUsageMiddleware
+from deerflow.agents.thread_state import ThreadState
+from deerflow.config.app_config import AppConfig
+from deerflow.config.loop_detection_config import LoopDetectionConfig
+from deerflow.config.memory_config import MemoryConfig
+from deerflow.config.sandbox_config import SandboxConfig
+from deerflow.config.summarization_config import SummarizationConfig
+from deerflow.config.token_budget_config import TokenBudgetConfig
+from deerflow.config.token_usage_config import TokenUsageConfig
+from deerflow.tools.builtins.task_tool import _task_result_command
+
+_ACCOUNTING_MIDDLEWARE_TYPES = (TokenBudgetMiddleware, TokenUsageMiddleware, CompletedSubagentUsageMiddleware)
+
+
+class _ToolCallingModel(FakeMessagesListChatModel):
+    def bind_tools(self, tools, **kwargs):
+        return self
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("async_graph", [False, True], ids=["sync", "async"])
+@pytest.mark.parametrize(
+    ("child_tokens", "usage_enabled"),
+    [(200, True), (500, True), (500, False)],
+    ids=["below-budget", "last-batch-exceeds-budget", "usage-disabled"],
+)
+async def test_lead_budget_counts_last_subagent_batch_across_same_run_continuation(async_graph, child_tokens, usage_enabled):
+    app_config = AppConfig(
+        sandbox=SandboxConfig(use="deerflow.sandbox.local:LocalSandboxProvider"),
+        memory=MemoryConfig(enabled=False),
+        summarization=SummarizationConfig(enabled=False),
+        token_budget=TokenBudgetConfig(enabled=True, max_tokens=1000),
+        token_usage=TokenUsageConfig(enabled=usage_enabled),
+    )
+    assembled = build_middlewares(
+        {"configurable": {"subagent_enabled": True}},
+        model_name=None,
+        app_config=app_config,
+        memory_enabled=False,
+        available_skills=set(),
+    )
+    # Preserve the production factory's accounting order; unrelated middleware
+    # and live child model execution are outside this offline regression.
+    accounting = [middleware for middleware in assembled if isinstance(middleware, _ACCOUNTING_MIDDLEWARE_TYPES)]
+    budget = next(middleware for middleware in accounting if isinstance(middleware, TokenBudgetMiddleware))
+
+    @tool("task")
+    def completed_task(tool_call_id: Annotated[str, InjectedToolCallId]):
+        """Return controlled terminal usage through the production serializer."""
+        return _task_result_command(
+            tool_call_id=tool_call_id,
+            status="completed",
+            result="Completed synthetic task.",
+            usage={"input_tokens": child_tokens, "output_tokens": 0, "total_tokens": child_tokens},
+        )
+
+    def response(message_id, content="", tool_calls=None):
+        return AIMessage(
+            id=message_id,
+            content=content,
+            tool_calls=tool_calls or [],
+            usage_metadata={"input_tokens": 100, "output_tokens": 0, "total_tokens": 100},
+        )
+
+    graph = create_agent(
+        model=_ToolCallingModel(
+            responses=[
+                response("dispatch", tool_calls=[{"name": "task", "id": f"child-{index}", "args": {}} for index in range(2)]),
+                response("final", "Completed answer."),
+                response("continuation", "Continued answer."),
+                response("new-run", "New answer."),
+            ]
+        ),
+        tools=[completed_task],
+        middleware=accounting,
+        state_schema=ThreadState,
+        context_schema=dict,
+        checkpointer=InMemorySaver(),
+    )
+    config = {"configurable": {"thread_id": "budget-thread"}}
+
+    async def invoke(text, run_id):
+        arguments = {"config": config, "context": {"thread_id": "budget-thread", "run_id": run_id}}
+        inputs = {"messages": [HumanMessage(content=text)]}
+        return await graph.ainvoke(inputs, **arguments) if async_graph else graph.invoke(inputs, **arguments)
+
+    expected = 200 + (2 * child_tokens if usage_enabled else 0)
+    first = await invoke("Complete two tasks.", "run-1")
+    assert budget._cumulative_usage["run-1"].total == expected
+    assert ("TOKEN BUDGET EXCEEDED" in first["messages"][-1].content) == (expected >= 1000)
+
+    # Rebuilding the seen-message baseline must neither discard the last child
+    # batch nor count checkpointed usage a second time.
+    await invoke("Continue within the same run.", "run-1")
+    assert budget._cumulative_usage["run-1"].total == expected + 100
+    assert budget.consume_stop_reason("run-1") == ("token_capped" if expected >= 1000 else None)
+
+    await invoke("Start another user run.", "run-2")
+    assert budget._cumulative_usage["run-2"].total == 100
+    assert budget.consume_stop_reason("run-2") is None
+
+
+def test_disabled_budget_preserves_attribution_after_subagent_call_limit():
+    """Default tracking must describe the calls retained by the real limit."""
+    app_config = AppConfig(
+        sandbox=SandboxConfig(use="deerflow.sandbox.local:LocalSandboxProvider"),
+        memory=MemoryConfig(enabled=False),
+        summarization=SummarizationConfig(enabled=False),
+        token_budget=TokenBudgetConfig(enabled=False),
+    )
+    assembled = build_middlewares(
+        {"configurable": {"subagent_enabled": True, "max_concurrent_subagents": 1}},
+        model_name=None,
+        app_config=app_config,
+        memory_enabled=False,
+        available_skills=set(),
+    )
+    middlewares = [middleware for middleware in assembled if isinstance(middleware, (SubagentLimitMiddleware, TokenUsageMiddleware))]
+    executed = []
+
+    @tool("task")
+    def completed_task(description: str) -> str:
+        """Record the child calls actually executed after limiting."""
+        executed.append(description)
+        return "Completed."
+
+    dispatch = AIMessage(id="dispatch", content="", tool_calls=[{"name": "task", "id": f"child-{index}", "args": {"description": str(index)}} for index in range(3)])
+    graph = create_agent(
+        model=_ToolCallingModel(responses=[dispatch, AIMessage(content="Done.")]),
+        tools=[completed_task],
+        middleware=middlewares,
+        state_schema=ThreadState,
+        context_schema=dict,
+    )
+    result = graph.invoke({"messages": [HumanMessage(content="Complete tasks.")]}, context={"run_id": "limited-run"})
+    recorded_dispatch = next(message for message in result["messages"] if message.id == "dispatch")
+    assert executed == ["0"]
+    assert recorded_dispatch.additional_kwargs[TOKEN_USAGE_ATTRIBUTION_KEY]["tool_call_ids"] == ["child-0"]
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("async_graph", [False, True], ids=["sync", "async"])
+@pytest.mark.parametrize("guard", ["concurrent", "total", "budget", "loop"])
+async def test_enabled_budget_attribution_matches_calls_retained_by_guards(async_graph, guard):
+    """Enabled budgeting must preserve attribution after actual call guards."""
+    app_config = AppConfig(
+        sandbox=SandboxConfig(use="deerflow.sandbox.local:LocalSandboxProvider"),
+        memory=MemoryConfig(enabled=False),
+        summarization=SummarizationConfig(enabled=False),
+        token_budget=TokenBudgetConfig(enabled=True, max_tokens=1000 if guard == "budget" else 10000),
+        loop_detection=LoopDetectionConfig(warn_threshold=1 if guard == "loop" else 100, hard_limit=1 if guard == "loop" else 100),
+    )
+    assembled = build_middlewares(
+        {"configurable": {"subagent_enabled": True, "max_concurrent_subagents": 1, "max_total_subagents": 1}},
+        model_name=None,
+        app_config=app_config,
+        memory_enabled=False,
+        available_skills=set(),
+    )
+    guard_classes = _ACCOUNTING_MIDDLEWARE_TYPES + (SubagentLimitMiddleware, LoopDetectionMiddleware)
+    middlewares = [middleware for middleware in assembled if isinstance(middleware, guard_classes)]
+    executed = []
+
+    @tool("task")
+    def completed_task(description: str, tool_call_id: Annotated[str, InjectedToolCallId]) -> str:
+        """Record the task calls that survive the production guards."""
+        executed.append(tool_call_id)
+        return "Completed."
+
+    tokens = 1000 if guard == "budget" else 100
+    dispatch = AIMessage(
+        id="dispatch",
+        content="",
+        tool_calls=[{"name": "task", "id": f"child-{index}", "args": {"description": str(index)}} for index in range(3)],
+        usage_metadata={"input_tokens": tokens, "output_tokens": 0, "total_tokens": tokens},
+    )
+    graph = create_agent(
+        model=_ToolCallingModel(responses=[dispatch, AIMessage(content="Done.")]),
+        tools=[completed_task],
+        middleware=middlewares,
+        state_schema=ThreadState,
+        context_schema=dict,
+    )
+    context = {"run_id": "guarded-run", "thread_id": "guarded-thread"}
+    inputs = {"messages": [HumanMessage(content="Complete tasks.")]}
+    if guard == "total":
+        inputs["delegations"] = [
+            {
+                "id": "prior-child",
+                "run_id": "guarded-run",
+                "description": "Prior completed task.",
+                "subagent_type": "general-purpose",
+                "status": "completed",
+                "created_at": "2026-01-01T00:00:00Z",
+            }
+        ]
+    result = await graph.ainvoke(inputs, context=context) if async_graph else graph.invoke(inputs, context=context)
+    recorded_dispatch = next(message for message in result["messages"] if message.id == "dispatch")
+    expected_calls = ["child-0"] if guard == "concurrent" else []
+    assert executed == expected_calls
+    assert [call["id"] for call in recorded_dispatch.tool_calls] == expected_calls
+    attribution = recorded_dispatch.additional_kwargs[TOKEN_USAGE_ATTRIBUTION_KEY]
+    assert attribution["tool_call_ids"] == expected_calls
+    assert [action["tool_call_id"] for action in attribution["actions"]] == expected_calls
+    assert attribution["shared_attribution"] is False