added streamablehttp mcp servers support

Author: kevalmahajan

mcpgateway/services/gateway_service.py

@@ -24,6 +24,7 @@
 from filelock import FileLock, Timeout
 from mcp import ClientSession
 from mcp.client.sse import sse_client
+from mcp.client.streamable_http import streamablehttp_client
 from sqlalchemy import select
 from sqlalchemy.orm import Session
 
@@ -186,7 +187,7 @@ async def register_gateway(self, db: Session, gateway: GatewayCreate) -> Gateway
             auth_type = getattr(gateway, "auth_type", None)
             auth_value = getattr(gateway, "auth_value", {})
 
-            capabilities, tools = await self._initialize_gateway(str(gateway.url), auth_value)
+            capabilities, tools = await self._initialize_gateway(str(gateway.url), auth_value, gateway.transport)
 
             all_names = [td.name for td in tools]
 
@@ -217,6 +218,7 @@ async def register_gateway(self, db: Session, gateway: GatewayCreate) -> Gateway
                 name=gateway.name,
                 url=str(gateway.url),
                 description=gateway.description,
+                transport=gateway.transport,
                 capabilities=capabilities,
                 last_seen=datetime.now(timezone.utc),
                 auth_type=auth_type,
@@ -311,6 +313,8 @@ async def update_gateway(self, db: Session, gateway_id: int, gateway_update: Gat
                 gateway.url = str(gateway_update.url)
             if gateway_update.description is not None:
                 gateway.description = gateway_update.description
+            if gateway_update.transport is not None:
+                gateway.transport = gateway_update.transport
 
             if getattr(gateway, "auth_type", None) is not None:
                 gateway.auth_type = gateway_update.auth_type
@@ -322,7 +326,7 @@ async def update_gateway(self, db: Session, gateway_id: int, gateway_update: Gat
             # Try to reinitialize connection if URL changed
             if gateway_update.url is not None:
                 try:
-                    capabilities, _ = await self._initialize_gateway(gateway.url, gateway.auth_value)
+                    capabilities, _ = await self._initialize_gateway(gateway.url, gateway.auth_value, gateway.transport)
                     gateway.capabilities = capabilities
                     gateway.last_seen = datetime.utcnow()
 
@@ -399,7 +403,7 @@ async def toggle_gateway_status(self, db: Session, gateway_id: int, activate: bo
                     self._active_gateways.add(gateway.url)
                     # Try to initialize if activating
                     try:
-                        capabilities, tools = await self._initialize_gateway(gateway.url, gateway.auth_value)
+                        capabilities, tools = await self._initialize_gateway(gateway.url, gateway.auth_value, gateway.transport)
                         gateway.capabilities = capabilities.dict()
                         gateway.last_seen = datetime.utcnow()
                     except Exception as e:
@@ -571,14 +575,20 @@ async def check_health_of_gateways(self, gateways: List[DbGateway]) -> bool:
                     headers = decode_auth(auth_data)
 
                     # Perform the GET and raise on 4xx/5xx
-                    async with client.stream("GET", gateway.url, headers=headers) as response:
-                        # This will raise immediately if status is 4xx/5xx
-                        response.raise_for_status()
+                    if (gateway.transport).lower() == "sse":
+                        async with client.stream("GET", gateway.url, headers=headers) as response:
+                            # This will raise immediately if status is 4xx/5xx
+                            response.raise_for_status()
+                    elif (gateway.transport).lower() == "streamablehttp":
+                        async with streamablehttp_client(url=gateway.url, headers=headers, timeout=5) as (read_stream, write_stream, get_session_id):
+                            async with ClientSession(read_stream, write_stream) as session:
+                                # Initialize the session
+                                response = await session.initialize()
 
                     # Mark successful check
                     gateway.last_seen = datetime.utcnow()
 
-                except Exception:
+                except Exception as e:
                     await self._handle_gateway_failure(gateway)
 
         # All gateways passed
@@ -629,7 +639,7 @@ async def subscribe_events(self) -> AsyncGenerator[Dict[str, Any], None]:
         finally:
             self._event_subscribers.remove(queue)
 
-    async def _initialize_gateway(self, url: str, authentication: Optional[Dict[str, str]] = None) -> Any:
+    async def _initialize_gateway(self, url: str, authentication: Optional[Dict[str, str]] = None, transport: str= "sse") -> Any:
         """Initialize connection to a gateway and retrieve its capabilities.
 
         Args:
@@ -676,7 +686,45 @@ async def connect_to_sse_server(server_url: str, authentication: Optional[Dict[s
 
                 return capabilities, tools
 
-            capabilities, tools = await connect_to_sse_server(url, authentication)
+            async def connect_to_streamablehttp_server(server_url: str, authentication: Optional[Dict[str, str]] = None):
+                """
+                Connect to an MCP server running with Streamable HTTP transport
+
+                Args:
+                    server_url: URL to connect to the server
+                    authentication: Authentication headers for connection to URL
+
+                Returns:
+                    list, list: List of capabilities and tools
+                """
+                if authentication is None:
+                    authentication = {}
+                # Store the context managers so they stay alive
+                decoded_auth = decode_auth(authentication)
+
+                # Use async with for both streamablehttp_client and ClientSession
+                async with streamablehttp_client(url=server_url, headers=decoded_auth) as (read_stream, write_stream, get_session_id):
+                    async with ClientSession(read_stream, write_stream) as session:
+                        # Initialize the session
+                        response = await session.initialize()
+                        # if get_session_id:
+                        #     session_id = get_session_id()
+                        #     if session_id:
+                        #         print(f"Session ID: {session_id}")
+                        capabilities = response.capabilities.model_dump(by_alias=True, exclude_none=True)
+                        response = await session.list_tools()
+                        tools = response.tools
+                        tools = [tool.model_dump(by_alias=True, exclude_none=True) for tool in tools]
+                        tools = [ToolCreate.model_validate(tool) for tool in tools]
+                        for tool in tools:
+                            tool.request_type = "STREAMABLEHTTP"
+
+                return capabilities, tools
+
+            if transport.lower() == "sse":
+                capabilities, tools = await connect_to_sse_server(url, authentication)
+            elif transport.lower() == "streamablehttp":
+                capabilities, tools = await connect_to_streamablehttp_server(url, authentication)
 
             return capabilities, tools
         except Exception as e:

mcpgateway/services/tool_service.py

@@ -25,6 +25,7 @@
 import httpx
 from mcp import ClientSession
 from mcp.client.sse import sse_client
+from mcp.client.streamable_http import streamablehttp_client
 from sqlalchemy import delete, func, not_, select
 from sqlalchemy.exc import IntegrityError
 from sqlalchemy.orm import Session
@@ -515,6 +516,7 @@ async def invoke_tool(self, db: Session, name: str, arguments: Dict[str, Any]) -
 
                 success = True
             elif tool.integration_type == "MCP":
+                transport = tool.request_type.lower()
                 gateway = db.execute(select(DbGateway).where(DbGateway.id == tool.gateway_id).where(DbGateway.is_active)).scalar_one_or_none()
                 if gateway.auth_type == "bearer":
                     headers = decode_auth(gateway.auth_value)
@@ -539,10 +541,31 @@ async def connect_to_sse_server(server_url: str) -> str:
                             tool_call_result = await session.call_tool(name, arguments)
                     return tool_call_result
 
+                async def connect_to_streamablehttp_server(server_url: str) -> str:
+                    """
+                    Connect to an MCP server running with Streamable HTTP transport
+
+                    Args:
+                        server_url (str): MCP Server URL
+
+                    Returns:
+                        str: Result of tool call
+                    """
+                    # Use async with directly to manage the context
+                    async with streamablehttp_client(url=server_url, headers=headers) as (read_stream, write_stream, get_session_id):
+                        async with ClientSession(read_stream, write_stream) as session:
+                            # Initialize the session
+                            await session.initialize()
+                            tool_call_result = await session.call_tool(name, arguments)
+                    return tool_call_result
+
                 tool_gateway_id = tool.gateway_id
                 tool_gateway = db.execute(select(DbGateway).where(DbGateway.id == tool_gateway_id).where(DbGateway.is_active)).scalar_one_or_none()
 
-                tool_call_result = await connect_to_sse_server(tool_gateway.url)
+                if transport == "sse":
+                    tool_call_result = await connect_to_sse_server(tool_gateway.url)
+                elif transport == "streamablehttp":
+                    tool_call_result = await connect_to_streamablehttp_server(tool_gateway.url)
                 content = tool_call_result.model_dump(by_alias=True).get("content", [])
 
                 success = True