FIX finish WebSocket conversation cleanup before cancellation

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:

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")