test

Author: ihrpr

src/mcp/server/sse.py

@@ -129,8 +129,20 @@ async def sse_writer():
                 logger.debug("Starting SSE response task")
                 tg.start_soon(response, scope, receive, send)
 
-                logger.debug("Yielding read and write streams")
-                yield (read_stream, write_stream, response)
+                try:
+                    logger.debug("Yielding read and write streams")
+                    yield (read_stream, write_stream, response)
+                finally:
+                    # Cleanup when connection closes
+                    logger.debug(f"Cleaning up SSE session {session_id}")
+                    try:
+                        # Remove session from tracking dictionary
+                        if session_id in self._read_stream_writers:
+                            del self._read_stream_writers[session_id]
+                        # Cancel any remaining tasks in the task group
+                        tg.cancel_scope.cancel()
+                    except Exception as e:
+                        logger.error(f"Error during SSE cleanup: {e}")
 
     async def handle_post_message(
         self, scope: Scope, receive: Receive, send: Send

tests/shared/test_sse.py

@@ -142,12 +142,19 @@ def server(server_port: int) -> Generator[None, None, None]:
 
     yield
 
-    print("killing server")
-    # Signal the server to stop
-    proc.kill()
-    proc.join(timeout=2)
+    print("shutting down server gracefully")
+    # Try graceful shutdown first
+    proc.terminate()
+    try:
+        proc.join(timeout=5)
+    except Exception:
+        print("Graceful shutdown failed, forcing kill")
+        proc.kill()
+        proc.join(timeout=2)
+    
     if proc.is_alive():
         print("server process failed to terminate")
+        proc.kill()  # Force kill as last resort
 
 
 
@@ -180,9 +187,7 @@ async def test_raw_sse_connection(server, server_url) -> None:
         pytest.fail(f"{e}")
 
 @pytest.mark.anyio
-@pytest.mark.skip(
-    "fails in CI, but works locally. Need to investigate why."
-)
+
 async def test_sse_client_basic_connection(server: None, server_url: str) -> None:
     async with sse_client(server_url + "/sse") as streams:
         async with ClientSession(*streams) as session: