portfolio-optimiser/src/portfolio_optimiser/budget.py

"""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.

S3.4 (F10) adds the cross-project half. ``PortfolioBudget`` + ``PortfolioMeter`` carry ONE token
ledger over a whole portfolio pass — and, seeded from ``read_spend``, across passes — while the
per-run ``Budget``/``TokenMeter`` pair stays exactly what it was. Two teeth follow from that:
the middleware now refuses a call BEFORE making it once the binding cap is exhausted (money not
spent, not money spent and regretted), and ``run_portfolio`` refuses to START a project it cannot
fund. Both are enforcement, never repair — nothing is trimmed, retried, or scaled down.
"""

from __future__ import annotations

import json
from collections.abc import Awaitable, Callable
from dataclasses import dataclass
from pathlib import Path

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


class BudgetRefused(RuntimeError):
    """Raised at STARTUP when the portfolio's remaining tokens cannot fund one run (S3.4).

    Distinct from ``BudgetExceeded`` on purpose: nothing was crossed and nothing was spent — the
    pass is refused before it begins, because a pass that can afford zero projects is a caller
    mistake, not a resource event. Named for the repo's fail-closed gate family
    (``PromotionRefused``, ``IngestStampError``): refusal, never repair.
    """

    def __init__(self, remaining: int, required: int) -> None:
        self.remaining = remaining
        self.required = required
        super().__init__(
            f"portfolio budget refused: remaining={remaining} cannot fund one run "
            f"(requires {required}); raise the global cap or lower min_run_reserve"
        )


@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}")


@dataclass(frozen=True)
class PortfolioBudget:
    """The GLOBAL token cap over a whole portfolio pass, plus the per-run cap it funds (S3.4/F10).

    ``max_total_tokens`` bounds everything a pass may spend — and, when a ``PortfolioMeter`` is
    seeded from persisted spend, everything a SERIES of passes may spend. ``max_tokens_per_run``
    is the cap each individual run is given. ``min_run_reserve`` (defaulting to the per-run cap)
    is the remainder a run must be able to claim before it is allowed to START: without it a pass
    would happily launch a project it can only half pay for, spending real tokens on a run that
    cannot finish.

    Fail-fast on configurations that cannot mean what they say: a per-run cap above the global one
    would let a single run cross the pass's own ceiling, and a reserve above the per-run cap would
    demand more than any run can ever spend, refusing every pass forever.
    """

    max_total_tokens: int
    max_tokens_per_run: int
    min_run_reserve: int | None = None

    def __post_init__(self) -> None:
        if self.max_total_tokens <= 0:
            raise ValueError(f"max_total_tokens must be positive, got {self.max_total_tokens}")
        if self.max_tokens_per_run <= 0:
            raise ValueError(f"max_tokens_per_run must be positive, got {self.max_tokens_per_run}")
        if self.max_tokens_per_run > self.max_total_tokens:
            raise ValueError(
                f"max_tokens_per_run ({self.max_tokens_per_run}) exceeds max_total_tokens "
                f"({self.max_total_tokens}): one run could cross the portfolio cap on its own"
            )
        if self.min_run_reserve is not None:
            if self.min_run_reserve <= 0:
                raise ValueError(f"min_run_reserve must be positive, got {self.min_run_reserve}")
            if self.min_run_reserve > self.max_tokens_per_run:
                raise ValueError(
                    f"min_run_reserve ({self.min_run_reserve}) exceeds max_tokens_per_run "
                    f"({self.max_tokens_per_run}): no run could ever meet it"
                )

    @property
    def required_per_run(self) -> int:
        """The remainder one run must be able to claim to be allowed to start."""
        return self.min_run_reserve if self.min_run_reserve is not None else self.max_tokens_per_run


class PortfolioMeter:
    """The ONE token ledger for a portfolio pass, shared by every run's ``TokenMeter`` (S3.4).

    ``spent`` seeds from an earlier pass (``read_spend``), which is what makes the cap hold ACROSS
    passes and not merely within one.

    ``record`` and ``check`` are deliberately SPLIT rather than fused into one ``charge``. Tokens
    the provider already billed must reach this ledger even when the charge simultaneously breaks
    the RUN's own cap and raises there — a fused charge that returned early on the run cap would
    silently forget that spend and hand the next project a budget that was never really there.
    """

    def __init__(self, budget: PortfolioBudget, *, spent: int = 0) -> None:
        if spent < 0:
            raise ValueError(f"spent must be non-negative, got {spent}")
        self.budget = budget
        self.spent = spent

    @property
    def required_per_run(self) -> int:
        return self.budget.required_per_run

    def remaining(self) -> int:
        """Tokens left in the global cap (never negative — a crossed cap reads as zero left)."""
        return max(0, self.budget.max_total_tokens - self.spent)

    def can_fund_run(self, *, reserved: int = 0) -> bool:
        """Whether one more run can be started. ``reserved`` is what the caller has already
        committed to runs it admitted but that have not spent yet (the wave case): every member of
        a concurrent wave is funded off the same pre-wave remainder, so admitting them without
        reserving would over-commit the cap by exactly the wave size."""
        return self.remaining() - reserved >= self.required_per_run

    def record(self, tokens: int) -> int:
        """Accumulate ``tokens`` into the global ledger. Never raises — see the class docstring."""
        self.spent += tokens
        return self.spent

    def check(self) -> None:
        """Raise ``BudgetExceeded`` if the global cap has been crossed."""
        if self.spent > self.budget.max_total_tokens:
            raise BudgetExceeded("portfolio_tokens", self.budget.max_total_tokens, self.spent)


class TokenMeter:
    """Accumulates token and round usage against a ``Budget``; raises the moment a cap is
    crossed. When ``portfolio`` is supplied (S3.4) every charge also lands in the shared portfolio
    ledger, so the run is bounded by BOTH its own cap and the pass's global one."""

    def __init__(self, budget: Budget, *, portfolio: PortfolioMeter | None = None) -> None:
        self.budget = budget
        self.tokens = 0
        self.rounds = 0
        self.portfolio = portfolio

    def charge(self, tokens: int) -> int:
        """Add ``tokens`` to the running total; raise ``BudgetExceeded`` if over cap.

        Both ledgers are credited BEFORE either cap is tested: the spend happened regardless of
        which cap it broke, so recording must not depend on the outcome of a check."""
        self.tokens += tokens
        if self.portfolio is not None:
            self.portfolio.record(tokens)
        if self.tokens > self.budget.max_tokens:
            raise BudgetExceeded("tokens", self.budget.max_tokens, self.tokens)
        if self.portfolio is not None:
            self.portfolio.check()
        return self.tokens

    def remaining(self) -> int:
        """Tokens left under whichever cap BINDS — the run's own, or the portfolio's."""
        own = self.budget.max_tokens - self.tokens
        if self.portfolio is None:
            return own
        return min(own, self.portfolio.remaining())

    def exhausted(self) -> BudgetExceeded | None:
        """The structured error to refuse a call with, or ``None`` while budget remains.

        Exhaustion is ``remaining() <= 0``, not ``< 0``: a chat call that costs zero tokens does
        not exist, so at exactly-zero the next call can only overspend. Naming WHICH cap binds is
        the point — a run refused because a sibling drained the pass reads as ``portfolio_tokens``,
        not as its own overrun."""
        if self.remaining() > 0:
            return None
        if self.portfolio is not None and self.portfolio.remaining() <= 0:
            return BudgetExceeded(
                "portfolio_tokens", self.portfolio.budget.max_total_tokens, self.portfolio.spent
            )
        return BudgetExceeded("tokens", self.budget.max_tokens, 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.

    Two teeth, not one. The PRE-call guard (S3.4) refuses to make a call the budget cannot pay
    for — that is money not spent. The post-charge check is what it always was: the real usage is
    only knowable after the response, so a call that crosses the cap can only be caught behind it.
    The guard does not replace the check; it stops the NEXT call, never the one in flight."""

    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:
        exhausted = self._meter.exhausted()
        if exhausted is not None:
            raise exhausted  # BEFORE call_next: the call must never be made
        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


def write_spend(path: str | Path, spent: int, *, stamp: str) -> Path:
    """Persist a portfolio pass's accumulated ``spent`` tokens as deterministic JSON.

    ``stamp`` is a REQUIRED keyword with no wall-clock default (mirroring ``promote_verdict`` and
    the outbox's ``run_id``): the file must be byte-identical for the same spend, so a diff means
    the spend changed and nothing else. This is what carries the global cap ACROSS passes — the
    next pass seeds its ``PortfolioMeter`` from ``read_spend``."""
    target = Path(path)
    target.parent.mkdir(parents=True, exist_ok=True)
    payload = {"spent_tokens": int(spent), "stamp": stamp}
    target.write_text(json.dumps(payload, sort_keys=True, indent=2), encoding="utf-8")
    return target


def read_spend(path: str | Path) -> int:
    """Read persisted spend; a MISSING file reads as ``0`` (there was no earlier pass).

    Malformed content RAISES — deliberately unlike the tolerant raw verdict inbox
    (``load_verdicts_from_dir``). That folder is written out of band by other parties, so skipping
    junk is correct there; this file is our OWN accounting state, and reading a corrupt one as zero
    would silently hand back a budget that had already been spent."""
    target = Path(path)
    if not target.exists():
        return 0
    try:
        payload = json.loads(target.read_text(encoding="utf-8"))
    except json.JSONDecodeError as exc:
        raise ValueError(f"malformed spend file {str(target)!r}: {exc}") from exc
    if not isinstance(payload, dict) or not isinstance(payload.get("spent_tokens"), int):
        raise ValueError(f"spend file {str(target)!r} lacks an integer 'spent_tokens'")
    spent = int(payload["spent_tokens"])
    if spent < 0:
        raise ValueError(f"spend file {str(target)!r} carries negative spent_tokens: {spent}")
    return spent