Request for guidance: preserving MCP schema and result semantics across toolkit adapters

from __future__ import annotations

import asyncio
import json
from contextlib import asynccontextmanager
from contextlib import AbstractAsyncContextManager
from typing import Any

from fastmcp import FastMCP as StandaloneFastMCP
from mcp.server.fastmcp import FastMCP as SDKFastMCP
from mcp.types import CallToolResult
from mcp.types import ListToolsResult
from mcp.types import TextContent
from mcp.types import Tool
from pydantic import BaseModel
from pydantic import ConfigDict
from pydantic import Field
from pydantic import create_model

from nat.builder.function import LambdaFunction
from nat.builder.function_info import FunctionInfo
from nat.data_models.function import FunctionBaseConfig
from nat.plugins.fastmcp.server.tool_converter import create_function_wrapper as create_fastmcp_wrapper
from nat.plugins.mcp.client.client_base import MCPBaseClient
from nat.plugins.mcp.client.client_base import MCPToolClient
from nat.plugins.mcp.server.tool_converter import create_function_wrapper as create_mcp_wrapper
from nat.plugins.mcp.server.tool_converter import register_function_with_mcp
from nat.plugins.mcp.utils import model_from_mcp_schema


class EchoInput(BaseModel):
    message: str = Field(description="Text to return.")


class EchoOutput(BaseModel):
    message: str
    label_count: int


class SchemaInput(BaseModel):
    model_config = ConfigDict(extra="forbid")

    query: str = Field(min_length=1, description="Text to return.")
    labels: list[str] = Field(default_factory=list)


class DefaultInput(BaseModel):
    limit: int = 10


class _Runner(AbstractAsyncContextManager):
    def __init__(self, result: Any):
        self._result = result

    async def __aenter__(self):
        return self

    async def __aexit__(self, exc_type, exc, traceback):
        return False

    async def result(self):
        return self._result


class _SessionManager:
    def __init__(self, result: Any):
        self._result = result
        self.payloads: list[Any] = []

    def run(self, payload: Any):
        self.payloads.append(payload)
        return _Runner(self._result)


class _MCPParent:
    server_name = "minimal-server"

    def __init__(self, result: CallToolResult):
        self._result = result

    async def call_tool(self, name: str, arguments: dict[str, Any]):
        return self._result


class _DiscoverySession:
    async def list_tools(self):
        return ListToolsResult(tools=[
            Tool(
                name="echo",
                description="Echo a value.",
                inputSchema={
                    "type": "object",
                    "properties": {"message": {"type": "string"}},
                    "required": ["message"],
                },
                outputSchema=EchoOutput.model_json_schema(),
            )
        ])


class _DiscoveryClient(MCPBaseClient):
    @asynccontextmanager
    async def connect_to_server(self):
        yield _DiscoverySession()


def show(case: str, actual: Any, expected: str) -> None:
    print(f"\n[{case}]")
    print("actual:  " + json.dumps(actual, default=str, sort_keys=True))
    print("expected: " + expected)


async def main() -> None:
    # 1. Both MCP server adapters stringify typed toolkit results.
    typed_result = EchoOutput(message="hello", label_count=0)

    async def echo(value: EchoInput) -> EchoOutput:
        return EchoOutput(message=value.message, label_count=0)

    echo_function = LambdaFunction.from_info(
        config=FunctionBaseConfig(),
        info=FunctionInfo.create(single_fn=echo),
        instance_name="echo",
    )

    mcp_manager = _SessionManager(typed_result)
    mcp_manager.workflow = echo_function
    mcp_wrapper = create_mcp_wrapper("echo", mcp_manager, EchoInput)
    mcp_wrapper_result = await mcp_wrapper(message="hello")
    show(
        "nvidia_nat_mcp typed server result",
        {"python_type": type(mcp_wrapper_result).__name__, "value": mcp_wrapper_result},
        "A structured EchoOutput value, not a string.",
    )

    sdk_server = SDKFastMCP("minimal-server")
    register_function_with_mcp(sdk_server, "echo", mcp_manager, function=echo_function)
    sdk_tool = (await sdk_server.list_tools())[0]
    show(
        "nvidia_nat_mcp advertised output schema",
        sdk_tool.outputSchema,
        "The EchoOutput object schema.",
    )
    sdk_result = await sdk_server.call_tool("echo", {"message": "hello"})
    show(
        "nvidia_nat_mcp tools/call result",
        {
            "python_type": type(sdk_result).__name__,
            "content": sdk_result[0],
            "structured_content": sdk_result[1],
        },
        "structuredContent containing the EchoOutput object, plus backward-compatible text content.",
    )

    fastmcp_manager = _SessionManager(typed_result)
    fastmcp_wrapper = create_fastmcp_wrapper("echo", fastmcp_manager, EchoInput)
    fastmcp_wrapper_result = await fastmcp_wrapper(message="hello")
    show(
        "nvidia_nat_fastmcp typed server result",
        {"python_type": type(fastmcp_wrapper_result).__name__, "value": fastmcp_wrapper_result},
        "A structured EchoOutput value, not a string.",
    )

    native_error = CallToolResult(
        content=[TextContent(type="text", text="invalid input")],
        structuredContent={"code": "invalid_input"},
        isError=True,
    )
    native_manager = _SessionManager(native_error)
    native_wrapper = create_mcp_wrapper("native_error", native_manager, EchoInput)
    native_wrapper_result = await native_wrapper(message="hello")
    show(
        "native CallToolResult through server adapter",
        {"python_type": type(native_wrapper_result).__name__, "value": native_wrapper_result},
        "A CallToolResult retaining content, structuredContent, and isError.",
    )

    # 2. The MCP client returns text and discards structured result fields.
    successful_result = CallToolResult(
        content=[TextContent(type="text", text='{"message":"hello","label_count":0}')],
        structuredContent={"message": "hello", "label_count": 0},
        isError=False,
    )
    client = MCPToolClient(
        session=object(),
        parent_client=_MCPParent(successful_result),
        tool_name="echo",
        tool_description="Echo a value.",
        tool_input_schema={"type": "object", "properties": {}},
    )
    client_result = await client.acall({})
    show(
        "MCP client structured success",
        {"python_type": type(client_result).__name__, "value": client_result},
        "The structuredContent object, or an API that makes it available.",
    )
    discovery_client = _DiscoveryClient(reconnect_enabled=False)
    discovery_client._session = _DiscoverySession()
    discovered_tool = (await discovery_client.get_tools())["echo"]
    show(
        "MCP client outputSchema discovery",
        {
            "server_output_schema": EchoOutput.model_json_schema(),
            "client_attributes": sorted(vars(discovered_tool)),
            "client_has_output_schema": hasattr(discovered_tool, "output_schema"),
        },
        "The discovered client tool retains and exposes the server outputSchema.",
    )

    error_client = MCPToolClient(
        session=object(),
        parent_client=_MCPParent(native_error),
        tool_name="fail",
        tool_description="Return a structured failure.",
        tool_input_schema={"type": "object", "properties": {}},
    )
    error_client_result = await error_client.acall({})
    show(
        "MCP client structured error",
        {"python_type": type(error_client_result).__name__, "value": error_client_result},
        "An error representation retaining structuredContent and isError.",
    )

    # 3a. Local references are not resolved by the MCP-schema converter.
    referenced_schema = {
        "type": "object",
        "$defs": {
            "Filter": {
                "type": "object",
                "properties": {"limit": {"type": "integer"}},
                "additionalProperties": False,
            }
        },
        "properties": {"filter": {"$ref": "#/$defs/Filter"}},
        "required": ["filter"],
    }
    referenced_model = model_from_mcp_schema("referenced", referenced_schema)
    referenced_value = referenced_model.model_validate({"filter": {"unexpected": True}})
    show(
        "local $defs/$ref",
        referenced_value.model_dump(),
        "ValidationError for the unexpected nested property.",
    )

    # 3b. An omittable string is converted into a nullable string.
    omittable_schema = {
        "type": "object",
        "properties": {"query": {"type": "string"}},
        "required": [],
    }
    omittable_model = model_from_mcp_schema("omittable", omittable_schema)
    explicit_null = omittable_model.model_validate({"query": None})
    show(
        "omittable but non-null property",
        explicit_null.model_dump(),
        "ValidationError for explicit null; omission should remain valid.",
    )

    # 3c. A concrete default is materialized before the toolkit handler.
    default_manager = _SessionManager("ok")
    default_wrapper = create_mcp_wrapper("defaulted", default_manager, DefaultInput)
    default_server = SDKFastMCP("default-server")
    default_server.tool(name="defaulted")(default_wrapper)
    await default_server.call_tool("defaulted", {})
    default_payload = default_manager.payloads[-1]
    show(
        "omitted concrete default",
        {
            "payload": default_payload.model_dump(),
            "model_fields_set": sorted(default_payload.model_fields_set),
        },
        'If caller-presence preservation is intended: omitted input has model_fields_set=[]; explicit {"limit":10} sets the field.',
    )

    # 3d-f. The standalone FastMCP frontend derives its wire schema from the
    # sanitized Python signature rather than preserving all Pydantic semantics.
    schema_manager = _SessionManager("ok")
    schema_wrapper = create_fastmcp_wrapper("schema_echo", schema_manager, SchemaInput)
    standalone_server = StandaloneFastMCP("schema-server")
    standalone_server.tool(name="schema_echo")(schema_wrapper)
    schema_tool = (await standalone_server.list_tools())[0].to_mcp_tool()
    show(
        "nvidia_nat_fastmcp default_factory and field metadata",
        schema_tool.inputSchema,
        "labels omitted from required; query description/minLength preserved.",
    )

    keyword_schema = create_model("KeywordInput", **{"from": (str, ...)})
    keyword_manager = _SessionManager("ok")
    keyword_wrapper = create_fastmcp_wrapper("keyword", keyword_manager, keyword_schema)
    keyword_server = StandaloneFastMCP("keyword-server")
    keyword_server.tool(name="keyword")(keyword_wrapper)
    keyword_tool = (await keyword_server.list_tools())[0].to_mcp_tool()
    show(
        "nvidia_nat_fastmcp sanitized wire name",
        keyword_tool.inputSchema,
        'The protocol-facing property remains "from"; only the Python wrapper uses "from_".',
    )


if __name__ == "__main__":
    asyncio.run(main())