updated validate_gateway_url - progress

Signed-off-by: Satya <[email protected]>

Author: TS0713

.env.example

@@ -240,6 +240,13 @@ MAX_PROMPT_SIZE=102400
 # Timeout for rendering prompt templates (in seconds)
 PROMPT_RENDER_TIMEOUT=10
 
+#####################################
+# Gateway Validation
+#####################################
+
+# Timeout for gateway validation (in seconds)
+GATEWAY_VALIDATION_TIMEOUT=5
+
 #####################################
 # Health Checks
 #####################################

mcpgateway/config.py

@@ -280,6 +280,9 @@ def _parse_federation_peers(cls, v):
     health_check_timeout: int = 10  # seconds
     unhealthy_threshold: int = 5  # after this many failures, mark as Offline
 
+    # Validation Gateway URL
+    gateway_validation_timeout: int = 5  # seconds
+
     filelock_name: str = "gateway_service_leader.lock"
 
     # Default Roots

mcpgateway/services/gateway_service.py

@@ -244,47 +244,54 @@ def __init__(self) -> None:
         else:
             self._redis_client = None
 
-    async def _validate_gateway_url(self, url: str, headers: dict, transport_type: str, timeout=5):
+    async def _validate_gateway_url(self, url: str, headers: dict, transport_type: str, timeout: Optional[int] = None):
         """
         Validate if the given URL is a live Server-Sent Events (SSE) endpoint.
 
-        This function performs a GET request followed by a HEAD request to the provided URL
-        to ensure the endpoint is reachable and returns a valid `Content-Type` header indicating
-        Server-Sent Events (`text/event-stream`).
-
         Args:
             url (str): The full URL of the endpoint to validate.
             headers (dict): Headers to be included in the requests (e.g., Authorization).
             transport_type (str): SSE or STREAMABLEHTTP
-            timeout (int, optional): Timeout in seconds for both requests. Defaults to 5.
+            timeout (int, optional): Timeout in seconds. Defaults to settings.gateway_validation_timeout.
 
         Returns:
-            bool: True if the endpoint is reachable and supports SSE (Content-Type is
-                'text/event-stream'), otherwise False.
+            bool: True if the endpoint is reachable and supports SSE/StreamableHTTP, otherwise False.
         """
-        async with httpx.AsyncClient() as client:
-            timeout = httpx.Timeout(timeout)
-            try:
-                async with client.stream("GET", url, headers=headers, timeout=timeout) as response:
-                    response_headers = dict(response.headers)
-                    location = response_headers.get("location")
-                    content_type = response_headers.get("content-type")
-                    if transport_type == "STREAMABLEHTTP":
-                        if location:
-                            async with client.stream("GET", location, headers=headers, timeout=timeout) as response_redirect:
-                                response_headers = dict(response_redirect.headers)
-                                mcp_session_id = response_headers.get("mcp-session-id")
-                                content_type = response_headers.get("content-type")
-                                if mcp_session_id is not None and mcp_session_id != "":
-                                    if content_type is not None and content_type != "" and "application/json" in content_type:
-                                        return True
-
-                    elif transport_type == "SSE":
-                        if content_type is not None and content_type != "" and "text/event-stream" in content_type:
-                            return True
+        if timeout is None:
+            timeout = settings.gateway_validation_timeout
+        validation_client = ResilientHttpClient(client_args={"timeout": settings.gateway_validation_timeout, "verify": not settings.skip_ssl_verify})
+        try:
+            async with validation_client.client.stream("GET", url, headers=headers, timeout=timeout) as response:
+                response_headers = dict(response.headers)
+                location = response_headers.get("location")
+                content_type = response_headers.get("content-type")
+                if response.status_code in (401, 403):
+                    logger.debug(f"Authentication failed for {url} with status {response.status_code}")
                     return False
-            except Exception:
+
+                if transport_type == "STREAMABLEHTTP":
+                    if location:
+                        async with validation_client.client.stream("GET", location, headers=headers, timeout=timeout) as response_redirect:
+                            response_headers = dict(response_redirect.headers)
+                            mcp_session_id = response_headers.get("mcp-session-id")
+                            content_type = response_headers.get("content-type")
+                            if response_redirect.status_code in (401, 403):
+                                logger.debug(f"Authentication failed at redirect location {location}")
+                                return False
+                            if mcp_session_id is not None and mcp_session_id != "":
+                                if content_type is not None and content_type != "" and "application/json" in content_type:
+                                    return True
+
+                elif transport_type == "SSE":
+                    if content_type is not None and content_type != "" and "text/event-stream" in content_type:
+                        return True
                 return False
+        except Exception as e:
+            print(str(e))
+            logger.debug(f"Gateway validation failed for {url}: {str(e)}", exc_info=True)
+            return False
+        finally:
+            await validation_client.aclose()
 
     async def initialize(self) -> None:
         """Initialize the service and start health check if this instance is the leader.
@@ -1260,7 +1267,8 @@ async def connect_to_streamablehttp_server(server_url: str, authentication: Opti
                 capabilities, tools = await connect_to_streamablehttp_server(url, authentication)
 
             return capabilities, tools
-        except Exception:
+        except Exception as e:
+            logger.debug(f"Gateway initialization failed for {url}: {str(e)}", exc_info=True)
             raise GatewayConnectionError(f"Failed to initialize gateway at {url}")
 
     def _get_gateways(self, include_inactive: bool = True) -> list[DbGateway]:

tests/unit/mcpgateway/services/test_gateway_service.py

@@ -20,6 +20,7 @@
 from unittest.mock import AsyncMock, MagicMock, Mock
 
 # Third-Party
+import httpx
 import pytest
 
 # First-Party
@@ -144,7 +145,6 @@ class TestGatewayService:
     # ────────────────────────────────────────────────────────────────────
     # REGISTER
     # ────────────────────────────────────────────────────────────────────
-
     @pytest.mark.asyncio
     async def test_register_gateway(self, gateway_service, test_db):
         """Successful gateway registration populates DB and returns data."""
@@ -231,6 +231,146 @@ async def test_register_gateway_connection_error(self, gateway_service, test_db)
 
         assert "Failed to connect" in str(exc_info.value)
 
+    # ────────────────────────────────────────────────────────────────────
+    # Validate Gateway URL Timeout
+    # ────────────────────────────────────────────────────────────────────
+    @pytest.mark.asyncio
+    async def test_gateway_validate_timeout(self, gateway_service, monkeypatch):
+        # creating a mock with a timeout error
+        mock_stream = AsyncMock(side_effect=httpx.ReadTimeout("Timeout"))
+
+        mock_aclose = AsyncMock()
+
+        # Step 3: Mock client with .stream and .aclose
+        mock_client_instance = MagicMock()
+        mock_client_instance.stream = mock_stream
+        mock_client_instance.aclose = mock_aclose
+
+        mock_http_client = MagicMock()
+        mock_http_client.client = mock_client_instance
+        mock_http_client.aclose = mock_aclose
+
+        monkeypatch.setattr("mcpgateway.services.gateway_service.ResilientHttpClient", MagicMock(return_value=mock_http_client))
+
+        result = await gateway_service._validate_gateway_url(url="http://example.com", headers={}, transport_type="SSE", timeout=2)
+
+        assert result is False
+
+    # ────────────────────────────────────────────────────────────────────
+    # Validate Gateway URL SSL Verification
+    # ────────────────────────────────────────────────────────────────────
+    @pytest.mark.asyncio
+    async def test_ssl_verification_bypass(self, gateway_service, monkeypatch):
+        # TODO
+        pass
+
+    # ────────────────────────────────────────────────────────────────────
+    # Validate Gateway URL Auth Failure
+    # ────────────────────────────────────────────────────────────────────
+    @pytest.mark.asyncio
+    async def test_validate_auth_failure(self, gateway_service, monkeypatch):
+        # Mock the response object to be returned inside the async with block
+        response_mock = MagicMock()
+        response_mock.status_code = 401
+        response_mock.headers = {"content-type": "text/event-stream"}
+
+        # Create an async context manager mock that returns response_mock
+        stream_context = MagicMock()
+        stream_context.__aenter__ = AsyncMock(return_value=response_mock)
+        stream_context.__aexit__ = AsyncMock(return_value=None)
+
+        # Mock the AsyncClient to return this context manager from .stream()
+        client_mock = MagicMock()
+        client_mock.stream = AsyncMock(return_value=stream_context)
+        client_mock.aclose = AsyncMock()
+
+        # Mock ResilientHttpClient to return this client
+        resilient_client_mock = MagicMock()
+        resilient_client_mock.client = client_mock
+        resilient_client_mock.aclose = AsyncMock()
+
+        monkeypatch.setattr("mcpgateway.services.gateway_service.ResilientHttpClient", MagicMock(return_value=resilient_client_mock))
+
+        # Run the method
+        result = await gateway_service._validate_gateway_url(url="http://example.com", headers={}, transport_type="SSE")
+
+        # Expect False due to 401
+        assert result is False
+
+    # ────────────────────────────────────────────────────────────────────
+    # Validate Gateway URL Connection Error
+    # ────────────────────────────────────────────────────────────────────
+    @pytest.mark.asyncio
+    async def test_validate_connectivity_failure(self, gateway_service, monkeypatch):
+        # Create an async context manager mock that raises ConnectError
+        stream_context = AsyncMock()
+        stream_context.__aenter__.side_effect = httpx.ConnectError("connection error")
+        stream_context.__aexit__.return_value = AsyncMock()
+
+        # Mock client with .stream() and .aclose()
+        mock_client = MagicMock()
+        mock_client.stream.return_value = stream_context
+        mock_client.aclose = AsyncMock()
+
+        # Patch ResilientHttpClient to return this mock client
+        resilient_client_mock = MagicMock()
+        resilient_client_mock.client = mock_client
+        resilient_client_mock.aclose = AsyncMock()
+
+        monkeypatch.setattr("mcpgateway.services.gateway_service.ResilientHttpClient", MagicMock(return_value=resilient_client_mock))
+
+        # Call the method and assert result
+        result = await gateway_service._validate_gateway_url(url="http://example.com", headers={}, transport_type="SSE")
+
+        assert result is False
+
+    # ────────────────────────────────────────────────────────────────────
+    # Validate Gateway URL Bulk Connections Validation
+    # ────────────────────────────────────────────────────────────────────
+    @pytest.mark.asyncio
+    async def test_bulk_concurrent_validation(self, gateway_service, monkeypatch):
+        # TODO
+        pass
+
+    # ───────────────────────────────────────────────────────────────────────────
+    # Validate Gateway - StreamableHTTP with mcp-session-id & redirected-url
+    # ───────────────────────────────────────────────────────────────────────────
+    @pytest.mark.asyncio
+    async def test_streamablehttp_redirect(self, gateway_service, monkeypatch):
+        # Mock first response (redirect)
+        first_response = MagicMock()
+        first_response.status_code = 200
+        first_response.headers = {"Location": "http://sampleredirected.com"}
+
+        first_cm = AsyncMock()
+        first_cm.__aenter__.return_value = first_response
+        first_cm.__aexit__.return_value = None
+
+        # Mock redirected response (final)
+        redirected_response = MagicMock()
+        redirected_response.status_code = 200
+        redirected_response.headers = {"Mcp-Session-Id": "sample123", "Content-Type": "application/json"}
+
+        second_cm = AsyncMock()
+        second_cm.__aenter__.return_value = redirected_response
+        second_cm.__aexit__.return_value = None
+
+        # Mock ResilientHttpClient client.stream to return redirect chain
+        client_mock = MagicMock()
+        client_mock.stream = AsyncMock(side_effect=[first_cm, second_cm])
+        client_mock.aclose = AsyncMock()
+
+        resilient_http_mock = MagicMock()
+        resilient_http_mock.client = client_mock
+        resilient_http_mock.aclose = AsyncMock()
+
+        monkeypatch.setattr("mcpgateway.services.gateway_service.ResilientHttpClient", MagicMock(return_value=resilient_http_mock))
+
+        result = await gateway_service._validate_gateway_url(url="http://example.com", headers={}, transport_type="STREAMABLEHTTP")
+        # TODO
+        # assert result is True
+        pass
+
     # ────────────────────────────────────────────────────────────────────
     # LIST / GET
     # ────────────────────────────────────────────────────────────────────