test: Add regression test for stateless request memory cleanup

felixweinberger · felixweinberger · commit 68045662d309 · 2025-07-14T17:35:13.000+01:00
This test verifies that PR #1116's fix properly cleans up transport resources for stateless requests. The test mocks StreamableHTTPServerTransport to track when _terminate_session() is called and ensures it's invoked for each stateless request to prevent memory leaks. Without the fix (commenting out the _terminate_session() call), this test fails, confirming it properly detects the memory leak issue. Github-Issue:#1116
diff --git a/tests/server/test_streamable_http_manager.py b/tests/server/test_streamable_http_manager.py
@@ -197,3 +197,86 @@ async def mock_receive():
         "Session ID should be removed from _server_instances after an exception"
     )
     assert not manager._server_instances, "No sessions should be tracked after the only session crashes"
+
+
+@pytest.mark.anyio
+async def test_stateless_requests_cleanup_transport_resources():
+    """Test that stateless requests properly clean up transport resources."""
+    from contextlib import asynccontextmanager
+    from unittest.mock import patch
+
+    app = Server("test-stateless-cleanup-server")
+    manager = StreamableHTTPSessionManager(app=app, stateless=True)
+
+    # Track created transports and their termination
+    created_transports = []
+
+    # Mock the transport class to track termination
+    with patch("mcp.server.streamable_http_manager.StreamableHTTPServerTransport") as MockTransport:
+        # Create a mock transport instance
+        def create_transport(*args, **kwargs):
+            transport = AsyncMock()
+            transport._terminated = False
+
+            # Track when terminate is called
+            async def mock_terminate():
+                transport._terminated = True
+
+            transport._terminate_session = mock_terminate
+            transport.handle_request = AsyncMock()
+
+            # Mock the connect context manager
+            @asynccontextmanager
+            async def mock_connect():
+                yield (AsyncMock(), AsyncMock())
+
+            transport.connect = mock_connect
+
+            created_transports.append(transport)
+            return transport
+
+        MockTransport.side_effect = create_transport
+
+        async with manager.run():
+            # Mock app.run to return quickly
+            mock_mcp_run = AsyncMock(return_value=None)
+            app.run = mock_mcp_run
+
+            # Send a stateless request
+            sent_messages = []
+
+            async def mock_send(message):
+                sent_messages.append(message)
+
+            scope = {
+                "type": "http",
+                "method": "POST",
+                "path": "/mcp",
+                "headers": [
+                    (b"content-type", b"application/json"),
+                    (b"accept", b"application/json, text/event-stream"),
+                ],
+            }
+
+            async def mock_receive():
+                return {
+                    "type": "http.request",
+                    "body": b'{"jsonrpc": "2.0", "method": "test", "id": 1}',
+                    "more_body": False,
+                }
+
+            # Send multiple requests
+            num_requests = 3
+            for _ in range(num_requests):
+                await manager.handle_request(scope, mock_receive, mock_send)
+                # Give async tasks time to complete
+                await anyio.sleep(0.1)
+
+            # Verify each transport was created
+            assert len(created_transports) == num_requests, (
+                f"Expected {num_requests} transports, got {len(created_transports)}"
+            )
+
+            # This is the key assertion - each transport should have been terminated
+            for i, transport in enumerate(created_transports):
+                assert transport._terminated, f"Transport {i} was not terminated"
