portfolio-optimiser/tests/test_step5_refine_loadbearing.py

"""Step 5 load-bearing seam (målbilde §5/§7): informed refinement MUST feed the previous
attempt's falsification into the next hypothesis prompt.

The gap (fot-i-bakken, verified in code): ``generate_via_llm``'s outer ``max_attempts`` loop
built the prompt ONCE and re-sent it identically — a BLIND retry. The validator's per-attempt
``Rejection.reason`` was captured in ``last`` but never reached the next prompt, so the model
re-answered the same question with no knowledge of *why* the prior proposal failed. Fase 4
routes that reason into the next attempt's prompt, under the EXISTING meter/``max_attempts``
cap (no new loop; §6).

Scope note: the only PER-ATTEMPT falsifier inside this loop is the validator. The checker is a
run-level, one-shot signal (run.py, before generation); seeding generation with the checker
critique is separately scoped and NOT exercised here.

These two tests are a load-bearing PAIR (per the Fase-2 green-but-dead trap):
- the positive test keys its proposer's flip on the validator REASON PAYLOAD (the rejected
  claim value), not on a wrapper phrase or call-count, and asserts the reason reached the
  next prompt VERBATIM. It goes RED if ``_build_messages`` stops injecting the reason (detach):
  the flip token never arrives, so the outcome never flips AND the verbatim assertion fails.
- the control proves the loop stays BOUNDED: a proposer that never fixes its claim exhausts
  exactly ``max_attempts`` and returns a ``Rejection`` — so the positive test's flip is caused
  by the fed-back reason, not by mere retrying.
"""

from __future__ import annotations

from collections.abc import Sequence
from typing import Any

from agent_framework import ChatResponse, ChatResponseUpdate, Message
from spikes._harness import FakeChatClient, message_texts

from portfolio_optimiser.budget import Budget, TokenMeter
from portfolio_optimiser.generate import generate_via_llm
from portfolio_optimiser.reference_domain import load_reference_projects
from portfolio_optimiser.validator import (
    Rejection,
    ValidatedProposal,
    proposal_for,
    validate_proposal,
)

# KONTOR-IT-E1 cost codes 05.2 + 03.1 -> affected total 1_482_500 -> degenerate Monte Carlo
# P90 = 0.30 x 1_482_500 = 444_750. A claim of 800_000 is Pydantic-parseable (< affected total)
# but > P90 -> the validator REJECTS it (exercising the OUTER max_attempts path). A claim of
# 200_000 (<= P90) validates. Both proposals are built via proposal_for + model_dump_json, so
# the SUT parses EXACTLY what the test built -> the recomputed reason is byte-identical.
_CODES = ["05.2", "03.1"]
_BAD_CLAIM = 800_000
_CORRECTED_CLAIM = 200_000


def _meter() -> TokenMeter:
    # max_rounds well above max_attempts so max_attempts -- not BudgetExceeded -- is the bound.
    return TokenMeter(Budget(max_tokens=10**9, max_rounds=20))


class _ReasonAwareChatClient(FakeChatClient):
    """A proposer whose reply depends on the prompt: if the prompt carries ``flip_key`` (the
    rejected claim value, present only once the validator's reason is fed back), it returns the
    CORRECTED proposal; otherwise the BAD one. Records ``received_texts`` per call (inherited)
    so the test can assert exactly what reached each attempt's prompt."""

    def __init__(self, flip_key: str, bad_reply: str, corrected_reply: str) -> None:
        super().__init__()
        self._flip_key = flip_key
        self._bad = bad_reply
        self._corrected = corrected_reply

    def _inner_get_response(
        self, *, messages: Sequence[Message], stream: bool, options: Any, **kwargs: Any
    ) -> Any:
        received = message_texts(messages)
        self.received_texts.append(received)
        self.call_count += 1
        reply = self._corrected if self._flip_key in " ".join(received) else self._bad

        if stream:

            async def _agen() -> Any:
                yield ChatResponseUpdate(
                    role="assistant", contents=[{"type": "text", "text": reply}]
                )

            return self._build_response_stream(_agen())

        async def _coro() -> ChatResponse:
            return ChatResponse(
                messages=[Message(role="assistant", contents=[reply])], response_id="fake"
            )

        return _coro()


async def test_refinement_feeds_prior_falsification_into_next_prompt() -> None:
    """LOAD-BEARING: the validator's rejection from attempt 1 reaches attempt 2's prompt, so the
    proposer can correct. Goes RED if ``_build_messages`` stops injecting ``prior_rejection.reason``
    (the flip token never arrives -> no validated outcome AND the verbatim assertion fails)."""
    project = load_reference_projects()[0]  # KONTOR-IT-E1
    bad = proposal_for(project, _CODES, claimed_saving_nok=_BAD_CLAIM)
    corrected = proposal_for(project, _CODES, claimed_saving_nok=_CORRECTED_CLAIM)

    # The reason is computed via the SAME validator the SUT uses, on the SAME proposal -> the
    # reason string is byte-identical to the one the loop will feed back. flip_key is the claim
    # value, guaranteed to appear in that reason.
    rej = validate_proposal(bad)
    assert isinstance(rej, Rejection), "fixture invariant: the BAD claim must reject"
    flip_key = f"{bad.claimed_saving_nok:.0f}"

    client = _ReasonAwareChatClient(flip_key, bad.model_dump_json(), corrected.model_dump_json())
    # context="" so the flip token cannot pre-exist in attempt 1's prompt.
    result = await generate_via_llm(client, project, "", _meter(), max_attempts=3)

    assert isinstance(result.outcome, ValidatedProposal), (
        "the proposer corrected on attempt 2 but the loop did not validate -- the falsification "
        "never reached the next prompt"
    )
    # Green-but-dead guard: the validator's reason reached attempt 2's prompt VERBATIM. This is
    # what makes the seam load-bearing -- an impl that injects a wrapper but drops the reason
    # leaves this RED.
    assert rej.reason in " ".join(client.received_texts[1]), (
        "the validator's rejection reason did not reach the next attempt's prompt"
    )
    assert client.call_count == 2


async def test_refinement_loop_stays_bounded_when_never_fixed() -> None:
    """CAUSALITY + BOUND CONTROL (§6): a proposer that never fixes its claim exhausts exactly
    ``max_attempts`` and returns a Rejection -- proving the loop is bounded and that the positive
    test's flip is caused by the fed-back reason, not by retrying alone."""
    project = load_reference_projects()[0]
    bad = proposal_for(project, _CODES, claimed_saving_nok=_BAD_CLAIM)
    client = FakeChatClient(default_reply=bad.model_dump_json())

    result = await generate_via_llm(client, project, "", _meter(), max_attempts=3)

    assert isinstance(result.outcome, Rejection)
    assert client.call_count == 3