feat(budget): enforce a global portfolio token cap before the call, not after it (S3.4/F10)

PortfolioBudget + PortfolioMeter carry ONE token ledger over a whole portfolio
pass -- and, seeded from a persisted spend file, across passes -- while the
per-run Budget/TokenMeter pair is untouched. Three enforcement points, each
doing a different job:

- startup: a remainder that cannot fund one run raises BudgetRefused before
  anything loads (a pass that can afford zero projects is a caller mistake,
  not a result);
- wave assembly: an unfundable project is NEVER STARTED and the pass stops
  structurally (budget_stop + stopped_early, completed runs preserved).
  Because every member of a wave is funded against the SAME pre-wave
  remainder, admission RESERVES each member's requirement -- otherwise a wave
  of k over-commits the cap by up to k runs;
- pre-call: BudgetMiddleware refuses a call the remainder cannot pay for
  instead of making it. The post-charge check stays: real usage is only
  knowable after the response, so the guard stops the NEXT call, never the
  one in flight.

budget_stop is its own field rather than a widened stop_reason -- a goal-stop
is success, this is resource exhaustion, and fusing them would make "we
stopped" unreadable. PortfolioMeter splits record/check so tokens the provider
already billed reach the ledger even when the same charge breaks the run's own
cap. read_spend raises on corrupt content (our own accounting state, unlike
the tolerant RAW inbox layer); write_spend takes a REQUIRED stamp with no
wall-clock default, mirroring promote_verdict.

Load-bearing MEASURED, not asserted -- 6 mutations, all red: detach the wave
check; detach the pre-call guard; detach the wave reservation; check the run
cap before crediting the global ledger; detach the startup refusal; make
read_spend tolerant. Files restored from shasum-verified copies after each.

Two findings worth keeping: the pre-call guard MASKS a detached wave check if
the test asserts on overspend (spend stays under the cap either way), so the
load-bearing assertion had to become failures == () plus never-started; and
the token arithmetic is probed (32 tokens/run at tokens=8), not guessed.

537 -> 553 tests, ruff + mypy green.

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_015EaxFnaDAbMQkmTeX4u7sd
This commit is contained in:
Kjell Tore Guttormsen 2026-07-31 21:34:48 +02:00
commit a831aa1e3b
6 changed files with 733 additions and 11 deletions

View file

@ -8,12 +8,21 @@ 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
@ -37,6 +46,24 @@ class UsageUnavailable(RuntimeError):
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).
@ -54,22 +81,142 @@ class Budget:
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."""
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) -> None:
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."""
"""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
@ -80,13 +227,21 @@ class TokenMeter:
class BudgetMiddleware(ChatMiddleware):
"""Chat middleware that charges a ``TokenMeter`` from each response's real
``UsageDetails`` and short-circuits when the cap is crossed."""
``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
@ -98,3 +253,39 @@ class BudgetMiddleware(ChatMiddleware):
)
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

View file

@ -35,7 +35,13 @@ from agent_framework import BaseChatClient, SessionContext
from pydantic import ValidationError
from portfolio_optimiser.backends import Profile, get_backend, resolve_model
from portfolio_optimiser.budget import Budget, BudgetMiddleware, TokenMeter
from portfolio_optimiser.budget import (
Budget,
BudgetMiddleware,
BudgetRefused,
PortfolioMeter,
TokenMeter,
)
from portfolio_optimiser.contracts import GoalConfig, GoalContract, load_contracts, load_goal_config
from portfolio_optimiser.ledger import SavingsLedger
from portfolio_optimiser.datasource import (
@ -141,6 +147,25 @@ class GoalReached:
observed_ore: int
@dataclass(frozen=True)
class BudgetStop:
"""A GLOBAL token-cap stop signal VALUE (S3.4/F10) — NOT an exception, and NOT a goal.
Structured like ``GoalReached`` and carried the same way (a ``stop_reason``-shaped value plus a
loop ``break``), but it is kept as its OWN field rather than widening ``stop_reason``: the two
stops mean opposite things. A goal-stop is success (the savings target was met); this is
resource exhaustion (the pass ran out of tokens). Folding them into one field would let a
caller read "we stopped" without being able to tell which happened.
``required_tokens`` is what one more run would have needed; ``remaining_tokens`` is what was
actually left. Both are recorded because their DIFFERENCE is the operator's next decision."""
limit_tokens: int
spent_tokens: int
remaining_tokens: int
required_tokens: int
@dataclass(frozen=True)
class PortfolioResult:
"""The outcome of a sequential fan-out over N projects (SC2).
@ -150,9 +175,11 @@ class PortfolioResult:
The remaining fields are a thin aggregate over ``runs``: ``validated_count`` /
``rejected_count`` partition the outcomes; ``sum_claimed_saving_nok`` totals the claimed
saving of the validated proposals only; ``sum_token_usage`` totals every run's
provenance token usage. ``stopped_early`` / ``stop_reason`` record a Step-8 goal-stop, and
``failures`` records the projects that RAISED (S3.3 collect-and-continue): all three default so
the frozen aggregate and every existing constructor call are unaffected.
provenance token usage. ``stopped_early`` / ``stop_reason`` record a Step-8 goal-stop,
``failures`` records the projects that RAISED (S3.3 collect-and-continue), and ``budget_stop``
records a S3.4 global-token-cap stop (which also sets ``stopped_early``, but is kept apart from
``stop_reason`` because exhaustion is not success): all four default so the frozen aggregate and
every existing constructor call are unaffected.
``runs`` and ``failures`` PARTITION the projects that were actually submitted a project
appears in exactly one of them, never both, and the counts do not overlap. ``validated_count`` /
@ -168,6 +195,7 @@ class PortfolioResult:
stopped_early: bool = False
stop_reason: GoalReached | None = None
failures: tuple[RunFailure, ...] = ()
budget_stop: BudgetStop | None = None
def _authored_texts(result: Any, name: str) -> list[str]:
@ -647,6 +675,30 @@ def _waves(ids: list[str], k: int) -> list[list[str]]:
return [ids[i : i + k] for i in range(0, len(ids), k)]
def _run_meter(
meter_factory: Callable[[], TokenMeter] | None,
portfolio_meter: PortfolioMeter | None,
max_rounds: int,
) -> TokenMeter | None:
"""The per-project meter ``run_portfolio`` hands to one run: the injected test seam, a meter
BOUND to the portfolio ledger, or ``None`` (letting ``run_project`` build its own, unchanged).
The bound meter mirrors ``run_project``'s own construction (``max(max_rounds * 4, 4)``) so the
only difference the portfolio cap introduces is the token ceiling and the shared ledger the
round budget is not quietly redefined along the way."""
if meter_factory is not None:
return meter_factory()
if portfolio_meter is None:
return None
return TokenMeter(
Budget(
max_tokens=portfolio_meter.budget.max_tokens_per_run,
max_rounds=max(max_rounds * 4, 4),
),
portfolio=portfolio_meter,
)
async def run_portfolio(
project_ids: Sequence[str] | None = None,
profile: Profile | str = Profile.LOCAL,
@ -661,6 +713,7 @@ async def run_portfolio(
top_k: int = 3,
concurrency: int = 1,
meter_factory: Callable[[], TokenMeter] | None = None,
portfolio_meter: PortfolioMeter | None = None,
semantic_retrieval: bool = False,
embedder: Embedder | None = None,
) -> PortfolioResult:
@ -727,12 +780,46 @@ async def run_portfolio(
property of the design, so it is pinned on a bundle-backed pair where the fold does fire
(``test_intra_wave_visibility_is_the_documented_semantic_difference``): same fixture, same
sentinel, only the wave boundary moves. Store content stays identical across ``k`` the
difference is confined to what each project READ, never to what the pass produced or persisted."""
difference is confined to what each project READ, never to what the pass produced or persisted.
``portfolio_meter`` (S3.4, F10) installs the GLOBAL token cap. Without it nothing changes: each
run is bounded only by its own ``max_tokens``, so N projects can cost N times that with no
ceiling over the pass. With it, one ``PortfolioMeter`` is shared by every run (each run's meter
is BOUND to it, which is why ``meter_factory`` whose meters are unbound is refused
alongside it: accepting both would run a pass that looks capped and is not), and the cap is
enforced in three places that are deliberately different:
- **At startup**, a remainder that cannot fund one run raises ``BudgetRefused`` a pass that
can afford zero projects is a caller mistake, not a result.
- **At wave assembly**, a project that cannot be funded is NEVER STARTED, and the pass stops
with ``budget_stop`` + ``stopped_early``, every completed run preserved. Never-started is the
property that matters: an unfunded project that is merely interrupted mid-run has already
cost model round-trips. Because every member of a wave is checked against the SAME pre-wave
remainder, admission RESERVES each member's requirement as it goes — otherwise a wave of k
would over-commit the cap by up to k runs. This makes membership identical at every
``concurrency``, matching the goal-stop's guarantee above.
- **Before each chat call** (``BudgetMiddleware``'s pre-call guard), a call the remainder cannot
pay for is refused rather than made. This is what bounds a run that was funded at admission
but whose siblings drained the pass while it was in flight; such a run surfaces as a
``RunFailure``, not as overspend.
Spend is carried ACROSS passes by seeding the meter from ``budget.read_spend`` and writing
``budget.write_spend`` afterwards. Those are the caller's calls, not this function's the pass
reads its ledger, it does not own the file (mirroring the Steg-7 role split)."""
if concurrency < 1:
raise ValueError(
f"concurrency must be >= 1, got {concurrency}: a non-positive wave size would run no "
"projects at all and read as an empty portfolio. Pass 1 for the sequential pass."
)
if portfolio_meter is not None and meter_factory is not None:
raise ValueError(
"portfolio_meter and meter_factory are mutually exclusive: a meter_factory meter is "
"not bound to the portfolio ledger, so the global cap would be silently unenforced"
)
# Startup refusal (fail-fast, before anything loads): a pass that cannot fund its first run
# must not begin. Raised, not returned — an empty PortfolioResult would read as "nothing to do".
if portfolio_meter is not None and not portfolio_meter.can_fund_run():
raise BudgetRefused(portfolio_meter.remaining(), portfolio_meter.required_per_run)
projects = {p.id: p for p in load_reference_projects()}
ids = list(project_ids) if project_ids is not None else list(projects)
store = store if store is not None else VerdictStore(verdicts=[])
@ -747,8 +834,13 @@ async def run_portfolio(
failures: list[RunFailure] = []
stopped_early = False
stop_reason: GoalReached | None = None
budget_stop: BudgetStop | None = None
for wave_ids in _waves(ids, concurrency):
members: list[str] = []
# Tokens committed to members already admitted to THIS wave but that have not spent yet.
# Every member reads the same pre-wave remainder, so without this a wave of k would admit
# k projects off one project's worth of budget.
reserved = 0
for pid in wave_ids:
if pid not in projects:
raise ValueError(f"unknown project_id: {pid!r}")
@ -777,6 +869,21 @@ async def run_portfolio(
if per_project_goal.mode == "hard":
continue # skip THIS pid; the rest of the pass proceeds
# S3.4 funding check, placed AFTER the goal checks: a reached goal is success and owns
# the stop when both apply, and a pid a hard per-project goal already skipped costs no
# budget, so it must not consume a reservation.
if portfolio_meter is not None and not portfolio_meter.can_fund_run(reserved=reserved):
stopped_early = True
budget_stop = BudgetStop(
limit_tokens=portfolio_meter.budget.max_total_tokens,
spent_tokens=portfolio_meter.spent,
remaining_tokens=portfolio_meter.remaining(),
required_tokens=portfolio_meter.required_per_run,
)
break
if portfolio_meter is not None:
reserved += portfolio_meter.required_per_run
members.append(pid)
# Every project in the wave reads the SAME wave-start state and writes only its own copy,
@ -805,7 +912,7 @@ async def run_portfolio(
top_k=top_k,
semantic_retrieval=semantic_retrieval,
embedder=embedder,
meter=meter_factory() if meter_factory is not None else None,
meter=_run_meter(meter_factory, portfolio_meter, max_rounds),
)
for pid, snapshot in snapshots
),
@ -836,12 +943,13 @@ async def run_portfolio(
break
base = _aggregate(tuple(runs), store)
if stopped_early or stop_reason is not None or failures:
if stopped_early or stop_reason is not None or failures or budget_stop is not None:
return replace(
base,
stopped_early=stopped_early,
stop_reason=stop_reason,
failures=tuple(failures),
budget_stop=budget_stop,
)
return base