streamable http server wise connection

Author: kevalmahajan

mcpgateway/main.py

@@ -105,8 +105,8 @@
 )
 from mcpgateway.transports.sse_transport import SSETransport
 from mcpgateway.transports.streamablehttp_transport import (
-    JWTAuthMiddlewareStreamableHttp,
     SessionManagerWrapper,
+    streamable_http_auth,
 )
 from mcpgateway.types import (
     InitializeRequest,
@@ -197,7 +197,7 @@ async def lifespan(_app: FastAPI) -> AsyncIterator[None]:
         await logging_service.initialize()
         await sampling_handler.initialize()
         await resource_cache.initialize()
-        await streamable_http_session.start()
+        await streamable_http_session.initialize()
 
         logger.info("All services initialized successfully")
         yield
@@ -262,6 +262,54 @@ async def dispatch(self, request: Request, call_next):
         return await call_next(request)
 
 
+class MCPPathRewriteMiddleware:
+    """
+    Supports requests like '/servers/<server_id>/mcp' by rewriting the path to '/mcp'.
+
+    - Only rewrites paths ending with '/mcp' but not exactly '/mcp'.
+    - Performs authentication before rewriting.
+    - Passes rewritten requests to `streamable_http_session`.
+    - All other requests are passed through unchanged.
+    """
+
+    def __init__(self, app):
+        """
+        Initialize the middleware with the ASGI application.
+
+        Args:
+            app (Callable): The next ASGI application in the middleware stack.
+        """
+        self.app = app
+
+    async def __call__(self, scope, receive, send):
+        """
+        Intercept and potentially rewrite the incoming HTTP request path.
+
+        Args:
+            scope (dict): The ASGI connection scope.
+            receive (Callable): Awaitable that yields events from the client.
+            send (Callable): Awaitable used to send events to the client.
+        """
+        # Only handle HTTP requests, HTTPS uses scope["type"] == "http" in ASGI
+        if scope["type"] != "http":
+            await self.app(scope, receive, send)
+            return
+
+        # Call auth check first
+        auth_ok = await streamable_http_auth(scope, receive, send)
+        if not auth_ok:
+            return
+
+        original_path = scope.get("path", "")
+        scope["modified_path"] = original_path
+        if (original_path.endswith("/mcp") and original_path != "/mcp") or (original_path.endswith("/mcp/") and original_path != "/mcp/"):
+            # Rewrite path so mounted app at /mcp handles it
+            scope["path"] = "/mcp"
+            await streamable_http_session.handle_streamable_http(scope, receive, send)
+            return
+        await self.app(scope, receive, send)
+
+
 # Configure CORS
 app.add_middleware(
     CORSMiddleware,
@@ -276,8 +324,9 @@ async def dispatch(self, request: Request, call_next):
 # Add custom DocsAuthMiddleware
 app.add_middleware(DocsAuthMiddleware)
 
-# Add streamable HTTP middleware for JWT auth
-app.add_middleware(JWTAuthMiddlewareStreamableHttp)
+# Add streamable HTTP middleware for /mcp routes
+app.add_middleware(MCPPathRewriteMiddleware)
+
 
 # Set up Jinja2 templates and store in app state for later use
 templates = Jinja2Templates(directory=str(settings.templates_dir))

mcpgateway/transports/streamablehttp_transport.py

@@ -9,15 +9,16 @@
 
 Key components include:
 - SessionManagerWrapper: Manages the lifecycle of streamable HTTP sessions
-- JWTAuthMiddlewareStreamableHttp: Middleware for JWT authentication
 - Configuration options for:
         1. stateful/stateless operation
         2. JSON response mode or SSE streams
 - InMemoryEventStore: A simple in-memory event storage system for maintaining session state
 
 """
 
+import contextvars
 import logging
+import re
 from collections import deque
 from contextlib import AsyncExitStack, asynccontextmanager
 from dataclasses import dataclass
@@ -38,10 +39,9 @@
 from mcp.types import JSONRPCMessage
 from starlette.datastructures import Headers
 from starlette.middleware.base import BaseHTTPMiddleware
-from starlette.requests import Request
 from starlette.responses import JSONResponse
 from starlette.status import HTTP_401_UNAUTHORIZED
-from starlette.types import ASGIApp, Receive, Scope, Send
+from starlette.types import Receive, Scope, Send
 
 from mcpgateway.config import settings
 from mcpgateway.db import SessionLocal
@@ -55,6 +55,8 @@
 tool_service = ToolService()
 mcp_app = Server("mcp-streamable-http-stateless")
 
+server_id_var: contextvars.ContextVar[str] = contextvars.ContextVar("server_id", default=None)
+
 ## ------------------------------ Event store ------------------------------
 
 
@@ -191,13 +193,24 @@ async def list_tools() -> List[types.Tool]:
         A list of Tool objects containing metadata such as name, description, and input schema.
     Logs and returns an empty list on failure.
     """
-    try:
-        async with get_db() as db:
-            tools = await tool_service.list_tools(db)
-            return [types.Tool(name=tool.name, description=tool.description, inputSchema=tool.input_schema) for tool in tools]
-    except Exception as e:
-        logger.exception("Error listing tools")
-        return []
+    server_id = server_id_var.get()
+
+    if server_id:
+        try:
+            async with get_db() as db:
+                tools = await tool_service.list_server_tools(db, server_id)
+                return [types.Tool(name=tool.name, description=tool.description, inputSchema=tool.input_schema) for tool in tools]
+        except Exception as e:
+            logger.exception("Error listing tools")
+            return []
+    else:
+        try:
+            async with get_db() as db:
+                tools = await tool_service.list_tools(db)
+                return [types.Tool(name=tool.name, description=tool.description, inputSchema=tool.input_schema) for tool in tools]
+        except Exception as e:
+            logger.exception("Error listing tools")
+            return []
 
 
 class SessionManagerWrapper:
@@ -226,7 +239,7 @@ def __init__(self) -> None:
         )
         self.stack = AsyncExitStack()
 
-    async def start(self) -> None:
+    async def initialize(self) -> None:
         """
         Starts the Streamable HTTP session manager context.
         """
@@ -250,80 +263,122 @@ async def handle_streamable_http(self, scope: Scope, receive: Receive, send: Sen
             send (Send): ASGI send callable.
         Logs any exceptions that occur during request handling.
         """
+
+        path = scope["modified_path"]
+        match = re.search(r"/servers/(?P<server_id>\d+)/mcp", path)
+
+        if match:
+            server_id = match.group("server_id")
+            server_id_var.set(server_id)
+
         try:
             await self.session_manager.handle_request(scope, receive, send)
         except Exception as e:
             logger.exception("Error handling streamable HTTP request")
             raise
 
 
-## ------------------------- FastAPI Middleware for Authentication ------------------------------
+## ------------------------- Authentication for /mcp routes ------------------------------
 
+# async def streamable_http_auth(scope, receive, send):
+#     """
+#     Perform authentication check in middleware context (ASGI scope).
 
-class JWTAuthMiddlewareStreamableHttp(BaseHTTPMiddleware):
-    """
-    Middleware for handling JWT authentication in an ASGI application.
-    This middleware checks for JWT tokens in the authorization header or cookies
-    and verifies the credentials before allowing access to protected routes.
+#     If path does not end with "/mcp", just continue (return True).
+
+#     If auth fails, sends 401 JSONResponse and returns False.
+
+#     If auth succeeds or not required, returns True.
+#     """
+
+#     path = scope.get("path", "")
+#     if not path.endswith("/mcp"):
+#         # No auth needed for other paths in this middleware usage
+#         return True
+
+#     headers = Headers(scope=scope)
+#     authorization = headers.get("authorization")
+#     cookie_header = headers.get("cookie", "")
+
+#     token = None
+#     if authorization:
+#         scheme, credentials = get_authorization_scheme_param(authorization)
+#         if scheme.lower() == "bearer" and credentials:
+#             token = credentials
+
+#     if not token:
+#         # parse cookie header manually
+#         for cookie in cookie_header.split(";"):
+#             if cookie.strip().startswith("jwt_token="):
+#                 token = cookie.strip().split("=", 1)[1]
+#                 break
+
+#     if settings.auth_required and not token:
+#         response = JSONResponse(
+#             {"detail": "Not authenticated"},
+#             status_code=HTTP_401_UNAUTHORIZED,
+#             headers={"WWW-Authenticate": "Bearer"},
+#         )
+#         await response(scope, receive, send)
+#         return False
+
+#     if token:
+#         try:
+#             await verify_credentials(token)
+#         except Exception:
+#             response = JSONResponse(
+#                 {"detail": "Authentication failed"},
+#                 status_code=HTTP_401_UNAUTHORIZED,
+#                 headers={"WWW-Authenticate": "Bearer"},
+#             )
+#             await response(scope, receive, send)
+#             return False
+
+#     return True
+
+
+async def streamable_http_auth(scope, receive, send):
     """
+    Perform authentication check in middleware context (ASGI scope).
 
-    def __init__(self, app: ASGIApp):
-        """
-        Initialize the middleware with the given ASGI application.
+    If path does not end with "/mcp", just continue (return True).
 
-        Args:
-            app (ASGIApp): The ASGI application to wrap.
-        """
-        super().__init__(app)
+    Only check Authorization header for Bearer token.
+    If no Bearer token provided, allow (return True).
 
-    async def dispatch(self, request: Request, call_next):
-        """
-        Dispatch the request to the appropriate handler after performing JWT authentication.
+    If auth_required is True and Bearer token provided, verify it.
+    If verification fails, send 401 JSONResponse and return False.
+    """
 
-        Args:
-            request (Request): The incoming request.
-            call_next: The next middleware or route handler in the chain.
+    path = scope.get("path", "")
+    if not path.endswith("/mcp") and not path.endswith("/mcp/"):
+        # No auth needed for other paths in this middleware usage
+        return True
 
-        Returns:
-            JSONResponse: A response indicating authentication failure if the token is invalid or missing.
-            Response: The response from the next middleware or route handler if authentication is successful.
-        """
-        # Only apply auth to /mcp path
-        if not request.url.path.startswith("/mcp"):
-            return await call_next(request)
-
-        headers = Headers(scope=request.scope)
-        authorization = headers.get("authorization")
-        cookie_header = headers.get("cookie", "")
-
-        token = None
-        if authorization:
-            scheme, credentials = get_authorization_scheme_param(authorization)
-            if scheme.lower() == "bearer" and credentials:
-                token = credentials
-
-        if not token:
-            for cookie in cookie_header.split(";"):
-                if cookie.strip().startswith("jwt_token="):
-                    token = cookie.strip().split("=", 1)[1]
-                    break
+    headers = Headers(scope=scope)
+    authorization = headers.get("authorization")
 
-        try:
-            if settings.auth_required and not token:
-                return JSONResponse(
-                    {"detail": "Not authenticated"},
-                    status_code=HTTP_401_UNAUTHORIZED,
-                    headers={"WWW-Authenticate": "Bearer"},
-                )
+    token = None
+    if authorization:
+        scheme, credentials = get_authorization_scheme_param(authorization)
+        if scheme.lower() == "bearer" and credentials:
+            token = credentials
 
-            if token:
-                await verify_credentials(token)
+    # # If no Bearer token in Authorization header, just allow (no auth)
+    # if not token:
+    #     return True
 
-            return await call_next(request)
+    # If token is present, verify it
+    print("TOKEN::::", token)
+    try:
+        await verify_credentials(token)
+    except Exception:
+        response = JSONResponse(
+            {"detail": "Authentication failed"},
+            status_code=HTTP_401_UNAUTHORIZED,
+            headers={"WWW-Authenticate": "Bearer"},
+        )
+        await response(scope, receive, send)
+        return False
 
-        except Exception as e:
-            return JSONResponse(
-                {"detail": "Authentication failed"},
-                status_code=HTTP_401_UNAUTHORIZED,
-                headers={"WWW-Authenticate": "Bearer"},
-            )
+    return True

tests/unit/mcpgateway/test_admin.py

@@ -109,18 +109,15 @@ def mock_request():
     )
 
     # Basic template rendering stub
-    request.app = MagicMock()           # ensure .app exists
-    request.app.state = MagicMock()     # ensure .app.state exists
+    request.app = MagicMock()  # ensure .app exists
+    request.app.state = MagicMock()  # ensure .app.state exists
     request.app.state.templates = MagicMock()
-    request.app.state.templates.TemplateResponse.return_value = HTMLResponse(
-        content="<html></html>"
-    )
+    request.app.state.templates.TemplateResponse.return_value = HTMLResponse(content="<html></html>")
 
     request.query_params = {"include_inactive": "false"}
     return request
 
 
-
 class TestAdminServerRoutes:
     """Test admin routes for server management."""
 
@@ -217,6 +214,7 @@ async def test_admin_delete_server(self, mock_delete_server, mock_request, mock_
         assert result.status_code == 303
         assert "/admin#catalog" in result.headers["location"]
 
+
 class TestAdminToolRoutes:
     """Test admin routes for tool management."""
 
@@ -602,6 +600,7 @@ async def test_admin_delete_gateway(self, mock_delete_gateway, mock_request, moc
         assert result.status_code == 303
         assert "/admin#gateways" in result.headers["location"]
 
+
 class TestAdminRootRoutes:
     """Test admin routes for root management."""
 
@@ -626,6 +625,7 @@ async def test_admin_delete_root(self, mock_remove_root, mock_request):
         assert result.status_code == 303
         assert "/admin#roots" in result.headers["location"]
 
+
 class TestAdminMetricsRoutes:
     """Test admin routes for metrics management."""