fix: bucket subagent token usage into parent run totals (#2838)

* fix: bucket subagent token usage into RunRow.subagent_tokens

Add caller-bucketed token tracking to RunJournal so subagent and
middleware LLM calls are written to the correct RunRow columns instead
of all falling into lead_agent_tokens (default 0).

- RunJournal: accumulate _lead_agent_tokens / _subagent_tokens /
  _middleware_tokens in on_llm_end, deduped by langchain run_id.
  Add record_external_llm_usage_records() for external sources
  (respects track_token_usage flag). Return caller buckets from
  get_completion_data().
- SubagentTokenCollector: new lightweight callback handler that
  collects LLM usage within subagent execution.
- SubagentExecutor: wire collector into subagent run_config and sync
  records to SubagentResult on every chunk (timeout/cancel safe).
- SubagentResult: add token_usage_records and usage_reported fields.
- task_tool: report subagent usage to parent RunJournal on every
  terminal status (COMPLETED/FAILED/CANCELLED/TIMED_OUT), including
  the CancelledError path, guarded against double-reporting.

No DB migration needed — RunRow columns already exist.

* Potential fix for pull request finding

Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com>

* fix: address token usage review feedback

* Address review follow-ups

---------

Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com>
This commit is contained in:
YuJitang
2026-05-10 22:47:30 +08:00
committed by GitHub
parent 94da8f67d7
commit 9892a7d468
8 changed files with 843 additions and 77 deletions
@@ -26,6 +26,7 @@ from deerflow.models import create_chat_model
from deerflow.skills.tool_policy import filter_tools_by_skill_allowed_tools
from deerflow.skills.types import Skill
from deerflow.subagents.config import SubagentConfig, resolve_subagent_model_name
from deerflow.subagents.token_collector import SubagentTokenCollector
logger = logging.getLogger(__name__)
@@ -70,6 +71,8 @@ class SubagentResult:
started_at: datetime | None = None
completed_at: datetime | None = None
ai_messages: list[dict[str, Any]] | None = None
token_usage_records: list[dict[str, int | str]] = field(default_factory=list)
usage_reported: bool = False
cancel_event: threading.Event = field(default_factory=threading.Event, repr=False)
def __post_init__(self):
@@ -412,13 +415,20 @@ class SubagentExecutor:
ai_messages = []
result.ai_messages = ai_messages
collector: SubagentTokenCollector | None = None
try:
state, filtered_tools = await self._build_initial_state(task)
agent = self._create_agent(filtered_tools)
# Token collector for subagent LLM calls
collector_caller = f"subagent:{self.config.name}"
collector = SubagentTokenCollector(caller=collector_caller)
# Build config with thread_id for sandbox access and recursion limit
run_config: RunnableConfig = {
"recursion_limit": self.config.max_turns,
"callbacks": [collector],
"tags": [collector_caller],
}
context: dict[str, Any] = {}
if self.thread_id:
@@ -441,6 +451,8 @@ class SubagentExecutor:
result.status = SubagentStatus.CANCELLED
result.error = "Cancelled by user"
result.completed_at = datetime.now()
if collector is not None:
result.token_usage_records = collector.snapshot_records()
return result
async for chunk in agent.astream(state, config=run_config, context=context, stream_mode="values"): # type: ignore[arg-type]
@@ -455,6 +467,7 @@ class SubagentExecutor:
result.status = SubagentStatus.CANCELLED
result.error = "Cancelled by user"
result.completed_at = datetime.now()
result.token_usage_records = collector.snapshot_records()
return result
final_state = chunk
@@ -481,6 +494,7 @@ class SubagentExecutor:
logger.info(f"[trace={self.trace_id}] Subagent {self.config.name} captured AI message #{len(ai_messages)}")
logger.info(f"[trace={self.trace_id}] Subagent {self.config.name} completed async execution")
result.token_usage_records = collector.snapshot_records()
if final_state is None:
logger.warning(f"[trace={self.trace_id}] Subagent {self.config.name} no final state")
@@ -560,6 +574,8 @@ class SubagentExecutor:
result.status = SubagentStatus.FAILED
result.error = str(e)
result.completed_at = datetime.now()
if collector is not None:
result.token_usage_records = collector.snapshot_records()
return result