feat(fase2): real-UsageDetails token meter + budget middleware

src/portfolio_optimiser/budget.py

@@ -0,0 +1,102 @@
+"""Token meter + budget ChatMiddleware fed from REAL UsageDetails (B4 / D6).
+
+The hard token cap is enforced off the provider's own ``UsageDetails`` — **never** a
+word-count proxy (research 03 Rec 3; the Fase 1 ``_word_tokens`` theatre is retired). A
+``BudgetMiddleware`` reads ``response.usage_details["total_token_count"]`` after each chat
+call and short-circuits with ``BudgetExceeded`` the moment the cap is crossed.
+
+``strict_usage`` (default ``True``) makes a missing usage a HARD FAIL (``UsageUnavailable``):
+a usage regression must never silently disable the cap. Test doubles that legitimately supply
+a synthetic ``UsageDetails`` do not trip it.
+"""
+
+from __future__ import annotations
+
+from collections.abc import Awaitable, Callable
+from dataclasses import dataclass
+
+from agent_framework import ChatContext, ChatMiddleware
+
+
+class BudgetExceeded(RuntimeError):
+    """Raised the moment a runtime cap (tokens or rounds) is crossed (B4).
+
+    Carries the breached ``kind`` ("tokens" | "rounds"), the ``limit`` set, and the
+    ``observed`` value that crossed it — a structured stop event, never a silent hang.
+    """
+
+    def __init__(self, kind: str, limit: int, observed: int) -> None:
+        self.kind = kind
+        self.limit = limit
+        self.observed = observed
+        super().__init__(f"budget exceeded: {kind} limit={limit} observed={observed}")
+
+
+class UsageUnavailable(RuntimeError):
+    """Raised when ``strict_usage`` is on and a response carries no usable usage — the cap
+    must fail closed rather than silently stop counting (research 03 Rec 3)."""
+
+
+@dataclass(frozen=True)
+class Budget:
+    """Hard token + round/iteration caps, required at startup (A4 / D6).
+
+    Refuses to construct without positive caps — fail-fast, never an unbounded loop.
+    """
+
+    max_tokens: int
+    max_rounds: int
+
+    def __post_init__(self) -> None:
+        if self.max_tokens <= 0:
+            raise ValueError(f"max_tokens must be positive, got {self.max_tokens}")
+        if self.max_rounds <= 0:
+            raise ValueError(f"max_rounds must be positive, got {self.max_rounds}")
+
+
+class TokenMeter:
+    """Accumulates token and round usage against a ``Budget``; raises the moment a cap is
+    crossed."""
+
+    def __init__(self, budget: Budget) -> None:
+        self.budget = budget
+        self.tokens = 0
+        self.rounds = 0
+
+    def charge(self, tokens: int) -> int:
+        """Add ``tokens`` to the running total; raise ``BudgetExceeded`` if over cap."""
+        self.tokens += tokens
+        if self.tokens > self.budget.max_tokens:
+            raise BudgetExceeded("tokens", self.budget.max_tokens, self.tokens)
+        return self.tokens
+
+    def tick_round(self) -> int:
+        """Increment the round counter; raise ``BudgetExceeded`` if over cap."""
+        self.rounds += 1
+        if self.rounds > self.budget.max_rounds:
+            raise BudgetExceeded("rounds", self.budget.max_rounds, self.rounds)
+        return self.rounds
+
+
+class BudgetMiddleware(ChatMiddleware):
+    """Chat middleware that charges a ``TokenMeter`` from each response's real
+    ``UsageDetails`` and short-circuits when the cap is crossed."""
+
+    def __init__(self, meter: TokenMeter, *, strict_usage: bool = True) -> None:
+        self._meter = meter
+        self._strict = strict_usage
+
+    async def process(
+        self, context: ChatContext, call_next: Callable[[], Awaitable[None]]
+    ) -> None:
+        await call_next()
+        usage = getattr(context.result, "usage_details", None)
+        total = usage.get("total_token_count") if usage is not None else None
+        if total is None:
+            if self._strict:
+                raise UsageUnavailable(
+                    "usage_details missing total_token_count; the cap must fail closed "
+                    "(set strict_usage=False only for non-counting paths)"
+                )
+            return
+        self._meter.charge(int(total))  # raises BudgetExceeded if over cap

tests/test_budget.py

@@ -0,0 +1,83 @@
+"""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,
+    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
+
+
+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]
+
+
+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}"