distinct_mcp/common.py

"""Plumbing shared by every server in this package.

Four things live here because each one would otherwise be copied into five
files and drift in four of them.

**The result envelope.** MCP returns a call result as a list of typed content
blocks, and :func:`distinct_tools.mcp._result_payload` reads
``structuredContent`` first when a server supplies it. Every tool here supplies
it. The alternative, letting the in-process transport wrap a plain return value
in a text block, encodes the payload as JSON and then encodes that string as
JSON again, so every quote and newline is escaped twice and a text-heavy result
can double in size before the broker measures it against
``JobPolicy.max_output_bytes``. The MCP specification says a server that
returns structured content should *also* serialise it into a text block for
clients that cannot read the structured field. This client can read it, and the
duplicate copy would cost the bytes the tool's own caps are sized against, so
these servers deliberately do not send one. A server published to unknown
clients should make the opposite choice.

**Argument checking.** Unknown arguments are refused rather than ignored, for
the reason :mod:`distinct_tools.local` gives: ignoring one is always the
permissive outcome and it hides the model misreading the schema, which is the
failure this library exists to make visible.

**Refusals are :class:`~distinct_tools.core.ToolInputError`.** Raising it from
an in-process MCP handler reaches the broker unchanged, so the model gets
``invalid_input`` and the reason rather than a generic tool failure. That
matters more than it looks: below 7B the single most common recovery from a
refusal is to call the tool again unchanged, and a message that names the
argument and what would be accepted is the only thing that redirects it.

**Publishing an existing spec.** :func:`definition_from_spec` republishes a
:class:`~distinct_tools.core.ToolSpec` over MCP with the same description and
the same schema the direct path advertises, so the text the model reads cannot
diverge between the two routes.
"""

from __future__ import annotations

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

from distinct_tools.core import ToolInputError, ToolSpec
from distinct_tools.mcp import McpToolDefinition


def structured(value: Mapping[str, Any]) -> dict[str, Any]:
    """Return an MCP call result carrying ``value`` as structured content."""

    return {"content": [], "structuredContent": dict(value)}


def reject_unknown(arguments: Mapping[str, Any], known: frozenset[str]) -> None:
    """Refuse any argument the tool does not understand."""

    if not isinstance(arguments, Mapping):
        raise ToolInputError("arguments must be an object")
    unexpected = sorted(str(key) for key in arguments if str(key) not in known)
    if unexpected:
        accepted = ", ".join(sorted(known))
        raise ToolInputError(
            f"unknown argument(s): {', '.join(unexpected)}. This tool accepts: {accepted}"
        )


def text_argument(
    arguments: Mapping[str, Any],
    key: str,
    limit: int,
    *,
    required: bool = True,
    default: str = "",
) -> str:
    """Return a bounded string argument, or refuse with the reason."""

    value = arguments.get(key)
    if value is None and not required:
        return default
    if not isinstance(value, str):
        raise ToolInputError(f"{key} must be a string")
    if required and not value.strip():
        raise ToolInputError(f"{key} must be a non-empty string")
    if len(value) > limit:
        raise ToolInputError(f"{key} must be at most {limit:,} characters")
    return value


def choice_argument(
    arguments: Mapping[str, Any],
    key: str,
    allowed: Sequence[str],
    *,
    default: str | None = None,
) -> str:
    """Return one of ``allowed``, or refuse with the full list of choices."""

    value = arguments.get(key)
    if value is None and default is not None:
        return default
    if not isinstance(value, str):
        raise ToolInputError(f"{key} must be a string")
    lowered = value.strip().lower()
    if lowered not in allowed:
        raise ToolInputError(f"{key} must be one of: {', '.join(allowed)}")
    return lowered


def flag_argument(arguments: Mapping[str, Any], key: str, *, default: bool = False) -> bool:
    """Return a strictly boolean argument.

    A string ``"true"`` is refused rather than coerced. Coercing it would make
    the tool's behaviour depend on how the model happened to spell a value the
    schema said was a boolean, and the refusal names the problem.
    """

    value = arguments.get(key)
    if value is None:
        return default
    if not isinstance(value, bool):
        raise ToolInputError(f"{key} must be true or false")
    return value


def definition_from_spec(spec: ToolSpec, *, name: str | None = None) -> McpToolDefinition:
    """Build an MCP definition from a spec the direct path already publishes.

    ``ToolSpec`` deep-freezes its schema into mapping proxies and tuples, which
    are not JSON types, so the published dictionary is taken from
    :meth:`ToolSpec.to_dict` rather than from the attribute. Handing the frozen
    form to the JSON encoder is exactly what the digest path refuses.
    """

    published = spec.to_dict()
    return McpToolDefinition(
        name=name or spec.tool_id,
        description=str(published["description"]),
        input_schema=published["input_schema"],
    )


__all__ = [
    "choice_argument",
    "definition_from_spec",
    "flag_argument",
    "reject_unknown",
    "structured",
    "text_argument",
]