fix: proxy resource management and test cleanup

- Remove stream closing from proxy (streams are owned by transport context managers)

- Add proper test stream cleanup in all test cases

- Fix test_proxy_cleans_up_streams to match actual behavior

- Fix formatting (add missing newlines)

Author: dgenio

examples/servers/simple-proxy/mcp_simple_proxy/__init__.py

@@ -3,5 +3,3 @@
 from .server import main
 
 __all__ = ["main"]
-
-

examples/servers/simple-proxy/mcp_simple_proxy/__main__.py

@@ -6,5 +6,3 @@
 
 if __name__ == "__main__":
     sys.exit(main())
-
-

examples/servers/simple-proxy/mcp_simple_proxy/server.py

@@ -109,5 +109,3 @@ async def run_proxy():
 
 if __name__ == "__main__":
     sys.exit(main())
-
-

src/mcp/shared/proxy.py

@@ -135,33 +135,26 @@ async def forward_messages(
                 onerror(exc)
 
     async with anyio.create_task_group() as tg:
-        try:
-            # Start forwarding tasks in both directions
-            tg.start_soon(
-                forward_messages,
-                client_read,
-                server_write,
-                on_client_message,
-                "client→server",
-            )
-            tg.start_soon(
-                forward_messages,
-                server_read,
-                client_write,
-                on_server_message,
-                "server→client",
-            )
+        # Start forwarding tasks in both directions
+        tg.start_soon(
+            forward_messages,
+            client_read,
+            server_write,
+            on_client_message,
+            "client→server",
+        )
+        tg.start_soon(
+            forward_messages,
+            server_read,
+            client_write,
+            on_server_message,
+            "server→client",
+        )
 
+        try:
             yield
-
         finally:
-            # Cancel all forwarding tasks
+            # Cancel all forwarding tasks when the context exits
+            # Note: We don't close the streams here because they're owned by the caller
+            # (the transport context managers that created them)
             tg.cancel_scope.cancel()
-
-            # Close all streams to ensure proper cleanup
-            # Note: We close streams even if they might already be closed
-            for stream in [client_read, client_write, server_read, server_write]:
-                try:
-                    await stream.aclose()
-                except Exception as exc:
-                    logger.debug(f"Error closing stream during cleanup: {exc}")

tests/shared/test_proxy.py

@@ -22,9 +22,7 @@ def _create():
         server_write, server_write_reader = anyio.create_memory_object_stream[SessionMessage](10)
 
         # Track all streams for cleanup
-        streams_to_cleanup.extend(
-            [client_read_writer, client_write_reader, server_read_writer, server_write_reader]
-        )
+        streams_to_cleanup.extend([client_read_writer, client_write_reader, server_read_writer, server_write_reader])
 
         return (
             (client_read, client_write),
@@ -62,6 +60,10 @@ async def test_proxy_forwards_client_to_server(create_streams):
             assert received.message.root.id == "1"
             assert received.message.root.method == "test_method"
 
+    # Clean up test streams
+    await client_read_writer.aclose()
+    await server_write_reader.aclose()
+
 
 @pytest.mark.anyio
 async def test_proxy_forwards_server_to_client(create_streams):
@@ -82,6 +84,10 @@ async def test_proxy_forwards_server_to_client(create_streams):
             assert received.message.root.id == "2"
             assert received.message.root.method == "server_method"
 
+    # Clean up test streams
+    await server_read_writer.aclose()
+    await client_write_reader.aclose()
+
 
 @pytest.mark.anyio
 async def test_proxy_bidirectional_forwarding(create_streams):
@@ -118,6 +124,12 @@ async def test_proxy_bidirectional_forwarding(create_streams):
             received_at_client = await client_write_reader.receive()
             assert received_at_client.message.root.id == "server_1"
 
+    # Clean up test streams
+    await client_read_writer.aclose()
+    await client_write_reader.aclose()
+    await server_read_writer.aclose()
+    await server_write_reader.aclose()
+
 
 @pytest.mark.anyio
 async def test_proxy_message_transformation(create_streams):
@@ -146,6 +158,10 @@ async def transform_client_message(msg: SessionMessage) -> SessionMessage | None
             assert received.message.root.id == "transformed_original"
             assert "client" in transformed
 
+    # Clean up test streams
+    await client_read_writer.aclose()
+    await server_write_reader.aclose()
+
 
 @pytest.mark.anyio
 async def test_proxy_message_dropping(create_streams):
@@ -168,11 +184,15 @@ async def drop_message(msg: SessionMessage) -> SessionMessage | None:
             # If we get here, the message was not dropped
             pytest.fail("Message should have been dropped")
 
+    # Clean up test streams
+    await client_read_writer.aclose()
+    await server_write_reader.aclose()
+
 
 @pytest.mark.anyio
 async def test_proxy_error_handling(create_streams):
     """Test that errors are caught and onerror callback is invoked."""
-    client_streams, server_streams, (client_read_writer, _), (_, _) = create_streams()
+    client_streams, server_streams, (client_read_writer, _), (_, server_write_reader) = create_streams()
 
     errors = []
 
@@ -194,6 +214,10 @@ def error_handler(error: Exception) -> None:
         assert isinstance(errors[0], ValueError)
         assert str(errors[0]) == "Test error"
 
+    # Clean up test streams
+    await client_read_writer.aclose()
+    await server_write_reader.aclose()
+
 
 @pytest.mark.anyio
 async def test_proxy_continues_after_error(create_streams):
@@ -222,6 +246,10 @@ def error_handler(error: Exception) -> None:
         # Error should have been captured
         assert len(errors) == 1
 
+    # Clean up test streams
+    await client_read_writer.aclose()
+    await server_write_reader.aclose()
+
 
 @pytest.mark.anyio
 async def test_proxy_transform_error_handling(create_streams):
@@ -255,6 +283,10 @@ async def failing_transform(msg: SessionMessage) -> SessionMessage | None:
             received = await server_write_reader.receive()
             pytest.fail("Message should not have been forwarded after transform error")
 
+    # Clean up test streams
+    await client_read_writer.aclose()
+    await server_write_reader.aclose()
+
 
 @pytest.mark.anyio
 async def test_proxy_cleans_up_streams(create_streams):
@@ -306,4 +338,6 @@ async def test_proxy_multiple_messages(create_streams):
                 assert received.message.root.id == str(i)
                 assert received.message.root.method == f"method_{i}"
 
-
+    # Clean up test streams
+    await client_read_writer.aclose()
+    await server_write_reader.aclose()