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