portfolio-optimiser/tests/test_portfolio_failure_accounting_loadbearing.py

"""S3.3 load-bearing — what a FAILED project does to the wave handler and to the token ledger.

Two properties on the same seam (``run.py`` wave-result handling), sharing one fixture because they
share one code path: a project that raises mid-wave.

1. **The wave handler catches ``BaseException``, not ``Exception``.** The existing collect-and-
   continue test (``test_one_project_failure_does_not_cancel_its_siblings``) raises ``RuntimeError``,
   which is an ``Exception`` — so it stays GREEN if the handler is narrowed, and cannot gate the
   width of the catch. ``asyncio.CancelledError`` derives from ``BaseException`` (3.8+) and is the
   one realistic vector that separates the two.

   MEASURED, not assumed, before this test was written: ``asyncio.gather(return_exceptions=True)``
   COLLECTS a ``CancelledError`` raised by a member coroutine into the result list, so the
   ``BaseException`` branch is genuinely reachable. ``KeyboardInterrupt`` is NOT a usable vector —
   asyncio special-cases it and propagates out of ``gather`` regardless of ``return_exceptions``,
   so widening the catch could never help there. That measurement is why this file tests exactly
   one BaseException subclass and does not pretend to cover the rest.

2. **A failed project's tokens still count against the GLOBAL cap.** ``PortfolioResult.
   sum_token_usage`` sums ``provenance.token_usage`` over ``runs``, and a failed project has no
   ``RunResult`` — so the aggregate necessarily UNDER-reports what the pass actually spent. That is
   the honest answer rather than a bug to be patched: a run that died before producing a stamp has
   no provenance, and inventing one would fabricate exactly the thing this repo's provenance rules
   exist to prevent (the same reasoning that made ``RunFailure`` a distinct type).

   What must NOT be true is that those tokens vanish from the ledger the cap is enforced against.
   If they did, a project that fails repeatedly could burn unbounded budget while the meter read
   clean. This pins the split: ``meter.spent`` is the pass's real cost, ``sum_token_usage`` is the
   completed-run subtotal, and the difference is exactly the failed project's spend.
"""

from __future__ import annotations

import asyncio
from typing import Any, Mapping, Sequence

from portfolio_optimiser.budget import PortfolioBudget, PortfolioMeter
from portfolio_optimiser.run import RunFailure, run_portfolio
from portfolio_optimiser.simulation import ScriptedChatClient

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

# The MIDDLE project. Middle is deliberate (mirroring the Step-4 test): a failure at either end can
# be dropped by a truncation bug and still leave the survivors in the right relative order.
_FAILING_PID = "RV13-RAS-TP"

_DEFAULT_REPLY = (
    '{"measure":"Reduce scope","affected_items":'
    '[{"code":"01.1","quantity":1,"unit_cost":100000}],"claimed_saving_nok":20000}'
)

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}'
    ),
}

# Measured in tests/test_portfolio_budget_loadbearing.py: 4 chat calls x ``tokens`` per reply, so a
# completed run costs a flat 32 tokens at tokens=8.
_TOKENS_PER_REPLY = 8
_PER_RUN_SPEND = 32
_SURVIVOR_COUNT = 2


class _FailAfterNCallsClient(ScriptedChatClient):
    """Project-aware scripted client that raises a chosen exception for ONE project, after that
    project has already completed ``fail_after`` calls.

    ``fail_after`` is what makes property 2 measurable: at 0 the project dies having spent nothing,
    which would make the ledger assertion trivially equal. The counter is SHARED across the clients
    the factory produces, because ``client_factory`` is called per ROLE and a run uses several — a
    per-instance counter would restart on every role and never reach the threshold.

    The failure is raised from inside a coroutine, not synchronously: ``_inner_get_response`` is a
    sync method RETURNING an awaitable, so raising synchronously would blow up during ``gather``'s
    argument construction — before any concurrency exists — and would test a different thing than a
    mid-flight failure.
    """

    def __init__(
        self,
        replies: dict[str, str],
        *,
        failing_pid: str,
        error: BaseException,
        counter: dict[str, int],
        fail_after: int = 0,
    ) -> None:
        table = dict(replies)

        def _select(blob: str, _role: str) -> str:
            return next((r for pid, r in table.items() if pid in blob), _DEFAULT_REPLY)

        super().__init__(
            reply_selector=_select,
            default_reply=_DEFAULT_REPLY,
            tokens_per_reply=_TOKENS_PER_REPLY,
        )
        self._failing_pid = failing_pid
        self._error = error
        self._counter = counter
        self._fail_after = fail_after

    def _inner_get_response(
        self,
        *,
        messages: Sequence[Any],
        options: Mapping[str, Any],
        stream: bool = False,
        **kwargs: Any,
    ) -> Any:
        blob = " ".join(getattr(m, "text", "") or "" for m in messages)
        if not stream and self._failing_pid in blob:
            seen = self._counter.get("n", 0)
            if seen >= self._fail_after:
                error = self._error

                async def _boom() -> Any:
                    raise error

                return _boom()
            self._counter["n"] = seen + 1
        return super()._inner_get_response(
            messages=messages, options=options, stream=stream, **kwargs
        )


def _factory(error: BaseException, *, fail_after: int = 0) -> Any:
    counter: dict[str, int] = {}

    def factory(_role: str) -> Any:
        return _FailAfterNCallsClient(
            REPLIES,
            failing_pid=_FAILING_PID,
            error=error,
            counter=counter,
            fail_after=fail_after,
        )

    return factory


async def test_a_cancelled_project_is_collected_not_propagated(fresh_store) -> None:
    """RED 1: a wave member raising ``asyncio.CancelledError`` is recorded as a ``RunFailure``
    while its siblings complete — the pass does not raise and does not lose their results.

    **Detach point: narrow ``run.py``'s ``isinstance(outcome, BaseException)`` to ``Exception``
    -> this test goes RED (measured).** ``CancelledError`` then misses the failure branch and falls
    through to ``cast(RunResult, outcome)``, which puts a live exception object into ``runs``; the
    pass dies in ``_aggregate`` reaching for ``.provenance`` on it. Note the failure mode is a
    crash in a DIFFERENT function than the defect — collect-and-continue is negated for cancelled
    members, and the traceback points away from the cause. That is precisely why the width of this
    catch is worth pinning rather than leaving to the ``RuntimeError`` test.

    Why ``CancelledError`` specifically is documented in this module's docstring: it is the only
    ``BaseException`` subclass ``gather(return_exceptions=True)`` actually hands back.
    """
    result = await run_portfolio(
        _PORTFOLIO_IDS,
        "local",
        store=fresh_store,
        client_factory=_factory(asyncio.CancelledError("synthetic cancellation")),
    )

    # Reaching this line at all is half the contract: the pass RETURNED rather than propagating.
    assert len(result.failures) == 1, (
        f"expected exactly one failure slot, got {result.failures} — a CancelledError that is not "
        "collected either propagates out of run_portfolio or is cast into runs as a fake RunResult"
    )
    failure = result.failures[0]
    assert isinstance(failure, RunFailure)
    assert failure.project_id == _FAILING_PID
    assert failure.error_type == "CancelledError"

    # The siblings survived and are real RunResults — the assertion that catches an exception
    # object smuggled into ``runs`` by a narrowed catch.
    assert len(result.runs) == _SURVIVOR_COUNT
    assert [r.verdict.proposal_features.claimed_saving_nok for r in result.runs] == [
        200_000.0,
        210_000.0,
    ]
    assert result.validated_count + result.rejected_count == _SURVIVOR_COUNT


async def test_a_failed_projects_tokens_still_count_against_the_global_cap(fresh_store) -> None:
    """RED 2: tokens spent by a project that later failed remain in the ledger the global cap is
    enforced against, even though they are absent from ``sum_token_usage``.

    The split is the point, and both halves are asserted:

    - ``meter.spent`` — the pass's REAL cost — includes the failed project's partial spend. If it
      did not, a project failing on every attempt would burn budget invisibly and the S3.4 cap
      would not bound the pass at all.
    - ``sum_token_usage`` — the completed-run subtotal — excludes it, because a run that never
      produced provenance has no honest token figure to contribute. Pinned so that a later
      "fix" cannot quietly fabricate one to make the two numbers agree.

    The difference between them is asserted to be EXACTLY the failed project's spend, so this stays
    a statement about accounting rather than a loose inequality that would hold for many wrong
    reasons.
    """
    # Two calls complete for the failing project before it dies, so its partial spend is non-zero
    # and known: without that, both sides of the ledger assertion would be equal by accident.
    fail_after = 2
    partial_spend = fail_after * _TOKENS_PER_REPLY

    meter = PortfolioMeter(PortfolioBudget(max_total_tokens=10_000, max_tokens_per_run=1_000))
    result = await run_portfolio(
        _PORTFOLIO_IDS,
        "local",
        store=fresh_store,
        client_factory=_factory(RuntimeError("synthetic backend failure"), fail_after=fail_after),
        portfolio_meter=meter,
    )

    assert len(result.runs) == _SURVIVOR_COUNT
    assert len(result.failures) == 1
    assert result.failures[0].project_id == _FAILING_PID

    completed_spend = _SURVIVOR_COUNT * _PER_RUN_SPEND
    assert result.sum_token_usage == completed_spend, (
        "sum_token_usage must total the COMPLETED runs' provenance and nothing else; a failed "
        "project has no provenance to contribute and must not be given a fabricated one"
    )
    assert meter.spent == completed_spend + partial_spend, (
        f"the global ledger recorded {meter.spent}, expected {completed_spend + partial_spend} — "
        f"the failed project's {partial_spend} tokens were really spent, so a cap that does not "
        "see them does not bound the pass"
    )
    # The under-report is real, bounded, and exactly the failed project's spend.
    assert meter.spent - result.sum_token_usage == partial_spend