portfolio-optimiser/src/portfolio_optimiser/mcp_tools.py

"""Concrete MCP servers as tools on the RUN path (Trekk B1, krav 3).

MAF already ships the MCP client (``MCPStdioTool`` / ``MCPStreamableHTTPTool``, verified present in
the pinned 1.9.0 alongside ``allowed_tools`` / ``request_timeout``), so this module owns only what
MAF cannot decide for us: **which** servers a run may contact, **which** of their tools it may call,
how long it waits, and where the credential comes from.

MAF-SPECIFIC by design, and therefore NOT in ``_MAF_FREE_MODULES``: it builds framework client
objects. That is the deliberate split from ``mandate.py``, which stays framework-neutral and
D7-portable. The D7 sibling will need its own transport module against the same config shape.

**This is a different seam from ``ingest_mcp.py``, on purpose.** That one pulls SOURCE DOCUMENTS
into a bundle before a run and speaks to null-argument tools; this one hands live tools to the
agents DURING the debate. They share a protocol, not a job.

The transport set is CLOSED (``stdio`` | ``http``): a config file can never name arbitrary
machinery to load — the same rule ``--embedder-config`` follows, for the same reason.

Every refusal below exists because its absence is a live hazard:

* an EMPTY ``allowed_tools`` hands the agents whatever the far end chooses to serve — authority
  granted by the other party. You must name what you allow;
* a non-positive timeout is an unbounded wait against a third party, which is the fail-fast
  invariant this repo applies to every loop and cap;
* an UNKNOWN field is refused rather than ignored, which is also what keeps a literal secret from
  being parked in the file and carried along unnoticed — there is no field for one;
* a named credential env var that is NOT set refuses the run rather than calling the service
  anonymously: an unauthenticated call can still succeed, with the wrong scope, which is worse
  than not calling at all.
"""

from __future__ import annotations

import os
from collections.abc import Awaitable, Callable
from pathlib import Path
from typing import Any, Literal

from agent_framework import (
    FunctionInvocationContext,
    FunctionMiddleware,
    MCPStdioTool,
    MCPStreamableHTTPTool,
)
from pydantic import BaseModel, ConfigDict, Field, model_validator

from portfolio_optimiser.provenance import ExternalCall


class McpServerConfig(BaseModel):
    """One external MCP server a run is permitted to contact.

    ``extra="forbid"`` is load-bearing, not tidiness: it is what turns a stray ``"credential":
    "sk-..."`` into a refusal instead of an ignored key sitting in a config file.
    """

    model_config = ConfigDict(extra="forbid", frozen=True)

    name: str = Field(min_length=1)
    transport: Literal["stdio", "http"]
    allowed_tools: tuple[str, ...] = Field(min_length=1)
    timeout_seconds: int = Field(gt=0)
    command: str | None = None
    args: tuple[str, ...] = ()
    url: str | None = None
    #: The NAME of an environment variable holding the credential — never the credential itself.
    credential_env: str | None = None

    @model_validator(mode="after")
    def _coordinates_match_the_transport(self) -> McpServerConfig:
        if self.transport == "stdio":
            if not self.command:
                raise ValueError(f"{self.name}: stdio transport requires 'command'")
            if self.url:
                raise ValueError(
                    f"{self.name}: stdio transport must not carry 'url' — it is ambiguous which "
                    "endpoint would actually be contacted"
                )
        else:
            if not self.url:
                raise ValueError(f"{self.name}: http transport requires 'url'")
            if self.command:
                raise ValueError(
                    f"{self.name}: http transport must not carry 'command' — it is ambiguous "
                    "which endpoint would actually be contacted"
                )
        return self


class McpConfig(BaseModel):
    """The full set of servers a run may contact."""

    model_config = ConfigDict(extra="forbid")

    servers: tuple[McpServerConfig, ...] = ()

    @model_validator(mode="after")
    def _names_are_unique(self) -> McpConfig:
        seen: set[str] = set()
        for server in self.servers:
            if server.name in seen:
                raise ValueError(f"duplicate MCP server name: {server.name!r}")
            seen.add(server.name)
        return self


def load_mcp_config(path: str | Path) -> tuple[McpServerConfig, ...]:
    """Fail-fast loader (mirrors ``load_dimension`` / ``load_mandate``).

    Authoritative startup input: missing -> ``FileNotFoundError``, malformed -> ``ValidationError``.
    A tolerant read here would be the worst of the three loaders in this repo — degrading a broken
    egress config to "no external services" would leave the announcement telling the truth about a
    run that was never configured, while a degraded one that *partially* parsed could contact a
    subset nobody chose.
    """
    p = Path(path)
    if not p.is_file():
        raise FileNotFoundError(f"MCP config not found: {str(path)!r}")
    return McpConfig.model_validate_json(p.read_text(encoding="utf-8")).servers


def _credential(config: McpServerConfig) -> str | None:
    """The credential value, read from the environment at BUILD time. A named-but-unset variable
    refuses: calling an external service anonymously can succeed with the wrong scope."""
    if config.credential_env is None:
        return None
    value = os.environ.get(config.credential_env)
    if not value:
        raise ValueError(
            f"{config.name}: credential env var {config.credential_env!r} is not set — refusing "
            "to contact the service unauthenticated"
        )
    return value


def build_mcp_tools(configs: tuple[McpServerConfig, ...]) -> list[Any]:
    """Build one MAF MCP client per configured server, with the allowlist and timeout applied.

    NOTE for the caller: an ``MCPTool`` is an async context manager. These objects are constructed
    here but NOT connected — the run path must enter them around the debate and exit afterwards.
    """
    tools: list[Any] = []
    for config in configs:
        credential = _credential(config)
        if config.transport == "stdio":
            # ``credential`` is non-None only when ``credential_env`` named a variable, so the key
            # is a str here — narrowed explicitly rather than asserted, so mypy sees it too.
            env = (
                {config.credential_env: credential}
                if credential and config.credential_env
                else None
            )
            tools.append(
                MCPStdioTool(
                    name=config.name,
                    command=config.command or "",
                    args=list(config.args),
                    env=env,
                    allowed_tools=config.allowed_tools,
                    request_timeout=config.timeout_seconds,
                )
            )
        else:
            header_provider = (
                (lambda _ctx, token=credential: {"Authorization": f"Bearer {token}"})
                if credential
                else None
            )
            tools.append(
                MCPStreamableHTTPTool(
                    name=config.name,
                    url=config.url or "",
                    allowed_tools=config.allowed_tools,
                    request_timeout=config.timeout_seconds,
                    header_provider=header_provider,
                )
            )
    return tools


def tool_server_index(configs: tuple[McpServerConfig, ...]) -> dict[str, str]:
    """Map each allowed tool name to the server that serves it — the attribution table B4 needs.

    A name allowed by MORE THAN ONE configured server maps to ``""`` (unattributed) rather than to
    one of them. MEASURED: MAF hands function middleware the BARE tool name, with no server prefix,
    so the two really are indistinguishable at that seam. Picking the first match would be a guess
    written into a provenance record, which is the one place a guess must never go — validation,
    never repair, applied to attribution.
    """
    index: dict[str, str] = {}
    for config in configs:
        for name in config.allowed_tools:
            index[name] = "" if name in index else config.name
    return index


class ToolCallRecorder(FunctionMiddleware):
    """Function middleware that records WHICH configured external tool a run actually called (B4).

    Mirrors ``budget.BudgetMiddleware(ChatMiddleware)`` one layer down: that one observes the
    debate's chat calls, this one its tool calls. It observes only — ``call_next`` is always
    awaited, and nothing here can block, alter or short-circuit an invocation. A trace that changed
    the run it traces would not be a trace.

    **Only CONFIGURED tools are recorded.** The middleware fires for every function the agents
    invoke, including the in-process ``retrieve_cost_docs`` on the road path. Logging those would
    turn the record into a false egress claim — the whole value of the record is that its entries
    mean "something outside this process was contacted".

    Recording is idempotent per ``(server, tool)`` and returned SORTED: the record answers *what was
    contacted*, and it is stamped into a byte-deterministic artefact, so it must not vary with how
    many times an agent happened to ask.
    """

    def __init__(self, index: dict[str, str]) -> None:
        self._index = index
        self._seen: set[tuple[str, str]] = set()

    def note(self, tool_name: str) -> None:
        """Record one invocation by tool name. Unconfigured names are IGNORED, not recorded as
        unattributed calls — an in-process tool is not an external service with a missing label."""
        if tool_name in self._index:
            self._seen.add((self._index[tool_name], tool_name))

    def calls(self) -> list[ExternalCall]:
        return [ExternalCall(server=s, tool=t) for s, t in sorted(self._seen)]

    async def process(
        self, context: FunctionInvocationContext, call_next: Callable[[], Awaitable[None]]
    ) -> None:
        name = getattr(getattr(context, "function", None), "name", None)
        if isinstance(name, str):
            self.note(name)
        await call_next()


def service_labels(configs: tuple[McpServerConfig, ...]) -> tuple[str, ...]:
    """The egress declaration, one label per server: what will be contacted and which tools are
    permitted. Feeds ``mandate.announce(external_services=...)``, so an operator sees the full
    reach of a run BEFORE the first call — a run never contacts a service it did not announce."""
    return tuple(f"{c.name} ({', '.join(c.allowed_tools)})" for c in configs)