portfolio-optimiser/tests/test_budget.py

"""Step 4 tests — token meter + budget middleware off REAL UsageDetails (no LLM).

Usage is hand-built (synthetic ``UsageDetails``), so these run in CI without an endpoint.
The strict-usage guard hard-fails on a missing usage; the meta-guard proves the meter is
fed from ``UsageDetails`` and never a word-count proxy. Pattern: tests/spikes/test_harness.py.
"""

import re
from pathlib import Path

import pytest
from agent_framework import ChatResponse, Message, UsageDetails

from portfolio_optimiser.budget import (
    Budget,
    BudgetExceeded,
    BudgetMiddleware,
    PortfolioBudget,
    PortfolioMeter,
    TokenMeter,
    UsageUnavailable,
)


class _Ctx:
    """Minimal ChatContext stand-in carrying the post-call ``result``."""

    def __init__(self, result: object) -> None:
        self.result = result


async def _noop() -> None:
    return None


def _resp(total: int | None) -> ChatResponse:
    usage = UsageDetails(total_token_count=total) if total is not None else None
    return ChatResponse(messages=[Message(role="assistant", contents=["x"])], usage_details=usage)


async def test_meter_reads_total_token_count_and_accumulates() -> None:
    meter = TokenMeter(Budget(max_tokens=1000, max_rounds=10))
    mw = BudgetMiddleware(meter)
    await mw.process(_Ctx(_resp(40)), _noop)  # type: ignore[arg-type]
    await mw.process(_Ctx(_resp(50)), _noop)  # type: ignore[arg-type]
    assert meter.tokens == 90  # read from UsageDetails, accumulated across calls


async def test_cap_crossed_raises_budget_exceeded() -> None:
    meter = TokenMeter(Budget(max_tokens=50, max_rounds=10))
    mw = BudgetMiddleware(meter)
    with pytest.raises(BudgetExceeded) as exc:
        await mw.process(_Ctx(_resp(60)), _noop)  # type: ignore[arg-type]
    assert exc.value.kind == "tokens"
    assert exc.value.limit == 50
    assert exc.value.observed == 60  # the spend that crossed it, not merely that something did


def test_round_cap_crossed_carries_the_round_ledger() -> None:
    """The rounds arm of the same stop event (kø-(y)).

    The shipped ``tick_round`` had NO direct test at all — only ``spikes/_harness.py``'s separate
    copy of ``TokenMeter`` did, and that file is not the one that runs. Its entire error contract
    was riding on a module the framework never imports."""
    meter = TokenMeter(Budget(max_tokens=1000, max_rounds=2))
    assert meter.tick_round() == 1
    assert meter.tick_round() == 2  # exactly at the cap is still within it
    with pytest.raises(BudgetExceeded) as exc:
        meter.tick_round()
    assert (exc.value.kind, exc.value.limit, exc.value.observed) == ("rounds", 2, 3)


def test_non_positive_cap_rejected() -> None:
    with pytest.raises(ValueError):
        Budget(max_tokens=0, max_rounds=5)
    with pytest.raises(ValueError):
        Budget(max_tokens=5, max_rounds=0)


async def test_strict_usage_none_hard_fails() -> None:
    meter = TokenMeter(Budget(max_tokens=100, max_rounds=10))
    mw = BudgetMiddleware(meter, strict_usage=True)
    with pytest.raises(UsageUnavailable):
        await mw.process(_Ctx(_resp(None)), _noop)  # type: ignore[arg-type]


async def test_budget_middleware_fires_on_real_agent_chat(make_client_factory) -> None:
    """F8 (real-client half): registering BudgetMiddleware on a real Agent (layered chat client
    that carries ChatMiddlewareLayer) and running an actual chat call short-circuits with
    BudgetExceeded — the middleware<->client integration the prior suite never exercised (the
    minimal-base stand-in silently no-ops the middleware)."""
    from agent_framework import Agent

    client = make_client_factory("ok", tokens=8)("proposer")  # 8 tokens/reply > cap 5
    meter = TokenMeter(Budget(max_tokens=5, max_rounds=10))
    agent = Agent(
        client, "propose a measure", name="proposer", middleware=[BudgetMiddleware(meter)]
    )
    with pytest.raises(BudgetExceeded) as exc:
        await agent.run("hi")
    assert exc.value.kind == "tokens"
    assert meter.tokens == 8  # charged from the synthetic UsageDetails via the middleware


async def test_pre_call_guard_does_not_await_call_next_when_exhausted() -> None:
    """S3.4 unit arm of the pre-call guard: the middleware refuses BEFORE ``call_next``. Counting
    the awaits is the whole measurement — the post-charge arm raises either way."""
    calls = {"n": 0}

    async def _counted() -> None:
        calls["n"] += 1

    meter = TokenMeter(Budget(max_tokens=10, max_rounds=10))
    meter.charge(10)  # exactly at the cap -> nothing left to fund a call with
    mw = BudgetMiddleware(meter)
    with pytest.raises(BudgetExceeded):
        await mw.process(_Ctx(_resp(5)), _counted)  # type: ignore[arg-type]
    assert calls["n"] == 0


async def test_below_the_cap_still_calls_through() -> None:
    """Control for the guard: with budget left, the call goes through exactly as before — the
    guard must gate exhaustion, not traffic."""
    calls = {"n": 0}

    async def _counted() -> None:
        calls["n"] += 1

    meter = TokenMeter(Budget(max_tokens=10, max_rounds=10))
    mw = BudgetMiddleware(meter)
    await mw.process(_Ctx(_resp(4)), _counted)  # type: ignore[arg-type]
    assert calls["n"] == 1 and meter.tokens == 4 and meter.remaining() == 6


def test_portfolio_budget_rejects_unusable_configs() -> None:
    """Fail-fast at construction (mirroring ``Budget``): non-positive caps, a per-run cap larger
    than the global one (one run could then cross the pass's own ceiling), and a reserve larger
    than a run can ever spend (which would refuse every pass forever)."""
    with pytest.raises(ValueError):
        PortfolioBudget(max_total_tokens=0, max_tokens_per_run=10)
    with pytest.raises(ValueError):
        PortfolioBudget(max_total_tokens=100, max_tokens_per_run=0)
    with pytest.raises(ValueError):
        PortfolioBudget(max_total_tokens=100, max_tokens_per_run=200)
    with pytest.raises(ValueError):
        PortfolioBudget(max_total_tokens=100, max_tokens_per_run=50, min_run_reserve=60)
    with pytest.raises(ValueError):
        PortfolioMeter(PortfolioBudget(max_total_tokens=100, max_tokens_per_run=50), spent=-1)


def test_portfolio_meter_accumulates_and_crosses() -> None:
    budget = PortfolioBudget(max_total_tokens=100, max_tokens_per_run=50)
    meter = PortfolioMeter(budget, spent=40)
    assert meter.remaining() == 60 and meter.required_per_run == 50
    meter.record(60)
    meter.check()  # exactly at the cap is still within it (mirrors TokenMeter's `>` boundary)
    meter.record(1)
    with pytest.raises(BudgetExceeded) as exc:
        meter.check()
    assert exc.value.kind == "portfolio_tokens"
    assert exc.value.limit == 100 and exc.value.observed == 101


async def test_run_meter_bound_to_portfolio_charges_both_ledgers() -> None:
    """A bound run meter charges the global ledger on every call, and ``remaining`` reads whichever
    cap binds — that is what lets one project's spend refuse another's next call."""
    portfolio = PortfolioMeter(PortfolioBudget(max_total_tokens=100, max_tokens_per_run=80))
    meter = TokenMeter(Budget(max_tokens=80, max_rounds=10), portfolio=portfolio)
    mw = BudgetMiddleware(meter)
    await mw.process(_Ctx(_resp(30)), _noop)  # type: ignore[arg-type]
    assert meter.tokens == 30 and portfolio.spent == 30
    assert meter.remaining() == 50  # run has 50 left, portfolio 70 -> the RUN binds
    await mw.process(_Ctx(_resp(45)), _noop)  # type: ignore[arg-type]
    assert meter.remaining() == 5  # run 5, portfolio 25 -> still the run


async def test_spend_is_ledgered_even_when_the_per_run_cap_raises() -> None:
    """Order matters: tokens the provider already billed must reach the global ledger even though
    the RUN's own cap raises on the same charge. Checking the run cap first and returning early
    would lose that spend and hand the next project a budget that was never really there."""
    portfolio = PortfolioMeter(PortfolioBudget(max_total_tokens=1000, max_tokens_per_run=50))
    meter = TokenMeter(Budget(max_tokens=50, max_rounds=10), portfolio=portfolio)
    with pytest.raises(BudgetExceeded) as exc:
        meter.charge(60)
    assert exc.value.kind == "tokens"  # the RUN's cap is the one that broke
    assert portfolio.spent == 60  # ...and the global ledger still saw the spend


def test_exhausted_names_one_ledger_in_all_three_fields() -> None:
    """kø-(y): the pre-call guard's refusal must describe ONE ledger consistently.

    ``exhausted`` is where the S3.4 guard decides what to refuse a call with, and it is the only
    raise site that CHOOSES between two ledgers. Naming the binding cap in ``kind`` is already
    tested; this pins the other two fields to that same choice — a refusal that says
    ``portfolio_tokens`` while reporting the run's own spend would misdirect every reader of it.

    Both arms are built with ``observed != limit`` on purpose. At exactly-exhausted the two
    coincide, so a test written at that point cannot tell them apart and would pass on an
    implementation that echoed the cap back as the spend. Overrunning first is also the honest
    case: a run that crossed its cap and had the error caught is precisely when the next call
    must be refused."""
    meter = TokenMeter(Budget(max_tokens=10, max_rounds=10))
    with pytest.raises(BudgetExceeded):
        meter.charge(15)  # the overrun is real: the meter now stands at 15 against a cap of 10
    own = meter.exhausted()
    assert own is not None
    assert (own.kind, own.limit, own.observed) == ("tokens", 10, 15)

    # The global arm: this run has spent NOTHING — a sibling drained the pass.
    portfolio = PortfolioMeter(
        PortfolioBudget(max_total_tokens=100, max_tokens_per_run=100), spent=150
    )
    bound = TokenMeter(Budget(max_tokens=100, max_rounds=10), portfolio=portfolio)
    assert bound.tokens == 0
    shared = bound.exhausted()
    assert shared is not None
    assert (shared.kind, shared.limit, shared.observed) == ("portfolio_tokens", 100, 150)


def test_no_word_count_token_proxy_in_src() -> None:
    # The meter is fed from real UsageDetails, never a len(text.split()) word-count proxy
    # (research 03 Rec 3 — the Fase 1 _word_tokens anti-pattern is retired). NOTE: a bare
    # `.split()` grep is wrong — retrieval.py legitimately uses .split() for KEYWORD scoring
    # (a [0,1] overlap ratio, not a token count). So guard the specific proxy signature.
    proxy = re.compile(r"len\(\s*[^)]*\.split\(\)\s*\)")
    offenders = [
        py.name
        for py in Path("src/portfolio_optimiser").rglob("*.py")
        if proxy.search(py.read_text(encoding="utf-8"))
    ]
    assert offenders == [], f"word-count token proxy found in: {offenders}"