#!/usr/bin/env python3 """Simple MCP client example with OAuth authentication support. This client connects to an MCP server using streamable HTTP transport with OAuth. """ from __future__ import annotations as _annotations import asyncio import os import socketserver import threading import time import webbrowser from http.server import BaseHTTPRequestHandler, HTTPServer from typing import Any from urllib.parse import parse_qs, urlparse import httpx from mcp.client._transport import ReadStream, WriteStream from mcp.client.auth import AuthorizationCodeResult, OAuthClientProvider, TokenStorage from mcp.client.session import ClientSession from mcp.client.sse import sse_client from mcp.client.streamable_http import streamable_http_client from mcp.shared.auth import OAuthClientInformationFull, OAuthClientMetadata, OAuthToken from mcp.shared.message import SessionMessage class InMemoryTokenStorage(TokenStorage): """Simple in-memory token storage implementation.""" def __init__(self): self._tokens: OAuthToken | None = None self._client_info: OAuthClientInformationFull | None = None async def get_tokens(self) -> OAuthToken | None: return self._tokens async def set_tokens(self, tokens: OAuthToken) -> None: self._tokens = tokens async def get_client_info(self) -> OAuthClientInformationFull | None: return self._client_info async def set_client_info(self, client_info: OAuthClientInformationFull) -> None: self._client_info = client_info class CallbackHandler(BaseHTTPRequestHandler): """Simple HTTP handler to capture OAuth callback.""" def __init__( self, request: Any, client_address: tuple[str, int], server: socketserver.BaseServer, callback_data: dict[str, Any], ): """Initialize with callback data storage.""" self.callback_data = callback_data super().__init__(request, client_address, server) def do_GET(self): """Handle GET request from OAuth redirect.""" parsed = urlparse(self.path) query_params = parse_qs(parsed.query) if "code" in query_params: self.callback_data["authorization_code"] = query_params["code"][0] self.callback_data["state"] = query_params.get("state", [None])[0] self.callback_data["iss"] = query_params.get("iss", [None])[0] self.send_response(200) self.send_header("Content-type", "text/html") self.end_headers() self.wfile.write(b"""

Authorization Successful!

You can close this window and return to the terminal.

""") elif "error" in query_params: self.callback_data["error"] = query_params["error"][0] self.send_response(400) self.send_header("Content-type", "text/html") self.end_headers() self.wfile.write( f"""

Authorization Failed

Error: {query_params["error"][0]}

You can close this window and return to the terminal.

""".encode() ) else: self.send_response(404) self.end_headers() def log_message(self, format: str, *args: Any): """Suppress default logging.""" class CallbackServer: """Simple server to handle OAuth callbacks.""" def __init__(self, port: int = 3000): self.port = port self.server = None self.thread = None self.callback_data = {"authorization_code": None, "state": None, "iss": None, "error": None} def _create_handler_with_data(self): """Create a handler class with access to callback data.""" callback_data = self.callback_data class DataCallbackHandler(CallbackHandler): def __init__( self, request: BaseHTTPRequestHandler, client_address: tuple[str, int], server: socketserver.BaseServer, ): super().__init__(request, client_address, server, callback_data) return DataCallbackHandler def start(self): """Start the callback server in a background thread.""" handler_class = self._create_handler_with_data() self.server = HTTPServer(("localhost", self.port), handler_class) self.thread = threading.Thread(target=self.server.serve_forever, daemon=True) self.thread.start() print(f"šŸ–„ļø Started callback server on http://localhost:{self.port}") def stop(self): """Stop the callback server.""" if self.server: self.server.shutdown() self.server.server_close() if self.thread: self.thread.join(timeout=1) def wait_for_callback(self, timeout: int = 300): """Wait for OAuth callback with timeout.""" start_time = time.time() while time.time() - start_time < timeout: if self.callback_data["authorization_code"]: return self.callback_data["authorization_code"] elif self.callback_data["error"]: raise Exception(f"OAuth error: {self.callback_data['error']}") time.sleep(0.1) raise Exception("Timeout waiting for OAuth callback") @property def state(self): """The received state parameter.""" return self.callback_data["state"] @property def iss(self): """The received iss parameter.""" return self.callback_data["iss"] class SimpleAuthClient: """Simple MCP client with auth support.""" def __init__( self, server_url: str, transport_type: str = "streamable-http", client_metadata_url: str | None = None, ): self.server_url = server_url self.transport_type = transport_type self.client_metadata_url = client_metadata_url self.session: ClientSession | None = None async def connect(self): """Connect to the MCP server.""" print(f"šŸ”— Attempting to connect to {self.server_url}...") try: callback_server = CallbackServer(port=3030) callback_server.start() async def callback_handler() -> AuthorizationCodeResult: """Wait for OAuth callback and return auth code, state, and iss.""" print("ā³ Waiting for authorization callback...") try: auth_code = callback_server.wait_for_callback(timeout=300) return AuthorizationCodeResult(code=auth_code, state=callback_server.state, iss=callback_server.iss) finally: callback_server.stop() client_metadata_dict = { "client_name": "Simple Auth Client", "redirect_uris": ["http://localhost:3030/callback"], "grant_types": ["authorization_code", "refresh_token"], "response_types": ["code"], } async def _default_redirect_handler(authorization_url: str) -> None: """Default redirect handler that opens the URL in a browser.""" print(f"Opening browser for authorization: {authorization_url}") webbrowser.open(authorization_url) # Create OAuth authentication handler using the new interface # Use client_metadata_url to enable CIMD when the server supports it oauth_auth = OAuthClientProvider( server_url=self.server_url.replace("/mcp", ""), client_metadata=OAuthClientMetadata.model_validate(client_metadata_dict), storage=InMemoryTokenStorage(), redirect_handler=_default_redirect_handler, callback_handler=callback_handler, client_metadata_url=self.client_metadata_url, ) # Create transport with auth handler based on transport type if self.transport_type == "sse": print("šŸ“” Opening SSE transport connection with auth...") async with sse_client( url=self.server_url, auth=oauth_auth, timeout=60.0, ) as (read_stream, write_stream): await self._run_session(read_stream, write_stream) else: print("šŸ“” Opening StreamableHTTP transport connection with auth...") async with httpx.AsyncClient(auth=oauth_auth, follow_redirects=True) as custom_client: async with streamable_http_client(url=self.server_url, http_client=custom_client) as ( read_stream, write_stream, ): await self._run_session(read_stream, write_stream) except Exception as e: print(f"āŒ Failed to connect: {e}") import traceback traceback.print_exc() async def _run_session( self, read_stream: ReadStream[SessionMessage | Exception], write_stream: WriteStream[SessionMessage], ): """Run the MCP session with the given streams.""" print("šŸ¤ Initializing MCP session...") async with ClientSession(read_stream, write_stream) as session: self.session = session print("āš” Starting session initialization...") await session.initialize() print("āœØ Session initialization complete!") print(f"\nāœ… Connected to MCP server at {self.server_url}") # Run interactive loop await self.interactive_loop() async def list_tools(self): """List available tools from the server.""" if not self.session: print("āŒ Not connected to server") return try: result = await self.session.list_tools() if hasattr(result, "tools") and result.tools: print("\nšŸ“‹ Available tools:") for i, tool in enumerate(result.tools, 1): print(f"{i}. {tool.name}") if tool.description: print(f" Description: {tool.description}") print() else: print("No tools available") except Exception as e: print(f"āŒ Failed to list tools: {e}") async def call_tool(self, tool_name: str, arguments: dict[str, Any] | None = None): """Call a specific tool.""" if not self.session: print("āŒ Not connected to server") return try: result = await self.session.call_tool(tool_name, arguments or {}) print(f"\nšŸ”§ Tool '{tool_name}' result:") if hasattr(result, "content"): for content in result.content: if content.type == "text": print(content.text) else: print(content) else: print(result) except Exception as e: print(f"āŒ Failed to call tool '{tool_name}': {e}") async def interactive_loop(self): """Run interactive command loop.""" print("\nšŸŽÆ Interactive MCP Client") print("Commands:") print(" list - List available tools") print(" call [args] - Call a tool") print(" quit - Exit the client") print() while True: try: command = input("mcp> ").strip() if not command: continue if command == "quit": break elif command == "list": await self.list_tools() elif command.startswith("call "): parts = command.split(maxsplit=2) tool_name = parts[1] if len(parts) > 1 else "" if not tool_name: print("āŒ Please specify a tool name") continue # Parse arguments (simple JSON-like format) arguments: dict[str, Any] = {} if len(parts) > 2: import json try: arguments = json.loads(parts[2]) except json.JSONDecodeError: print("āŒ Invalid arguments format (expected JSON)") continue await self.call_tool(tool_name, arguments) else: print("āŒ Unknown command. Try 'list', 'call ', or 'quit'") except KeyboardInterrupt: print("\n\nšŸ‘‹ Goodbye!") break except EOFError: break async def main(): """Main entry point.""" # Default server URL - can be overridden with environment variable # Most MCP streamable HTTP servers use /mcp as the endpoint server_url = os.getenv("MCP_SERVER_PORT", 8000) transport_type = os.getenv("MCP_TRANSPORT_TYPE", "streamable-http") client_metadata_url = os.getenv("MCP_CLIENT_METADATA_URL") server_url = ( f"http://localhost:{server_url}/mcp" if transport_type == "streamable-http" else f"http://localhost:{server_url}/sse" ) print("šŸš€ Simple MCP Auth Client") print(f"Connecting to: {server_url}") print(f"Transport type: {transport_type}") if client_metadata_url: print(f"Client metadata URL: {client_metadata_url}") # Start connection flow - OAuth will be handled automatically client = SimpleAuthClient(server_url, transport_type, client_metadata_url) await client.connect() def cli(): """CLI entry point for uv script.""" asyncio.run(main()) if __name__ == "__main__": cli()