Merge pull request #44 from IBM/fix-test-tool

Fix test tool and improve stability

Author: crivetimihai

mcpgateway/cache/session_registry.py

@@ -129,9 +129,6 @@ def __init__(
             database_url: Database connection URL (required for database backend)
             session_ttl: Session time-to-live in seconds
             message_ttl: Message time-to-live in seconds
-
-        Raises:
-            ValueError: If backend is invalid or required URL is missing
         """
         super().__init__(backend=backend, redis_url=redis_url, database_url=database_url, session_ttl=session_ttl, message_ttl=message_ttl)
         self._sessions: Dict[str, Any] = {}  # Local transport cache
@@ -407,13 +404,15 @@ async def respond(
         server_id: Optional[str],
         user: json,
         session_id: str,
+        base_url: str,
     ) -> None:
         """Respond to broadcast message is transport relevant to session_id is found locally
 
         Args:
             server_id: Server ID
             session_id: Session ID
             user: User information
+            base_url: Base URL for the FastAPI request
 
         """
 
@@ -425,7 +424,7 @@ async def respond(
             transport = self.get_session_sync(session_id)
             if transport:
                 message = json.loads(self._session_message.get("message"))
-                await self.generate_response(message=message, transport=transport, server_id=server_id, user=user)
+                await self.generate_response(message=message, transport=transport, server_id=server_id, user=user, base_url=base_url)
 
         elif self._backend == "redis":
             await self._pubsub.subscribe(session_id)
@@ -440,7 +439,7 @@ async def respond(
                         message = json.loads(message)
                     transport = self.get_session_sync(session_id)
                     if transport:
-                        await self.generate_response(message=message, transport=transport, server_id=server_id, user=user)
+                        await self.generate_response(message=message, transport=transport, server_id=server_id, user=user, base_url=base_url)
             except asyncio.CancelledError:
                 logger.info(f"PubSub listener for session {session_id} cancelled")
             finally:
@@ -494,7 +493,7 @@ async def message_check_loop(session_id):
                         transport = self.get_session_sync(session_id)
                         if transport:
                             logger.info("Ready to respond")
-                            await self.generate_response(message=message, transport=transport, server_id=server_id, user=user)
+                            await self.generate_response(message=message, transport=transport, server_id=server_id, user=user, base_url=base_url)
 
                             await asyncio.to_thread(_db_remove, session_id, record.message)
 
@@ -671,7 +670,7 @@ async def handle_initialize_logic(self, body: dict) -> InitializeResult:
             instructions=("MCP Gateway providing federated tools, resources and prompts. Use /admin interface for configuration."),
         )
 
-    async def generate_response(self, message: json, transport: SSETransport, server_id: Optional[str], user: dict):
+    async def generate_response(self, message: json, transport: SSETransport, server_id: Optional[str], user: dict, base_url: str):
         """
         Generates response according to SSE specifications
 
@@ -680,6 +679,7 @@ async def generate_response(self, message: json, transport: SSETransport, server
             transport: Transport where message should be responded in
             server_id: Server ID
             user: User information
+            base_url: Base URL for the FastAPI request
 
         """
         result = {}
@@ -745,9 +745,10 @@ async def generate_response(self, message: json, transport: SSETransport, server
                     "id": 1,
                 }
                 headers = {"Authorization": f"Bearer {user['token']}", "Content-Type": "application/json"}
+                rpc_url = base_url + "/rpc"
                 async with httpx.AsyncClient(timeout=settings.federation_timeout, verify=not settings.skip_ssl_verify) as client:
                     rpc_response = await client.post(
-                        f"http://localhost:{settings.port}/rpc",
+                        url=rpc_url,
                         json=rpc_input,
                         headers=headers,
                     )

mcpgateway/db.py

@@ -1086,3 +1086,7 @@ def init_db():
         Base.metadata.create_all(bind=engine)
     except SQLAlchemyError as e:
         raise Exception(f"Failed to initialize database: {str(e)}")
+
+
+if __name__ == "__main__":
+    init_db()

mcpgateway/handlers/sampling.py

@@ -150,9 +150,9 @@ async def _add_context(self, _db: Session, messages: List[Dict[str, Any]], _cont
         """Add context to messages.
 
         Args:
-            db: Database session
+            _db: Database session
             messages: Message list
-            context_type: Context inclusion type
+            _context_type: Context inclusion type
 
         Returns:
             Messages with added context

mcpgateway/main.py

@@ -157,12 +157,71 @@
 # Initialize cache
 resource_cache = ResourceCache(max_size=settings.resource_cache_size, ttl=settings.resource_cache_ttl)
 
+
+####################
+# Startup/Shutdown #
+####################
+@asynccontextmanager
+async def lifespan(_app: FastAPI) -> AsyncIterator[None]:
+    """
+    Manage the application's startup and shutdown lifecycle.
+
+    The function initialises every core service on entry and then
+    shuts them down in reverse order on exit.
+
+    Args:
+        app (FastAPI): FastAPI app
+
+    Yields:
+        None
+
+    Raises:
+        Exception: Any unhandled error that occurs during service
+            initialisation or shutdown is re-raised to the caller.
+    """
+    logger.info("Starting MCP Gateway services")
+    try:
+        await tool_service.initialize()
+        await resource_service.initialize()
+        await prompt_service.initialize()
+        await gateway_service.initialize()
+        await root_service.initialize()
+        await completion_service.initialize()
+        await logging_service.initialize()
+        await sampling_handler.initialize()
+        await resource_cache.initialize()
+        logger.info("All services initialized successfully")
+        yield
+    except Exception as e:
+        logger.error(f"Error during startup: {str(e)}")
+        raise
+    finally:
+        logger.info("Shutting down MCP Gateway services")
+        for service in [
+            resource_cache,
+            sampling_handler,
+            logging_service,
+            completion_service,
+            root_service,
+            gateway_service,
+            prompt_service,
+            resource_service,
+            tool_service,
+        ]:
+            try:
+                await service.shutdown()
+            except Exception as e:
+                logger.error(f"Error shutting down {service.__class__.__name__}: {str(e)}")
+        logger.info("Shutdown complete")
+
+
 # Initialize FastAPI app
 app = FastAPI(
     title=settings.app_name,
     version=__version__,
     description="A FastAPI-based MCP Gateway with federation support",
     root_path=settings.app_root_path,
+    lifespan=lifespan,
 )
 
 
@@ -620,7 +679,7 @@ async def sse_endpoint(request: Request, server_id: int, user: str = Depends(req
         await session_registry.add_session(transport.session_id, transport)
         response = await transport.create_sse_response(request)
 
-        asyncio.create_task(session_registry.respond(server_id, user, session_id=transport.session_id))
+        asyncio.create_task(session_registry.respond(server_id, user, session_id=transport.session_id, base_url=base_url))
 
         tasks = BackgroundTasks()
         tasks.add_task(session_registry.remove_session, transport.session_id)
@@ -1744,7 +1803,7 @@ async def utility_sse_endpoint(request: Request, user: str = Depends(require_aut
         await transport.connect()
         await session_registry.add_session(transport.session_id, transport)
 
-        asyncio.create_task(session_registry.respond(None, user, session_id=transport.session_id))
+        asyncio.create_task(session_registry.respond(None, user, session_id=transport.session_id, base_url=base_url))
 
         response = await transport.create_sse_response(request)
         tasks = BackgroundTasks()
@@ -1909,60 +1968,6 @@ async def healthcheck(db: Session = Depends(get_db)):
     return {"status": "healthy"}
 
 
-####################
-# Startup/Shutdown #
-####################
-@asynccontextmanager
-async def lifespan() -> AsyncIterator[None]:
-    """
-    Manage the application's startup and shutdown lifecycle.
-
-    The function initialises every core service on entry and then
-    shuts them down in reverse order on exit.
-
-    Yields:
-        None
-
-    Raises:
-        Exception: Any unhandled error that occurs during service
-            initialisation or shutdown is re-raised to the caller.
-    """
-    logger.info("Starting MCP Gateway services")
-    try:
-        await tool_service.initialize()
-        await resource_service.initialize()
-        await prompt_service.initialize()
-        await gateway_service.initialize()
-        await root_service.initialize()
-        await completion_service.initialize()
-        await logging_service.initialize()
-        await sampling_handler.initialize()
-        await resource_cache.initialize()
-        logger.info("All services initialized successfully")
-        yield
-    except Exception as e:
-        logger.error(f"Error during startup: {str(e)}")
-        raise
-    finally:
-        logger.info("Shutting down MCP Gateway services")
-        for service in [
-            resource_cache,
-            sampling_handler,
-            logging_service,
-            completion_service,
-            root_service,
-            gateway_service,
-            prompt_service,
-            resource_service,
-            tool_service,
-        ]:
-            try:
-                await service.shutdown()
-            except Exception as e:
-                logger.error(f"Error shutting down {service.__class__.__name__}: {str(e)}")
-        logger.info("Shutdown complete")
-
-
 # Mount static files
 app.mount("/static", StaticFiles(directory=str(settings.static_dir)), name="static")
 

mcpgateway/services/gateway_service.py

@@ -28,6 +28,7 @@
 
 from mcpgateway.config import settings
 from mcpgateway.db import Gateway as DbGateway
+from mcpgateway.db import SessionLocal
 from mcpgateway.db import Tool as DbTool
 from mcpgateway.schemas import GatewayCreate, GatewayRead, GatewayUpdate, ToolCreate
 from mcpgateway.services.tool_service import ToolService
@@ -455,28 +456,29 @@ async def forward_request(self, gateway: DbGateway, method: str, params: Optiona
         except Exception as e:
             raise GatewayConnectionError(f"Failed to forward request to {gateway.name}: {str(e)}")
 
-    async def check_gateway_health(self, gateway: DbGateway) -> bool:
-        """Check if a gateway is healthy.
+    async def check_health_of_gateways(self, gateways: List[DbGateway]) -> bool:
+        """Health check for gateways
 
         Args:
-            gateway: Gateway to check
+            gateways: Gateways to check
 
         Returns:
             True if gateway is healthy
         """
-        if not gateway.is_active:
-            return False
+        for gateway in gateways:
+            if not gateway.is_active:
+                return False
 
-        try:
-            # Try to initialize connection
-            await self._initialize_gateway(gateway.url, gateway.auth_value)
+            try:
+                # Try to initialize connection
+                await self._initialize_gateway(gateway.url, gateway.auth_value)
 
-            # Update last seen
-            gateway.last_seen = datetime.utcnow()
-            return True
+                # Update last seen
+                gateway.last_seen = datetime.utcnow()
+                return True
 
-        except Exception:
-            return False
+            except Exception:
+                return False
 
     async def aggregate_capabilities(self, db: Session) -> Dict[str, Any]:
         """Aggregate capabilities from all gateways.
@@ -576,29 +578,24 @@ async def connect_to_sse_server(server_url: str, authentication: Optional[Dict[s
         except Exception as e:
             raise GatewayConnectionError(f"Failed to initialize gateway at {url}: {str(e)}")
 
+    def _get_active_gateways(self) -> list[DbGateway]:
+        """Sync function for database operations (runs in thread)."""
+        with SessionLocal() as db:
+            return db.execute(select(DbGateway).where(DbGateway.is_active)).scalars().all()
+
     async def _run_health_checks(self) -> None:
-        """Run periodic health checks on all gateways."""
+        """Run health checks with sync Session in async code."""
         while True:
             try:
-                async with Session() as db:
-                    # Get active gateways
-                    gateways = db.execute(select(DbGateway).where(DbGateway.is_active)).scalars().all()
-
-                    # Check each gateway
-                    for gateway in gateways:
-                        try:
-                            is_healthy = await self.check_gateway_health(gateway)
-                            if not is_healthy:
-                                logger.warning(f"Gateway {gateway.name} is unhealthy")
-                        except Exception as e:
-                            logger.error(f"Health check failed for {gateway.name}: {str(e)}")
+                # Run sync database code in a thread
+                gateways = await asyncio.to_thread(self._get_active_gateways)
 
-                    db.commit()
+                # Async health checks (non-blocking)
+                await self.check_health_of_gateways(gateways)
 
             except Exception as e:
                 logger.error(f"Health check run failed: {str(e)}")
 
-            # Wait for next check
             await asyncio.sleep(self._health_check_interval)
 
     def _get_auth_headers(self) -> Dict[str, str]:

mcpgateway/services/resource_service.py

@@ -222,8 +222,6 @@ async def list_resources(self, db: Session, include_inactive: bool = False) -> L
             db (Session): The SQLAlchemy database session.
             include_inactive (bool): If True, include inactive resources in the result.
                 Defaults to False.
-            cursor (Optional[str], optional): An opaque cursor token for pagination. Currently,
-                this parameter is ignored. Defaults to None.
 
         Returns:
             List[ResourceRead]: A list of resources represented as ResourceRead objects.
@@ -249,8 +247,6 @@ async def list_server_resources(self, db: Session, server_id: int, include_inact
             server_id (int): Server ID
             include_inactive (bool): If True, include inactive resources in the result.
                 Defaults to False.
-            cursor (Optional[str], optional): An opaque cursor token for pagination. Currently,
-                this parameter is ignored. Defaults to None.
 
         Returns:
             List[ResourceRead]: A list of resources represented as ResourceRead objects.

mcpgateway/services/tool_service.py

@@ -521,6 +521,9 @@ async def invoke_tool(self, db: Session, name: str, arguments: Dict[str, Any]) -
                 async def connect_to_sse_server(server_url: str):
                     """
                     Connect to an MCP server running with SSE transport
+
+                    Args:
+                        server_url (str): MCP Server SSE URL
                     """
                     # Use async with directly to manage the context
                     async with sse_client(url=server_url, headers=headers) as streams: