Cleaner fix for RPC issue

stepup/core/rpc.py

@@ -27,6 +27,7 @@
 import pickle
 import socket
 import subprocess
+import sys
 import traceback
 from collections.abc import Awaitable, Callable, Collection
 from functools import partial
@@ -221,31 +222,25 @@ async def _serve_rpc_send_loop(
 
 async def _handle_connection(
     handler,
-    stop_events: list[asyncio.Event],
     reader: asyncio.StreamReader,
     writer: asyncio.StreamWriter,
 ):
-    stop_event = asyncio.Event()
-    stop_events.append(stop_event)
-    await serve_rpc(handler, reader, writer, stop_event)
+    await serve_rpc(handler, reader, writer)
     await writer.drain()
     writer.close()
     await writer.wait_closed()
 
 
-async def serve_socket_rpc(handler, path, stop_event):
-    # Keep a list of stop_events, including one for each handler.
-    # This works around an apparent (but difficult to isolate) issue in Python 3.11:
-    # The server context handler exits before all open requests are handled,
-    # resulting in lost connection errors.
-    # (This is not needed for Python 3.12)
-    stop_events = [stop_event]
-    server = await asyncio.start_unix_server(
-        partial(_handle_connection, handler, stop_events), path
-    )
+async def serve_socket_rpc(handler, path: str, stop_event: asyncio.Event):
+    server = await asyncio.start_unix_server(partial(_handle_connection, handler), path)
     async with server:
-        while len(stop_events) > 0:
-            await stop_events.pop().wait()
+        await stop_event.wait()
+    if sys.version_info < (3, 12, 1) and server._waiters is not None:
+        # Workaround for server.wait_closed() issue fixed in Python 3.12.1
+        # See https://github.com/python/cpython/issues/120866
+        waiter = server.get_loop().create_future()
+        server._waiters.append(waiter)
+        await waiter
 
 
 async def serve_stdio_rpc(handler):