From ae20e48406b310e19bcd1668dbb822a7d30990da Mon Sep 17 00:00:00 2001 From: Copilot <223556219+Copilot@users.noreply.github.com> Date: Wed, 19 Aug 2026 01:40:43 -0700 Subject: [PATCH 1/3] FIX preserve WebSocket cleanup and reference lock Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> --- doc/references.bib | 3 +- pyrit/prompt_target/websocket_target.py | 10 ++++++- .../target/test_websocket_target.py | 28 +++++++++++++++++++ 3 files changed, 39 insertions(+), 2 deletions(-) diff --git a/doc/references.bib b/doc/references.bib index 5d9475d354..3046e3a725 100644 --- a/doc/references.bib +++ b/doc/references.bib @@ -78,7 +78,8 @@ @inproceedings{mazeika2023tdc author = {Mantas Mazeika and Andy Zou and Norman Mu and Long Phan and Zifan Wang and Chunru Yu and Adam Khoja and Fengqing Jiang and Aidan O'Gara and Zhen Xiang and Arezoo Rajabi and Dan Hendrycks and Radha Poovendran and Bo Li and David Forsyth}, booktitle = {NeurIPS Competition Track}, year = {2023}, - url = {https://neurips.cc/virtual/2023/competition/66583}, + url = {https://proceedings.mlr.press/v220/mazeika23a.html}, + note = {NeurIPS Trojan Detection Challenge series. Official challenge site may be unavailable; using the proceedings page for archival access.}, } @misc{promptfoo2025ccp, diff --git a/pyrit/prompt_target/websocket_target.py b/pyrit/prompt_target/websocket_target.py index b2ab2be695..9705f165d0 100644 --- a/pyrit/prompt_target/websocket_target.py +++ b/pyrit/prompt_target/websocket_target.py @@ -188,13 +188,21 @@ async def cleanup_conversation_async(self, conversation_id: str) -> None: Args: conversation_id (str): PyRIT conversation ID. + + Raises: + asyncio.CancelledError: If cleanup is cancelled after the connection has finished closing. """ conversation_lock = self._conversation_locks.setdefault(conversation_id, asyncio.Lock()) async with conversation_lock: websocket = self._existing_conversation.pop(conversation_id, None) if websocket is None: return - await websocket.close() + close_future = asyncio.ensure_future(websocket.close()) + try: + await asyncio.shield(close_future) + except asyncio.CancelledError: + await close_future + raise logger.info("Disconnected WebSocket conversation: %s", conversation_id) async def cleanup_target_async(self) -> None: diff --git a/tests/unit/prompt_target/target/test_websocket_target.py b/tests/unit/prompt_target/target/test_websocket_target.py index 281dc5d5cd..f853a0e476 100644 --- a/tests/unit/prompt_target/target/test_websocket_target.py +++ b/tests/unit/prompt_target/target/test_websocket_target.py @@ -588,6 +588,34 @@ async def test_cleanup_conversation_async_does_not_retain_unknown_lock(websocket assert "missing" not in websocket_target._conversation_locks +async def test_cleanup_conversation_async_cancellation_finishes_closing_connection( + websocket_target: WebsocketTarget, +) -> None: + connection = AsyncMock(spec=ClientConnection) + websocket_target._existing_conversation["conversation"] = connection + close_started = asyncio.Event() + finish_close = asyncio.Event() + + async def close_connection() -> None: + close_started.set() + await finish_close.wait() + + connection.close.side_effect = close_connection + cleanup_task = asyncio.create_task(websocket_target.cleanup_conversation_async("conversation")) + await close_started.wait() + + cleanup_task.cancel() + await asyncio.sleep(0) + assert not cleanup_task.done() + + finish_close.set() + with pytest.raises(asyncio.CancelledError): + await cleanup_task + + connection.close.assert_awaited_once() + assert websocket_target._existing_conversation == {} + + async def test_cleanup_target_async_attempts_every_connection(websocket_target: WebsocketTarget) -> None: failing_connection = AsyncMock(spec=ClientConnection) failing_connection.close.side_effect = RuntimeError("close failed") From 3b25a90ae4da7d183efc47991635a17a96a4efa4 Mon Sep 17 00:00:00 2001 From: Copilot <223556219+Copilot@users.noreply.github.com> Date: Wed, 19 Aug 2026 05:58:17 -0700 Subject: [PATCH 2/3] DOC retain NeurIPS TDC link Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Copilot-Session: 5cd848f4-78c8-4b36-83c6-70f07c323896 --- doc/references.bib | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/doc/references.bib b/doc/references.bib index 3046e3a725..5d9475d354 100644 --- a/doc/references.bib +++ b/doc/references.bib @@ -78,8 +78,7 @@ @inproceedings{mazeika2023tdc author = {Mantas Mazeika and Andy Zou and Norman Mu and Long Phan and Zifan Wang and Chunru Yu and Adam Khoja and Fengqing Jiang and Aidan O'Gara and Zhen Xiang and Arezoo Rajabi and Dan Hendrycks and Radha Poovendran and Bo Li and David Forsyth}, booktitle = {NeurIPS Competition Track}, year = {2023}, - url = {https://proceedings.mlr.press/v220/mazeika23a.html}, - note = {NeurIPS Trojan Detection Challenge series. Official challenge site may be unavailable; using the proceedings page for archival access.}, + url = {https://neurips.cc/virtual/2023/competition/66583}, } @misc{promptfoo2025ccp, From 390966eac0ae2aeb15eac91cd16eec889f291f47 Mon Sep 17 00:00:00 2001 From: Copilot <223556219+Copilot@users.noreply.github.com> Date: Thu, 20 Aug 2026 15:44:57 -0700 Subject: [PATCH 3/3] FIX preserve cancellation when WebSocket close fails Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Copilot-Session: 5cd848f4-78c8-4b36-83c6-70f07c323896 --- pyrit/prompt_target/websocket_target.py | 7 +++-- .../target/test_websocket_target.py | 29 +++++++++++++++++++ 2 files changed, 34 insertions(+), 2 deletions(-) diff --git a/pyrit/prompt_target/websocket_target.py b/pyrit/prompt_target/websocket_target.py index 9705f165d0..d8fb94a1b4 100644 --- a/pyrit/prompt_target/websocket_target.py +++ b/pyrit/prompt_target/websocket_target.py @@ -200,8 +200,11 @@ async def cleanup_conversation_async(self, conversation_id: str) -> None: close_future = asyncio.ensure_future(websocket.close()) try: await asyncio.shield(close_future) - except asyncio.CancelledError: - await close_future + except asyncio.CancelledError as cancellation_error: + try: + await close_future + except BaseException as close_error: + raise cancellation_error from close_error raise logger.info("Disconnected WebSocket conversation: %s", conversation_id) diff --git a/tests/unit/prompt_target/target/test_websocket_target.py b/tests/unit/prompt_target/target/test_websocket_target.py index f853a0e476..0b762a93ab 100644 --- a/tests/unit/prompt_target/target/test_websocket_target.py +++ b/tests/unit/prompt_target/target/test_websocket_target.py @@ -616,6 +616,35 @@ async def close_connection() -> None: assert websocket_target._existing_conversation == {} +async def test_cleanup_conversation_async_cancellation_preserved_when_close_fails( + websocket_target: WebsocketTarget, +) -> None: + connection = AsyncMock(spec=ClientConnection) + websocket_target._existing_conversation["conversation"] = connection + close_started = asyncio.Event() + finish_close = asyncio.Event() + close_error = ConnectionError("close failed") + + async def close_connection() -> None: + close_started.set() + await finish_close.wait() + raise close_error + + connection.close.side_effect = close_connection + cleanup_task = asyncio.create_task(websocket_target.cleanup_conversation_async("conversation")) + await close_started.wait() + + cleanup_task.cancel() + finish_close.set() + + with pytest.raises(asyncio.CancelledError) as exc_info: + await cleanup_task + + assert exc_info.value.__cause__ is close_error + connection.close.assert_awaited_once() + assert websocket_target._existing_conversation == {} + + async def test_cleanup_target_async_attempts_every_connection(websocket_target: WebsocketTarget) -> None: failing_connection = AsyncMock(spec=ClientConnection) failing_connection.close.side_effect = RuntimeError("close failed")