github/deer-flow
1
0
Fork 0
mirror of https://github.com/bytedance/deer-flow.git synced 2026-07-22 06:00:21 +00:00
Code Issues Packages Projects Releases Wiki Activity

fix(gateway): attribute token usage to actual models (#3658)

Browse Source 
* fix(gateway): attribute token usage to actual models

Capture per-call model names from LLM response metadata for lead, middleware, and subagent calls.

Persist a per-run token_usage_by_model breakdown and aggregate by that map in both SQL and memory stores, with legacy fallback to the run-level model_name for older rows.

Add regression coverage for by_model totals, caller consistency, active progress snapshots, store parity, and SubagentTokenCollector model propagation.

* fix(gateway): harden by-model token aggregation

Use usage.get("total_tokens", 0) when reducing per-model token usage maps so aggregation tolerates partially written or manually edited JSON blobs without changing behavior for journal-written rows.

* docs(gateway): clarify by-model run count semantics

Document that by_model[*].runs counts the number of runs in which a model appeared, so multi-model runs can increment multiple model buckets.
This commit is contained in:
AnoobFeng 2026-06-19 21:42:42 +08:00 committed by GitHub
parent 8cde305fe4
commit e7a03e5243
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
12 changed files with 631 additions and 47 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

5
backend/app/gateway/routers/thread_runs.py
View File

@ -79,7 +79,10 @@ class RunResponse(BaseModel):
class ThreadTokenUsageModelBreakdown(BaseModel):
tokens: int = 0
runs: int = 0
runs: int = Field(
default=0,
description="Number of runs in which this model appeared; counts are non-exclusive for runs that used multiple models.",
)
class ThreadTokenUsageCallerBreakdown(BaseModel):

2
backend/packages/harness/deerflow/persistence/run/model.py
View File

@ -39,6 +39,8 @@ class RunRow(Base):
lead_agent_tokens: Mapped[int] = mapped_column(default=0)
subagent_tokens: Mapped[int] = mapped_column(default=0)
middleware_tokens: Mapped[int] = mapped_column(default=0)
# Per-model token breakdown
token_usage_by_model: Mapped[dict] = mapped_column(JSON, default=dict)
# Follow-up association
follow_up_to_run_id: Mapped[str | None] = mapped_column(String(64))

70
backend/packages/harness/deerflow/persistence/run/sql.py
View File

@ -11,7 +11,7 @@ import json
from datetime import UTC, datetime
from typing import Any
from sqlalchemy import func, select, update
from sqlalchemy import select, update
from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker
from deerflow.persistence.run.model import RunRow
@ -230,6 +230,7 @@ class RunRepository(RunStore):
lead_agent_tokens: int = 0,
subagent_tokens: int = 0,
middleware_tokens: int = 0,
token_usage_by_model: dict[str, dict[str, int]] | None = None,
message_count: int = 0,
last_ai_message: str | None = None,
first_human_message: str | None = None,
@ -248,6 +249,7 @@ class RunRepository(RunStore):
"lead_agent_tokens": lead_agent_tokens,
"subagent_tokens": subagent_tokens,
"middleware_tokens": middleware_tokens,
"token_usage_by_model": self._safe_json(token_usage_by_model) or {},
"message_count": message_count,
"updated_at": datetime.now(UTC),
}
@ -273,6 +275,7 @@ class RunRepository(RunStore):
lead_agent_tokens: int | None = None,
subagent_tokens: int | None = None,
middleware_tokens: int | None = None,
token_usage_by_model: dict[str, dict[str, int]] | None = None,
message_count: int | None = None,
last_ai_message: str | None = None,
first_human_message: str | None = None,
@ -292,6 +295,8 @@ class RunRepository(RunStore):
for key, value in optional_counters.items():
if value is not None:
values[key] = value
if token_usage_by_model is not None:
values["token_usage_by_model"] = self._safe_json(token_usage_by_model) or {}
if last_ai_message is not None:
values["last_ai_message"] = last_ai_message[:2000]
if first_human_message is not None:
@ -301,26 +306,33 @@ class RunRepository(RunStore):
await session.commit()
async def aggregate_tokens_by_thread(self, thread_id: str, *, include_active: bool = False) -> dict[str, Any]:
"""Aggregate token usage via a single SQL GROUP BY query."""
"""Aggregate token usage for a thread.
``by_model`` is reduced in Python from each row's ``token_usage_by_model``
JSON column so subagent / middleware tokens land on the model that
actually produced them (issue #3645). Rows written before that column
existed fall back to ``RunRow.model_name`` + ``RunRow.total_tokens``,
preserving the legacy lead-only behavior instead of dropping the data.
Headline totals (``total_tokens``, ``total_input_tokens``,
``total_output_tokens``) and the ``by_caller`` bucket are summed from
their own columns and are therefore unaffected by the JSON column being
empty.
"""
statuses = ("success", "error", "running") if include_active else ("success", "error")
_completed = RunRow.status.in_(statuses)
_thread = RunRow.thread_id == thread_id
model_name = func.coalesce(RunRow.model_name, "unknown")
stmt = (
select(
model_name.label("model"),
func.count().label("runs"),
func.coalesce(func.sum(RunRow.total_tokens), 0).label("total_tokens"),
func.coalesce(func.sum(RunRow.total_input_tokens), 0).label("total_input_tokens"),
func.coalesce(func.sum(RunRow.total_output_tokens), 0).label("total_output_tokens"),
func.coalesce(func.sum(RunRow.lead_agent_tokens), 0).label("lead_agent"),
func.coalesce(func.sum(RunRow.subagent_tokens), 0).label("subagent"),
func.coalesce(func.sum(RunRow.middleware_tokens), 0).label("middleware"),
)
.where(_thread, _completed)
.group_by(model_name)
)
stmt = select(
RunRow.model_name,
RunRow.total_tokens,
RunRow.total_input_tokens,
RunRow.total_output_tokens,
RunRow.lead_agent_tokens,
RunRow.subagent_tokens,
RunRow.middleware_tokens,
RunRow.token_usage_by_model,
).where(_thread, _completed)
async with self._sf() as session:
rows = (await session.execute(stmt)).all()
@ -329,14 +341,28 @@ class RunRepository(RunStore):
lead_agent = subagent = middleware = 0
by_model: dict[str, dict] = {}
for r in rows:
by_model[r.model] = {"tokens": r.total_tokens, "runs": r.runs}
total_runs += 1
total_tokens += r.total_tokens
total_input += r.total_input_tokens
total_output += r.total_output_tokens
total_runs += r.runs
lead_agent += r.lead_agent
subagent += r.subagent
middleware += r.middleware
lead_agent += r.lead_agent_tokens
subagent += r.subagent_tokens
middleware += r.middleware_tokens
# ``or {}`` covers rows written before ``token_usage_by_model``
# existed (the column is NULL on a manual ALTER ADD COLUMN without
# backfill); fresh rows always carry the journal-produced dict.
usage_by_model = r.token_usage_by_model or {}
if usage_by_model:
for model, usage in usage_by_model.items():
entry = by_model.setdefault(model, {"tokens": 0, "runs": 0})
entry["tokens"] += usage.get("total_tokens", 0)
entry["runs"] += 1
else:
model = r.model_name or "unknown"
entry = by_model.setdefault(model, {"tokens": 0, "runs": 0})
entry["tokens"] += r.total_tokens
entry["runs"] += 1
return {
"total_tokens": total_tokens,

47
backend/packages/harness/deerflow/runtime/journal.py
View File

@ -77,6 +77,9 @@ class RunJournal(BaseCallbackHandler):
self._subagent_tokens = 0
self._middleware_tokens = 0
# Per-model token accumulator
self._tokens_by_model: dict[str, dict[str, int]] = {}
# Dedup: LangChain may fire on_llm_end multiple times for the same run_id
self._counted_llm_run_ids: set[str] = set()
self._counted_external_source_ids: set[str] = set()
@ -327,6 +330,13 @@ class RunJournal(BaseCallbackHandler):
else:
self._lead_agent_tokens += total_tk
# Per-model bucket
response_metadata = getattr(message, "response_metadata", None) or {}
per_call_model: str | None = None
if isinstance(response_metadata, Mapping):
per_call_model = response_metadata.get("model_name") or response_metadata.get("model")
self._record_model_usage(per_call_model, input_tk, output_tk, total_tk)
self._schedule_progress_flush()
if messages:
@ -435,17 +445,42 @@ class RunJournal(BaseCallbackHandler):
# themselves.
return "lead_agent"
def _record_model_usage(
self,
model_name: str | None,
input_tokens: int,
output_tokens: int,
total_tokens: int,
) -> None:
"""Add a single LLM call's token usage to the per-model accumulator.
Missing / empty ``model_name`` collapses into a shared ``"unknown"``
bucket so the breakdown stays usable when a provider doesn't surface
``response_metadata.model_name``.
"""
if total_tokens <= 0:
return
bucket = self._tokens_by_model.setdefault(
model_name or "unknown",
{"input_tokens": 0, "output_tokens": 0, "total_tokens": 0},
)
bucket["input_tokens"] += int(input_tokens or 0)
bucket["output_tokens"] += int(output_tokens or 0)
bucket["total_tokens"] += int(total_tokens)
# -- Public methods (called by worker) --
def record_external_llm_usage_records(
self,
records: list[dict[str, int | str]],
records: list[dict[str, int | str | None]],
) -> None:
"""Record token usage from external sources (e.g., subagents).
Each record should contain:
source_run_id: Unique identifier to prevent double-counting
caller: Caller tag (e.g. "subagent:general-purpose")
model_name: Real per-call model name (str or None; falls back to
``"unknown"`` bucket when missing)
input_tokens: Input token count
output_tokens: Output token count
total_tokens: Total token count (computed from input+output if 0/missing)
@ -467,9 +502,12 @@ class RunJournal(BaseCallbackHandler):
if total_tk <= 0:
continue
input_tk = record.get("input_tokens", 0) or 0
output_tk = record.get("output_tokens", 0) or 0
self._counted_external_source_ids.add(source_id)
self._total_input_tokens += record.get("input_tokens", 0) or 0
self._total_output_tokens += record.get("output_tokens", 0) or 0
self._total_input_tokens += input_tk
self._total_output_tokens += output_tk
self._total_tokens += total_tk
caller = str(record.get("caller", ""))
@ -480,6 +518,8 @@ class RunJournal(BaseCallbackHandler):
else:
self._lead_agent_tokens += total_tk
self._record_model_usage(record.get("model_name"), input_tk, output_tk, total_tk)
self._schedule_progress_flush()
def set_first_human_message(self, content: str) -> None:
@ -590,6 +630,7 @@ class RunJournal(BaseCallbackHandler):
"lead_agent_tokens": self._lead_agent_tokens,
"subagent_tokens": self._subagent_tokens,
"middleware_tokens": self._middleware_tokens,
"token_usage_by_model": {model: dict(usage) for model, usage in self._tokens_by_model.items()},
"message_count": self._msg_count,
"last_ai_message": self._last_ai_msg,
"first_human_message": self._first_human_msg,

3
backend/packages/harness/deerflow/runtime/runs/manager.py
View File

@ -99,6 +99,8 @@ class RunRecord:
lead_agent_tokens: int = 0
subagent_tokens: int = 0
middleware_tokens: int = 0
# Per-model token breakdown
token_usage_by_model: dict[str, dict[str, int]] = field(default_factory=dict)
message_count: int = 0
last_ai_message: str | None = None
first_human_message: str | None = None
@ -291,6 +293,7 @@ class RunManager:
lead_agent_tokens=row.get("lead_agent_tokens") or 0,
subagent_tokens=row.get("subagent_tokens") or 0,
middleware_tokens=row.get("middleware_tokens") or 0,
token_usage_by_model=row.get("token_usage_by_model") or {},
message_count=row.get("message_count") or 0,
last_ai_message=row.get("last_ai_message"),
first_human_message=row.get("first_human_message"),

2
backend/packages/harness/deerflow/runtime/runs/store/base.py
View File

@ -93,6 +93,7 @@ class RunStore(abc.ABC):
lead_agent_tokens: int = 0,
subagent_tokens: int = 0,
middleware_tokens: int = 0,
token_usage_by_model: dict[str, dict[str, int]] | None = None,
message_count: int = 0,
last_ai_message: str | None = None,
first_human_message: str | None = None,
@ -115,6 +116,7 @@ class RunStore(abc.ABC):
lead_agent_tokens: int | None = None,
subagent_tokens: int | None = None,
middleware_tokens: int | None = None,
token_usage_by_model: dict[str, dict[str, int]] | None = None,
message_count: int | None = None,
last_ai_message: str | None = None,
first_human_message: str | None = None,

19
backend/packages/harness/deerflow/runtime/runs/store/memory.py
View File

@ -110,10 +110,21 @@ class MemoryRunStore(RunStore):
completed = [r for r in self._runs.values() if r["thread_id"] == thread_id and r.get("status") in statuses]
by_model: dict[str, dict] = {}
for r in completed:
model = r.get("model_name") or "unknown"
entry = by_model.setdefault(model, {"tokens": 0, "runs": 0})
entry["tokens"] += r.get("total_tokens", 0)
entry["runs"] += 1
usage_by_model = r.get("token_usage_by_model") or {}
if usage_by_model:
for model, usage in usage_by_model.items():
entry = by_model.setdefault(model, {"tokens": 0, "runs": 0})
entry["tokens"] += usage.get("total_tokens", 0)
entry["runs"] += 1
else:
# Fallback for rows written before per-model accounting landed:
# attribute the whole run to its single ``model_name``. Keeps
# the legacy lead-only behavior for old data instead of
# silently dropping it.
model = r.get("model_name") or "unknown"
entry = by_model.setdefault(model, {"tokens": 0, "runs": 0})
entry["tokens"] += r.get("total_tokens", 0)
entry["runs"] += 1
return {
"total_tokens": sum(r.get("total_tokens", 0) for r in completed),
"total_input_tokens": sum(r.get("total_input_tokens", 0) for r in completed),

4
backend/packages/harness/deerflow/subagents/executor.py
View File

@ -89,7 +89,7 @@ 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)
token_usage_records: list[dict[str, int | str | None]] = field(default_factory=list)
usage_reported: bool = False
cancel_event: threading.Event = field(default_factory=threading.Event, repr=False)
_state_lock: threading.Lock = field(default_factory=threading.Lock, init=False, repr=False)
@ -107,7 +107,7 @@ class SubagentResult:
error: str | None = None,
completed_at: datetime | None = None,
ai_messages: list[dict[str, Any]] | None = None,
token_usage_records: list[dict[str, int | str]] | None = None,
token_usage_records: list[dict[str, int | str | None]] | None = None,
) -> bool:
"""Set a terminal status exactly once.

13
backend/packages/harness/deerflow/subagents/token_collector.py
View File

@ -7,6 +7,7 @@ via :meth:`RunJournal.record_external_llm_usage_records`.
from __future__ import annotations
from collections.abc import Mapping
from typing import Any
from langchain_core.callbacks import BaseCallbackHandler
@ -18,7 +19,7 @@ class SubagentTokenCollector(BaseCallbackHandler):
def __init__(self, caller: str):
super().__init__()
self.caller = caller
self._records: list[dict[str, int | str]] = []
self._records: list[dict[str, int | str | None]] = []
self._counted_run_ids: set[str] = set()
def on_llm_end(
@ -46,11 +47,19 @@ class SubagentTokenCollector(BaseCallbackHandler):
total_tk = input_tk + output_tk
if total_tk <= 0:
continue
# Capture the model that actually produced this response so the
# parent journal can bucket tokens by real model rather than the
# lead agent's resolved model
response_metadata = getattr(gen.message, "response_metadata", None) or {}
model_name: str | None = None
if isinstance(response_metadata, Mapping):
model_name = response_metadata.get("model_name") or response_metadata.get("model")
self._counted_run_ids.add(rid)
self._records.append(
{
"source_run_id": rid,
"caller": self.caller,
"model_name": model_name,
"input_tokens": input_tk,
"output_tokens": output_tk,
"total_tokens": total_tk,
@ -58,6 +67,6 @@ class SubagentTokenCollector(BaseCallbackHandler):
)
return
def snapshot_records(self) -> list[dict[str, int | str]]:
def snapshot_records(self) -> list[dict[str, int | str | None]]:
"""Return a copy of the accumulated usage records."""
return list(self._records)

53
backend/tests/test_run_repository.py
View File

@ -3,8 +3,6 @@
Uses a temp SQLite DB to test ORM-backed CRUD operations.
"""
import re
import pytest
from sqlalchemy.dialects import postgresql
@ -451,7 +449,11 @@ class TestRunRepository:
await _cleanup()
@pytest.mark.anyio
async def test_aggregate_tokens_by_thread_reuses_shared_model_name_expression(self):
async def test_aggregate_tokens_by_thread_returns_zeros_when_no_rows(self):
"""Empty thread aggregates to all-zero totals, no model buckets, and a
single query — replaces the older test that pinned the now-removed
``GROUP BY coalesce(model_name)`` shape (issue #3645 reduces by_model
in Python from each row's per-model JSON column instead)."""
captured = []
class FakeResult:
@ -483,17 +485,44 @@ class TestRunRepository:
}
assert len(captured) == 1
stmt = captured[0]
compiled_sql = str(stmt.compile(dialect=postgresql.dialect()))
select_sql, group_by_sql = compiled_sql.split(" GROUP BY ", maxsplit=1)
model_expr_pattern = r"coalesce\(runs\.model_name, %\(([^)]+)\)s\)"
@pytest.mark.anyio
async def test_aggregate_tokens_by_thread_compiles_on_postgres_dialect(self):
"""Compile-smoke the new SELECT on the postgres dialect.
select_match = re.search(model_expr_pattern + r" AS model", select_sql)
group_by_match = re.fullmatch(model_expr_pattern, group_by_sql.strip())
The project ships both SQLite and Postgres backends. The new aggregation
projects ``RunRow.token_usage_by_model`` (a JSON column) directly into
the row set instead of grouping on a scalar, so the SQL needs to compile
cleanly under PG's JSON/JSONB binding too. Pins:
* the JSON column is selected by name (PG would otherwise need a
``::jsonb`` cast or coalesce around it)
* there is no GROUP BY / aggregate function left (the per-model
reduction now happens in Python — see issue #3645)
"""
assert select_match is not None
assert group_by_match is not None
assert select_match.group(1) == group_by_match.group(1)
captured = []
class FakeResult:
def all(self):
return []
class FakeSession:
async def execute(self, stmt):
captured.append(stmt)
return FakeResult()
class FakeSessionContext:
async def __aenter__(self):
return FakeSession()
async def __aexit__(self, exc_type, exc, tb):
return None
repo = RunRepository(lambda: FakeSessionContext())
await repo.aggregate_tokens_by_thread("t1")
compiled = str(captured[0].compile(dialect=postgresql.dialect()))
assert "token_usage_by_model" in compiled
assert "GROUP BY" not in compiled.upper()
@pytest.mark.anyio
async def test_run_manager_hydrates_store_only_run_from_sql(self, tmp_path):

29
backend/tests/test_subagent_token_collector.py
View File

@ -6,11 +6,12 @@ from uuid import uuid4
from deerflow.subagents.token_collector import SubagentTokenCollector
def _make_llm_response(content="Hello", usage=None):
def _make_llm_response(content="Hello", usage=None, response_metadata=None):
"""Create a mock LLM response with a message."""
msg = MagicMock()
msg.content = content
msg.usage_metadata = usage
msg.response_metadata = response_metadata or {}
gen = MagicMock()
gen.message = msg
@ -50,6 +51,32 @@ class TestSubagentTokenCollector:
assert records[0]["total_tokens"] == 150
assert "source_run_id" in records[0]
def test_collects_model_name_from_response_metadata(self):
collector = SubagentTokenCollector(caller="subagent:test")
usage = {"input_tokens": 100, "output_tokens": 50, "total_tokens": 150}
collector.on_llm_end(
_make_llm_response("Hi", usage=usage, response_metadata={"model_name": "subagent-model"}),
run_id=uuid4(),
)
records = collector.snapshot_records()
assert len(records) == 1
assert records[0]["model_name"] == "subagent-model"
def test_collects_model_name_from_response_metadata_model_fallback(self):
collector = SubagentTokenCollector(caller="subagent:test")
usage = {"input_tokens": 100, "output_tokens": 50, "total_tokens": 150}
collector.on_llm_end(
_make_llm_response("Hi", usage=usage, response_metadata={"model": "provider-model"}),
run_id=uuid4(),
)
records = collector.snapshot_records()
assert len(records) == 1
assert records[0]["model_name"] == "provider-model"
def test_total_tokens_zero_uses_input_plus_output(self):
collector = SubagentTokenCollector(caller="subagent:test")
usage = {"input_tokens": 200, "output_tokens": 100, "total_tokens": 0}

431
backend/tests/test_token_usage_by_model.py Normal file
View File

@ -0,0 +1,431 @@
"""Per-model token usage regression tests (issue #3645).
Covers the full path that powers ``GET /api/threads/{id}/token-usage``'s
``by_model`` field:
* ``RunJournal`` capturing each LLM call's real ``response_metadata.model_name``
for both the lead agent / middleware path (``on_llm_end``) and the subagent
external-records path (``record_external_llm_usage_records``).
* ``RunJournal.get_completion_data`` exposing the per-model breakdown so it can
be threaded into the run store on completion.
* ``MemoryRunStore`` and ``RunRepository`` (SQLAlchemy) returning the same
``by_model`` shape from ``aggregate_tokens_by_thread``, with the invariant
``sum(by_model[*].tokens) == total_tokens``.
* Legacy rows written before this fix (``token_usage_by_model`` empty) falling
back to the old ``model_name + total_tokens`` attribution instead of being
silently dropped.
"""
from __future__ import annotations
from unittest.mock import MagicMock
from uuid import uuid4
import pytest
from deerflow.persistence.run import RunRepository
from deerflow.runtime.events.store.memory import MemoryRunEventStore
from deerflow.runtime.journal import RunJournal
from deerflow.runtime.runs.store.memory import MemoryRunStore
# ---------------------------------------------------------------------------
# Test doubles
# ---------------------------------------------------------------------------
def _make_llm_response(*, usage: dict | None, model_name: str | None = "lead-model"):
"""Build a minimal LLM response carrying the bits journal/collector read."""
msg = MagicMock()
msg.type = "ai"
msg.content = ""
msg.id = f"msg-{id(msg)}"
msg.tool_calls = []
msg.invalid_tool_calls = []
msg.response_metadata = {} if model_name is None else {"model_name": model_name}
msg.usage_metadata = usage
msg.additional_kwargs = {}
msg.name = None
msg.model_dump.return_value = {
"content": "",
"additional_kwargs": {},
"response_metadata": msg.response_metadata,
"type": "ai",
"name": None,
"id": msg.id,
"tool_calls": [],
"invalid_tool_calls": [],
"usage_metadata": usage,
}
gen = MagicMock()
gen.message = msg
response = MagicMock()
response.generations = [[gen]]
return response
def _journal() -> RunJournal:
return RunJournal("r1", "t1", MemoryRunEventStore(), flush_threshold=100)
# ---------------------------------------------------------------------------
# RunJournal: per-call model accounting
# ---------------------------------------------------------------------------
class TestJournalByModel:
def test_lead_agent_call_lands_on_real_model(self) -> None:
j = _journal()
j.on_llm_end(
_make_llm_response(usage={"input_tokens": 10, "output_tokens": 5, "total_tokens": 15}, model_name="lead-model"),
run_id=uuid4(),
parent_run_id=None,
tags=["lead_agent"],
)
data = j.get_completion_data()
assert data["token_usage_by_model"] == {
"lead-model": {"input_tokens": 10, "output_tokens": 5, "total_tokens": 15},
}
assert data["lead_agent_tokens"] == 15
assert data["total_tokens"] == 15
def test_middleware_call_lands_on_its_own_model(self) -> None:
"""A middleware (e.g. title/summarization) on a different model gets its own bucket."""
j = _journal()
j.on_llm_end(
_make_llm_response(usage={"input_tokens": 10, "output_tokens": 5, "total_tokens": 15}, model_name="lead-model"),
run_id=uuid4(),
parent_run_id=None,
tags=["lead_agent"],
)
j.on_llm_end(
_make_llm_response(usage={"input_tokens": 4, "output_tokens": 1, "total_tokens": 5}, model_name="title-model"),
run_id=uuid4(),
parent_run_id=None,
tags=["middleware:title"],
)
data = j.get_completion_data()
assert data["token_usage_by_model"] == {
"lead-model": {"input_tokens": 10, "output_tokens": 5, "total_tokens": 15},
"title-model": {"input_tokens": 4, "output_tokens": 1, "total_tokens": 5},
}
assert data["lead_agent_tokens"] == 15
assert data["middleware_tokens"] == 5
def test_missing_model_name_falls_back_to_unknown(self) -> None:
j = _journal()
j.on_llm_end(
_make_llm_response(usage={"input_tokens": 3, "output_tokens": 2, "total_tokens": 5}, model_name=None),
run_id=uuid4(),
parent_run_id=None,
tags=["lead_agent"],
)
data = j.get_completion_data()
assert data["token_usage_by_model"] == {
"unknown": {"input_tokens": 3, "output_tokens": 2, "total_tokens": 5},
}
def test_same_model_aggregates_across_calls(self) -> None:
j = _journal()
for _ in range(2):
j.on_llm_end(
_make_llm_response(usage={"input_tokens": 7, "output_tokens": 3, "total_tokens": 10}, model_name="lead-model"),
run_id=uuid4(),
parent_run_id=None,
tags=["lead_agent"],
)
data = j.get_completion_data()
assert data["token_usage_by_model"] == {
"lead-model": {"input_tokens": 14, "output_tokens": 6, "total_tokens": 20},
}
def test_subagent_external_records_attribute_to_real_model(self) -> None:
"""The fix's headline behavior: subagent on a different model no longer
steals tokens from the lead model bucket."""
j = _journal()
# Lead emits 10 tokens on lead-model.
j.on_llm_end(
_make_llm_response(usage={"input_tokens": 6, "output_tokens": 4, "total_tokens": 10}, model_name="lead-model"),
run_id=uuid4(),
parent_run_id=None,
tags=["lead_agent"],
)
# Subagent ran on subagent-model and reports 25 tokens via the
# external-records bridge (the path SubagentTokenCollector uses).
j.record_external_llm_usage_records(
[
{
"source_run_id": "sub-1",
"caller": "subagent:general-purpose",
"model_name": "subagent-model",
"input_tokens": 15,
"output_tokens": 10,
"total_tokens": 25,
},
],
)
data = j.get_completion_data()
assert data["token_usage_by_model"] == {
"lead-model": {"input_tokens": 6, "output_tokens": 4, "total_tokens": 10},
"subagent-model": {"input_tokens": 15, "output_tokens": 10, "total_tokens": 25},
}
assert data["total_tokens"] == 35
# by_caller stays accurate too.
assert data["lead_agent_tokens"] == 10
assert data["subagent_tokens"] == 25
# Invariant the issue calls out: by_model sums to total_tokens.
assert sum(b["total_tokens"] for b in data["token_usage_by_model"].values()) == data["total_tokens"]
def test_subagent_record_without_model_falls_back_to_unknown(self) -> None:
j = _journal()
j.record_external_llm_usage_records(
[
{
"source_run_id": "sub-1",
"caller": "subagent:bash",
"input_tokens": 5,
"output_tokens": 2,
"total_tokens": 7,
},
],
)
data = j.get_completion_data()
assert data["token_usage_by_model"] == {
"unknown": {"input_tokens": 5, "output_tokens": 2, "total_tokens": 7},
}
def test_on_llm_end_dedup_does_not_double_count_model(self) -> None:
j = _journal()
rid = uuid4()
usage = {"input_tokens": 10, "output_tokens": 5, "total_tokens": 15}
j.on_llm_end(_make_llm_response(usage=usage, model_name="lead-model"), run_id=rid, parent_run_id=None, tags=["lead_agent"])
# Same langchain run_id firing twice (real callbacks do this) must
# not inflate either total_tokens or the per-model bucket.
j.on_llm_end(_make_llm_response(usage=usage, model_name="lead-model"), run_id=rid, parent_run_id=None, tags=["lead_agent"])
data = j.get_completion_data()
assert data["total_tokens"] == 15
assert data["token_usage_by_model"] == {
"lead-model": {"input_tokens": 10, "output_tokens": 5, "total_tokens": 15},
}
def test_external_records_dedup_does_not_double_count_model(self) -> None:
j = _journal()
record = {
"source_run_id": "sub-1",
"caller": "subagent:general-purpose",
"model_name": "subagent-model",
"input_tokens": 15,
"output_tokens": 10,
"total_tokens": 25,
}
j.record_external_llm_usage_records([record])
j.record_external_llm_usage_records([record])
data = j.get_completion_data()
assert data["subagent_tokens"] == 25
assert data["token_usage_by_model"] == {
"subagent-model": {"input_tokens": 15, "output_tokens": 10, "total_tokens": 25},
}
def test_track_tokens_disabled_keeps_by_model_empty(self) -> None:
store = MemoryRunEventStore()
j = RunJournal("r1", "t1", store, track_token_usage=False, flush_threshold=100)
j.on_llm_end(
_make_llm_response(usage={"input_tokens": 10, "output_tokens": 5, "total_tokens": 15}, model_name="lead-model"),
run_id=uuid4(),
parent_run_id=None,
tags=["lead_agent"],
)
j.record_external_llm_usage_records(
[{"source_run_id": "sub", "caller": "subagent:x", "model_name": "sub-model", "input_tokens": 1, "output_tokens": 1, "total_tokens": 2}],
)
data = j.get_completion_data()
assert data["token_usage_by_model"] == {}
assert data["total_tokens"] == 0
# ---------------------------------------------------------------------------
# Store aggregation: invariants and parity across MemoryRunStore + RunRepository
# ---------------------------------------------------------------------------
_THREAD = "thread-by-model"
def _completed_run(
run_id: str,
*,
model_name: str | None,
total_tokens: int,
lead: int = 0,
sub: int = 0,
mw: int = 0,
by_model: dict | None = None,
) -> dict:
"""Shape that both stores accept for completion writes (kwargs to update_run_completion)."""
return {
"run_id": run_id,
"model_name": model_name,
"completion": {
"status": "success",
"total_input_tokens": 0,
"total_output_tokens": 0,
"total_tokens": total_tokens,
"llm_call_count": 1,
"lead_agent_tokens": lead,
"subagent_tokens": sub,
"middleware_tokens": mw,
"token_usage_by_model": by_model or {},
"message_count": 0,
},
}
async def _seed_run(store, *, run_id: str, model_name: str | None, completion: dict) -> None:
await store.put(run_id, thread_id=_THREAD, status="pending", model_name=model_name)
await store.update_run_completion(run_id, **completion)
_RUN_FIXTURES = [
# 1. Run where subagent and middleware ran on different models than lead.
_completed_run(
"run-1",
model_name="lead-model",
total_tokens=300,
lead=100,
sub=150,
mw=50,
by_model={
"lead-model": {"input_tokens": 60, "output_tokens": 40, "total_tokens": 100},
"subagent-model": {"input_tokens": 90, "output_tokens": 60, "total_tokens": 150},
"middleware-model": {"input_tokens": 30, "output_tokens": 20, "total_tokens": 50},
},
),
# 2. Another run, lead on a *different* lead model — exercises multi-run merge.
_completed_run(
"run-2",
model_name="lead-model-b",
total_tokens=80,
lead=80,
by_model={
"lead-model-b": {"input_tokens": 50, "output_tokens": 30, "total_tokens": 80},
},
),
# 3. Legacy row written before this fix: empty token_usage_by_model. Must
# fall back to (model_name, total_tokens) instead of disappearing from
# by_model entirely.
_completed_run(
"run-3",
model_name="legacy-model",
total_tokens=42,
lead=42,
by_model={},
),
]
async def _seed_all(store) -> None:
for fix in _RUN_FIXTURES:
await _seed_run(store, run_id=fix["run_id"], model_name=fix["model_name"], completion=fix["completion"])
def _assert_aggregate_shape(agg: dict) -> None:
"""Pin the contract that powers /api/threads/{id}/token-usage."""
# The headline totals stay the simple SUMs.
assert agg["total_tokens"] == 300 + 80 + 42
assert agg["total_runs"] == 3
assert agg["by_caller"] == {
"lead_agent": 100 + 80 + 42,
"subagent": 150,
"middleware": 50,
}
# The core fix: subagent / middleware models show up in by_model with their
# real tokens; the lead-model bucket is NOT inflated with subagent tokens.
assert agg["by_model"]["lead-model"] == {"tokens": 100, "runs": 1}
assert agg["by_model"]["subagent-model"] == {"tokens": 150, "runs": 1}
assert agg["by_model"]["middleware-model"] == {"tokens": 50, "runs": 1}
assert agg["by_model"]["lead-model-b"] == {"tokens": 80, "runs": 1}
# Legacy fallback path — empty token_usage_by_model maps to the row's
# ``model_name`` with the full total_tokens.
assert agg["by_model"]["legacy-model"] == {"tokens": 42, "runs": 1}
# Invariant from issue #3645.
assert sum(b["tokens"] for b in agg["by_model"].values()) == agg["total_tokens"]
@pytest.mark.anyio
async def test_memory_store_by_model_invariant_and_fallback():
store = MemoryRunStore()
await _seed_all(store)
agg = await store.aggregate_tokens_by_thread(_THREAD)
_assert_aggregate_shape(agg)
async def _make_sql_repo(tmp_path):
from deerflow.persistence.engine import get_session_factory, init_engine
url = f"sqlite+aiosqlite:///{tmp_path / 'by-model.db'}"
await init_engine("sqlite", url=url, sqlite_dir=str(tmp_path))
return RunRepository(get_session_factory())
async def _close_sql_engine() -> None:
from deerflow.persistence.engine import close_engine
await close_engine()
@pytest.mark.anyio
async def test_sql_store_by_model_invariant_and_fallback(tmp_path):
repo = await _make_sql_repo(tmp_path)
try:
await _seed_all(repo)
agg = await repo.aggregate_tokens_by_thread(_THREAD)
_assert_aggregate_shape(agg)
finally:
await _close_sql_engine()
@pytest.mark.anyio
async def test_memory_and_sql_stores_agree(tmp_path):
"""Memory and SQL stores must return byte-identical aggregations so
behavior does not silently diverge based on database.backend choice."""
mem = MemoryRunStore()
sql = await _make_sql_repo(tmp_path)
try:
await _seed_all(mem)
await _seed_all(sql)
mem_agg = await mem.aggregate_tokens_by_thread(_THREAD)
sql_agg = await sql.aggregate_tokens_by_thread(_THREAD)
assert mem_agg == sql_agg
finally:
await _close_sql_engine()
@pytest.mark.anyio
async def test_include_active_picks_up_running_progress_snapshot(tmp_path):
"""``update_run_progress`` must persist ``token_usage_by_model`` so the
``include_active=true`` view of /token-usage reflects in-flight tokens."""
repo = await _make_sql_repo(tmp_path)
try:
await repo.put("run-active", thread_id=_THREAD, status="pending")
# Transition to running so update_run_progress' status guard fires.
await repo.update_status("run-active", "running")
await repo.update_run_progress(
"run-active",
total_tokens=70,
total_input_tokens=40,
total_output_tokens=30,
lead_agent_tokens=70,
token_usage_by_model={
"lead-model": {"input_tokens": 40, "output_tokens": 30, "total_tokens": 70},
},
)
# Default (completed-only) excludes running runs.
completed_only = await repo.aggregate_tokens_by_thread(_THREAD)
assert completed_only["total_runs"] == 0
assert completed_only["by_model"] == {}
active = await repo.aggregate_tokens_by_thread(_THREAD, include_active=True)
assert active["total_runs"] == 1
assert active["by_model"] == {"lead-model": {"tokens": 70, "runs": 1}}
assert active["total_tokens"] == 70
finally:
await _close_sql_engine()