lda-wf/tests/wf_mcp/test_support.py

from __future__ import annotations

import os
import sys
from pathlib import Path
from typing import Any, overload
from warnings import deprecated

from pydantic import BaseModel, Field

from wf_authoring import NodeReturn, node
from wf_core import RuntimeContext
from wf_core.models.steps import InputPathBinding, OutputBinding
from wf_core.paths import GraphSourcePath, LocalPath, StatePath
from wf_mcp.capabilities import DiscoveredPrompt, DiscoveredResource, DiscoveredTool
from wf_mcp.models import ConnectionConfig
from wf_mcp.sdk import ToolCallResult


class EchoInput(BaseModel):
    text: str = Field(description="Text to echo")


class EchoOutput(BaseModel):
    echoed: str = Field(description="Echoed text")


class FinalizeInput(BaseModel):
    echoed: str


class FinalizeOutput(BaseModel):
    result: str


@node()
async def echo_tool(payload: EchoInput, ctx: RuntimeContext) -> EchoOutput:
    return EchoOutput(echoed=payload.text)


@node(outcomes=("done",))
def finalize_tool(
    payload: FinalizeInput, ctx: RuntimeContext
) -> NodeReturn[FinalizeOutput]:
    return NodeReturn(
        outcome="done",
        output=FinalizeOutput(result=f"final:{payload.echoed}"),
    )


@deprecated("Use pytests tmp_path fixture instead")
@overload
def local_temp_root() -> Path: ...


@overload
def local_temp_root(root_path: Path) -> Path: ...


def local_temp_root(root_path: Path | None = None) -> Path:
    root = root_path or (Path("test-artifacts") / "wf_mcp_store")
    root.mkdir(parents=True, exist_ok=True)
    return root


def fixture_server_path() -> str:
    return str(Path(__file__).resolve().parents[1] / "fixtures" / "mcp_echo_server.py")


def input_binding(path: str, target: str) -> dict[str, object]:
    """Return canonical JSON for a node input path binding in raw plans."""
    return InputPathBinding(
        path=GraphSourcePath.parse(path),
        target=LocalPath.parse(target),
    ).model_dump(mode="json")


def output_binding(source: str, target: str) -> dict[str, object]:
    """Return canonical JSON for a node output binding in raw plans."""
    return OutputBinding(
        source=LocalPath.parse(source),
        target=StatePath.parse(target),
    ).model_dump(mode="json")


def everything_server_connection() -> ConnectionConfig | None:
    transport = os.environ.get("MCP_EVERYTHING_TRANSPORT", "stdio")
    if transport == "stdio":
        command = os.environ.get("MCP_EVERYTHING_COMMAND")
        if not command:
            return None

        raw_args = os.environ.get("MCP_EVERYTHING_ARGS", "")
        args = [arg for arg in raw_args.split(" ") if arg]

        metadata: dict[str, Any] = {
            "transport": transport,
            "command": command,
            "args": args,
        }
    elif transport == "streamable_http":
        url = os.environ.get("MCP_EVERYTHING_URL")
        if not url:
            raise AssertionError(
                "MCP_EVERYTHING_URL must be set when MCP_EVERYTHING_TRANSPORT=streamable_http"
            )
        metadata = {
            "transport": "streamable_http",
            "url": url,
        }
    else:
        return None

    return ConnectionConfig(
        id="everything.default",
        server="everything",
        account="default",
        metadata=metadata,
    )


class FakeAdapter:
    async def list_tools(
        self,
        connection,
        auth,
    ) -> list[DiscoveredTool]:
        return [
            DiscoveredTool(
                name="echo_tool",
                title="Echo Tool",
                description="Echo text back",
                input_schema={
                    "type": "object",
                    "properties": {
                        "text": {
                            "type": "string",
                            "description": "Text to echo",
                        }
                    },
                    "required": ["text"],
                },
                output_schema={
                    "type": "object",
                    "properties": {
                        "echoed": {
                            "type": "string",
                            "description": "Echoed text",
                        }
                    },
                    "required": ["echoed"],
                },
            )
        ]

    async def list_resources(
        self,
        connection,
        auth,
    ) -> list[DiscoveredResource]:
        return [
            DiscoveredResource(
                uri="demo://docs/welcome",
                name="resource.welcome",
                title="Welcome Resource",
                description="Welcome resource",
                mime_type="text/plain",
                metadata={"kind": "static"},
            )
        ]

    async def list_prompts(
        self,
        connection,
        auth,
    ) -> list[DiscoveredPrompt]:
        return [
            DiscoveredPrompt(
                name="prompt.summarize",
                title="Summarize Prompt",
                description="Summarize text",
                arguments=[
                    {
                        "name": "text",
                        "required": True,
                        "description": "Text to summarize",
                    }
                ],
                metadata={"kind": "template"},
            )
        ]

    async def get_connection_metadata(
        self,
        connection,
        auth,
    ) -> dict[str, Any]:
        return {
            "server": getattr(
                connection, "provider", getattr(connection, "server", None)
            ),
            "account": connection.account,
            "auth_scheme": auth.scheme if auth is not None else None,
        }

    async def read_resource(
        self,
        connection,
        auth,
        uri: str,
    ) -> dict[str, Any]:
        if uri != "demo://docs/welcome":
            raise KeyError(uri)
        return {
            "contents": [
                {
                    "uri": uri,
                    "mimeType": "text/plain",
                    "text": "Welcome from the fake adapter resource.",
                }
            ]
        }

    async def get_prompt(
        self,
        connection,
        auth,
        prompt_name: str,
        arguments: dict[str, str] | None = None,
    ) -> dict[str, Any]:
        if prompt_name != "prompt.summarize":
            raise KeyError(prompt_name)
        text = (arguments or {}).get("text", "")
        return {
            "description": "Summarize text",
            "messages": [
                {
                    "role": "user",
                    "content": {
                        "type": "text",
                        "text": f"Summarize this text:\n\n{text}",
                    },
                }
            ],
        }

    async def invoke_method(
        self,
        connection,
        auth,
        method: str,
        params: dict[str, Any] | None = None,
    ) -> dict[str, Any]:
        if method == "ping":
            return {}
        if method == "demo.echo":
            return {"echoed": (params or {}).get("text", "")}
        raise KeyError(method)

    async def send_notification(
        self,
        connection,
        auth,
        method: str,
        params: dict[str, Any] | None = None,
    ) -> None:
        return None

    async def call_tool(
        self,
        connection,
        auth,
        tool_name: str,
        payload: dict[str, Any],
    ) -> ToolCallResult:
        if tool_name != "echo_tool":
            raise KeyError(tool_name)
        return ToolCallResult(
            outcome="ok",
            output={"echoed": str(payload["text"])},
        )


class FailingDiscoveryAdapter(FakeAdapter):
    async def list_tools(
        self,
        connection,
        auth,
    ) -> list[DiscoveredTool]:
        raise PermissionError("Access is denied")


__all__ = [
    "FailingDiscoveryAdapter",
    "FakeAdapter",
    "echo_tool",
    "everything_server_connection",
    "finalize_tool",
    "fixture_server_path",
    "local_temp_root",
    "sys",
]