diff --git a/pyrit/prompt_target/websocket_target.py b/pyrit/prompt_target/websocket_target.py index b2ab2be695..d8fb94a1b4 100644 --- a/pyrit/prompt_target/websocket_target.py +++ b/pyrit/prompt_target/websocket_target.py @@ -188,13 +188,24 @@ 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 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) 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..0b762a93ab 100644 --- a/tests/unit/prompt_target/target/test_websocket_target.py +++ b/tests/unit/prompt_target/target/test_websocket_target.py @@ -588,6 +588,63 @@ 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_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")