feat: Add method to extract session ID from headers with hyphen and underscore support (#601)

jirispilka · web-flow · commit 47737007cdd3 · 2025-12-22T12:23:23.000+01:00
diff --git a/templates/python-mcp-empty/requirements.txt b/templates/python-mcp-empty/requirements.txt
@@ -3,5 +3,5 @@
 
 apify < 4.0.0
 apify-client < 3.0.0
-fastmcp>=0.2.0
-mcp>=1.16.0
+fastmcp>=2.14.0
+mcp>=1.25.0
diff --git a/templates/python-mcp-empty/src/main.py b/templates/python-mcp-empty/src/main.py
@@ -1,41 +1,34 @@
 """Main entry point for the MCP Server Actor."""
 
+import asyncio
 import os
+import time
+from collections.abc import Mapping, MutableMapping
+from typing import Any
 
+import uvicorn
 from apify import Actor
 from fastmcp import FastMCP
-
-# Initialize the Apify Actor environment
-# This call configures the Actor for its environment and should be called at startup
+from starlette.requests import Request
+from starlette.types import Receive, Scope, Send
 
 
 def get_server() -> FastMCP:
-    """Create an MCP server with implementation details."""
+    """Create an MCP server with tools and resources."""
     server = FastMCP('python-mcp-empty', '1.0.0')
 
     @server.tool()  # type: ignore[misc]
     def add(a: float, b: float) -> dict:
-        """Add two numbers together and return the sum with structured output.
-
-        Args:
-            a: First number to add
-            b: Second number to add
-
-        Returns:
-            Dictionary with the sum result and structured output
-        """
-        # Note: We can't await here in sync context, so charging happens in async wrapper
-        sum_result = a + b
-        structured_content = {
-            'result': sum_result,
-            'operands': {'a': a, 'b': b},
-            'operation': 'addition',
-        }
-
+        """Add two numbers together and return the sum."""
+        result = a + b
         return {
             'type': 'text',
-            'text': f'The sum of {a} and {b} is {sum_result}',
-            'structuredContent': structured_content,
+            'text': f'The sum of {a} and {b} is {result}',
+            'structuredContent': {
+                'result': result,
+                'operands': {'a': a, 'b': b},
+                'operation': 'addition',
+            },
         }
 
     @server.resource(uri='https://example.com/calculator', name='calculator-info')  # type: ignore[misc]
@@ -46,39 +39,150 @@ def calculator_info() -> str:
     return server
 
 
+def get_session_id(headers: Mapping[str, str]) -> str | None:
+    """Extract session ID from request headers."""
+    for key in ('mcp-session-id', 'mcp_session_id'):
+        if value := headers.get(key):
+            return value
+    return None
+
+
+class SessionTrackingMiddleware:
+    """ASGI middleware that tracks MCP sessions and closes idle ones."""
+
+    def __init__(self, app: Any, port: int, timeout_secs: int) -> None:
+        self.app = app
+        self.port = port
+        self.timeout_secs = timeout_secs
+
+        # Session tracking state
+        self._last_activity: dict[str, float] = {}
+        self._timers: dict[str, asyncio.Task[None]] = {}
+
+    def _session_cleanup(self, sid: str) -> None:
+        self._last_activity.pop(sid, None)
+        if (timer := self._timers.pop(sid, None)) and not timer.done():
+            timer.cancel()
+
+    def _touch(self, sid: str) -> None:
+        self._last_activity[sid] = time.time()
+
+        # Cancel existing timer
+        if (timer := self._timers.get(sid)) and not timer.done():
+            timer.cancel()
+
+        async def close_if_idle() -> None:
+            try:
+                await asyncio.sleep(self.timeout_secs)
+
+                # Check if activity occurred during sleep
+                elapsed = time.time() - self._last_activity.get(sid, 0)
+                if elapsed < self.timeout_secs * 0.9:
+                    return
+
+                Actor.log.info(f'Closing idle session: {sid}')
+
+                # Send internal DELETE request to close session
+                scope: Scope = {
+                    'type': 'http',
+                    'http_version': '1.1',
+                    'method': 'DELETE',
+                    'scheme': 'http',
+                    'path': '/mcp',
+                    'raw_path': b'/mcp',
+                    'query_string': b'',
+                    'headers': [(b'mcp-session-id', sid.encode())],
+                    'server': ('127.0.0.1', self.port),
+                    'client': ('127.0.0.1', 0),
+                    '_idle_close': True,
+                }
+
+                async def noop_receive() -> MutableMapping[str, Any]:
+                    return {'type': 'http.request', 'body': b'', 'more_body': False}
+
+                async def noop_send(_: MutableMapping[str, Any]) -> None:
+                    pass
+
+                # Re-enter middleware with an internal DELETE; _idle_close will skip tracking
+                await self(scope, noop_receive, noop_send)
+                self._session_cleanup(sid)
+
+            except asyncio.CancelledError:
+                pass
+            except Exception as e:
+                Actor.log.exception(f'Failed to close idle session {sid}: {e}')
+
+        self._timers[sid] = asyncio.create_task(close_if_idle())
+
+    async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None:
+        """ASGI entry point that wraps the underlying app."""
+        # Pass through non-MCP requests
+        path = scope.get('path', '')
+        if scope.get('type') != 'http' or path not in ('/mcp', '/mcp/'):
+            await self.app(scope, receive, send)
+            return
+
+        # Skip tracking for internal idle-close requests
+        if scope.get('_idle_close'):
+            await self.app(scope, receive, send)
+            return
+
+        request = Request(scope, receive)
+        sid = get_session_id(request.headers)
+        is_delete = scope.get('method') == 'DELETE'
+
+        # Track activity for existing sessions (skip DELETE)
+        if sid and not is_delete:
+            self._touch(sid)
+
+        # Capture new session ID from response headers
+        new_sid: str | None = None
+
+        async def capture_send(msg: MutableMapping[str, Any]) -> None:
+            nonlocal new_sid
+            if msg.get('type') == 'http.response.start':
+                for k, v in msg.get('headers', []):
+                    if k.decode().lower() == 'mcp-session-id':
+                        new_sid = v.decode()
+                        break
+            await send(msg)
+
+        await self.app(scope, receive, capture_send)
+
+        # Track a newly created session
+        if not sid and new_sid:
+            Actor.log.info(f'New session: {new_sid}')
+            self._touch(new_sid)
+
+        # Cleanup on explicit DELETE
+        if is_delete and sid:
+            Actor.log.info(f'Session closed: {sid}')
+            self._session_cleanup(sid)
+
+
 async def main() -> None:
-    """Run the MCP Server Actor.
-
-    This function:
-    1. Initializes the Actor
-    2. Creates and configures the MCP server
-    3. Starts the HTTP server with Streamable HTTP transport
-    4. Handles MCP requests
-    """
+    """Run the MCP Server Actor with session timeout handling."""
     await Actor.init()
 
-    # Get port from environment or default to 3000
     port = int(os.environ.get('APIFY_CONTAINER_PORT', '3000'))
+    timeout_secs = int(os.environ.get('SESSION_TIMEOUT_SECS', '300'))
 
     server = get_server()
+    app = server.http_app(transport='streamable-http')
+
+    # Wrap the app with session tracking middleware to handle idle timeouts
+    app = SessionTrackingMiddleware(app=app, port=port, timeout_secs=timeout_secs)
 
     try:
-        Actor.log.info('Starting MCP server with FastMCP')
-
-        # Start the FastMCP server with HTTP transport
-        # This starts the server on the specified port and handles MCP protocol messages
-        await server.run_http_async(
-            host='0.0.0.0',  # noqa: S104 - Required for container networking
-            port=port,
-        )
+        Actor.log.info(f'Starting MCP server on port {port} (session timeout: {timeout_secs}s)')
+        config = uvicorn.Config(app, host='0.0.0.0', port=port, log_level='info')  # noqa: S104
+        await uvicorn.Server(config).serve()
     except KeyboardInterrupt:
-        Actor.log.info('Shutting down server...')
-    except Exception as error:
-        Actor.log.error(f'Server failed to start: {error}')
+        Actor.log.info('Shutting down...')
+    except Exception as e:
+        Actor.log.error(f'Server failed: {e}')
         raise
 
 
 if __name__ == '__main__':
-    import asyncio
-
     asyncio.run(main())
diff --git a/templates/python-mcp-proxy/requirements.txt b/templates/python-mcp-proxy/requirements.txt
@@ -6,7 +6,7 @@ apify-client < 3.0.0
 arxiv-mcp-server==0.3.1
 fastapi==0.119.1
 httpx>=0.24.0
-mcp==1.23.0
+mcp>=1.25.0
 pydantic>=2.0.0
 uv>=0.7.8
 uvicorn>=0.27.0
diff --git a/templates/python-mcp-proxy/src/server.py b/templates/python-mcp-proxy/src/server.py
@@ -263,6 +263,26 @@ async def capturing_send(message: dict[str, Any]) -> None:
 
         return capturing_send
 
+    @staticmethod
+    def _get_session_id_from_headers(headers: Any) -> str | None:
+        """Extract session ID from headers, trying both hyphen and underscore variants.
+
+        HTTP headers are case-insensitive per spec, and Starlette's Request.headers
+        handles this automatically. We only need to check for hyphen vs underscore variants.
+
+        Args:
+            headers: Either a Starlette Request.headers object or a dict
+
+        Returns:
+            Session ID string if found, None otherwise
+        """
+        # Try both hyphen and underscore variants
+        # Case is handled automatically by Starlette's case-insensitive headers
+        for key in ('mcp-session-id', 'mcp_session_id'):
+            if value := headers.get(key):
+                return value  # type: ignore[no-any-return]
+        return None
+
     async def create_starlette_app(self, mcp_server: Server) -> Starlette:
         """Create a Starlette app that exposes /mcp endpoint for Streamable HTTP transport."""
         event_store = InMemoryEventStore()
@@ -341,7 +361,7 @@ async def handle_streamable_http(scope: Scope, receive: Receive, send: Send) ->
 
             if scope['method'] == 'DELETE':
                 await session_manager.handle_request(scope, receive, send)
-                if req_sid := request.headers.get('mcp-session-id'):
+                if req_sid := self._get_session_id_from_headers(request.headers):
                     self._cleanup_session_last_activity(req_sid)
                     self._cleanup_session_timer(req_sid)
                 return
@@ -352,7 +372,7 @@ async def handle_streamable_http(scope: Scope, receive: Receive, send: Send) ->
             capturing_send = self._create_capturing_send(send, session_id_from_resp)
 
             # Log and touch existing session if present on request
-            if req_sid := request.headers.get('mcp-session-id'):
+            if req_sid := self._get_session_id_from_headers(request.headers):
                 self._touch_session(req_sid, session_manager)
 
             await session_manager.handle_request(scope, receive, capturing_send)  # type: ignore[arg-type]
