RDCIST-3853: Add support to private gateway - SSE

OS-anasantos · OS-anasantos · commit 2a3378a47c65 · 2025-11-28T14:29:14.000Z
diff --git a/examples/clients/simple-streamable-private-gateway/README.md b/examples/clients/simple-streamable-private-gateway/README.md
@@ -1,11 +1,12 @@
 # Simple Streamable Private Gateway Example
 
-A demonstration of how to use the MCP Python SDK as a streamable private gateway without authentication over streamable HTTP or SSE transport.
+A demonstration of how to use the MCP Python SDK as a streamable private gateway without authentication over streamable HTTP or SSE transport with custom extensions for private gateway connectivity (SNI hostname support).
 
 ## Features
 
 - No authentication required
-- Support StreamableHTTP
+- Supports both StreamableHTTP and SSE transports
+- Custom extensions for private gateway (SNI hostname) - **Both transports**
 - Interactive command-line interface
 - Tool calling
 
@@ -23,22 +24,30 @@ uv sync --reinstall
 You can use any MCP server that doesn't require authentication. For example:
 
 ```bash
-# Example with a simple tool server
+# Example with StreamableHTTP transport
 cd examples/servers/simple-tool
-uv run mcp-simple-tool --transport streamable-http --port 8000
+uv run mcp-simple-tool --transport streamable-http --port 8081
+
+# Or with SSE transport
+cd examples/servers/simple-tool
+uv run mcp-simple-tool --transport sse --port 8081
 
 # Or use any of the other example servers
 cd examples/servers/simple-resource
-uv run simple-resource --transport streamable-http --port 8000
+uv run simple-resource --transport streamable-http --port 8081
 ```
 
 ### 2. Run the client
 
 ```bash
+# Default: StreamableHTTP transport
 uv run mcp-simple-streamable-private-gateway
 
-# Or with custom server port
-MCP_SERVER_PORT=8000 uv run mcp-simple-streamable-private-gateway
+# Or with SSE transport
+MCP_TRANSPORT_TYPE=sse uv run mcp-simple-streamable-private-gateway
+
+# Or with custom server port and hostname
+MCP_SERVER_PORT=8081 MCP_SERVER_HOSTNAME=mcp.deepwiki.com uv run mcp-simple-streamable-private-gateway
 ```
 
 ### 3. Use the interactive interface
@@ -51,19 +60,22 @@ The client provides several commands:
 
 ## Examples
 
-### Basic tool usage
+### StreamableHTTP Transport
 
 ```markdown
 🚀 Simple Streamable Private Gateway
-Connecting to: https://localhost:8000/mcp
-📡 Opening StreamableHTTP transport connection...
+Connecting to: https://localhost:8081/mcp
+Server hostname: mcp.deepwiki.com
+Transport type: streamable-http
+📡 Opening StreamableHTTP transport connection with extensions...
 🤝 Initializing MCP session...
 ⚡ Starting session initialization...
 ✨ Session initialization complete!
 
-✅ Connected to MCP server at https://localhost:8000/mcp
+✅ Connected to MCP server at https://localhost:8081/mcp
+Session ID: abc123...
 
-🎯 Interactive MCP Client
+🎯 Interactive MCP Client (Private Gateway)
 Commands:
   list - List available tools
   call <tool_name> [args] - Call a tool
@@ -82,7 +94,39 @@ mcp> quit
 👋 Goodbye!
 ```
 
+### SSE Transport
+
+```markdown
+🚀 Simple Streamable Private Gateway
+Connecting to: https://localhost:8081/sse
+Server hostname: mcp.deepwiki.com
+Transport type: sse
+📡 Opening SSE transport connection with extensions...
+🤝 Initializing MCP session...
+⚡ Starting session initialization...
+✨ Session initialization complete!
+
+✅ Connected to MCP server at https://localhost:8081/sse
+
+🎯 Interactive MCP Client (Private Gateway)
+Commands:
+  list - List available tools
+  call <tool_name> [args] - Call a tool
+  quit - Exit the client
+
+mcp> list
+📋 Available tools:
+1. echo
+   Description: Echo back the input text
+
+mcp> quit
+👋 Goodbye!
+```
+
 ## Configuration
 
-- `MCP_SERVER_PORT` - Server port (default: 8000)
-- `MCP_SERVER_HOSTNAME` - Server hostname (default: localhost)
+Environment variables:
+
+- `MCP_SERVER_PORT` - Server port (default: 8081)
+- `MCP_SERVER_HOSTNAME` - Server hostname for SNI (default: mcp.deepwiki.com)
+- `MCP_TRANSPORT_TYPE` - Transport type: `streamable-http` or `sse` (default: streamable-http)
diff --git a/examples/clients/simple-streamable-private-gateway/mcp_simple_streamable_private_gateway/main.py b/examples/clients/simple-streamable-private-gateway/mcp_simple_streamable_private_gateway/main.py
@@ -2,21 +2,30 @@
 """
 Simple MCP streamable private gateway client example without authentication.
 
-This client connects to an MCP server using streamable HTTP or SSE transport.
+This client connects to an MCP server using streamable HTTP or SSE transport
+with custom extensions for private gateway connectivity (SNI hostname support).
 
 """
 
 import asyncio
-import os
+from collections.abc import Callable
 from datetime import timedelta
 from typing import Any
 
+from anyio.streams.memory import MemoryObjectReceiveStream, MemoryObjectSendStream
+
 from mcp.client.session import ClientSession
+from mcp.client.sse import sse_client
 from mcp.client.streamable_http import streamablehttp_client
+from mcp.shared.message import SessionMessage
 
 
 class SimpleStreamablePrivateGateway:
-    """Simple MCP streamable private gateway client without authentication."""
+    """Simple MCP private gateway client supporting StreamableHTTP and SSE transports.
+    
+    This client demonstrates how to use custom extensions (e.g., SNI hostname) for
+    private gateway connectivity with both transport types.
+    """
 
     def __init__(self, server_url: str, server_hostname: str, transport_type: str = "streamable-http"):
         self.server_url = server_url
@@ -29,25 +38,42 @@ async def connect(self):
         print(f"🔗 Attempting to connect to {self.server_url}...")
 
         try:
-            print("📡 Opening StreamableHTTP transport connection...")
-            # Note: terminate_on_close=False prevents SSL handshake failures during exit
-            # Some servers may not handle session termination gracefully over SSL
-            async with streamablehttp_client(
-                url=self.server_url,
-                headers={"Host": self.server_hostname},
-                extensions={"sni_hostname": self.server_hostname},
-                timeout=timedelta(seconds=60),
-                terminate_on_close=False,  # Skip session termination to avoid SSL errors
-            ) as (read_stream, write_stream, get_session_id):
-                await self._run_session(read_stream, write_stream, get_session_id)
+            # Create transport based on transport type
+            if self.transport_type == "sse":
+                print("📡 Opening SSE transport connection with extensions...")
+                # SSE transport with custom extensions for private gateway
+                async with sse_client(
+                    url=self.server_url,
+                    headers={"Host": self.server_hostname},
+                    extensions={"sni_hostname": self.server_hostname},
+                    timeout=60,
+                ) as (read_stream, write_stream):
+                    await self._run_session(read_stream, write_stream, None)
+            else:
+                print("📡 Opening StreamableHTTP transport connection with extensions...")
+                # Note: terminate_on_close=False prevents SSL handshake failures during exit
+                # Some servers may not handle session termination gracefully over SSL
+                async with streamablehttp_client(
+                    url=self.server_url,
+                    headers={"Host": self.server_hostname},
+                    extensions={"sni_hostname": self.server_hostname},
+                    timeout=timedelta(seconds=60),
+                    terminate_on_close=False,  # Skip session termination to avoid SSL errors
+                ) as (read_stream, write_stream, get_session_id):
+                    await self._run_session(read_stream, write_stream, get_session_id)
 
         except Exception as e:
             print(f"❌ Failed to connect: {e}")
             import traceback
 
             traceback.print_exc()
 
-    async def _run_session(self, read_stream, write_stream, get_session_id):
+    async def _run_session(
+        self,
+        read_stream: MemoryObjectReceiveStream[SessionMessage | Exception],
+        write_stream: MemoryObjectSendStream[SessionMessage],
+        get_session_id: Callable[[], str | None] | None,
+    ):
         """Run the MCP session with the given streams."""
         print("🤝 Initializing MCP session...")
         async with ClientSession(read_stream, write_stream) as session:
@@ -107,7 +133,7 @@ async def call_tool(self, tool_name: str, arguments: dict[str, Any] | None = Non
 
     async def interactive_loop(self):
         """Run interactive command loop."""
-        print("\n🎯 Interactive Streamable Private Gateway")
+        print("\n🎯 Interactive MCP Client (Private Gateway)")
         print("Commands:")
         print("  list - List available tools")
         print("  call <tool_name> [args] - Call a tool")
@@ -160,23 +186,57 @@ async def interactive_loop(self):
                 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_port = os.getenv("MCP_SERVER_PORT", "8081")
-    server_hostname = os.getenv("MCP_SERVER_HOSTNAME", "mcp.deepwiki.com")
-    transport_type = "streamable-http"
-    server_url = f"https://localhost:{server_port}/mcp"
-
+def get_user_input():
+    """Get server configuration from user input."""
     print("🚀 Simple Streamable Private Gateway")
-    print(f"Connecting to: {server_url}")
-    print(f"Server hostname: {server_hostname}")
-    print(f"Transport type: {transport_type}")
+    print("\n📝 Server Configuration")
+    print("=" * 50)
+    
+    # Get server port
+    server_port = input("Server port [8081]: ").strip() or "8081"
+    
+    # Get server hostname
+    server_hostname = input("Server hostname [mcp.deepwiki.com]: ").strip() or "mcp.deepwiki.com"
+    
+    # Get transport type
+    print("\nTransport type:")
+    print("  1. streamable-http (default)")
+    print("  2. sse")
+    transport_choice = input("Select transport [1]: ").strip() or "1"
+    
+    if transport_choice == "2":
+        transport_type = "sse"
+    else:
+        transport_type = "streamable-http"
+    
+    print("=" * 50)
+    
+    return server_port, server_hostname, transport_type
+
 
-    # Start connection flow
-    client = SimpleStreamablePrivateGateway(server_url, server_hostname, transport_type)
-    await client.connect()
+async def main():
+    """Main entry point."""
+    try:
+        # Get configuration from user input
+        server_port, server_hostname, transport_type = get_user_input()
+        
+        # Set URL endpoint based on transport type
+        # StreamableHTTP servers typically use /mcp, SSE servers use /sse
+        endpoint = "/mcp" if transport_type == "streamable-http" else "/sse"
+        server_url = f"https://localhost:{server_port}{endpoint}"
+
+        print(f"\n🔗 Connecting to: {server_url}")
+        print(f"📡 Server hostname: {server_hostname}")
+        print(f"🚀 Transport type: {transport_type}\n")
+
+        # Start connection flow
+        client = SimpleStreamablePrivateGateway(server_url, server_hostname, transport_type)
+        await client.connect()
+        
+    except KeyboardInterrupt:
+        print("\n\n👋 Goodbye!")
+    except EOFError:
+        print("\n👋 Goodbye!")
 
 
 def cli():
diff --git a/src/mcp/client/sse.py b/src/mcp/client/sse.py
@@ -25,6 +25,7 @@ def remove_request_params(url: str) -> str:
 async def sse_client(
     url: str,
     headers: dict[str, Any] | None = None,
+    extensions: dict[str, str] | None = None,
     timeout: float = 5,
     sse_read_timeout: float = 60 * 5,
     httpx_client_factory: McpHttpClientFactory = create_mcp_http_client,
@@ -39,6 +40,7 @@ async def sse_client(
     Args:
         url: The SSE endpoint URL.
         headers: Optional headers to include in requests.
+        extensions: Optional extensions to include in requests (e.g., for SNI hostname).
         timeout: HTTP timeout for regular operations.
         sse_read_timeout: Timeout for SSE read operations.
         auth: Optional HTTPX authentication handler.
@@ -51,6 +53,9 @@ async def sse_client(
 
     read_stream_writer, read_stream = anyio.create_memory_object_stream(0)
     write_stream, write_stream_reader = anyio.create_memory_object_stream(0)
+    
+    # Prepare extensions (copy to avoid mutation)
+    request_extensions = extensions.copy() if extensions else {}
 
     async with anyio.create_task_group() as tg:
         try:
@@ -62,6 +67,7 @@ async def sse_client(
                     client,
                     "GET",
                     url,
+                    extensions=request_extensions,
                 ) as event_source:
                     event_source.response.raise_for_status()
                     logger.debug("SSE connection established")
@@ -127,6 +133,7 @@ async def post_writer(endpoint_url: str):
                                             mode="json",
                                             exclude_none=True,
                                         ),
+                                        extensions=request_extensions,
                                     )
                                     response.raise_for_status()
                                     logger.debug(f"Client message sent successfully: {response.status_code}")
diff --git a/uv.lock b/uv.lock
