/backend/tests/test_token_usage_by_model.py

"""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()


# ---------------------------------------------------------------------------
# Prompt-cache-hit accounting (powers cache-aware cost estimation in
# /api/console): cache_read_tokens is a *sparse* bucket key — present only
# when a provider reported cache hits — so pre-existing bucket shapes and
# exact-equality assertions above stay valid.
# ---------------------------------------------------------------------------


class TestJournalCacheRead:
    def test_cache_read_accumulates_as_sparse_key(self) -> None:
        j = _journal()
        j.on_llm_end(
            _make_llm_response(
                usage={"input_tokens": 100, "output_tokens": 10, "total_tokens": 110, "input_token_details": {"cache_read": 80}},
                model_name="m",
            ),
            run_id=uuid4(),
            parent_run_id=None,
            tags=["lead_agent"],
        )
        # Second call without cache hits still accumulates into the same bucket.
        j.on_llm_end(
            _make_llm_response(usage={"input_tokens": 50, "output_tokens": 5, "total_tokens": 55}, model_name="m"),
            run_id=uuid4(),
            parent_run_id=None,
            tags=["lead_agent"],
        )
        data = j.get_completion_data()
        assert data["token_usage_by_model"]["m"] == {
            "input_tokens": 150,
            "output_tokens": 15,
            "total_tokens": 165,
            "cache_read_tokens": 80,
        }

    def test_bucket_without_cache_hits_keeps_legacy_shape(self) -> None:
        j = _journal()
        j.on_llm_end(
            _make_llm_response(usage={"input_tokens": 10, "output_tokens": 5, "total_tokens": 15}, model_name="m"),
            run_id=uuid4(),
            parent_run_id=None,
            tags=["lead_agent"],
        )
        assert j.get_completion_data()["token_usage_by_model"]["m"] == {
            "input_tokens": 10,
            "output_tokens": 5,
            "total_tokens": 15,
        }

    def test_external_records_carry_cache_read(self) -> None:
        j = _journal()
        j.record_external_llm_usage_records(
            [
                {
                    "source_run_id": "sub-1",
                    "caller": "subagent:general-purpose",
                    "model_name": "sub-m",
                    "input_tokens": 40,
                    "output_tokens": 10,
                    "total_tokens": 50,
                    "cache_read_tokens": 25,
                },
            ],
        )
        assert j.get_completion_data()["token_usage_by_model"]["sub-m"]["cache_read_tokens"] == 25

    def test_deepseek_raw_usage_normalizes_to_cache_read(self) -> None:
        """Pin the DeepSeek chat-completions shape end-to-end: the raw
        ``prompt_tokens_details.cached_tokens`` field is what langchain-openai's
        ``_create_usage_metadata`` normalizes into
        ``input_token_details.cache_read`` (DeepSeek's top-level
        ``prompt_cache_hit/miss_tokens`` are redundant aliases LangChain does
        not read), and the journal captures it. The derived cache-miss count
        (input − cache_read) must equal DeepSeek's own
        ``prompt_cache_miss_tokens``, which is what cache-aware pricing bills
        at the full input price."""
        from langchain_openai.chat_models.base import _create_usage_metadata

        raw = {
            "prompt_tokens": 106,
            "completion_tokens": 112,
            "total_tokens": 218,
            "prompt_tokens_details": {"cached_tokens": 64},
            "prompt_cache_hit_tokens": 64,
            "prompt_cache_miss_tokens": 42,
        }
        usage = _create_usage_metadata(raw)
        j = _journal()
        j.on_llm_end(
            _make_llm_response(usage=dict(usage), model_name="deepseek-chat"),
            run_id=uuid4(),
            parent_run_id=None,
            tags=["lead_agent"],
        )
        bucket = j.get_completion_data()["token_usage_by_model"]["deepseek-chat"]
        assert bucket == {
            "input_tokens": 106,
            "output_tokens": 112,
            "total_tokens": 218,
            "cache_read_tokens": 64,
        }
        assert bucket["input_tokens"] - bucket["cache_read_tokens"] == raw["prompt_cache_miss_tokens"]

    def test_collector_extracts_cache_read_from_usage_metadata(self) -> None:
        from deerflow.subagents.token_collector import SubagentTokenCollector

        collector = SubagentTokenCollector("subagent:general-purpose")
        collector.on_llm_end(
            _make_llm_response(
                usage={"input_tokens": 30, "output_tokens": 6, "total_tokens": 36, "input_token_details": {"cache_read": 20}},
                model_name="sub-m",
            ),
            run_id=uuid4(),
        )
        records = collector.snapshot_records()
        assert len(records) == 1
        assert records[0]["cache_read_tokens"] == 20

    def test_collector_omits_cache_read_key_when_no_cache_hits(self) -> None:
        from deerflow.subagents.token_collector import SubagentTokenCollector

        collector = SubagentTokenCollector("subagent:general-purpose")
        collector.on_llm_end(
            _make_llm_response(
                usage={"input_tokens": 30, "output_tokens": 6, "total_tokens": 36},
                model_name="sub-m",
            ),
            run_id=uuid4(),
        )
        records = collector.snapshot_records()
        assert len(records) == 1
        # Sparse record shape: no explicit 0 when the provider reported no
        # cache hits (record_external_llm_usage_records treats absent as 0).
        assert "cache_read_tokens" not in records[0]