portfolio-optimiser/tests/test_b4_mcp_call_trace_loadbearing.py

"""Load-bearing: a run records WHICH external service it actually called (Trekk B4).

Trekk B3 made a run declare, before its first call, which servers it MAY contact. That is the
permission. It is not the record: after the run, the egress declaration still cannot tell you
whether the agents used the price register or ignored it, and a proposal that rests on an external
service should be traceable to it.

``FunctionMiddleware`` + ``FunctionInvocationContext`` were verified present in the pinned MAF 1.9.0,
and the middleware was MEASURED to fire for a tool served over a REAL MCP stdio subprocess — with
``context.function.name`` carrying the bare tool name and NO server prefix. That measurement is why
attribution comes from our own config rather than from MAF: the framework cannot tell us which
server a tool came from, so two servers exposing one tool name are indistinguishable at this seam
and must be reported as unattributed rather than guessed.

**Honesty limit, pinned in the code it describes:** this records the CALL and its SOURCE. It is not
evidence that the service's answer reached the proposal, nor a verified rendering of that answer.

Detach points, each RED on its own:

* drop the recorder from the debate middleware -> a called service leaves no trace, and the run
  reports the same empty record as a run that called nothing;
* record every function invocation -> the reference path's IN-PROCESS retrieval tool is reported as an
  external service call, which turns the record into a false egress claim;
* attribute an ambiguous tool name to the first server that allows it -> the record names a service
  that may never have been contacted.

The control (``test_a_run_without_mcp_servers_records_no_calls``) proves the empty record is a real
statement rather than the only thing the seam can produce.
"""

from __future__ import annotations

from collections.abc import Mapping, Sequence
from pathlib import Path
from typing import Any

import pytest
from agent_framework import ChatResponse, Content, FunctionTool, Message, UsageDetails, tool

from portfolio_optimiser import run as run_module
from portfolio_optimiser.mcp_tools import McpServerConfig, ToolCallRecorder, tool_server_index
from portfolio_optimiser.provenance import ExternalCall
from portfolio_optimiser.reference_domain import load_reference_projects
from portfolio_optimiser.run import run_project
from portfolio_optimiser.simulation import ScriptedChatClient
from portfolio_optimiser.verdicts import VerdictStore

BUNDLE_DIR = Path(__file__).resolve().parents[1] / "shared" / "examples" / "bygg-energi-mikro"
_VERDICT_INPUT = {"decision": "approved", "rationale": "expert reviewed (sim)"}
_REPLY = (
    '{"measure":"LED-retrofit av kontorbelysning","affected_items":'
    '[{"code":"ENERGI-TOTAL-EL","quantity":300000,"unit_cost":1.0}],"claimed_saving_nok":30000}'
)

_SERVER = McpServerConfig(
    name="prisregister",
    transport="http",
    url="https://intern.example/mcp",
    allowed_tools=("lookup_unit_price",),
    timeout_seconds=15,
)


@tool(name="lookup_unit_price", description="Look up a unit price from the price register.")
def _lookup_unit_price(code: str) -> str:
    return f"unit price for {code}: 1.0"


def _as_context_manager(function_tool: FunctionTool) -> Any:
    """Wrap a real ``FunctionTool`` so it is ALSO an async context manager — the two halves a MAF
    ``MCPTool`` presents to the run path (``AsyncExitStack`` enters it; the agents invoke it).

    A stand-in rather than a live server on purpose: the subject here is what OUR middleware records,
    and the real MCP transport is already measured against a spawned server in
    ``test_ingest_golden_mcp.py``. What matters is that the agent genuinely INVOKES the tool, so the
    middleware fires for real rather than being called by the test directly.
    """
    cls = type(
        "_ContextManagedTool",
        (type(function_tool),),
        {
            "__aenter__": lambda self: _aenter(self),
            "__aexit__": lambda self, *exc: _aexit(),
        },
    )
    wrapped = cls.__new__(cls)
    wrapped.__dict__.update(function_tool.__dict__)
    return wrapped


async def _aenter(self: Any) -> Any:
    return self


async def _aexit() -> bool:
    return False


class _ToolCallingClient(ScriptedChatClient):
    """Emits ONE tool call on its first response, then delegates to the canonical scripted body.

    Registered in ``test_scripted_client_consolidation.py``'s ``_DELEGATING_OVERRIDES``: it cannot
    live in the ``reply_selector`` seam, which returns a reply STRING — this double's whole subject
    is a response that is a function CALL rather than text (the same reason the two ordering/failure
    probes are registered there).
    """

    def __init__(self, call_tool: str, arguments: dict[str, Any], reply: str) -> None:
        super().__init__(reply=reply, tokens_per_reply=8)
        self._call_tool = call_tool
        self._arguments = arguments
        self._called = False

    def _inner_get_response(
        self,
        *,
        messages: Sequence[Any],
        options: Mapping[str, Any],
        stream: bool = False,
        **kwargs: Any,
    ) -> Any:
        if self._called or stream:
            return super()._inner_get_response(
                messages=messages, options=options, stream=stream, **kwargs
            )
        self._called = True
        self.call_count += 1

        async def _coro() -> ChatResponse:
            return ChatResponse(
                messages=[
                    Message(
                        role="assistant",
                        contents=[
                            Content(
                                "function_call",
                                call_id="call-1",
                                name=self._call_tool,
                                arguments=self._arguments,
                            )
                        ],
                    )
                ],
                response_id="synthetic",
                usage_details=UsageDetails(total_token_count=8),
            )

        return _coro()


#: The argument name each tool actually declares. MEASURED, and load-bearing: a call whose argument
#: names do not match the tool's signature is rejected by MAF BEFORE invocation, so the middleware
#: never fires — the first version of the reference-path test below passed ``code`` to
#: ``retrieve_cost_docs(query)`` and was therefore VACUOUS. It asserted an empty record against a run
#: where no tool ran at all, and stayed green under the mutation it exists to catch.
_ARGUMENTS = {
    "lookup_unit_price": {"code": "ENERGI-TOTAL-EL"},
    "retrieve_cost_docs": {"query": "cost saving measure"},
}


def _factory(call_tool: str):
    def factory(role: str) -> ScriptedChatClient:
        # Only the proposer calls a tool; the checker stays textual, so the recorded call cannot be
        # an artefact of every agent calling everything.
        return (
            _ToolCallingClient(call_tool, _ARGUMENTS[call_tool], _REPLY)
            if role == "proposer"
            else ScriptedChatClient(reply="VERDICT: APPROVE", tokens_per_reply=8)
        )

    return factory


# --- The recorder's own contract ------------------------------------------------------------


def test_an_unconfigured_tool_name_is_not_recorded_as_egress() -> None:
    """Only a CONFIGURED MCP tool is an external call. RED when the recorder logs every function:
    the in-process retrieval tool would then be reported as contact with a third party."""
    recorder = ToolCallRecorder(tool_server_index((_SERVER,)))
    recorder.note("retrieve_cost_docs")
    assert recorder.calls() == []


def test_an_ambiguous_tool_name_is_recorded_without_a_server() -> None:
    """Two servers allowing one tool name cannot be told apart at this seam (MEASURED: MAF passes
    the bare tool name). RED when the recorder picks the first match — the record would then name a
    service that may never have been contacted."""
    other = _SERVER.model_copy(update={"name": "reservepris"})
    recorder = ToolCallRecorder(tool_server_index((_SERVER, other)))
    recorder.note("lookup_unit_price")
    assert recorder.calls() == [ExternalCall(server="", tool="lookup_unit_price")]


# --- The run path ---------------------------------------------------------------------------


async def test_a_called_mcp_tool_is_recorded_in_provenance(monkeypatch: pytest.MonkeyPatch) -> None:
    """The seam itself: the agents call the configured service, and the run's provenance says so.

    RED when the recorder is not in the debate middleware — the run then reports the same empty
    record whether it contacted the price register or ignored it.
    """
    monkeypatch.setattr(
        run_module, "build_mcp_tools", lambda _c: [_as_context_manager(_lookup_unit_price)]
    )
    result = await run_project(
        "BYGG-KONTOR-NORD",
        "local",
        docs_dir=str(BUNDLE_DIR),
        bundle_dir=str(BUNDLE_DIR),
        verdict_input=_VERDICT_INPUT,
        store=VerdictStore(verdicts=[]),
        client_factory=_factory("lookup_unit_price"),
        mcp_servers=(_SERVER,),
    )
    assert result.provenance.external_calls == [
        ExternalCall(server="prisregister", tool="lookup_unit_price")
    ]


async def test_a_local_tool_call_is_not_recorded_as_an_external_call(
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    """The reference path's ``retrieve_cost_docs`` runs IN PROCESS. Calling it is not egress, and a run
    that only called it must not claim it contacted the price register.

    RED when the recorder logs every function invocation rather than the configured ones.

    The ``observed`` spy is what keeps this test HONEST rather than merely green. An empty record is
    also what a run that invoked nothing produces, so without proof that the invocation actually
    reached the recorder this asserts nothing — which is exactly what happened to the first version
    of it (the tool call named an argument the tool does not declare, MAF rejected it before
    invocation, and the test stayed green under the very mutation it exists to catch).
    """
    monkeypatch.setattr(
        run_module, "build_mcp_tools", lambda _c: [_as_context_manager(_lookup_unit_price)]
    )
    observed: list[str] = []
    original = ToolCallRecorder.note

    def spy(self: ToolCallRecorder, tool_name: str) -> None:
        observed.append(tool_name)
        original(self, tool_name)

    monkeypatch.setattr(ToolCallRecorder, "note", spy)

    nett = {p.id: p for p in load_reference_projects()}["NETT-SIKR-TP"]
    result = await run_project(
        "NETT-SIKR-TP",
        "local",
        docs_dir=nett.docs_dir,
        verdict_input=nett.verdict_input,
        store=VerdictStore(verdicts=[]),
        client_factory=_factory("retrieve_cost_docs"),
        mcp_servers=(_SERVER,),
    )
    assert observed == ["retrieve_cost_docs"], "the local tool must actually have been invoked"
    assert result.provenance.external_calls == []


async def test_a_run_without_mcp_servers_records_no_calls() -> None:
    """Control: no configured servers -> an empty record, on the pre-B4 path, unchanged.

    Without this the assertions above could all pass against a seam that can only ever produce one
    answer.
    """
    result = await run_project(
        "BYGG-KONTOR-NORD",
        "local",
        docs_dir=str(BUNDLE_DIR),
        bundle_dir=str(BUNDLE_DIR),
        verdict_input=_VERDICT_INPUT,
        store=VerdictStore(verdicts=[]),
        client_factory=_factory("lookup_unit_price"),
    )
    assert result.provenance.external_calls == []