fix(runtime): charge a steered subagent's usage to the token budget (#7776)

* fix(runtime): charge a steered subagent's usage to the token budget

A task that finishes after a Steer reports through a receipt HumanMessage, so its usage never reached the dispatching AIMessage the token budget reads. Carry the usage over when the receipt is delivered.

* fix(runtime): report deferred usage with no AIMessage to charge

Log the rare case instead of dropping the usage silently, and scan for AIMessages once per delivery.

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

@@ -277,6 +277,36 @@ def _build_attribution(message: AIMessage, todos: list[Todo]) -> dict[str, Any]:
     }
 
 
+def _add_subagent_usage(message: AIMessage, usage: dict[str, int]) -> AIMessage:
+    previous = message.usage_metadata or {}
+    merged = {
+        **previous,
+        "input_tokens": previous.get("input_tokens", 0) + usage["input_tokens"],
+        "output_tokens": previous.get("output_tokens", 0) + usage["output_tokens"],
+        "total_tokens": previous.get("total_tokens", 0) + usage["total_tokens"],
+    }
+    return message.model_copy(update={"usage_metadata": merged})
+
+
+def attribute_deferred_subagent_usage(messages: list[Any], deferred: list[tuple[str, dict[str, int]]]) -> list[AIMessage]:
+    """Return AIMessage copies carrying usage of subagents that outlived their ToolMessage.
+
+    A steered task's pending ToolMessage is no longer next to the latest
+    AIMessage and its receipt is a HumanMessage, so the terminal-message walk
+    below cannot see this usage. The caller owns delivering each entry once.
+    """
+    updates: dict[int, AIMessage] = {}
+    indexes = [idx for idx in range(len(messages) - 1, -1, -1) if isinstance(messages[idx], AIMessage)]
+    for tool_call_id, usage in deferred:
+        # Compaction can drop the dispatch, so fall back to the newest AIMessage.
+        target = next((idx for idx in indexes if _has_tool_call(messages[idx], tool_call_id)), indexes[0] if indexes else None)
+        if target is None:
+            logger.warning("Token budget missed subagent usage for %s: no AIMessage to charge", tool_call_id)
+            continue
+        updates[target] = _add_subagent_usage(updates.get(target, messages[target]), usage)
+    return [updates[idx] for idx in sorted(updates)]
+
+
 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.
@@ -306,14 +336,7 @@ def _attribute_completed_subagent_usage(messages: list[Any]) -> dict[int, AIMess
                         # 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})
+                        state_updates[dispatch_idx] = _add_subagent_usage(existing_update if isinstance(existing_update, AIMessage) else candidate, subagent_usage)
                         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})

backend/packages/harness/deerflow/runtime/AGENTS.md

@@ -23,6 +23,8 @@ owns cooperative pause, timeout and cancellation. A steer withdrawn before
 delivery resumes the children it paused, except those a delivered steer still
 holds. Only direct-user input exposes
 current-work control. Preserve permissions, pool caps and token/loop stops.
+A task finishing after Steer charges usage to its dispatch `AIMessage` when its
+receipt is delivered (newest `AIMessage` if compacted); budgets read only those.
 On restart retain unapplied text for explicit retry; never replay side effects
 automatically. Regressions: `test_run_inputs.py`, `test_run_input_graph.py`.
 

backend/packages/harness/deerflow/runtime/runs/input_control.py

@@ -78,6 +78,7 @@ class RunInputControl:
         self.children: dict[str, tuple[ChildInputGate, str | None]] = {}
         self.interrupted_children: set[str] = set()
         self.receipts: list[HumanMessage] = []
+        self.deferred_usage: list[tuple[str, dict[str, int]]] = []
         self.pending_calls: set[str] = set()
         self.finished_calls: set[str] = set()
         self.on_result = on_result
@@ -102,11 +103,14 @@ class RunInputControl:
         return [x for x in state["items"] if x["mode"] == "steer" and x["status"] in ("pending", "applying") and x["id"] not in self.delivered]
 
     async def deliver(self, state):
+        from deerflow.agents.middlewares.token_usage_middleware import attribute_deferred_subagent_usage
+
         entries = await self.pending()
         if entries:
             entries = await self.inputs.claim_steer([x["id"] for x in entries])
-        messages = self.receipts[:]
+        messages = [*attribute_deferred_subagent_usage(state.get("messages", []), self.deferred_usage), *self.receipts]
         self.receipts.clear()
+        self.deferred_usage.clear()
         if entries:
             self.abandoned = False
             self.steering_active = True
@@ -194,9 +198,13 @@ class RunInputControl:
         except Exception:
             message = self._failure(call_id, "failed", "The subagent result could not be collected")
         from deerflow.agents.middlewares.delegation_ledger import extract_delegations
+        from deerflow.subagents.status_contract import SUBAGENT_TOKEN_USAGE_KEY, normalize_token_usage
 
         self.finished_calls.add(call_id)
         self.delegations.extend({**entry, "run_id": self.inputs.run_id} for entry in extract_delegations([AIMessage(content="", tool_calls=[self.calls[call_id]]), message]))
+        usage = normalize_token_usage(message.additional_kwargs.get(SUBAGENT_TOKEN_USAGE_KEY))
+        if usage:
+            self.deferred_usage.append((call_id, usage))
         self.receipts.append(
             HumanMessage(
                 content=json.dumps({"source": "subagent", "task_call_id": call_id, "result": message.content}, ensure_ascii=False),

backend/tests/test_run_input_graph.py

@@ -2,21 +2,26 @@
 
 import asyncio
 from types import SimpleNamespace
+from typing import Annotated
 
 import pytest
 from langchain.agents import create_agent
 from langchain_core.language_models import BaseChatModel
 from langchain_core.messages import AIMessage, HumanMessage, RemoveMessage, ToolMessage
 from langchain_core.outputs import ChatGeneration, ChatResult
-from langchain_core.tools import tool
+from langchain_core.tools import InjectedToolCallId, tool
 from langgraph.checkpoint.memory import InMemorySaver
+from langgraph.types import Command
 from pydantic import ConfigDict
 
 from deerflow.agents.middlewares.run_input_middleware import RunInputAdmissionMiddleware, RunInputMiddleware
+from deerflow.agents.middlewares.token_budget_middleware import TokenBudgetMiddleware
 from deerflow.agents.thread_state import ThreadState
+from deerflow.config.token_budget_config import TokenBudgetConfig
 from deerflow.runtime.runs.input_control import ChildInputGate, RunInputControl
 from deerflow.runtime.runs.inputs import INPUT_CONTEXT_KEY, RunInputs
 from deerflow.runtime.runs.store.memory import MemoryRunStore
+from deerflow.subagents.status_contract import SUBAGENT_TOKEN_USAGE_KEY
 
 
 def test_input_control_reaches_the_factory_only_from_the_worker():
@@ -264,6 +269,135 @@ async def test_graph_keeps_a_pending_result_for_every_parallel_task():
         await control.close()
 
 
+class SpendingModel(InputModel):
+    def _generate(self, messages, stop=None, run_manager=None, **kwargs):
+        self.observed.append(messages)
+        step = len(self.observed)
+        calls = [{"name": "task", "args": {}, "id": "long"}] if step == 1 else [{"name": "follow_up", "args": {}, "id": f"follow-{step}"}] if step <= 4 else []
+        message = AIMessage(content=f"Step {step}", tool_calls=calls, usage_metadata={"input_tokens": 10, "output_tokens": 1, "total_tokens": 11})
+        return ChatResult(generations=[ChatGeneration(message=message)])
+
+
+def _task_result(tool_call_id, total):
+    usage = {"input_tokens": total - 100, "output_tokens": 100, "total_tokens": total}
+    return Command(update={"messages": [ToolMessage(content="Child finished", tool_call_id=tool_call_id, name="task", additional_kwargs={SUBAGENT_TOKEN_USAGE_KEY: usage})]})
+
+
+async def _run_controller():
+    store = MemoryRunStore()
+    await store.put("run", thread_id="thread", user_id="owner", status="running", owner_worker_id="worker")
+    inputs = RunInputs(store, "run", "owner", "worker")
+    saver = InMemorySaver()
+    return inputs, saver, RunInputControl(inputs, agent_runs=SimpleNamespace(), thread_id="thread", checkpointer=saver)
+
+
+@pytest.mark.asyncio
+async def test_budget_stops_on_usage_of_a_task_finished_after_a_steer():
+    inputs, saver, control = await _run_controller()
+    entered, release = asyncio.Event(), asyncio.Event()
+    followed = []
+
+    @tool
+    async def task(tool_call_id: Annotated[str, InjectedToolCallId]):
+        """A child that spends the whole budget after the lead stopped waiting."""
+        entered.set()
+        await release.wait()
+        return _task_result(tool_call_id, 1_000_000)
+
+    @tool
+    async def follow_up() -> str:
+        """Lead work that must stop once the child's tokens are charged."""
+        followed.append(True)
+        release.set()
+        async with asyncio.timeout(5):
+            while not control.receipts:
+                await asyncio.sleep(0.01)
+        return "Done"
+
+    budget = TokenBudgetMiddleware.from_config(TokenBudgetConfig(enabled=True, max_tokens=1000))
+    context = {INPUT_CONTEXT_KEY: control, "run_id": "run"}
+    graph = create_agent(SpendingModel(observed=[]), tools=[task, follow_up], middleware=[RunInputAdmissionMiddleware(), budget, RunInputMiddleware()], state_schema=ThreadState, context_schema=dict, checkpointer=saver)
+    run = asyncio.create_task(graph.ainvoke({"messages": [HumanMessage(content="Original task")]}, {"configurable": {"thread_id": "thread"}}, context=context))
+    try:
+        await asyncio.wait_for(entered.wait(), 5)
+        await inputs.submit("update", "Handle this too", "steer")
+        result = await asyncio.wait_for(run, 10)
+        assert context.get("stop_reason") == "token_capped"
+        assert len(followed) == 1
+        assert not result["messages"][-1].tool_calls
+    finally:
+        release.set()
+        await asyncio.wait_for(asyncio.gather(run, return_exceptions=True), 5)
+        await control.close()
+
+
+@pytest.mark.asyncio
+async def test_parallel_deferred_usage_is_charged_once_to_its_dispatch():
+    inputs, saver, control = await _run_controller()
+    entered, release = asyncio.Event(), asyncio.Event()
+    started = []
+
+    @tool
+    async def task(tool_call_id: Annotated[str, InjectedToolCallId]):
+        """Both parallel calls finish after the Steer released their waiters."""
+        started.append(tool_call_id)
+        if len(started) == 2:
+            entered.set()
+        await release.wait()
+        return _task_result(tool_call_id, 200 if tool_call_id == "first" else 300)
+
+    model = ParallelTaskModel(observed=[])
+    budget = TokenBudgetMiddleware.from_config(TokenBudgetConfig(enabled=True, max_tokens=10_000))
+    config = {"configurable": {"thread_id": "thread"}}
+    graph = create_agent(model, tools=[task], middleware=[RunInputAdmissionMiddleware(), budget, RunInputMiddleware()], state_schema=ThreadState, context_schema=dict, checkpointer=saver)
+    run = asyncio.create_task(graph.ainvoke({"messages": [HumanMessage(content="Original task")]}, config, context={INPUT_CONTEXT_KEY: control, "run_id": "run"}))
+    try:
+        await asyncio.wait_for(entered.wait(), 5)
+        await inputs.submit("update", "Handle this first", "steer")
+        async with asyncio.timeout(5):
+            while len(model.observed) < 2:
+                await asyncio.sleep(0.01)
+        release.set()
+        result = await asyncio.wait_for(run, 10)
+        dispatch = next(message for message in result["messages"] if isinstance(message, AIMessage) and message.tool_calls)
+        assert dispatch.usage_metadata == {"input_tokens": 300, "output_tokens": 200, "total_tokens": 500}
+        assert budget._cumulative_usage["run"].total == 500
+    finally:
+        release.set()
+        await asyncio.wait_for(asyncio.gather(run, return_exceptions=True), 5)
+        await control.close()
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize(("usage", "charged"), [({"input_tokens": 3, "output_tokens": 4, "total_tokens": 7}, 7), ({"total_tokens": "many"}, None)])
+async def test_deferred_usage_survives_a_compacted_dispatch(usage, charged):
+    _, _, control = await _run_controller()
+    control.calls["compacted"] = {"name": "task", "args": {}, "id": "compacted"}
+
+    async def finished():
+        return ToolMessage(content="Child finished", tool_call_id="compacted", name="task", additional_kwargs={SUBAGENT_TOKEN_USAGE_KEY: usage})
+
+    await control._observe(asyncio.create_task(finished()), "compacted")
+    latest = AIMessage(id="latest", content="Summarized history", usage_metadata={"input_tokens": 1, "output_tokens": 1, "total_tokens": 2})
+    update = await control.deliver({"messages": [AIMessage(id="older", content="Before"), HumanMessage(content="Next"), latest]})
+
+    rewritten = [message for message in update["messages"] if isinstance(message, AIMessage)]
+    if charged is None:
+        assert not rewritten
+    else:
+        assert [(message.id, message.usage_metadata["total_tokens"]) for message in rewritten] == [("latest", 2 + charged)]
+    assert await control.deliver({"messages": [latest]}) is None
+
+
+def test_deferred_usage_with_no_ai_message_to_charge_is_reported(caplog):
+    from deerflow.agents.middlewares.token_usage_middleware import attribute_deferred_subagent_usage
+
+    usage = {"input_tokens": 3, "output_tokens": 4, "total_tokens": 7}
+    with caplog.at_level("WARNING", logger="deerflow.agents.middlewares.token_usage_middleware"):
+        assert attribute_deferred_subagent_usage([HumanMessage(content="Only input")], [("lost", usage)]) == []
+    assert "lost" in caplog.text
+
+
 @pytest.mark.asyncio
 @pytest.mark.parametrize("capped", [False, True])
 async def test_input_during_inference_defers_old_actions_but_never_overrides_budget(capped):