diff --git a/doc/code/targets/websocket_target.ipynb b/doc/code/targets/websocket_target.ipynb new file mode 100644 index 0000000000..5698408565 --- /dev/null +++ b/doc/code/targets/websocket_target.ipynb @@ -0,0 +1,174 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "id": "0", + "metadata": {}, + "source": [ + "# WebSocket Target\n", + "\n", + "`WebsocketTarget` connects PyRIT to services that use a custom WebSocket protocol.\n", + "Supply the service-specific initialization messages, prompt builder, and response parser.\n", + "The `protocol_identifier` is a non-secret name for this complete protocol configuration.\n", + "\n", + "This example starts a local PyRIT WebSocket service. It exercises the real WebSocket\n", + "transport without credentials or an external endpoint." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "1", + "metadata": {}, + "outputs": [], + "source": [ + "import json\n", + "import uuid\n", + "\n", + "from websockets.asyncio.client import ClientConnection\n", + "from websockets.asyncio.server import ServerConnection, serve\n", + "\n", + "from pyrit.models import Message, MessagePiece\n", + "from pyrit.prompt_target import WebsocketTarget\n", + "from pyrit.setup import IN_MEMORY, initialize_pyrit_async\n", + "\n", + "await initialize_pyrit_async(memory_db_type=IN_MEMORY, load_defaults=False, silent=True) # type: ignore\n", + "\n", + "\n", + "async def pyrit_websocket_handler(websocket: ServerConnection) -> None:\n", + " initialization = json.loads(await websocket.recv())\n", + " if initialization != {\"type\": \"initialize\", \"client\": \"PyRIT\"}:\n", + " await websocket.close(code=1002, reason=\"Invalid initialization message\")\n", + " return\n", + "\n", + " await websocket.send(json.dumps({\"message\": \"PyRIT WebSocket target ready\"}))\n", + "\n", + " async for raw_message in websocket:\n", + " request = json.loads(raw_message)\n", + " if request[\"type\"] == \"restore\":\n", + " await websocket.send(json.dumps({\"type\": \"restored\"}))\n", + " continue\n", + "\n", + " await websocket.send(json.dumps({\"event\": \"processing\"}))\n", + " await websocket.send(json.dumps({\"message\": f\"PyRIT received: {request['prompt']}\"}))\n", + "\n", + "\n", + "def response_parser(message: str | bytes) -> str | None:\n", + " if isinstance(message, bytes):\n", + " message = message.decode()\n", + " return json.loads(message).get(\"message\")\n", + "\n", + "\n", + "def message_builder(prompt: str) -> str:\n", + " return json.dumps({\"type\": \"prompt\", \"prompt\": prompt})\n", + "\n", + "\n", + "async def restore_conversation_async(\n", + " websocket: ClientConnection,\n", + " conversation_history: list[Message],\n", + ") -> None:\n", + " history = [\n", + " {\n", + " \"role\": message.message_pieces[0].role,\n", + " \"content\": message.get_value(),\n", + " }\n", + " for message in conversation_history\n", + " ]\n", + " await websocket.send(json.dumps({\"type\": \"restore\", \"history\": history}))\n", + " acknowledgement = json.loads(await websocket.recv())\n", + " if acknowledgement != {\"type\": \"restored\"}:\n", + " raise ConnectionError(\"The WebSocket service did not restore the conversation.\")" + ] + }, + { + "cell_type": "markdown", + "id": "2", + "metadata": {}, + "source": [ + "Start the local service on an available loopback port, then configure the target for its protocol.\n", + "\n", + "The restore callback is service-specific. PyRIT calls it when a multi-turn conversation needs a\n", + "replacement connection. Without this callback, the target fails instead of silently losing history." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "3", + "metadata": {}, + "outputs": [], + "source": [ + "server = await serve(pyrit_websocket_handler, \"127.0.0.1\", 0) # type: ignore\n", + "port = server.sockets[0].getsockname()[1]\n", + "\n", + "target = WebsocketTarget(\n", + " endpoint=f\"ws://127.0.0.1:{port}\",\n", + " protocol_identifier=\"local-pyrit-echo-v1\",\n", + " initialization_strings=[json.dumps({\"type\": \"initialize\", \"client\": \"PyRIT\"})],\n", + " response_parser=response_parser,\n", + " message_builder=message_builder,\n", + " conversation_restore_callback=restore_conversation_async,\n", + " discard_initial_messages=1,\n", + ")" + ] + }, + { + "cell_type": "markdown", + "id": "4", + "metadata": {}, + "source": [ + "Send a prompt through the target and close both sides of the connection.\n", + "\n", + "Cleanup is terminal for this target instance. If a connection fails while a prompt is in\n", + "progress, the target discards that connection and raises the error instead of retrying the prompt." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "5", + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "PyRIT received: Hello\n" + ] + } + ], + "source": [ + "request = MessagePiece(\n", + " role=\"user\",\n", + " original_value=\"Hello\",\n", + " original_value_data_type=\"text\",\n", + " conversation_id=str(uuid.uuid4()),\n", + ").to_message()\n", + "\n", + "try:\n", + " response = await target.send_prompt_async(message=request) # type: ignore\n", + " print(response[0].get_value())\n", + "finally:\n", + " await target.cleanup_target_async() # type: ignore\n", + " server.close()\n", + " await server.wait_closed() # type: ignore" + ] + } + ], + "metadata": { + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.13.5" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} diff --git a/doc/code/targets/websocket_target.py b/doc/code/targets/websocket_target.py new file mode 100644 index 0000000000..095fa97895 --- /dev/null +++ b/doc/code/targets/websocket_target.py @@ -0,0 +1,120 @@ +# --- +# jupyter: +# jupytext: +# text_representation: +# extension: .py +# format_name: percent +# format_version: '1.3' +# jupytext_version: 1.19.0 +# --- + +# %% [markdown] +# # WebSocket Target +# +# `WebsocketTarget` connects PyRIT to services that use a custom WebSocket protocol. +# Supply the service-specific initialization messages, prompt builder, and response parser. +# The `protocol_identifier` is a non-secret name for this complete protocol configuration. +# +# This example starts a local PyRIT WebSocket service. It exercises the real WebSocket +# transport without credentials or an external endpoint. + +# %% +import json +import uuid + +from websockets.asyncio.client import ClientConnection +from websockets.asyncio.server import ServerConnection, serve + +from pyrit.models import Message, MessagePiece +from pyrit.prompt_target import WebsocketTarget +from pyrit.setup import IN_MEMORY, initialize_pyrit_async + +await initialize_pyrit_async(memory_db_type=IN_MEMORY, load_defaults=False, silent=True) # type: ignore + + +async def pyrit_websocket_handler(websocket: ServerConnection) -> None: + initialization = json.loads(await websocket.recv()) + if initialization != {"type": "initialize", "client": "PyRIT"}: + await websocket.close(code=1002, reason="Invalid initialization message") + return + + await websocket.send(json.dumps({"message": "PyRIT WebSocket target ready"})) + + async for raw_message in websocket: + request = json.loads(raw_message) + if request["type"] == "restore": + await websocket.send(json.dumps({"type": "restored"})) + continue + + await websocket.send(json.dumps({"event": "processing"})) + await websocket.send(json.dumps({"message": f"PyRIT received: {request['prompt']}"})) + + +def response_parser(message: str | bytes) -> str | None: + if isinstance(message, bytes): + message = message.decode() + return json.loads(message).get("message") + + +def message_builder(prompt: str) -> str: + return json.dumps({"type": "prompt", "prompt": prompt}) + + +async def restore_conversation_async( + websocket: ClientConnection, + conversation_history: list[Message], +) -> None: + history = [ + { + "role": message.message_pieces[0].role, + "content": message.get_value(), + } + for message in conversation_history + ] + await websocket.send(json.dumps({"type": "restore", "history": history})) + acknowledgement = json.loads(await websocket.recv()) + if acknowledgement != {"type": "restored"}: + raise ConnectionError("The WebSocket service did not restore the conversation.") + + +# %% [markdown] +# Start the local service on an available loopback port, then configure the target for its protocol. +# +# The restore callback is service-specific. PyRIT calls it when a multi-turn conversation needs a +# replacement connection. Without this callback, the target fails instead of silently losing history. + +# %% +server = await serve(pyrit_websocket_handler, "127.0.0.1", 0) # type: ignore +port = server.sockets[0].getsockname()[1] + +target = WebsocketTarget( + endpoint=f"ws://127.0.0.1:{port}", + protocol_identifier="local-pyrit-echo-v1", + initialization_strings=[json.dumps({"type": "initialize", "client": "PyRIT"})], + response_parser=response_parser, + message_builder=message_builder, + conversation_restore_callback=restore_conversation_async, + discard_initial_messages=1, +) + +# %% [markdown] +# Send a prompt through the target and close both sides of the connection. +# +# Cleanup is terminal for this target instance. If a connection fails while a prompt is in +# progress, the target discards that connection and raises the error instead of retrying the prompt. + +# %% +request = MessagePiece( + role="user", + original_value="Hello", + original_value_data_type="text", + conversation_id=str(uuid.uuid4()), +).to_message() + +try: + response = await target.send_prompt_async(message=request) # type: ignore + print(response[0].get_value()) +finally: + await target.cleanup_target_async() # type: ignore + server.close() + await server.wait_closed() # type: ignore diff --git a/doc/myst.yml b/doc/myst.yml index 5b192d30b0..7c2f509974 100644 --- a/doc/myst.yml +++ b/doc/myst.yml @@ -134,6 +134,7 @@ project: - file: code/targets/prompt_shield_target.ipynb - file: code/targets/realtime_target.ipynb - file: code/targets/use_huggingface_chat_target.ipynb + - file: code/targets/websocket_target.ipynb - file: code/targets/round_robin_target.ipynb - file: code/converters/0_converters.ipynb children: diff --git a/pyrit/prompt_target/__init__.py b/pyrit/prompt_target/__init__.py index 91b6cce7f8..1d5608d691 100644 --- a/pyrit/prompt_target/__init__.py +++ b/pyrit/prompt_target/__init__.py @@ -52,6 +52,7 @@ from pyrit.prompt_target.round_robin_target import RoundRobinTarget from pyrit.prompt_target.text_target import TextTarget from pyrit.prompt_target.websocket_copilot_target import WebSocketCopilotTarget +from pyrit.prompt_target.websocket_target import WebsocketTarget if TYPE_CHECKING: from pyrit.prompt_target.hugging_face.hugging_face_chat_target import HuggingFaceChatTarget @@ -109,6 +110,7 @@ def __getattr__(name: str) -> object: "TargetRequirements", "UnsupportedCapabilityBehavior", "TextTarget", + "WebsocketTarget", "discover_target_capabilities_async", "get_known_capabilities", "WebSocketCopilotTarget", diff --git a/pyrit/prompt_target/websocket_target.py b/pyrit/prompt_target/websocket_target.py new file mode 100644 index 0000000000..b2ab2be695 --- /dev/null +++ b/pyrit/prompt_target/websocket_target.py @@ -0,0 +1,397 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT license. + +import asyncio +import logging +from collections.abc import Awaitable, Callable +from typing import Any +from weakref import WeakValueDictionary + +import websockets +from websockets.asyncio.client import ClientConnection +from websockets.protocol import State + +from pyrit.exceptions import EmptyResponseException, pyrit_target_retry +from pyrit.models import ComponentIdentifier, Message, construct_response_from_request +from pyrit.prompt_target import PromptTarget, limit_requests_per_minute +from pyrit.prompt_target.common.target_capabilities import TargetCapabilities +from pyrit.prompt_target.common.target_configuration import TargetConfiguration + +logger = logging.getLogger(__name__) + +WebSocketMessage = str | bytes +ResponseParser = Callable[[WebSocketMessage], str | None] +MessageBuilder = Callable[[str], WebSocketMessage] +ConversationRestoreCallback = Callable[[ClientConnection, list[Message]], Awaitable[None]] + + +class WebsocketTarget(PromptTarget): + """ + Send text prompts to a configurable WebSocket service. + + The target keeps one initialized WebSocket connection for each PyRIT + conversation. Callers provide the service-specific initialization messages, + prompt builder, and response parser. + """ + + RESPONSE_TIMEOUT_SECONDS: float = 30.0 + _DEFAULT_CONFIGURATION: TargetConfiguration = TargetConfiguration( + capabilities=TargetCapabilities(supports_multi_turn=True) + ) + + def __init__( + self, + *, + endpoint: str, + protocol_identifier: str, + initialization_strings: list[WebSocketMessage], + response_parser: ResponseParser, + message_builder: MessageBuilder, + conversation_restore_callback: ConversationRestoreCallback | None = None, + discard_initial_messages: int = 1, + response_timeout_seconds: float = RESPONSE_TIMEOUT_SECONDS, + existing_convo: dict[str, ClientConnection] | None = None, + max_requests_per_minute: int | None = None, + custom_configuration: TargetConfiguration | None = None, + **websockets_kwargs: Any, + ) -> None: + """ + Initialize the WebSocket target. + + Args: + endpoint (str): WebSocket endpoint. Must use the ``ws://`` or ``wss://`` scheme. + protocol_identifier (str): Non-secret name that uniquely identifies the service protocol + and callback behavior. + initialization_strings (list[str | bytes]): Messages to send when a connection opens. + response_parser (ResponseParser): Function that returns response text or ``None`` for + a frame that the target should ignore. + message_builder (MessageBuilder): Function that converts prompt text to a WebSocket message. + conversation_restore_callback (ConversationRestoreCallback | None): Async function that restores + prior normalized messages on a replacement connection. A multi-turn conversation cannot + reconnect without this callback because the target cannot infer the service-specific protocol. + The callback must consume all frames produced by its restoration exchange. + discard_initial_messages (int): Number of parsed messages to discard after initialization. + response_timeout_seconds (float): Maximum time to wait for a parsed response. + existing_convo (dict[str, ClientConnection] | None): Pre-initialized connections by + PyRIT conversation ID. + max_requests_per_minute (int | None): Maximum number of requests per minute. + custom_configuration (TargetConfiguration | None): Override the default target configuration. + websockets_kwargs (Any): Additional keyword arguments for ``websockets.connect``. + + Raises: + ValueError: If endpoint or numeric arguments are invalid. + """ + if not endpoint.startswith(("ws://", "wss://")): + raise ValueError("endpoint must start with 'ws://' or 'wss://'.") + if not protocol_identifier.strip(): + raise ValueError("protocol_identifier must not be empty.") + if discard_initial_messages < 0: + raise ValueError("discard_initial_messages must be nonnegative.") + if response_timeout_seconds <= 0: + raise ValueError("response_timeout_seconds must be positive.") + + super().__init__( + endpoint=endpoint, + max_requests_per_minute=max_requests_per_minute, + custom_configuration=custom_configuration, + ) + + self._protocol_identifier = protocol_identifier + self._initialization_strings = initialization_strings + self._response_parser = response_parser + self._message_builder = message_builder + self._conversation_restore_callback = conversation_restore_callback + self._discard_initial_messages = discard_initial_messages + self._response_timeout_seconds = response_timeout_seconds + self._existing_conversation = existing_convo if existing_convo is not None else {} + self._conversation_locks: WeakValueDictionary[str, asyncio.Lock] = WeakValueDictionary() + self._is_closed = False + self._is_cleaning_up = False + self._websockets_kwargs = websockets_kwargs + + def _build_identifier(self) -> ComponentIdentifier: + """ + Build the identifier with the caller-defined protocol identity. + + Returns: + ComponentIdentifier: The identifier for this target instance. + """ + return self._create_identifier(params={"protocol_identifier": self._protocol_identifier}) + + async def _connect_async(self) -> ClientConnection: + """ + Open a connection to the configured WebSocket endpoint. + + Returns: + ClientConnection: The open WebSocket connection. + """ + logger.info("Connecting to WebSocket endpoint: %s", self._endpoint) + connection = await websockets.connect(uri=self._endpoint, **self._websockets_kwargs) + logger.info("Connected to WebSocket endpoint") + return connection + + async def _send_message_async(self, *, message: WebSocketMessage, conversation_id: str) -> None: + """ + Send one message on an existing conversation connection. + + Args: + message (str | bytes): Message to send. + conversation_id (str): PyRIT conversation ID. + """ + websocket = self._get_websocket(conversation_id=conversation_id) + await websocket.send(message) + + async def _receive_messages_async(self, conversation_id: str) -> str: + """ + Receive frames until the response parser returns text. + + Args: + conversation_id (str): PyRIT conversation ID. + + Returns: + str: Parsed response text. + """ + websocket = self._get_websocket(conversation_id=conversation_id) + return await self._receive_message_async(websocket=websocket) + + async def _send_text_async(self, *, text: str, conversation_id: str) -> str: + """ + Send a text prompt and wait for its response. + + Args: + text (str): Prompt text. + conversation_id (str): PyRIT conversation ID. + + Returns: + str: Parsed response text. + + Raises: + TimeoutError: If no parsed response arrives before the configured timeout. + """ + await self._send_message_async( + message=self._message_builder(text), + conversation_id=conversation_id, + ) + try: + return await asyncio.wait_for( + self._receive_messages_async(conversation_id), + timeout=self._response_timeout_seconds, + ) + except asyncio.TimeoutError: + raise TimeoutError( + f"Timed out waiting for a WebSocket response after {self._response_timeout_seconds} seconds." + ) from None + + async def cleanup_conversation_async(self, conversation_id: str) -> None: + """ + Close and remove one conversation connection. + + Args: + conversation_id (str): PyRIT conversation ID. + """ + 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() + logger.info("Disconnected WebSocket conversation: %s", conversation_id) + + async def cleanup_target_async(self) -> None: + """ + Close and remove all conversation connections. + + Raises: + ConnectionError: If one or more connections cannot be closed. + RuntimeError: If another cleanup operation is already in progress. + """ + if self._is_cleaning_up: + raise RuntimeError("WebsocketTarget cleanup is already in progress.") + + self._is_closed = True + self._is_cleaning_up = True + try: + await self._close_all_connections_async() + finally: + self._conversation_locks.clear() + self._is_cleaning_up = False + + @limit_requests_per_minute + @pyrit_target_retry + async def _send_prompt_to_target_async(self, *, normalized_conversation: list[Message]) -> list[Message]: + """ + Send the current normalized message to the WebSocket service. + + Args: + normalized_conversation (list[Message]): Normalized conversation with the current request last. + + Returns: + list[Message]: A list containing the target response. + + Raises: + ValueError: If the current message is not text. + """ + self._raise_if_closed() + request = normalized_conversation[-1].message_pieces[0] + if request.converted_value_data_type != "text": + raise ValueError(f"Unsupported response type: {request.converted_value_data_type}") + + conversation_id = request.conversation_id + if not conversation_id: + raise ValueError("WebsocketTarget requires a conversation_id on the message being sent.") + conversation_lock = self._conversation_locks.setdefault(conversation_id, asyncio.Lock()) + + async with conversation_lock: + try: + await self._get_or_create_connection_async( + conversation_id=conversation_id, + conversation_history=normalized_conversation[:-1], + ) + result = await self._send_text_async( + text=request.converted_value, + conversation_id=conversation_id, + ) + except BaseException: + await self._discard_connection_async(conversation_id=conversation_id) + raise + + response_piece = construct_response_from_request( + request=request, + response_text_pieces=[result], + response_type="text", + ).message_pieces[0] + return [Message(message_pieces=[response_piece])] + + async def _get_or_create_connection_async( + self, + *, + conversation_id: str, + conversation_history: list[Message], + ) -> ClientConnection: + self._raise_if_closed() + existing_connection = self._existing_conversation.get(conversation_id) + if existing_connection is not None and existing_connection.state is State.OPEN: + return existing_connection + + if existing_connection is not None: + logger.info("Replacing closed WebSocket conversation: %s", conversation_id) + await self._discard_connection_async(conversation_id=conversation_id) + + restore_callback = self._conversation_restore_callback + if conversation_history and restore_callback is None: + raise ConnectionError( + "The WebSocket connection must be replaced, but conversation history cannot be restored. " + "Configure conversation_restore_callback for multi-turn reconnection." + ) + + websocket = await self._connect_async() + try: + await self._initialize_connection_async(websocket=websocket) + if conversation_history and restore_callback is not None: + await self._restore_conversation_async( + websocket=websocket, + conversation_history=conversation_history, + restore_callback=restore_callback, + ) + self._raise_if_closed() + except BaseException: + try: + await websocket.close() + except Exception as error: + logger.warning("Failed to close an uninitialized WebSocket connection: %s", error) + raise + + self._existing_conversation[conversation_id] = websocket + return websocket + + async def _restore_conversation_async( + self, + *, + websocket: ClientConnection, + conversation_history: list[Message], + restore_callback: ConversationRestoreCallback, + ) -> None: + try: + await asyncio.wait_for( + restore_callback(websocket, conversation_history), + timeout=self._response_timeout_seconds, + ) + except asyncio.TimeoutError: + raise TimeoutError( + f"Timed out restoring WebSocket conversation history after {self._response_timeout_seconds} seconds." + ) from None + + async def _initialize_connection_async(self, *, websocket: ClientConnection) -> None: + for initialization_string in self._initialization_strings: + await websocket.send(initialization_string) + + for _ in range(self._discard_initial_messages): + try: + await asyncio.wait_for( + self._receive_message_async(websocket=websocket), + timeout=self._response_timeout_seconds, + ) + except asyncio.TimeoutError: + raise TimeoutError( + "Timed out waiting for an initial WebSocket message after " + f"{self._response_timeout_seconds} seconds." + ) from None + + async def _receive_message_async(self, *, websocket: ClientConnection) -> str: + async for message in websocket: + parsed_message = self._response_parser(message) + if parsed_message is None: + continue + if not parsed_message: + raise EmptyResponseException(message="The WebSocket target returned an empty response.") + return parsed_message + + raise ConnectionError("The WebSocket connection closed before a response was received.") + + async def _discard_connection_async(self, *, conversation_id: str) -> None: + websocket = self._existing_conversation.pop(conversation_id, None) + if websocket is None: + return + try: + await websocket.close() + except Exception as error: + logger.warning("Failed to close unusable WebSocket conversation %s: %s", conversation_id, error) + + def _raise_if_closed(self) -> None: + if self._is_closed: + raise RuntimeError("WebsocketTarget has been cleaned up and cannot send more prompts.") + + async def _close_all_connections_async(self) -> None: + connections = list(self._existing_conversation.items()) + close_future = asyncio.gather( + *(websocket.close() for _, websocket in connections), + return_exceptions=True, + ) + cancellation_error: asyncio.CancelledError | None = None + + try: + close_results = await asyncio.shield(close_future) + except asyncio.CancelledError as error: + cancellation_error = error + close_results = await close_future + + first_error: BaseException | None = None + for (conversation_id, websocket), close_result in zip(connections, close_results, strict=True): + if self._existing_conversation.get(conversation_id) is websocket: + del self._existing_conversation[conversation_id] + if isinstance(close_result, BaseException): + logger.error("Failed to close WebSocket conversation %s: %s", conversation_id, close_result) + first_error = first_error or close_result + continue + logger.info("Disconnected WebSocket conversation: %s", conversation_id) + + if cancellation_error is not None: + raise cancellation_error + if first_error is not None: + raise ConnectionError("Failed to close one or more WebSocket connections.") from first_error + + def _get_websocket(self, *, conversation_id: str) -> ClientConnection: + websocket = self._existing_conversation.get(conversation_id) + if websocket is None: + raise ConnectionError(f"WebSocket connection is not established for conversation {conversation_id}.") + return websocket diff --git a/tests/integration/targets/test_websocket_target_integration.py b/tests/integration/targets/test_websocket_target_integration.py new file mode 100644 index 0000000000..c23f2a550e --- /dev/null +++ b/tests/integration/targets/test_websocket_target_integration.py @@ -0,0 +1,179 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT license. + +import asyncio +import json +import uuid + +import pytest +from websockets.asyncio.client import ClientConnection +from websockets.asyncio.server import ServerConnection, serve + +from pyrit.memory import SQLiteMemory +from pyrit.models import Conversation, Message, MessagePiece +from pyrit.prompt_target import WebsocketTarget + + +@pytest.mark.run_only_if_all_tests +async def test_websocket_target_round_trip_with_local_pyrit_server(sqlite_instance: SQLiteMemory) -> None: + received_messages: list[dict[str, str]] = [] + + async def pyrit_websocket_handler(websocket: ServerConnection) -> None: + initialization = json.loads(await websocket.recv()) + received_messages.append(initialization) + await websocket.send(json.dumps({"message": "PyRIT WebSocket target ready"})) + + async for raw_message in websocket: + prompt_message = json.loads(raw_message) + received_messages.append(prompt_message) + await websocket.send(json.dumps({"event": "processing"})) + await websocket.send(json.dumps({"message": f"PyRIT received: {prompt_message['prompt']}"})) + + def response_parser(message: str | bytes) -> str | None: + if isinstance(message, bytes): + message = message.decode() + return json.loads(message).get("message") + + def message_builder(prompt: str) -> str: + return json.dumps({"type": "prompt", "prompt": prompt}) + + async with serve(pyrit_websocket_handler, "127.0.0.1", 0) as server: + port = server.sockets[0].getsockname()[1] + target = WebsocketTarget( + endpoint=f"ws://127.0.0.1:{port}", + protocol_identifier="local-pyrit-echo-v1", + initialization_strings=[json.dumps({"type": "initialize", "client": "PyRIT"})], + response_parser=response_parser, + message_builder=message_builder, + discard_initial_messages=1, + ) + + conversation_id = str(uuid.uuid4()) + request = MessagePiece( + role="user", + original_value="Hello", + original_value_data_type="text", + conversation_id=conversation_id, + ).to_message() + + try: + response = await target.send_prompt_async(message=request) + finally: + await target.cleanup_target_async() + + assert response[0].get_value() == "PyRIT received: Hello" + assert received_messages == [ + {"type": "initialize", "client": "PyRIT"}, + {"type": "prompt", "prompt": "Hello"}, + ] + + +@pytest.mark.run_only_if_all_tests +async def test_websocket_target_restores_history_after_server_disconnect(sqlite_instance: SQLiteMemory) -> None: + received_messages: list[dict[str, object]] = [] + first_connection_closed = asyncio.Event() + connection_count = 0 + + async def pyrit_websocket_handler(websocket: ServerConnection) -> None: + nonlocal connection_count + connection_count += 1 + current_connection = connection_count + + initialization = json.loads(await websocket.recv()) + received_messages.append(initialization) + await websocket.send(json.dumps({"message": "PyRIT WebSocket target ready"})) + + async for raw_message in websocket: + request = json.loads(raw_message) + received_messages.append(request) + + if request["type"] == "restore": + await websocket.send(json.dumps({"type": "restored"})) + continue + + await websocket.send(json.dumps({"message": f"PyRIT received: {request['prompt']}"})) + if current_connection == 1: + await websocket.close() + first_connection_closed.set() + return + + def response_parser(message: str | bytes) -> str | None: + if isinstance(message, bytes): + message = message.decode() + return json.loads(message).get("message") + + def message_builder(prompt: str) -> str: + return json.dumps({"type": "prompt", "prompt": prompt}) + + async def restore_conversation_async( + websocket: ClientConnection, + conversation_history: list[Message], + ) -> None: + history = [ + { + "role": message.message_pieces[0].role, + "content": message.get_value(), + } + for message in conversation_history + ] + await websocket.send(json.dumps({"type": "restore", "history": history})) + acknowledgement = json.loads(await websocket.recv()) + if acknowledgement != {"type": "restored"}: + raise ConnectionError("The local WebSocket server did not restore the conversation.") + + async with serve(pyrit_websocket_handler, "127.0.0.1", 0) as server: + port = server.sockets[0].getsockname()[1] + target = WebsocketTarget( + endpoint=f"ws://127.0.0.1:{port}", + protocol_identifier="local-pyrit-echo-v1", + initialization_strings=[json.dumps({"type": "initialize", "client": "PyRIT"})], + response_parser=response_parser, + message_builder=message_builder, + conversation_restore_callback=restore_conversation_async, + discard_initial_messages=1, + ) + conversation_id = str(uuid.uuid4()) + first_request = MessagePiece( + role="user", + original_value="First", + original_value_data_type="text", + conversation_id=conversation_id, + ).to_message() + + try: + first_response = await target.send_prompt_async(message=first_request) + await first_connection_closed.wait() + + sqlite_instance.add_conversation_to_memory( + conversation=Conversation( + conversation_id=conversation_id, + target_identifier=target.get_identifier(), + ) + ) + sqlite_instance.add_message_to_memory(request=first_request) + sqlite_instance.add_message_to_memory(request=first_response[0]) + + second_request = MessagePiece( + role="user", + original_value="Second", + original_value_data_type="text", + conversation_id=conversation_id, + ).to_message() + second_response = await target.send_prompt_async(message=second_request) + finally: + await target.cleanup_target_async() + + assert second_response[0].get_value() == "PyRIT received: Second" + assert received_messages == [ + {"type": "initialize", "client": "PyRIT"}, + {"type": "prompt", "prompt": "First"}, + {"type": "initialize", "client": "PyRIT"}, + { + "type": "restore", + "history": [ + {"role": "user", "content": "First"}, + {"role": "assistant", "content": "PyRIT received: First"}, + ], + }, + {"type": "prompt", "prompt": "Second"}, + ] diff --git a/tests/unit/prompt_target/target/test_websocket_target.py b/tests/unit/prompt_target/target/test_websocket_target.py new file mode 100644 index 0000000000..281dc5d5cd --- /dev/null +++ b/tests/unit/prompt_target/target/test_websocket_target.py @@ -0,0 +1,690 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT license. + +import asyncio +import json +from collections.abc import Callable +from unittest.mock import AsyncMock, patch + +import pytest +from websockets.asyncio.client import ClientConnection +from websockets.exceptions import ConnectionClosed +from websockets.frames import Close +from websockets.protocol import State + +from pyrit.exceptions import EmptyResponseException +from pyrit.memory import SQLiteMemory +from pyrit.models import Message, MessagePiece +from pyrit.prompt_target import WebsocketTarget + + +@pytest.fixture +def response_parser() -> Callable[[str | bytes], str | None]: + def parse_response(message: str | bytes) -> str | None: + if isinstance(message, bytes): + message = message.decode() + return json.loads(message).get("message") + + return parse_response + + +@pytest.fixture +def message_builder() -> Callable[[str], str | bytes]: + def build_message(prompt: str) -> str: + return json.dumps({"message": prompt}) + + return build_message + + +@pytest.fixture +def websocket_target( + response_parser: Callable[[str | bytes], str | None], + message_builder: Callable[[str], str | bytes], + sqlite_instance: SQLiteMemory, +) -> WebsocketTarget: + return WebsocketTarget( + endpoint="wss://example.com", + protocol_identifier="test-protocol", + initialization_strings=["connect", "authenticate"], + response_parser=response_parser, + message_builder=message_builder, + discard_initial_messages=0, + ) + + +def create_message(*, value: str = "Hello", conversation_id: str = "conversation") -> Message: + return MessagePiece( + original_value=value, + original_value_data_type="text", + converted_value=value, + converted_value_data_type="text", + role="user", + conversation_id=conversation_id, + ).to_message() + + +def test_init_invalid_endpoint_raises( + response_parser: Callable[[str | bytes], str | None], + message_builder: Callable[[str], str | bytes], + sqlite_instance: SQLiteMemory, +) -> None: + with pytest.raises(ValueError, match="endpoint must start"): + WebsocketTarget( + endpoint="https://example.com", + protocol_identifier="test-protocol", + initialization_strings=[], + response_parser=response_parser, + message_builder=message_builder, + ) + + +def test_init_empty_protocol_identifier_raises( + response_parser: Callable[[str | bytes], str | None], + message_builder: Callable[[str], str | bytes], + sqlite_instance: SQLiteMemory, +) -> None: + with pytest.raises(ValueError, match="protocol_identifier must not be empty"): + WebsocketTarget( + endpoint="wss://example.com", + protocol_identifier=" ", + initialization_strings=[], + response_parser=response_parser, + message_builder=message_builder, + ) + + +def test_init_invalid_discard_count_raises( + response_parser: Callable[[str | bytes], str | None], + message_builder: Callable[[str], str | bytes], + sqlite_instance: SQLiteMemory, +) -> None: + with pytest.raises(ValueError, match="must be nonnegative"): + WebsocketTarget( + endpoint="wss://example.com", + protocol_identifier="test-protocol", + initialization_strings=[], + response_parser=response_parser, + message_builder=message_builder, + discard_initial_messages=-1, + ) + + +def test_init_invalid_timeout_raises( + response_parser: Callable[[str | bytes], str | None], + message_builder: Callable[[str], str | bytes], + sqlite_instance: SQLiteMemory, +) -> None: + with pytest.raises(ValueError, match="must be positive"): + WebsocketTarget( + endpoint="wss://example.com", + protocol_identifier="test-protocol", + initialization_strings=[], + response_parser=response_parser, + message_builder=message_builder, + response_timeout_seconds=0, + ) + + +def test_identifier_distinguishes_protocols( + websocket_target: WebsocketTarget, + response_parser: Callable[[str | bytes], str | None], + message_builder: Callable[[str], str | bytes], +) -> None: + other_target = WebsocketTarget( + endpoint="wss://example.com", + protocol_identifier="other-protocol", + initialization_strings=[], + response_parser=response_parser, + message_builder=message_builder, + discard_initial_messages=0, + ) + + assert websocket_target.get_identifier() != other_target.get_identifier() + + +async def test_connect_async_passes_websocket_arguments( + response_parser: Callable[[str | bytes], str | None], + message_builder: Callable[[str], str | bytes], + sqlite_instance: SQLiteMemory, +) -> None: + target = WebsocketTarget( + endpoint="wss://example.com", + protocol_identifier="test-protocol", + initialization_strings=[], + response_parser=response_parser, + message_builder=message_builder, + proxy="http://proxy.example.com", + ) + connection = AsyncMock(spec=ClientConnection) + + with patch( + "pyrit.prompt_target.websocket_target.websockets.connect", + new_callable=AsyncMock, + return_value=connection, + ) as mock_connect: + result = await target._connect_async() + + assert result is connection + mock_connect.assert_awaited_once_with(uri="wss://example.com", proxy="http://proxy.example.com") + + +async def test_send_prompt_async_initializes_connection_once(websocket_target: WebsocketTarget) -> None: + connection = AsyncMock(spec=ClientConnection) + connection.state = State.OPEN + + with ( + patch.object(websocket_target, "_connect_async", new_callable=AsyncMock, return_value=connection) as connect, + patch.object( + websocket_target, + "_send_text_async", + new_callable=AsyncMock, + side_effect=["First response", "Second response"], + ) as send_text, + ): + first_response = await websocket_target.send_prompt_async( + message=create_message(value="First", conversation_id="shared") + ) + second_response = await websocket_target.send_prompt_async( + message=create_message(value="Second", conversation_id="shared") + ) + + connect.assert_awaited_once() + assert connection.send.await_count == 2 + assert [call.args[0] for call in connection.send.await_args_list] == ["connect", "authenticate"] + assert send_text.await_count == 2 + assert first_response[0].get_value() == "First response" + assert second_response[0].get_value() == "Second response" + + await websocket_target.cleanup_target_async() + + +async def test_send_prompt_async_serializes_same_conversation(websocket_target: WebsocketTarget) -> None: + connection = AsyncMock(spec=ClientConnection) + connection.state = State.OPEN + active_requests = 0 + maximum_active_requests = 0 + + async def send_text(*, text: str, conversation_id: str) -> str: + nonlocal active_requests, maximum_active_requests + active_requests += 1 + maximum_active_requests = max(maximum_active_requests, active_requests) + await asyncio.sleep(0) + active_requests -= 1 + return text + + with ( + patch.object(websocket_target, "_connect_async", new_callable=AsyncMock, return_value=connection) as connect, + patch.object(websocket_target, "_send_text_async", side_effect=send_text), + ): + await asyncio.gather( + websocket_target.send_prompt_async(message=create_message(value="First", conversation_id="shared")), + websocket_target.send_prompt_async(message=create_message(value="Second", conversation_id="shared")), + ) + + connect.assert_awaited_once() + assert maximum_active_requests == 1 + + await websocket_target.cleanup_target_async() + + +async def test_send_prompt_async_failure_discards_connection(websocket_target: WebsocketTarget) -> None: + connection = AsyncMock(spec=ClientConnection) + connection.state = State.OPEN + websocket_target._existing_conversation["conversation"] = connection + + with ( + patch.object(websocket_target, "_connect_async", new_callable=AsyncMock) as connect, + patch.object( + websocket_target, + "_send_text_async", + new_callable=AsyncMock, + side_effect=ConnectionError("connection failed"), + ), + ): + with pytest.raises(ConnectionError, match="connection failed"): + await websocket_target.send_prompt_async(message=create_message()) + + connect.assert_not_awaited() + connection.close.assert_awaited_once() + assert "conversation" not in websocket_target._existing_conversation + + +async def test_get_or_create_connection_async_restores_history_on_reconnect( + response_parser: Callable[[str | bytes], str | None], + message_builder: Callable[[str], str | bytes], + sqlite_instance: SQLiteMemory, +) -> None: + stale_connection = AsyncMock(spec=ClientConnection) + stale_connection.state = State.CLOSED + replacement_connection = AsyncMock(spec=ClientConnection) + replacement_connection.state = State.OPEN + restore_callback = AsyncMock() + target = WebsocketTarget( + endpoint="wss://example.com", + protocol_identifier="test-protocol", + initialization_strings=[], + response_parser=response_parser, + message_builder=message_builder, + conversation_restore_callback=restore_callback, + discard_initial_messages=0, + existing_convo={"conversation": stale_connection}, + ) + history = [create_message(value="Prior")] + + with patch.object( + target, + "_connect_async", + new_callable=AsyncMock, + return_value=replacement_connection, + ) as connect: + result = await target._get_or_create_connection_async( + conversation_id="conversation", + conversation_history=history, + ) + + assert result is replacement_connection + stale_connection.close.assert_awaited_once() + connect.assert_awaited_once() + restore_callback.assert_awaited_once_with(replacement_connection, history) + assert target._existing_conversation == {"conversation": replacement_connection} + + +async def test_get_or_create_connection_async_fails_when_history_cannot_be_restored( + websocket_target: WebsocketTarget, +) -> None: + stale_connection = AsyncMock(spec=ClientConnection) + stale_connection.state = State.CLOSED + websocket_target._existing_conversation["conversation"] = stale_connection + + with ( + patch.object(websocket_target, "_connect_async", new_callable=AsyncMock) as connect, + pytest.raises(ConnectionError, match="Configure conversation_restore_callback"), + ): + await websocket_target._get_or_create_connection_async( + conversation_id="conversation", + conversation_history=[create_message(value="Prior")], + ) + + stale_connection.close.assert_awaited_once() + connect.assert_not_awaited() + assert websocket_target._existing_conversation == {} + + +async def test_send_prompt_async_retry_restores_history( + response_parser: Callable[[str | bytes], str | None], + message_builder: Callable[[str], str | bytes], + sqlite_instance: SQLiteMemory, +) -> None: + existing_connection = AsyncMock(spec=ClientConnection) + existing_connection.state = State.OPEN + replacement_connection = AsyncMock(spec=ClientConnection) + replacement_connection.state = State.OPEN + restore_callback = AsyncMock() + target = WebsocketTarget( + endpoint="wss://example.com", + protocol_identifier="test-protocol", + initialization_strings=[], + response_parser=response_parser, + message_builder=message_builder, + conversation_restore_callback=restore_callback, + discard_initial_messages=0, + existing_convo={"conversation": existing_connection}, + ) + prior_message = create_message(value="Prior") + current_message = create_message(value="Current") + + with ( + patch.object( + target, + "_connect_async", + new_callable=AsyncMock, + return_value=replacement_connection, + ) as connect, + patch.object( + target, + "_send_text_async", + new_callable=AsyncMock, + side_effect=[ + EmptyResponseException(message="empty response"), + "Recovered response", + ], + ), + ): + response = await target._send_prompt_to_target_async( + normalized_conversation=[prior_message, current_message], + ) + + existing_connection.close.assert_awaited_once() + connect.assert_awaited_once() + restore_callback.assert_awaited_once_with(replacement_connection, [prior_message]) + assert response[0].get_value() == "Recovered response" + + +async def test_get_or_create_connection_async_closes_connection_when_initialization_fails( + websocket_target: WebsocketTarget, +) -> None: + connection = AsyncMock(spec=ClientConnection) + + with ( + patch.object( + websocket_target, + "_connect_async", + new_callable=AsyncMock, + return_value=connection, + ), + patch.object( + websocket_target, + "_initialize_connection_async", + new_callable=AsyncMock, + side_effect=ConnectionError("initialization failed"), + ), + pytest.raises(ConnectionError, match="initialization failed"), + ): + await websocket_target._get_or_create_connection_async( + conversation_id="conversation", + conversation_history=[], + ) + + connection.close.assert_awaited_once() + assert websocket_target._existing_conversation == {} + + +async def test_get_or_create_connection_async_closes_connection_when_restore_times_out( + response_parser: Callable[[str | bytes], str | None], + message_builder: Callable[[str], str | bytes], + sqlite_instance: SQLiteMemory, +) -> None: + async def restore_forever( + websocket: ClientConnection, + conversation_history: list[Message], + ) -> None: + await asyncio.sleep(1) + + target = WebsocketTarget( + endpoint="wss://example.com", + protocol_identifier="test-protocol", + initialization_strings=[], + response_parser=response_parser, + message_builder=message_builder, + conversation_restore_callback=restore_forever, + discard_initial_messages=0, + response_timeout_seconds=0.001, + ) + connection = AsyncMock(spec=ClientConnection) + + with ( + patch.object( + target, + "_connect_async", + new_callable=AsyncMock, + return_value=connection, + ), + pytest.raises(TimeoutError, match="Timed out restoring WebSocket conversation history"), + ): + await target._get_or_create_connection_async( + conversation_id="conversation", + conversation_history=[create_message(value="Prior")], + ) + + connection.close.assert_awaited_once() + assert target._existing_conversation == {} + + +def test_validate_request_invalid_type_raises(websocket_target: WebsocketTarget) -> None: + message = MessagePiece( + original_value="image.png", + original_value_data_type="image_path", + converted_value="image.png", + converted_value_data_type="image_path", + role="user", + ).to_message() + + with pytest.raises(ValueError, match="supports only the following data types: text"): + websocket_target._validate_request(normalized_conversation=[message]) + + +async def test_send_prompt_async_without_conversation_id_raises(websocket_target: WebsocketTarget) -> None: + message = create_message() + message.message_pieces[0].conversation_id = None + + with pytest.raises(ValueError, match="requires a conversation_id"): + await websocket_target.send_prompt_async(message=message) + + +async def test_receive_messages_async_ignores_unparsed_frames(websocket_target: WebsocketTarget) -> None: + connection = AsyncMock(spec=ClientConnection) + connection.__aiter__.return_value = [ + json.dumps({"event": "progress"}), + json.dumps({"message": "response"}), + ] + websocket_target._existing_conversation["conversation"] = connection + + result = await websocket_target._receive_messages_async("conversation") + + assert result == "response" + + +async def test_receive_messages_async_propagates_parser_error(websocket_target: WebsocketTarget) -> None: + connection = AsyncMock(spec=ClientConnection) + connection.__aiter__.return_value = ["not-json"] + websocket_target._existing_conversation["conversation"] = connection + + with pytest.raises(json.JSONDecodeError): + await websocket_target._receive_messages_async("conversation") + + +async def test_receive_messages_async_propagates_connection_closed(websocket_target: WebsocketTarget) -> None: + connection = AsyncMock(spec=ClientConnection) + close_frame = Close(1000, "Normal closure") + + class FailingAsyncIterator: + def __aiter__(self) -> "FailingAsyncIterator": + return self + + async def __anext__(self) -> str: + raise ConnectionClosed(rcvd=close_frame, sent=None) + + connection.__aiter__.side_effect = lambda: FailingAsyncIterator() + websocket_target._existing_conversation["conversation"] = connection + + with pytest.raises(ConnectionClosed): + await websocket_target._receive_messages_async("conversation") + + +async def test_receive_messages_async_accepts_binary_frame(websocket_target: WebsocketTarget) -> None: + connection = AsyncMock(spec=ClientConnection) + connection.__aiter__.return_value = [b'{"message": "response"}'] + websocket_target._existing_conversation["conversation"] = connection + + result = await websocket_target._receive_messages_async("conversation") + + assert result == "response" + + +async def test_receive_messages_async_without_connection_raises(websocket_target: WebsocketTarget) -> None: + with pytest.raises(ConnectionError, match="not established"): + await websocket_target._receive_messages_async("missing") + + +async def test_send_text_async_timeout_raises(websocket_target: WebsocketTarget) -> None: + connection = AsyncMock(spec=ClientConnection) + websocket_target._existing_conversation["conversation"] = connection + websocket_target._response_timeout_seconds = 0.001 + + async def wait_forever(conversation_id: str) -> str: + await asyncio.sleep(1) + return "unreachable" + + with patch.object(websocket_target, "_receive_messages_async", side_effect=wait_forever): + with pytest.raises(TimeoutError, match="Timed out waiting for a WebSocket response"): + await websocket_target._send_text_async(text="Hello", conversation_id="conversation") + + +async def test_initialize_connection_async_discards_configured_messages( + response_parser: Callable[[str | bytes], str | None], + message_builder: Callable[[str], str | bytes], + sqlite_instance: SQLiteMemory, +) -> None: + target = WebsocketTarget( + endpoint="wss://example.com", + protocol_identifier="test-protocol", + initialization_strings=["connect", "authenticate"], + response_parser=response_parser, + message_builder=message_builder, + discard_initial_messages=2, + ) + connection = AsyncMock(spec=ClientConnection) + + with patch.object( + target, + "_receive_message_async", + new_callable=AsyncMock, + side_effect=["first", "second"], + ) as receive: + await target._initialize_connection_async(websocket=connection) + + assert [call.args[0] for call in connection.send.await_args_list] == ["connect", "authenticate"] + assert receive.await_count == 2 + + +async def test_initialize_connection_async_timeout_raises( + response_parser: Callable[[str | bytes], str | None], + message_builder: Callable[[str], str | bytes], + sqlite_instance: SQLiteMemory, +) -> None: + target = WebsocketTarget( + endpoint="wss://example.com", + protocol_identifier="test-protocol", + initialization_strings=[], + response_parser=response_parser, + message_builder=message_builder, + discard_initial_messages=1, + response_timeout_seconds=0.001, + ) + connection = AsyncMock(spec=ClientConnection) + + async def wait_forever(*, websocket: ClientConnection) -> str: + await asyncio.sleep(1) + return "unreachable" + + with patch.object(target, "_receive_message_async", side_effect=wait_forever): + with pytest.raises(TimeoutError, match="Timed out waiting for an initial WebSocket message"): + await target._initialize_connection_async(websocket=connection) + + +async def test_cleanup_conversation_async_removes_connection(websocket_target: WebsocketTarget) -> None: + connection = AsyncMock(spec=ClientConnection) + websocket_target._existing_conversation["conversation"] = connection + + await websocket_target.cleanup_conversation_async("conversation") + + connection.close.assert_awaited_once() + assert websocket_target._existing_conversation == {} + + +async def test_cleanup_conversation_async_does_not_retain_unknown_lock(websocket_target: WebsocketTarget) -> None: + await websocket_target.cleanup_conversation_async("missing") + + assert "missing" not in websocket_target._conversation_locks + + +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") + successful_connection = AsyncMock(spec=ClientConnection) + websocket_target._existing_conversation = { + "failing": failing_connection, + "successful": successful_connection, + } + conversation_lock = asyncio.Lock() + websocket_target._conversation_locks["failing"] = conversation_lock + + with pytest.raises(ConnectionError, match="one or more"): + await websocket_target.cleanup_target_async() + + failing_connection.close.assert_awaited_once() + successful_connection.close.assert_awaited_once() + assert websocket_target._existing_conversation == {} + assert websocket_target._conversation_locks == {} + + +async def test_cleanup_target_async_makes_target_terminal(websocket_target: WebsocketTarget) -> None: + connection = AsyncMock(spec=ClientConnection) + connection.state = State.OPEN + websocket_target._existing_conversation["conversation"] = connection + + await websocket_target.cleanup_target_async() + + with ( + patch.object(websocket_target, "_connect_async", new_callable=AsyncMock) as connect, + pytest.raises(RuntimeError, match="has been cleaned up"), + ): + await websocket_target.send_prompt_async(message=create_message()) + + connect.assert_not_awaited() + connection.close.assert_awaited_once() + + +async def test_cleanup_target_async_blocks_rate_limited_send( + response_parser: Callable[[str | bytes], str | None], + message_builder: Callable[[str], str | bytes], + sqlite_instance: SQLiteMemory, +) -> None: + target = WebsocketTarget( + endpoint="wss://example.com", + protocol_identifier="test-protocol", + initialization_strings=[], + response_parser=response_parser, + message_builder=message_builder, + discard_initial_messages=0, + max_requests_per_minute=1, + ) + rate_limit_started = asyncio.Event() + release_rate_limit = asyncio.Event() + + async def wait_for_rate_limit(delay: float) -> None: + rate_limit_started.set() + await release_rate_limit.wait() + + with ( + patch("pyrit.prompt_target.common.utils.asyncio.sleep", side_effect=wait_for_rate_limit), + patch.object(target, "_connect_async", new_callable=AsyncMock) as connect, + ): + send_task = asyncio.create_task(target.send_prompt_async(message=create_message())) + await rate_limit_started.wait() + await target.cleanup_target_async() + release_rate_limit.set() + + with pytest.raises(RuntimeError, match="has been cleaned up"): + await send_task + + connect.assert_not_awaited() + + +async def test_cleanup_target_async_cancellation_finishes_closing_connections( + 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_target_async()) + 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 + + assert websocket_target._existing_conversation == {} + with pytest.raises(RuntimeError, match="has been cleaned up"): + await websocket_target.send_prompt_async(message=create_message())