rixdagen/mcp_server/test_mcp_client.py

"""
Test script for the RiksdagenTools MCP server (HTTP transport).
This script connects to the MCP server via Streamable HTTP and tests your main tools.
Ensure the MCP server is running and that the environment variable MCP_SERVER_TOKEN is set.
Also adjust MCP_SERVER_URL if needed.
"""

import os
import asyncio
from typing import Any, Dict, List, Optional, Sequence

from mcp.client.streamable_http import streamablehttp_client
from mcp.client.session import ClientSession   # adjusted import per SDK version

TOKEN: str = os.environ.get("MCP_SERVER_TOKEN", "2q89rwpfaiukdjshp298n3qw")
SERVER_URL: str = os.environ.get("MCP_SERVER_URL", "https://api.rixdagen.se/mcp")  # use the public HTTPS endpoint by default so tests target the nginx proxy

async def run_tests() -> None:
    """
    Attempt to connect to SERVER_URL and run the tool tests. On SSL certificate
    verification failures, optionally retry against the backend IP if
    MCP_FALLBACK_TO_IP=1 is set (or a custom MCP_FALLBACK_URL is provided).
    """
    async def run_with_url(url: str) -> None:
        print(f"Connecting to server URL: {url}")
        async with streamablehttp_client(
            url=url,
            headers={ "Authorization": f"Bearer {TOKEN}" }
        ) as (read_stream, write_stream, get_session_id):
            async with ClientSession(read_stream, write_stream) as session:
                # initialize the session (if needed)
                init_result = await session.initialize()
                print("Initialized session:", init_result)

                print("\nListing available tools...")
                tools_info = await session.list_tools()
                print("Tools:", [ tool.name for tool in tools_info.tools ])

                # Test aql_query
                print("\n== Testing aql_query ==")
                result1 = await session.call_tool(
                    "aql_query",
                    arguments={
                        "token": TOKEN,
                        "query": "FOR doc IN talks LIMIT 2 RETURN { _id: doc._id, talare: doc.talare }"
                    }
                )
                print("aql_query result:", result1)

                # Test search_documents
                print("\n== Testing search_documents ==")
                result2 = await session.call_tool(
                    "search_documents",
                    arguments={
                        "token": TOKEN,
                        "aql_query": "FOR doc IN talks LIMIT 2 RETURN { _id: doc._id, talare: doc.talare }"
                    }
                )
                print("search_documents result:", result2)

                # Test vector_search_talks
                print("\n== Testing vector_search_talks ==")
                result3 = await session.call_tool(
                    "vector_search_talks",
                    arguments={
                        "token": TOKEN,
                        "query": "klimatförändringar",
                        "limit": 2
                    }
                )
                print("vector_search_talks result:", result3)

                # Test fetch_documents
                print("\n== Testing fetch_documents ==")
                # try to pull out IDs from result3 if available
                doc_ids: List[str]
                maybe = result3
                if hasattr(maybe, "output") and isinstance(maybe.output, list) and maybe.output:
                    doc_ids = [ maybe.output[0].get("_id", "") ]
                else:
                    doc_ids = ["talks/1"]
                result4 = await session.call_tool(
                    "fetch_documents",
                    arguments={
                        "token": TOKEN,
                        "document_ids": doc_ids
                    }
                )
                print("fetch_documents result:", result4)

                # Test arango_search
                print("\n== Testing arango_search ==")
                result5 = await session.call_tool(
                    "arango_search",
                    arguments={
                        "token": TOKEN,
                        "query": "klimat",
                        "limit": 2
                    }
                )
                print("arango_search result:", result5)

    # try primary URL first
    try:
        await run_with_url(SERVER_URL)
    except Exception as e:  # capture failures from streamablehttp_client / httpx
        err_str = str(e).lower()
        ssl_fail = "certificate_verify_failed" in err_str or "hostname mismatch" in err_str or "certificate verify failed" in err_str
        gateway_fail = "502" in err_str or "bad gateway" in err_str or "502 bad gateway" in err_str

        if gateway_fail:
            print("Received 502 Bad Gateway from the proxy when connecting to the server URL.")
            # If user explicitly set fallback env var, retry against backend IP or custom fallback
            fallback_flag = os.environ.get("MCP_FALLBACK_TO_IP", "0").lower() in ("1", "true", "yes")
            if fallback_flag:
                fallback_url = os.environ.get("MCP_FALLBACK_URL", "http://127.0.0.1:8010/mcp")
                print(f"Retrying with fallback URL (MCP_FALLBACK_URL or default backend): {fallback_url}")
                await run_with_url(fallback_url)
                return
            print("")
            print("Possible causes:")
            print("- The proxy (nginx) couldn't reach the backend (backend down, wrong proxy_pass or path).")
            print("- Proxy buffering or HTTP version issues interfering with streaming transport.")
            print("")
            print("Options:")
            print("- Bypass the proxy and target the backend directly:")
            print("    export MCP_SERVER_URL='http://127.0.0.1:8010/mcp'")
            print("- Or enable automatic fallback to the backend (insecure) for testing:")
            print("    export MCP_FALLBACK_TO_IP=1")
            print("    # optionally override the fallback target")
            print("    export MCP_FALLBACK_URL='http://127.0.0.1:8010/mcp'")
            print("- Check the proxy's error log (e.g. /var/log/nginx/error.log) for upstream errors.")
            print("")
            # re-raise so caller still sees the error if they don't follow guidance
            raise

        if ssl_fail:
            print("SSL certificate verification failed while connecting to the server URL.")
            # If user explicitly set fallback env var, retry against backend IP or custom fallback
            fallback_flag = os.environ.get("MCP_FALLBACK_TO_IP", "0").lower() in ("1", "true", "yes")
            if fallback_flag:
                fallback_url = os.environ.get("MCP_FALLBACK_URL", "http://192.168.1.10:8010/mcp")
                print(f"Retrying with fallback URL (MCP_FALLBACK_URL or default backend IP): {fallback_url}")
                await run_with_url(fallback_url)
                return
            # Otherwise give actionable guidance
            print("")
            print("Possible causes:")
            print("- The TLS certificate served for api.rixdagen.se does not match that hostname on this machine.")
            print("")
            print("Options:")
            print("- Set MCP_SERVER_URL to the backend HTTP address to bypass TLS: e.g.")
            print("    export MCP_SERVER_URL='http://192.168.1.10:8010/mcp'")
            print("- Or enable automatic fallback to the backend IP for testing (insecure):")
            print("    export MCP_FALLBACK_TO_IP=1")
            print("    # optionally override the fallback target")
            print("    export MCP_FALLBACK_URL='http://192.168.1.10:8010/mcp'")
            print("")
            # re-raise so caller still sees the error if they don't follow guidance
            raise
        # Not an SSL or gateway failure: re-raise
        raise

def main() -> None:
    asyncio.run(run_tests())

if __name__ == "__main__":
    main()