Fix type annotations and rename reproduction script

- Add proper type annotations to all mock classes and methods
- Fix Pyright type checking errors
- Rename test_resource_leak_reproduction.py to resource_leak_reproduction.py
  to prevent pytest from treating it as a test file
- Remove unused AsyncIterator import

Author: certainly-param

test_resource_leak_reproduction.py renamed to resource_leak_reproduction.py

@@ -13,6 +13,8 @@
 This can cause connection pool exhaustion in production.
 """
 
+from typing import Any
+
 import pytest
 
 from mcp.client.streamable_http import StreamableHTTPTransport
@@ -21,43 +23,49 @@
 class MockResponse:
     """Simple mock to track if aclose() gets called"""
 
-    def __init__(self):
+    closed: bool
+    close_count: int
+    _is_closed: bool
+
+    def __init__(self) -> None:
         self.closed = False
         self.close_count = 0
         self._is_closed = False
 
-    async def aclose(self):
+    async def aclose(self) -> None:
         self.closed = True
         self.close_count += 1
         self._is_closed = True
 
     @property
-    def is_closed(self):
+    def is_closed(self) -> bool:
         return self._is_closed
 
 
 class MockEventSource:
     """Mock that throws an exception to simulate broken SSE"""
 
-    def __init__(self, response):
+    def __init__(self, response: MockResponse) -> None:
         self.response = response
 
-    def __aiter__(self):
+    def __aiter__(self) -> "MockEventSource":
         return self
 
-    async def __anext__(self):
+    async def __anext__(self) -> Any:
         # Simulate what happens when SSE parsing fails
         raise Exception("SSE parsing failed - connection broken")
 
 
 class MockTransport(StreamableHTTPTransport):
     """Mock that shows the same bug as the real code"""
 
-    def __init__(self):
+    def __init__(self) -> None:
         super().__init__("http://test")
         self.mock_response = MockResponse()
 
-    async def _handle_sse_response(self, response, ctx, is_initialization=False):
+    async def _handle_sse_response(
+        self, response: MockResponse, ctx: Any, is_initialization: bool = False
+    ) -> None:
         """
         This mimics the actual bug in the real code.
 
@@ -66,7 +74,7 @@ async def _handle_sse_response(self, response, ctx, is_initialization=False):
         """
         try:
             event_source = MockEventSource(response)
-            async for sse in event_source:
+            async for _sse in event_source:
                 # This never runs because the exception happens first
                 is_complete = False  # Simulate event processing
                 if is_complete:
@@ -81,13 +89,13 @@ class TestStreamableHTTPResourceLeak:
     """Tests for the resource leak I found in streamable HTTP"""
 
     @pytest.mark.anyio
-    async def test_handle_sse_response_resource_leak(self):
+    async def test_handle_sse_response_resource_leak(self) -> None:
         """Test that _handle_sse_response leaks resources when SSE fails"""
         transport = MockTransport()
 
         # Create mock context
         class MockContext:
-            def __init__(self):
+            def __init__(self) -> None:
                 self.read_stream_writer = None
                 self.metadata = None
 
@@ -106,33 +114,33 @@ def __init__(self):
         )
 
     @pytest.mark.anyio
-    async def test_handle_resumption_request_resource_leak(self):
+    async def test_handle_resumption_request_resource_leak(self) -> None:
         """Test that _handle_resumption_request leaks resources when SSE fails"""
         transport = MockTransport()
 
         # Override the method to reproduce the bug
-        async def mock_handle_resumption_request(ctx):
+        async def mock_handle_resumption_request(ctx: Any) -> None:
             try:
                 # Mock aconnect_sse context manager
                 class MockEventSourceWithResponse:
-                    def __init__(self, response):
+                    def __init__(self, response: MockResponse) -> None:
                         self.response = response
 
-                    async def __aenter__(self):
+                    async def __aenter__(self) -> "MockEventSourceWithResponse":
                         return self
 
-                    async def __aexit__(self, exc_type, exc_val, exc_tb):
+                    async def __aexit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None:
                         # Even if context manager exits, the response might not be closed
                         pass
 
-                    def __aiter__(self):
+                    def __aiter__(self) -> "MockEventSourceWithResponse":
                         return self
 
-                    async def __anext__(self):
+                    async def __anext__(self) -> Any:
                         raise Exception("Resumption SSE parsing failed")
 
                 async with MockEventSourceWithResponse(transport.mock_response) as event_source:
-                    async for sse in event_source:
+                    async for _sse in event_source:
                         # This code will never be reached due to the exception
                         is_complete = False
                         if is_complete:
@@ -144,7 +152,7 @@ async def __anext__(self):
 
         # Create mock context with resumption token
         class MockResumptionContext:
-            def __init__(self):
+            def __init__(self) -> None:
                 self.read_stream_writer = None
                 self.metadata = type("obj", (object,), {"resumption_token": "test-token"})()
                 self.session_message = type(
@@ -168,23 +176,23 @@ def __init__(self):
         )
 
     @pytest.mark.anyio
-    async def test_resource_leak_fix_verification(self):
+    async def test_resource_leak_fix_verification(self) -> None:
         """Test that shows how the fix should work"""
         transport = MockTransport()
 
         # Create mock context
         class MockContext:
-            def __init__(self):
+            def __init__(self) -> None:
                 self.read_stream_writer = None
                 self.metadata = None
 
         ctx = MockContext()
 
         # Simulate the FIXED version with finally block
-        async def fixed_handle_sse_response(response, ctx, is_initialization=False):
+        async def fixed_handle_sse_response(response: MockResponse, ctx: Any, is_initialization: bool = False) -> None:
             try:
                 event_source = MockEventSource(response)
-                async for sse in event_source:
+                async for _sse in event_source:
                     # This code will never be reached due to the exception
                     is_complete = False  # Simulate event processing
                     if is_complete: