/tests/fleet_monitoring/test_token_rate.py

"""Tests for the per-PID sliding-window token rate tracker (#2023)."""

from __future__ import annotations

import threading
import time

import pytest

from tools.system.fleet_monitoring.meters import TokenUsage
from tools.system.fleet_monitoring.token_rate import TokenRateTracker


@pytest.fixture
def tracker() -> TokenRateTracker:
    """Fresh tracker per test so module-level singleton state never leaks."""
    return TokenRateTracker()


class TestTokensPerMin:
    def test_unknown_pid_returns_none(self, tracker: TokenRateTracker) -> None:
        # Distinct from "known but idle". ``None`` lets the dashboard
        # render ``-`` instead of ``0`` for PIDs that were never
        # observed (e.g. claude-code agent on a TTY whose JSONL was
        # never resolved).
        assert tracker.tokens_per_min(9999) is None

    def test_known_idle_pid_returns_zero(self, tracker: TokenRateTracker) -> None:
        tracker.record(1234, tokens=0)
        # PID has been recorded — even with zero — so the dashboard
        # should distinguish it from never-seen.
        assert tracker.tokens_per_min(1234) == pytest.approx(0.0)

    def test_single_recent_record(self, tracker: TokenRateTracker) -> None:
        tracker.record(1234, tokens=150)
        assert tracker.tokens_per_min(1234) == pytest.approx(150.0)

    def test_multiple_records_in_window_sum(self, tracker: TokenRateTracker) -> None:
        tracker.record(1234, tokens=100)
        tracker.record(1234, tokens=200)
        tracker.record(1234, tokens=50)
        assert tracker.tokens_per_min(1234) == pytest.approx(350.0)

    def test_negative_tokens_clamped_to_zero(self, tracker: TokenRateTracker) -> None:
        # Defensive: a misbehaving meter should not bring the cell
        # negative. Clamp at insertion so the window math stays sane.
        tracker.record(1234, tokens=-10)
        tracker.record(1234, tokens=100)
        assert tracker.tokens_per_min(1234) == pytest.approx(100.0)

    def test_structured_usage_visible_total(self, tracker: TokenRateTracker) -> None:
        tracker.record(
            1234,
            usage=TokenUsage(
                input_tokens=100,
                output_tokens=50,
                cached_input_tokens=90,
                cache_read_input_tokens=200,
                cache_creation_input_tokens=25,
            ),
        )
        assert tracker.tokens_per_min(1234) == pytest.approx(375.0)

    def test_usage_per_min_returns_bucket_totals(self, tracker: TokenRateTracker) -> None:
        tracker.record(1234, usage=TokenUsage(input_tokens=100, cached_input_tokens=25))
        tracker.record(
            1234,
            usage=TokenUsage(
                output_tokens=50,
                cache_read_input_tokens=200,
                cache_creation_input_tokens=10,
            ),
        )
        usage = tracker.usage_per_min(1234)
        assert usage is not None
        assert usage.input_tokens == pytest.approx(100.0)
        assert usage.cached_input_tokens == pytest.approx(25.0)
        assert usage.output_tokens == pytest.approx(50.0)
        assert usage.cache_read_input_tokens == pytest.approx(200.0)
        assert usage.cache_creation_input_tokens == pytest.approx(10.0)


class TestWindowEviction:
    def test_old_entries_are_evicted(
        self, tracker: TokenRateTracker, monkeypatch: pytest.MonkeyPatch
    ) -> None:
        # Shrink the window to 1s so we can prove eviction without
        # racing wall-clock time. Class-level constant + class attr
        # patching keeps the lock behavior intact.
        monkeypatch.setattr(TokenRateTracker, "WINDOW_SECONDS", 1.0)

        tracker.record(1234, tokens=100)
        time.sleep(1.2)
        tracker.record(1234, tokens=50)

        # 100-token entry aged out, 50-token entry remains. Result is
        # scaled by the 60/WINDOW factor, so 50 × 60 = 3000.
        assert tracker.tokens_per_min(1234) == pytest.approx(3000.0)

    def test_all_entries_aged_out_returns_zero_not_none(
        self, tracker: TokenRateTracker, monkeypatch: pytest.MonkeyPatch
    ) -> None:
        # Once a PID has been recorded, it stays known forever (until
        # ``forget`` is called) — even when every entry has aged out.
        monkeypatch.setattr(TokenRateTracker, "WINDOW_SECONDS", 0.5)
        tracker.record(1234, tokens=100)
        time.sleep(0.7)
        assert tracker.tokens_per_min(1234) == pytest.approx(0.0)


class TestLatestModel:
    def test_unknown_pid_returns_none(self, tracker: TokenRateTracker) -> None:
        assert tracker.latest_model(9999) is None

    def test_returns_most_recent_non_none_model(self, tracker: TokenRateTracker) -> None:
        tracker.record(1234, tokens=10, model="claude-sonnet-4-5")
        tracker.record(1234, tokens=20, model="claude-opus-4-1")
        # Pricing should follow the active model — the most recent
        # one observed. A mid-session ``/model`` swap should be
        # reflected immediately.
        assert tracker.latest_model(1234) == "claude-opus-4-1"

    def test_skips_none_model_entries(self, tracker: TokenRateTracker) -> None:
        tracker.record(1234, tokens=10, model="claude-sonnet-4-5")
        # Subsequent record without a model hint must not erase the
        # previously known one — chunks that contain no assistant
        # event (system init, plain prose, etc.) still need to apply
        # a known price.
        tracker.record(1234, tokens=5, model=None)
        assert tracker.latest_model(1234) == "claude-sonnet-4-5"

    def test_returns_none_when_no_entry_carries_model(self, tracker: TokenRateTracker) -> None:
        tracker.record(1234, tokens=10, model=None)
        assert tracker.latest_model(1234) is None

    def test_survives_window_eviction(
        self, tracker: TokenRateTracker, monkeypatch: pytest.MonkeyPatch
    ) -> None:
        # A session that emitted a model once but then went idle for
        # > 60 s loses every token entry to the rolling window. The
        # model identity must persist so pricing keeps working.
        monkeypatch.setattr(TokenRateTracker, "WINDOW_SECONDS", 0.3)
        tracker.record(1234, tokens=10, model="claude-sonnet-4-5")
        time.sleep(0.5)
        assert tracker.tokens_per_min(1234) == pytest.approx(0.0)
        assert tracker.latest_model(1234) == "claude-sonnet-4-5"

    def test_codex_model_hint_survives_later_usage_after_window_eviction(
        self, tracker: TokenRateTracker, monkeypatch: pytest.MonkeyPatch
    ) -> None:
        # Codex emits the model in ``turn_context`` at turn start, but
        # the later ``token_count`` event has usage without a model.
        monkeypatch.setattr(TokenRateTracker, "WINDOW_SECONDS", 0.3)
        tracker.record(1234, tokens=0, model="gpt-5-codex")
        time.sleep(0.5)
        tracker.record(1234, tokens=225, model=None)

        assert tracker.tokens_per_min(1234) == pytest.approx(45000.0)
        assert tracker.latest_model(1234) == "gpt-5-codex"


class TestForgetAndKnownPids:
    def test_forget_removes_pid(self, tracker: TokenRateTracker) -> None:
        tracker.record(1234, tokens=100)
        tracker.forget(1234)
        # After forget, the PID is fully unknown — distinct from idle.
        assert tracker.tokens_per_min(1234) is None
        assert tracker.latest_model(1234) is None

    def test_forget_unknown_pid_is_silent(self, tracker: TokenRateTracker) -> None:
        # Sampler GC calls ``forget`` defensively on every disappeared
        # PID; it must not raise for PIDs the tracker never saw.
        tracker.forget(9999)

    def test_known_pids_lists_recorded(self, tracker: TokenRateTracker) -> None:
        tracker.record(1, tokens=10)
        tracker.record(2, tokens=20)
        # Sampler iterates the snapshot to compute the GC set; order
        # does not matter, presence does.
        assert set(tracker.known_pids()) == {1, 2}

    def test_known_pids_excludes_forgotten(self, tracker: TokenRateTracker) -> None:
        tracker.record(1, tokens=10)
        tracker.record(2, tokens=20)
        tracker.forget(1)
        assert set(tracker.known_pids()) == {2}


class TestThreadSafety:
    def test_concurrent_record_and_read_does_not_corrupt(self, tracker: TokenRateTracker) -> None:
        # The sampler writes from an asyncio executor while the view
        # reads from a worker thread. The lock must guarantee no
        # ``RuntimeError: deque mutated during iteration`` (which is
        # exactly the failure mode an unlocked deque + sum() race
        # would produce). 100 iterations is enough to make an
        # unlocked version fail reliably on macOS/Linux.
        stop = threading.Event()
        errors: list[Exception] = []

        def writer() -> None:
            for i in range(1000):
                if stop.is_set():
                    return
                try:
                    tracker.record(42, tokens=i % 10, model="claude-sonnet-4-5")
                except Exception as exc:  # noqa: BLE001
                    errors.append(exc)
                    return

        def reader() -> None:
            for _ in range(1000):
                if stop.is_set():
                    return
                try:
                    tracker.tokens_per_min(42)
                    tracker.latest_model(42)
                except Exception as exc:  # noqa: BLE001
                    errors.append(exc)
                    return

        threads = [
            threading.Thread(target=writer),
            threading.Thread(target=reader),
            threading.Thread(target=reader),
        ]
        for t in threads:
            t.start()
        for t in threads:
            t.join(timeout=10.0)
        stop.set()
        for t in threads:
            assert not t.is_alive(), f"thread {t.name} hung"
        assert errors == []