portfolio-optimiser/tests/test_portfolio_budget_loadbearing.py

"""S3.4 load-bearing — the GLOBAL token cap across a portfolio pass (F10).

Two properties the session plan names as RED conditions, plus the wave-size and control arms
that keep them honest:

1. **The global cap stops the pass mid-way, with completed runs preserved.** Detach the
   wave-assembly funding check and the third project is STARTED anyway — it then dies on the
   meter mid-run, so the pass reports a failure instead of a structured stop. The property is
   never-started, not merely never-overspent: a project that cannot be funded must cost zero
   model round-trips.
2. **The pre-call guard refuses the call BEFORE it is made.** The client double counts its own
   invocations, so "the call never happened" is measured, not inferred. Today's post-charge
   middleware would let the call through and only then raise — which is exactly the money the
   guard exists to not spend.

Token arithmetic is MEASURED, not guessed: every project in the reference portfolio spends
4 chat calls x ``tokens`` per reply. At ``tokens=8`` that is a flat 32 tokens per run (probed
against ``provenance.token_usage`` before these tests were written), which is what makes the
80/64/32 budget below land the stop between run 2 and run 3.
"""

from __future__ import annotations

import json
from pathlib import Path

import pytest
from agent_framework import Agent

from portfolio_optimiser.budget import (
    Budget,
    BudgetExceeded,
    BudgetMiddleware,
    BudgetRefused,
    PortfolioBudget,
    PortfolioMeter,
    TokenMeter,
    read_spend,
    write_spend,
)
from portfolio_optimiser.run import run_portfolio

_PORTFOLIO_IDS = ["FV42-GSV-E1", "RV13-RAS-TP", "BRU-LAKS-REHAB"]

# Per-project replies (the tested constants of tests/test_portfolio.py, unchanged): all three
# validate, and their claimed savings are DISTINCT — which is how a run is keyed back to its
# project here, since ``RunResult`` carries no project id. 200000 + 130000 = the first two.
REPLIES = {
    "FV42-GSV-E1": (
        '{"measure":"Reduce scope","affected_items":['
        '{"code":"05.2","quantity":4300,"unit_cost":215},'
        '{"code":"03.1","quantity":1800,"unit_cost":310}],"claimed_saving_nok":200000}'
    ),
    "RV13-RAS-TP": (
        '{"measure":"Material substitution","affected_items":['
        '{"code":"88.2","quantity":180,"unit_cost":4200}],"claimed_saving_nok":130000}'
    ),
    "BRU-LAKS-REHAB": (
        '{"measure":"Reduce scope","affected_items":['
        '{"code":"05.2","quantity":4300,"unit_cost":215},'
        '{"code":"07.4","quantity":2400,"unit_cost":690}],"claimed_saving_nok":210000}'
    ),
}
_FIRST_TWO_SAVING = 330000

# Measured: 32 tokens per run at tokens=8. 80 total funds two runs (32 + 32 = 64) and leaves 16,
# which is below the 32-token reserve one more run requires -> the third is never started.
_TOTAL = 80
_PER_RUN = 64
_RESERVE = 32
_PER_RUN_SPEND = 32


def _budget() -> PortfolioBudget:
    return PortfolioBudget(
        max_total_tokens=_TOTAL, max_tokens_per_run=_PER_RUN, min_run_reserve=_RESERVE
    )


async def test_global_cap_stops_the_pass_mid_way_with_completed_runs_preserved(
    make_portfolio_client_factory, fresh_store
) -> None:
    """RED 1: the pass stops the moment the REMAINING global budget cannot fund another run, and
    every completed run survives the stop.

    Detach the wave-assembly funding check and this goes red three ways at once: project 3 is
    started (so it lands in ``failures`` when its charge crosses the cap), ``stopped_early`` stays
    False, and ``budget_stop`` stays None. The strongest of the three is ``failures == ()``: an
    unfunded project that is merely *interrupted* has already cost model round-trips, which is the
    precise thing a global cap is for."""
    meter = PortfolioMeter(_budget())
    result = await run_portfolio(
        _PORTFOLIO_IDS,
        "local",
        store=fresh_store,
        client_factory=make_portfolio_client_factory(REPLIES, tokens=8),
        portfolio_meter=meter,
    )

    assert result.stopped_early is True
    assert result.budget_stop is not None
    assert result.budget_stop.limit_tokens == _TOTAL
    assert result.budget_stop.spent_tokens == 2 * _PER_RUN_SPEND
    assert result.budget_stop.remaining_tokens == _TOTAL - 2 * _PER_RUN_SPEND
    assert result.budget_stop.required_tokens == _RESERVE

    # Completed runs preserved, and the unfunded project NEVER STARTED: absent from BOTH sides of
    # the runs/failures partition, not merely absent from runs. A project that had been started and
    # then killed by the meter would land in ``failures`` — which is what makes ``== ()`` the
    # sharpest assertion here. The claimed-saving sum names WHICH two ran (the replies are distinct).
    assert len(result.runs) == 2
    assert result.failures == ()
    assert result.validated_count == 2
    assert result.sum_claimed_saving_nok == _FIRST_TWO_SAVING

    # The cap itself held: the pass spent what two runs cost and no more.
    assert meter.spent == 2 * _PER_RUN_SPEND
    assert meter.spent <= _TOTAL
    assert meter.spent == result.sum_token_usage


async def test_generous_budget_runs_the_whole_pass(
    make_portfolio_client_factory, fresh_store
) -> None:
    """Control: the stop above is CAUSED by the cap, not by wiring a meter in at all. With a global
    budget that funds every project the pass completes untouched — a gate that can only fire proves
    nothing about what fires it."""
    meter = PortfolioMeter(PortfolioBudget(max_total_tokens=10_000, max_tokens_per_run=_PER_RUN))
    result = await run_portfolio(
        _PORTFOLIO_IDS,
        "local",
        store=fresh_store,
        client_factory=make_portfolio_client_factory({}, tokens=8),
        portfolio_meter=meter,
    )
    assert result.stopped_early is False
    assert result.budget_stop is None
    assert len(result.runs) == 3
    assert meter.spent == 3 * _PER_RUN_SPEND


async def test_global_cap_holds_at_wave_size_three(
    make_portfolio_client_factory, fresh_store
) -> None:
    """The cap is a property of the PASS, not of the sequential schedule: at ``concurrency=3`` the
    whole portfolio is assembled into ONE wave, and every member's funding check reads the same
    pre-wave ``remaining``. Admitting all three off that shared reading would overspend by design.
    The wave therefore RESERVES each admitted member's requirement as it admits it, so the same 80
    tokens fund the same two projects at k=3 as at k=1."""
    meter = PortfolioMeter(_budget())
    result = await run_portfolio(
        _PORTFOLIO_IDS,
        "local",
        store=fresh_store,
        client_factory=make_portfolio_client_factory({}, tokens=8),
        concurrency=3,
        portfolio_meter=meter,
    )
    assert len(result.runs) == 2
    assert result.failures == ()
    assert result.stopped_early is True
    assert meter.spent <= _TOTAL


async def test_pre_call_guard_refuses_the_call_before_it_is_made(make_client_factory) -> None:
    """RED 2: with the budget exhausted, the chat call is REFUSED — the client is never invoked.

    ``call_count`` on the scripted client is the measurement: post-charge middleware alone would
    make the call, read its usage, and only then raise, so ``pytest.raises`` would pass either way.
    Only the counter separates 'refused' from 'made and then regretted'."""
    client = make_client_factory("ok", tokens=8)("proposer")
    meter = TokenMeter(Budget(max_tokens=10, max_rounds=10))
    meter.charge(10)  # exactly at the cap: within budget, but nothing left to fund a call with
    assert meter.remaining() == 0

    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 client.call_count == 0, "the call must NEVER be made once the budget is exhausted"


async def test_pre_call_guard_fires_on_the_GLOBAL_remainder_too(make_client_factory) -> None:
    """The guard reads whichever cap binds. Here the run's OWN budget is untouched and only the
    portfolio remainder is exhausted — a run that starts funded can still be refused mid-flight
    because a sibling spent the rest of the global budget."""
    client = make_client_factory("ok", tokens=8)("proposer")
    portfolio = PortfolioMeter(
        PortfolioBudget(max_total_tokens=100, max_tokens_per_run=100), spent=100
    )
    meter = TokenMeter(Budget(max_tokens=100, max_rounds=10), portfolio=portfolio)
    assert meter.tokens == 0  # the RUN has spent nothing
    assert meter.remaining() == 0  # ...but the portfolio has nothing left

    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 == "portfolio_tokens"
    assert client.call_count == 0


async def test_startup_refusal_when_the_remainder_cannot_fund_one_run(
    make_portfolio_client_factory, fresh_store
) -> None:
    """Fail-fast at startup (never a silently truncated pass): a meter whose remainder is already
    below one run's requirement refuses BEFORE anything runs. This is the across-passes arm — the
    remainder here is what an earlier pass left behind."""
    meter = PortfolioMeter(_budget(), spent=_TOTAL - (_RESERVE - 1))
    with pytest.raises(BudgetRefused) as exc:
        await run_portfolio(
            _PORTFOLIO_IDS,
            "local",
            store=fresh_store,
            client_factory=make_portfolio_client_factory({}, tokens=8),
            portfolio_meter=meter,
        )
    assert exc.value.remaining == _RESERVE - 1
    assert exc.value.required == _RESERVE
    assert meter.spent == _TOTAL - (_RESERVE - 1)  # nothing ran, nothing charged


async def test_spend_persists_across_passes(
    make_portfolio_client_factory, fresh_store, tmp_path: Path
) -> None:
    """The cap holds ACROSS passes, which is what the persisted spend buys: pass A writes its
    spend, pass B is seeded from that file and stops one run earlier than it otherwise would."""
    spend_file = tmp_path / "spend.json"
    assert read_spend(spend_file) == 0  # no prior pass -> no prior spend

    budget = PortfolioBudget(
        max_total_tokens=3 * _PER_RUN_SPEND, max_tokens_per_run=_PER_RUN, min_run_reserve=_RESERVE
    )
    meter_a = PortfolioMeter(budget)
    result_a = await run_portfolio(
        _PORTFOLIO_IDS[:1],
        "local",
        store=fresh_store,
        client_factory=make_portfolio_client_factory({}, tokens=8),
        portfolio_meter=meter_a,
    )
    assert len(result_a.runs) == 1
    write_spend(spend_file, meter_a.spent, stamp="pass-a")

    meter_b = PortfolioMeter(budget, spent=read_spend(spend_file))
    assert meter_b.spent == _PER_RUN_SPEND  # pass B starts where pass A stopped
    result_b = await run_portfolio(
        _PORTFOLIO_IDS,
        "local",
        store=fresh_store,
        client_factory=make_portfolio_client_factory({}, tokens=8),
        portfolio_meter=meter_b,
    )
    # 96 total funds three runs; one is already spent, so pass B affords two and stops before the
    # third — the earlier pass's spend is what moved the stop.
    assert len(result_b.runs) == 2
    assert result_b.stopped_early is True
    assert meter_b.spent == 3 * _PER_RUN_SPEND


def test_spend_file_is_deterministic_and_stamped(tmp_path: Path) -> None:
    """Byte-determinism with an EXPLICIT stamp (mirroring ``promote_verdict``): no wall-clock, so
    two writes of the same spend are byte-identical and a diff means the spend changed."""
    a, b = tmp_path / "a.json", tmp_path / "b.json"
    write_spend(a, 1234, stamp="run-7")
    write_spend(b, 1234, stamp="run-7")
    assert a.read_bytes() == b.read_bytes()
    assert json.loads(a.read_text(encoding="utf-8")) == {"spent_tokens": 1234, "stamp": "run-7"}
    assert read_spend(a) == 1234


def test_spend_file_malformed_fails_fast(tmp_path: Path) -> None:
    """The spend file is OUR OWN accounting state, not the tolerant raw inbox layer: a corrupt file
    raises rather than silently reading as zero spend, which would hand back the whole budget."""
    bad = tmp_path / "bad.json"
    bad.write_text("{not json", encoding="utf-8")
    with pytest.raises(ValueError):
        read_spend(bad)
    negative = tmp_path / "neg.json"
    negative.write_text('{"spent_tokens": -5, "stamp": "x"}', encoding="utf-8")
    with pytest.raises(ValueError):
        read_spend(negative)


async def test_meter_factory_and_portfolio_meter_are_mutually_exclusive(
    make_portfolio_client_factory, fresh_store
) -> None:
    """Fail-fast on the one wiring that would silently disable the global cap: a ``meter_factory``
    meter is not bound to the portfolio ledger, so accepting both would run a pass that LOOKS
    capped and is not."""
    with pytest.raises(ValueError):
        await run_portfolio(
            _PORTFOLIO_IDS,
            "local",
            store=fresh_store,
            client_factory=make_portfolio_client_factory({}, tokens=8),
            meter_factory=lambda: TokenMeter(Budget(max_tokens=1000, max_rounds=10)),
            portfolio_meter=PortfolioMeter(_budget()),
        )