Static typing (#725)

* Mypy admin.py and streamablehttp_transport.py

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

* Mypy admin.py

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

* Mypy admin.py

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

* Cherry-pick commits

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

* Mypy mcpgateway/transports/streamablehttp_transport.py

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

* Mypy mcpgateway/transports/streamablehttp_transport.py part2

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

* Small fixes to session_registry.py and streamablehttp_transport.py

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

* Mypy cache/session_registry.py

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

* Small fixes in mcpgateway/cache/session_registry.py

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

* Mypy for session_registry

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

* Mypy streamablehttp_transport after rebase

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

* Mypy admin.py after rebase

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

* Run linters

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

* Fix unit tests

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

* Pytest fixes

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

* Flake8

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

* Apply pre-commit formatting fixes

Signed-off-by: Mihai Criveti <[email protected]>

---------

Signed-off-by: Sebastian <[email protected]>
Signed-off-by: Mihai Criveti <[email protected]>
Co-authored-by: Sebastian <[email protected]>
Co-authored-by: Mihai Criveti <[email protected]>

Author: 3 people

docs/docs/manage/bulk-import.md

@@ -235,4 +235,4 @@ The endpoint provides detailed error information for each failed tool:
 - Use the bulk import for initial setup and migrations
 - Export existing tools first to understand the schema
 - Test with a small subset before importing hundreds of tools
-- Keep your import files in version control for reproducibility
+- Keep your import files in version control for reproducibility

mcpgateway/admin.py

@@ -69,12 +69,12 @@
 from mcpgateway.utils.retry_manager import ResilientHttpClient
 
 # Initialize logging service first
-logging_service = LoggingService()
+logging_service: LoggingService = LoggingService()
 logger = logging_service.get_logger(__name__)
 
-tool_service = ToolService()
-resource_service = ResourceService()
-prompt_service = PromptService()
+tool_service: ToolService = ToolService()
+resource_service: ResourceService = ResourceService()
+prompt_service: PromptService = PromptService()
 
 try:
     # Third-Party
@@ -423,7 +423,7 @@ async def add_session(self, session_id: str, transport: SSETransport) -> None:
             # Store session in database
             try:
 
-                def _db_add():
+                def _db_add() -> None:
                     """Store session record in the database.
 
                     Creates a new SessionRecord entry in the database for tracking
@@ -520,7 +520,7 @@ async def get_session(self, session_id: str) -> Any:
         elif self._backend == "database":
             try:
 
-                def _db_check():
+                def _db_check() -> bool:
                     """Check if a session exists in the database.
 
                     Queries the SessionRecord table to determine if a session with
@@ -614,7 +614,7 @@ async def remove_session(self, session_id: str) -> None:
         elif self._backend == "database":
             try:
 
-                def _db_remove():
+                def _db_remove() -> None:
                     """Delete session record from the database.
 
                     Removes the SessionRecord entry with the specified session_id
@@ -649,7 +649,7 @@ def _db_remove():
 
         logger.info(f"Removed session: {session_id}")
 
-    async def broadcast(self, session_id: str, message: dict) -> None:
+    async def broadcast(self, session_id: str, message: Dict[str, Any]) -> None:
         """Broadcast a message to a session.
 
         Sends a message to the specified session. The behavior depends on the backend:
@@ -691,7 +691,7 @@ async def broadcast(self, session_id: str, message: dict) -> None:
             else:
                 msg_json = json.dumps(str(message))
 
-            self._session_message = {"session_id": session_id, "message": msg_json}
+            self._session_message: Dict[str, Any] = {"session_id": session_id, "message": msg_json}
 
         elif self._backend == "redis":
             try:
@@ -710,7 +710,7 @@ async def broadcast(self, session_id: str, message: dict) -> None:
                 else:
                     msg_json = json.dumps(str(message))
 
-                def _db_add():
+                def _db_add() -> None:
                     """Store message in the database for inter-process communication.
 
                     Creates a new SessionMessageRecord entry containing the session_id
@@ -791,7 +791,7 @@ def get_session_sync(self, session_id: str) -> Any:
     async def respond(
         self,
         server_id: Optional[str],
-        user: json,
+        user: Dict[str, Any],
         session_id: str,
         base_url: str,
     ) -> None:
@@ -830,7 +830,7 @@ async def respond(
             # if self._session_message:
             transport = self.get_session_sync(session_id)
             if transport:
-                message = json.loads(self._session_message.get("message"))
+                message = json.loads(str(self._session_message.get("message")))
                 await self.generate_response(message=message, transport=transport, server_id=server_id, user=user, base_url=base_url)
 
         elif self._backend == "redis":
@@ -857,7 +857,7 @@ async def respond(
 
         elif self._backend == "database":
 
-            def _db_read_session(session_id):
+            def _db_read_session(session_id: str) -> SessionRecord:
                 """Check if session still exists in the database.
 
                 Queries the SessionRecord table to verify that the session
@@ -892,7 +892,7 @@ def _db_read_session(session_id):
                 finally:
                     db_session.close()
 
-            def _db_read(session_id):
+            def _db_read(session_id: str) -> SessionMessageRecord:
                 """Read pending message for a session from the database.
 
                 Retrieves the first (oldest) unprocessed message for the given
@@ -927,7 +927,7 @@ def _db_read(session_id):
                 finally:
                     db_session.close()
 
-            def _db_remove(session_id, message):
+            def _db_remove(session_id: str, message: str) -> None:
                 """Remove processed message from the database.
 
                 Deletes a specific message record after it has been successfully
@@ -960,7 +960,7 @@ def _db_remove(session_id, message):
                 finally:
                     db_session.close()
 
-            async def message_check_loop(session_id):
+            async def message_check_loop(session_id: str) -> None:
                 """Poll database for messages and deliver to local transport.
 
                 Continuously checks the database for new messages directed to
@@ -1042,7 +1042,7 @@ async def _db_cleanup_task(self) -> None:
         while True:
             try:
                 # Clean up expired sessions every 5 minutes
-                def _db_cleanup():
+                def _db_cleanup() -> int:
                     """Remove expired sessions from the database.
 
                     Deletes all SessionRecord entries that haven't been accessed
@@ -1093,7 +1093,7 @@ def _db_cleanup():
                             continue
 
                         # Refresh session in database
-                        def _refresh_session(session_id=session_id):
+                        def _refresh_session(session_id: str = session_id) -> bool:
                             """Update session's last accessed timestamp in the database.
 
                             Refreshes the last_accessed field for an active session to
@@ -1184,7 +1184,7 @@ async def _memory_cleanup_task(self) -> None:
                 await asyncio.sleep(300)  # Sleep longer on error
 
     # Handle initialize logic
-    async def handle_initialize_logic(self, body: dict) -> InitializeResult:
+    async def handle_initialize_logic(self, body: Dict[str, Any]) -> InitializeResult:
         """Process MCP protocol initialization request.
 
         Validates the protocol version and returns server capabilities and information.
@@ -1240,14 +1240,13 @@ async def handle_initialize_logic(self, body: dict) -> InitializeResult:
                 resources={"subscribe": True, "listChanged": True},
                 tools={"listChanged": True},
                 logging={},
-                roots={"listChanged": True},
-                sampling={},
+                # roots={"listChanged": True}
             ),
             serverInfo=Implementation(name=settings.app_name, version=__version__),
             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, base_url: str):
+    async def generate_response(self, message: Dict[str, Any], transport: SSETransport, server_id: Optional[str], user: Dict[str, Any], base_url: str) -> None:
         """Generate and send response for incoming MCP protocol message.
 
         Processes MCP protocol messages and generates appropriate responses based on

mcpgateway/cache/session_registry.py

@@ -23,7 +23,7 @@
 
 # Standard
 from datetime import datetime, timezone
-from typing import Any, Dict, List, Optional
+from typing import Any, Dict, Generator, List, Optional
 import uuid
 
 # Third-Party
@@ -32,7 +32,7 @@
 from sqlalchemy.event import listen
 from sqlalchemy.exc import SQLAlchemyError
 from sqlalchemy.ext.hybrid import hybrid_property
-from sqlalchemy.orm import DeclarativeBase, Mapped, mapped_column, relationship, sessionmaker
+from sqlalchemy.orm import DeclarativeBase, Mapped, mapped_column, relationship, Session, sessionmaker
 from sqlalchemy.orm.attributes import get_history
 
 # First-Party
@@ -1276,7 +1276,7 @@ def validate_prompt_schema(mapper, connection, target):
 listen(Prompt, "before_update", validate_prompt_schema)
 
 
-def get_db():
+def get_db() -> Generator[Session, Any, None]:
     """
     Dependency to get database session.
 

mcpgateway/db.py

@@ -94,7 +94,7 @@ class LoggingService:
     - Logger name tracking
     """
 
-    def __init__(self):
+    def __init__(self) -> None:
         """Initialize logging service."""
         self._level = LogLevel.INFO
         self._subscribers: List[asyncio.Queue] = []

mcpgateway/services/logging_service.py

@@ -109,7 +109,7 @@ class ResourceService:
     - Active/inactive status management
     """
 
-    def __init__(self):
+    def __init__(self) -> None:
         """Initialize the resource service."""
         self._event_subscribers: Dict[str, List[asyncio.Queue]] = {}
         self._template_cache: Dict[str, ResourceTemplate] = {}

mcpgateway/services/resource_service.py

@@ -39,7 +39,7 @@ class RootService:
     - Root permissions and access control
     """
 
-    def __init__(self):
+    def __init__(self) -> None:
         """Initialize root service."""
         self._roots: Dict[str, Root] = {}
         self._subscribers: List[asyncio.Queue] = []

mcpgateway/services/root_service.py

@@ -151,7 +151,7 @@ class ToolService:
     - Active/inactive tool management.
     """
 
-    def __init__(self):
+    def __init__(self) -> None:
         """Initialize the tool service.
 
         Examples:

mcpgateway/services/tool_service.py

@@ -35,7 +35,7 @@
 import contextvars
 from dataclasses import dataclass
 import re
-from typing import List, Union
+from typing import Any, AsyncGenerator, List, Union
 from uuid import uuid4
 
 # Third-Party
@@ -45,6 +45,7 @@
 from mcp.server.streamable_http import EventCallback, EventId, EventMessage, EventStore, StreamId
 from mcp.server.streamable_http_manager import StreamableHTTPSessionManager
 from mcp.types import JSONRPCMessage
+from sqlalchemy.orm import Session
 from starlette.datastructures import Headers
 from starlette.responses import JSONResponse
 from starlette.status import HTTP_401_UNAUTHORIZED
@@ -62,11 +63,11 @@
 logger = logging_service.get_logger(__name__)
 
 # Initialize ToolService and MCP Server
-tool_service = ToolService()
-mcp_app = Server("mcp-streamable-http-stateless")
+tool_service: ToolService = ToolService()
+mcp_app: Server[Any] = Server("mcp-streamable-http-stateless")
 
-server_id_var: contextvars.ContextVar[str] = contextvars.ContextVar("server_id", default=None)
-request_headers_var = contextvars.ContextVar("request_headers", default={})
+server_id_var: contextvars.ContextVar[str] = contextvars.ContextVar("server_id", default="default_server_id")
+request_headers_var: contextvars.ContextVar[dict[str, Any]] = contextvars.ContextVar("request_headers", default={})
 
 # ------------------------------ Event store ------------------------------
 
@@ -305,7 +306,7 @@ async def replay_events_after(
 
 
 @asynccontextmanager
-async def get_db():
+async def get_db() -> AsyncGenerator[Session, Any]:
     """
     Asynchronous context manager for database sessions.
 
@@ -536,7 +537,7 @@ async def handle_streamable_http(self, scope: Scope, receive: Receive, send: Sen
 # ------------------------- Authentication for /mcp routes ------------------------------
 
 
-async def streamable_http_auth(scope, receive, send):
+async def streamable_http_auth(scope: Any, receive: Any, send: Any) -> bool:
     """
     Perform authentication check in middleware context (ASGI scope).
 
@@ -584,6 +585,8 @@ async def streamable_http_auth(scope, receive, send):
         if scheme.lower() == "bearer" and credentials:
             token = credentials
     try:
+        if token is None:
+            raise Exception()
         await verify_credentials(token)
     except Exception:
         response = JSONResponse(