Add lifespan to FastAPI app

Author: madhav165

mcpgateway/main.py

@@ -157,12 +157,70 @@
 # 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,
 )
 
 
@@ -1909,60 +1967,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

@@ -29,6 +29,7 @@
 from mcpgateway.config import settings
 from mcpgateway.db import Gateway as DbGateway
 from mcpgateway.db import Tool as DbTool
+from mcpgateway.db import SessionLocal
 from mcpgateway.schemas import GatewayCreate, GatewayRead, GatewayUpdate, ToolCreate
 from mcpgateway.services.tool_service import ToolService
 from mcpgateway.utils.services_auth import decode_auth
@@ -579,7 +580,7 @@ async def connect_to_sse_server(server_url: str, authentication: Optional[Dict[s
 
     def _get_active_gateways(self) -> list[DbGateway]:
         """Sync function for database operations (runs in thread)."""
-        with Session() as db:
+        with SessionLocal() as db:
             return db.execute(select(DbGateway).where(DbGateway.is_active)).scalars().all()
 
     async def _run_health_checks(self) -> None: