diff --git a/backend/packages/harness/deerflow/agents/middlewares/token_usage_middleware.py b/backend/packages/harness/deerflow/agents/middlewares/token_usage_middleware.py index b72204db8f..9e34a7c53d 100644 --- a/backend/packages/harness/deerflow/agents/middlewares/token_usage_middleware.py +++ b/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}) diff --git a/backend/packages/harness/deerflow/runtime/AGENTS.md b/backend/packages/harness/deerflow/runtime/AGENTS.md index d5f1a78e64..6110e2a333 100644 --- a/backend/packages/harness/deerflow/runtime/AGENTS.md +++ b/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`. diff --git a/backend/packages/harness/deerflow/runtime/runs/input_control.py b/backend/packages/harness/deerflow/runtime/runs/input_control.py index 04549f5f66..b905097bce 100644 --- a/backend/packages/harness/deerflow/runtime/runs/input_control.py +++ b/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), diff --git a/backend/tests/test_run_input_graph.py b/backend/tests/test_run_input_graph.py index f8080f5e9a..b451a5d3fe 100644 --- a/backend/tests/test_run_input_graph.py +++ b/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):