portfolio-optimiser-claude/tests/test_step5_refine_loadbearing.py

"""Load-bearing: informed, bounded refinement (method-spec §3 Step 5, §11).

RED-conditions proved here: the prior rejection reason no longer appears
VERBATIM in the next prompt, or the outcome never flips. The flip proof uses a
prompt-SENSITIVE scripted stand-in (honesty rule §1): it returns a feasible
proposal ONLY when the full rejection reason is present in the prompt — so the
test goes red the moment the informed block is detached. Bounds proved: only
the MOST RECENT reason is carried (never accumulated history), never the prior
proposal JSON, attempt 1 uses the unchanged base prompt, and the loop stops at
``max_attempts``.
"""

from __future__ import annotations

import json

import pytest
from _scripted import ScriptedClient, reply

from portfolio_optimiser_claude.budget import BudgetMeter
from portfolio_optimiser_claude.contracts import TerminationContract
from portfolio_optimiser_claude.ir import SavingsProposal
from portfolio_optimiser_claude.loop import ModelReply, run_candidate_loop
from portfolio_optimiser_claude.validator import Rejection, ValidatedProposal, validate_proposal

# 100 × 1000 = 100_000 NOK total; nominal feasible 30_000 == p90 (no band).
GOOD = {
    "project_id": "p1",
    "measure": "led-retrofit",
    "affected_items": [{"code": "E01", "quantity": 100, "unit_cost": 1000}],
    "claimed_saving_nok": 25000,
}
BAD_90K = dict(GOOD, claimed_saving_nok=90000)  # IR-valid, above p90 → rejected
BAD_80K = dict(GOOD, claimed_saving_nok=80000)  # a SECOND distinct rejection reason


def _rejection_reason(raw: dict[str, object]) -> str:
    outcome = validate_proposal(SavingsProposal.model_validate(raw))
    assert isinstance(outcome, Rejection)
    return outcome.reason


def _meter() -> BudgetMeter:
    return BudgetMeter(TerminationContract(max_rounds=50, max_tokens=10_000))


class TestInformedRefinement:
    def test_the_outcome_flips_because_the_reason_reaches_the_prompt(self) -> None:
        # THE load-bearing flip: the stand-in addresses the falsification ONLY
        # if the full rejection reason appears verbatim — red at detach.
        reason_90k = _rejection_reason(BAD_90K)

        def script(role: str, prompt: str) -> ModelReply:
            if reason_90k in prompt:
                return reply(json.dumps(GOOD))
            return reply(json.dumps(BAD_90K))

        client = ScriptedClient(script=script)
        result = run_candidate_loop(client, "base prompt", meter=_meter(), max_attempts=3)
        assert isinstance(result.outcome, ValidatedProposal)
        assert result.attempts == 2
        assert reason_90k in client.prompts("proposer")[1]

    def test_attempt_one_uses_the_unchanged_base_prompt(self) -> None:
        client = ScriptedClient(replies=[reply(json.dumps(GOOD))])
        run_candidate_loop(client, "base prompt", meter=_meter(), max_attempts=3)
        assert client.prompts("proposer")[0] == "base prompt"

    def test_only_the_most_recent_reason_is_carried(self) -> None:
        # §3 Step 5: never an accumulated history — bounded prompt growth.
        reason_90k = _rejection_reason(BAD_90K)
        reason_80k = _rejection_reason(BAD_80K)
        client = ScriptedClient(
            replies=[
                reply(json.dumps(BAD_90K)),
                reply(json.dumps(BAD_80K)),
                reply(json.dumps(GOOD)),
            ]
        )
        result = run_candidate_loop(client, "base prompt", meter=_meter(), max_attempts=3)
        assert isinstance(result.outcome, ValidatedProposal)
        third_prompt = client.prompts("proposer")[2]
        assert reason_80k in third_prompt
        assert reason_90k not in third_prompt

    def test_the_prior_proposal_json_is_never_carried(self) -> None:
        # The model must address the falsification, not parrot the rejected
        # candidate — only the REASON crosses attempts.
        client = ScriptedClient(replies=[reply(json.dumps(BAD_90K)), reply(json.dumps(GOOD))])
        run_candidate_loop(client, "base prompt", meter=_meter(), max_attempts=3)
        second_prompt = client.prompts("proposer")[1]
        # Positive control: the REASON did cross into the retry prompt — so the two
        # negatives measure what was deliberately withheld, not an empty or misindexed
        # prompt that would carry nothing either way.
        assert _rejection_reason(BAD_90K) in second_prompt
        assert json.dumps(BAD_90K) not in second_prompt
        assert '"affected_items"' not in second_prompt

    def test_the_loop_stops_at_max_attempts(self) -> None:
        client = ScriptedClient(script=lambda role, prompt: reply(json.dumps(BAD_90K)))
        result = run_candidate_loop(client, "base prompt", meter=_meter(), max_attempts=3)
        assert isinstance(result.outcome, Rejection)
        assert result.attempts == 3
        assert len(client.prompts("proposer")) == 3

    def test_max_attempts_must_be_positive(self) -> None:
        client = ScriptedClient(replies=[reply(json.dumps(GOOD))])
        with pytest.raises(ValueError):
            run_candidate_loop(client, "base prompt", meter=_meter(), max_attempts=0)

    def test_round_ticks_are_charged_between_attempts(self) -> None:
        client = ScriptedClient(replies=[reply(json.dumps(BAD_90K)), reply(json.dumps(GOOD))])
        meter = _meter()
        run_candidate_loop(client, "base prompt", meter=meter, max_attempts=3)
        assert meter.rounds_used == 1

    def test_a_never_validating_run_yields_a_typed_rejection_not_a_bare_failure(self) -> None:
        # §3 Step 6: the outcome is either the validated proposal or a TYPED
        # rejection with its reason — never a bare failure.
        client = ScriptedClient(script=lambda role, prompt: reply(json.dumps(BAD_90K)))
        result = run_candidate_loop(client, "base prompt", meter=_meter(), max_attempts=2)
        assert isinstance(result.outcome, Rejection)
        assert "exceeds the optimistic feasible bound" in result.outcome.reason