fix(asyncio): abort protocol calls on task cancellation (#3144)

Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com>

Author: Skn0tt

playwright/_impl/_connection.py

@@ -142,13 +142,21 @@ async def _inner_send(
         callback = self._connection._send_message_to_server(
             self._object, method, augmented_params, timeout
         )
-        done, _ = await asyncio.wait(
-            {
-                self._connection._transport.on_error_future,
-                callback.future,
-            },
-            return_when=asyncio.FIRST_COMPLETED,
-        )
+        try:
+            done, _ = await asyncio.wait(
+                {
+                    self._connection._transport.on_error_future,
+                    callback.future,
+                },
+                return_when=asyncio.FIRST_COMPLETED,
+            )
+        except asyncio.CancelledError as exc:
+            await self._connection._abort(
+                self._object,
+                callback,
+                str(exc) or "Task was cancelled",
+            )
+            raise
         if not callback.future.done():
             callback.future.cancel()
         result = next(iter(done)).result()
@@ -240,7 +248,10 @@ def remove_listener(self, event: str, f: Any) -> None:
 
 
 class ProtocolCallback:
-    def __init__(self, loop: asyncio.AbstractEventLoop, no_reply: bool = False) -> None:
+    def __init__(
+        self, loop: asyncio.AbstractEventLoop, id: int, no_reply: bool = False
+    ) -> None:
+        self.id = id
         self.stack_trace: traceback.StackSummary
         self.no_reply = no_reply
         self.future = loop.create_future()
@@ -389,7 +400,7 @@ def _send_message_to_server(
             )
         self._last_id += 1
         id = self._last_id
-        callback = ProtocolCallback(self._loop, no_reply=no_reply)
+        callback = ProtocolCallback(self._loop, id, no_reply=no_reply)
         task = asyncio.current_task(self._loop)
         callback.stack_trace = cast(
             traceback.StackSummary,
@@ -433,6 +444,34 @@ def _send_message_to_server(
 
         return callback
 
+    async def _abort(
+        self, object: ChannelOwner, callback: ProtocolCallback, reason: str
+    ) -> None:
+        try:
+            self._transport.send(
+                {
+                    "guid": object._guid,
+                    "method": "__abort__",
+                    "params": {"id": callback.id, "reason": reason},
+                }
+            )
+        except (Error, OSError):
+            pass
+        try:
+            done, _ = await asyncio.wait(
+                {
+                    self._transport.on_error_future,
+                    callback.future,
+                },
+                return_when=asyncio.FIRST_COMPLETED,
+            )
+        finally:
+            if not callback.future.done():
+                callback.future.cancel()
+        for future in done:
+            if not future.cancelled():
+                future.exception()
+
     def dispatch(self, msg: ParsedMessagePayload) -> None:
         if self._closed_error:
             return

tests/async/test_asyncio.py

@@ -27,7 +27,8 @@
 
 
 async def test_should_cancel_underlying_protocol_calls(
-    browser_name: str, launch_arguments: Dict
+    browser_name: str,
+    launch_arguments: Dict,
 ) -> None:
     handler_exception = None
 
@@ -40,12 +41,23 @@ def exception_handlerdler(loop: asyncio.AbstractEventLoop, context: Dict) -> Non
     async with async_playwright() as p:
         browser = await p[browser_name].launch(**launch_arguments)
         page = await browser.new_page()
-        task = asyncio.create_task(page.wait_for_selector("will-never-find"))
-        # make sure that the wait_for_selector message was sent to the server (driver)
-        await asyncio.sleep(0.1)
+        await page.set_content(
+            """
+            <button disabled onclick="window.clicked = true">click me</button>
+            <script>window.clicked = false</script>
+            """
+        )
+        task = asyncio.create_task(page.locator("button").click(timeout=0))
+        await page.wait_for_timeout(100)
+        assert not task.done()
+
         task.cancel()
         with pytest.raises(asyncio.CancelledError):
             await task
+
+        await page.locator("button").evaluate("button => button.disabled = false")
+        await page.wait_for_timeout(700)
+        assert not await page.evaluate("window.clicked")
         await browser.close()
 
     # The actual 'Future exception was never retrieved' is logged inside the Future destructor (__del__).