pandaserver/pandamcp/mcp_test_client.py

0001 import argparse
0002 import asyncio
0003
0004 from fastmcp import Client
0005 from fastmcp.client.transports import SSETransport, StreamableHttpTransport
0006
0007 # argparse setup
0008 parser = argparse.ArgumentParser(description="MCP test client")
0009 parser.add_argument(
0010     "--transport",
0011     type=str,
0012     choices=("streamable-http", "sse"),
0013     default="streamable-http",
0014     help="Transport to use: streamable-http or sse (default: streamable-http)",
0015 )
0016 parser.add_argument("--host", type=str, default="localhost", help="PanDA Server host (default: localhost)")
0017 parser.add_argument("--port", type=int, default=25443, help="PanDA Server port (default: 25443)")
0018 parser.add_argument("--use_http", action=argparse.BooleanOptionalAction, default=False, help="Use HTTP instead of HTTPS (default: False)")
0019 parser.add_argument("--token", type=str, default=None, help="OIDC ID token for write-operations (default: None)")
0020 parser.add_argument("--vo", type=str, default=None, help="Virtual organization with ID token is given (default: None)")
0021 parser.add_argument("--tool", type=str, default="is_alive", help="A tool name to test (default: is_alive)")
0022 parser.add_argument(
0023     "--kv",
0024     dest="kv",
0025     metavar="KEY=VAL",
0026     action="append",
0027     help="Arguments to invoke the test function. Key=Value pairs repeat for multiple (e.g. --kv a=1 --kv b=2 --kv c=True)",
0028 )
0029 args = parser.parse_args()
0030
0031 # parse --kv KEY=VAL (repeatable) into a dict stored back on args.kv
0032 raw_kv = args.kv or []
0033 kv = {}
0034 for item in raw_kv:
0035     if "=" not in item:
0036         parser.error(f"invalid --kv value {item!r}, expected KEY=VALUE")
0037     k, v = item.split("=", 1)
0038     v = v.strip()
0039     if v == "True":
0040         v = True
0041     elif v == "False":
0042         v = False
0043     kv[k] = v
0044 args.kv = kv
0045
0046 # construct base URL
0047 if args.use_http:
0048     base_url = f"http://{args.host}:{args.port}/mcp/"
0049 else:
0050     base_url = f"https://{args.host}:{args.port}/mcp/"
0051
0052 headers = {"Origin": args.vo} if args.token else None
0053
0054 # select transport
0055 if args.transport == "streamable-http":
0056     transport = StreamableHttpTransport(url=base_url, auth=args.token, headers=headers)
0057 else:
0058     transport = SSETransport(url=base_url, auth=args.token, headers=headers)
0059
0060 # create client
0061 client = Client(transport)
0062
0063
0064 # test function
0065 async def cl():
0066     # Connection is established here
0067     async with client:
0068         if client.is_connected():
0069             print("Client connected")
0070         else:
0071             print("Client failed to connect")
0072             return
0073
0074         # Make MCP calls within the context
0075         tools = await client.list_tools()
0076         print(f"\nAvailable tools:")
0077         for tool in tools:
0078             print(f"- {tool.name} -\n")
0079             print(f" Description: {tool.description}")
0080
0081         print("\n" * 2)
0082         print(f"Testing {args.tool}:")
0083         result = await client.call_tool(args.tool, args.kv)
0084         print(f"Result: {result}\n")
0085
0086     # Connection is closed automatically here
0087     if not client.is_connected():
0088         print("Client disconnected")
0089     else:
0090         print("Client still connected")
0091     print("Done")
0092
0093
0094 if __name__ == "__main__":
0095     asyncio.run(cl())