diff --git a/.github/instructions/targets.instructions.md b/.github/instructions/targets.instructions.md index 7d21935b4e..36b985e736 100644 --- a/.github/instructions/targets.instructions.md +++ b/.github/instructions/targets.instructions.md @@ -43,6 +43,12 @@ class MyTarget(PromptTarget): ``send_prompt_async`` (the public entry point) is ``@final`` and MUST NOT be overridden. Override ``_send_prompt_to_target_async`` instead. +Targets that hold external state keyed by conversation (a websocket +connection, browser page, or upstream session) SHOULD override +``reset_conversation_async``. It must be safe to call more than once and for +unknown conversation IDs. Whole-target cleanup stays in +``cleanup_target_async``. + ## Keyword-only ``__init__`` is enforced Every ``PromptTarget`` subclass MUST make all ``__init__`` parameters diff --git a/pyrit/executor/attack/core/attack_strategy.py b/pyrit/executor/attack/core/attack_strategy.py index c1e260fc12..82879cdbf6 100644 --- a/pyrit/executor/attack/core/attack_strategy.py +++ b/pyrit/executor/attack/core/attack_strategy.py @@ -3,6 +3,7 @@ from __future__ import annotations +import asyncio import dataclasses import logging # noqa: TC003 import time @@ -47,6 +48,8 @@ from pyrit.prompt_target.common.target_requirements import TargetRequirements if TYPE_CHECKING: + from types import TracebackType + from pyrit.executor.attack.component.prepended_conversation_config import ( PrependedConversationConfig, ) @@ -72,6 +75,61 @@ class _NextMessageOverrideState(Enum): UNSET = "unset" +class _ObjectiveTargetConversationLifecycle: + """Track and release objective-target conversations for one attack execution.""" + + def __init__( + self, + *, + objective_target: PromptTarget, + logger: logging.Logger | logging.LoggerAdapter[logging.Logger], + ) -> None: + self._objective_target = objective_target + self._logger = logger + self._conversation_ids: set[str] = set() + + async def __aenter__(self) -> _ObjectiveTargetConversationLifecycle: + """ + Start tracking target invocations. + + Returns: + _ObjectiveTargetConversationLifecycle: This lifecycle instance. + """ + return self + + async def __aexit__( + self, + exc_type: type[BaseException] | None, + exc_value: BaseException | None, + traceback: TracebackType | None, + ) -> None: + """Release each conversation invoked during the attack.""" + pending_cancellation: asyncio.CancelledError | None = None + for conversation_id in self._conversation_ids: + try: + await self._objective_target.reset_conversation_async(conversation_id=conversation_id) + except asyncio.CancelledError as cancellation: + # Attempt every reset, then honor the first cancellation after the loop. + pending_cancellation = pending_cancellation or cancellation + except Exception as error: # noqa: BLE001 - cleanup must not replace the attack outcome + self._logger.warning( + "Failed to reset objective-target conversation %s: %s", + conversation_id, + error, + ) + if pending_cancellation is not None: + raise pending_cancellation + + def record_invocation(self, *, conversation_id: str) -> None: + """ + Record one objective-target invocation. + + Args: + conversation_id (str): The conversation ID used by the target. + """ + self._conversation_ids.add(conversation_id) + + @dataclass class AttackContext(StrategyContext, ABC, Generic[AttackParamsT]): """ @@ -101,6 +159,12 @@ class AttackContext(StrategyContext, ABC, Generic[AttackParamsT]): _prepended_conversation_override: list[Message] | None = None _memory_labels_override: dict[str, str] | None = None _error_result_persistence_error: Exception | None = field(default=None, init=False, repr=False) + _objective_target_conversation_lifecycle: _ObjectiveTargetConversationLifecycle | None = field( + default=None, + init=False, + repr=False, + compare=False, + ) # Per-execution prepended-history boundary and send lifecycle. Never persisted. prepended_history_send_context: PrependedHistorySendContext | None = field( @@ -167,6 +231,21 @@ def next_message(self, value: Message | None) -> None: """Set the next message (for attacks that generate internally).""" self._next_message_override = value + def _record_objective_target_invocation(self, *, conversation_id: str) -> None: + """ + Record an objective-target invocation for lifecycle cleanup. + + Recording is a no-op when no objective-target scope is active, so attack + helpers remain callable outside a full execution. + + Args: + conversation_id (str): The conversation ID used by the target. + """ + lifecycle = self._objective_target_conversation_lifecycle + if lifecycle is None: + return + lifecycle.record_invocation(conversation_id=conversation_id) + class _DefaultAttackStrategyEventHandler(StrategyEventHandler[AttackStrategyContextT, AttackStrategyResultT]): """ @@ -699,16 +778,25 @@ async def execute_with_context_async(self, *, context: AttackStrategyContextT) - ExceptionGroup: If attack execution and recording its error result both fail. """ context._error_result_persistence_error = None + lifecycle = _ObjectiveTargetConversationLifecycle( + objective_target=self._objective_target, + logger=self._logger, + ) + context._objective_target_conversation_lifecycle = lifecycle try: - result = await super().execute_with_context_async(context=context) - except Exception as attack_error: - persistence_error = context._error_result_persistence_error - if persistence_error is not None: - raise ExceptionGroup( - "Attack execution and error result persistence failed", - [attack_error, persistence_error], - ) from None - raise + async with lifecycle: + try: + result = await super().execute_with_context_async(context=context) + except Exception as attack_error: + persistence_error = context._error_result_persistence_error + if persistence_error is not None: + raise ExceptionGroup( + "Attack execution and error result persistence failed", + [attack_error, persistence_error], + ) from None + raise + finally: + context._objective_target_conversation_lifecycle = None self._default_event_handler._persist_result(result=result) return result diff --git a/pyrit/executor/attack/multi_turn/chunked_request.py b/pyrit/executor/attack/multi_turn/chunked_request.py index b90eca4d09..0c906b4157 100644 --- a/pyrit/executor/attack/multi_turn/chunked_request.py +++ b/pyrit/executor/attack/multi_turn/chunked_request.py @@ -287,6 +287,7 @@ async def _perform_async(self, *, context: ChunkedRequestAttackContext) -> Attac objective_target_conversation_id=context.session.conversation_id, objective=context.objective, ): + context._record_objective_target_invocation(conversation_id=context.session.conversation_id) response = await self._prompt_normalizer.send_prompt_async( message=message, target=self._objective_target, diff --git a/pyrit/executor/attack/multi_turn/crescendo.py b/pyrit/executor/attack/multi_turn/crescendo.py index b9fb98a677..a267ab50a2 100644 --- a/pyrit/executor/attack/multi_turn/crescendo.py +++ b/pyrit/executor/attack/multi_turn/crescendo.py @@ -635,6 +635,7 @@ async def _send_prompt_to_objective_target_async( objective_target_conversation_id=context.session.conversation_id, objective=context.objective, ): + context._record_objective_target_invocation(conversation_id=context.session.conversation_id) response = await self._prompt_normalizer.send_prompt_async( message=attack_message, target=self._objective_target, diff --git a/pyrit/executor/attack/multi_turn/multi_prompt_sending.py b/pyrit/executor/attack/multi_turn/multi_prompt_sending.py index eb065fb6d0..7ee10e06df 100644 --- a/pyrit/executor/attack/multi_turn/multi_prompt_sending.py +++ b/pyrit/executor/attack/multi_turn/multi_prompt_sending.py @@ -365,6 +365,7 @@ async def _send_prompt_to_objective_target_async( objective_target_conversation_id=context.session.conversation_id, objective=context.objective, ): + context._record_objective_target_invocation(conversation_id=context.session.conversation_id) return await self._prompt_normalizer.send_prompt_async( message=current_message, target=self._objective_target, diff --git a/pyrit/executor/attack/multi_turn/red_teaming.py b/pyrit/executor/attack/multi_turn/red_teaming.py index 44893bb755..0b225eff04 100644 --- a/pyrit/executor/attack/multi_turn/red_teaming.py +++ b/pyrit/executor/attack/multi_turn/red_teaming.py @@ -488,6 +488,7 @@ async def _send_prompt_to_objective_target_async( objective=context.objective, ): # Send the message to the target + context._record_objective_target_invocation(conversation_id=context.session.conversation_id) response = await self._prompt_normalizer.send_prompt_async( message=message, conversation_id=context.session.conversation_id, diff --git a/pyrit/executor/attack/multi_turn/tree_of_attacks.py b/pyrit/executor/attack/multi_turn/tree_of_attacks.py index 7e86386c4c..26ed278507 100644 --- a/pyrit/executor/attack/multi_turn/tree_of_attacks.py +++ b/pyrit/executor/attack/multi_turn/tree_of_attacks.py @@ -74,7 +74,7 @@ from pyrit.score.true_false.true_false_inverter_scorer import TrueFalseInverterScorer if TYPE_CHECKING: - from collections.abc import AsyncIterator + from collections.abc import AsyncIterator, Callable from pathlib import Path from pyrit.models.literals import PromptDataType @@ -377,6 +377,7 @@ def __init__( attack_id: ComponentIdentifier, attack_strategy_name: str, modality_router: _ModalityFeedbackRouter, + record_objective_conversation: Callable[..., None], use_score_as_feedback: bool = True, memory_labels: dict[str, str] | None = None, parent_id: str | None = None, @@ -405,6 +406,8 @@ def __init__( whether prior media should travel back to the adversarial chat or forward to the objective target, and fills adversarial-placeholder pieces in seed messages. Typically shared across all nodes of the same attack. + record_objective_conversation (Callable[..., None]): Records an objective-target + conversation ID for cleanup before each objective send. use_score_as_feedback (bool): Whether subsequent adversarial prompts include the objective score. Defaults to True. memory_labels (dict[str, str] | None): Labels for memory storage. @@ -432,6 +435,7 @@ def __init__( self._attack_strategy_name = attack_strategy_name self._memory_labels = memory_labels or {} self._modality_router = modality_router + self._record_objective_conversation = record_objective_conversation self._prepended_conversation_config = prepended_conversation_config or PrependedConversationConfig() self._use_score_as_feedback = use_score_as_feedback @@ -682,6 +686,7 @@ async def _send_prompt_to_target_async(self, prompt: str) -> Message: objective_target_conversation_id=self.objective_target_conversation_id, objective=self._objective, ): + self._record_objective_conversation(conversation_id=self.objective_target_conversation_id) response = await self._prompt_normalizer.send_prompt_async( message=message, request_converter_configurations=self._request_converters, @@ -761,6 +766,7 @@ async def _send_initial_prompt_to_target_async(self) -> Message: objective_target_conversation_id=self.objective_target_conversation_id, objective=self._objective, ): + self._record_objective_conversation(conversation_id=self.objective_target_conversation_id) response = await self._prompt_normalizer.send_prompt_async( message=message, request_converter_configurations=self._request_converters, @@ -967,6 +973,7 @@ def duplicate(self) -> _TreeOfAttacksNode: attack_id=self._attack_id, attack_strategy_name=self._attack_strategy_name, modality_router=self._modality_router, + record_objective_conversation=self._record_objective_conversation, use_score_as_feedback=self._use_score_as_feedback, memory_labels=self._memory_labels, desired_response_prefix=self._desired_response_prefix, @@ -1392,9 +1399,9 @@ async def execute_nodes_async( """ Execute nodes in ordered batches and yield each completed batch. - Node instances own all branch-specific mutable state. This executor only - schedules their existing execution protocol, so failures and cancellation - retain ``asyncio.gather`` semantics. + Node instances own all branch-specific mutable state. If one node fails, + the executor cancels and awaits the other nodes before it propagates the + error. Args: nodes (list[_TreeOfAttacksNode]): Nodes to execute. @@ -1408,7 +1415,14 @@ async def execute_nodes_async( batch_nodes = nodes[batch_start : batch_start + self._batch_size] self._log_batch_start(batch_start=batch_start, batch_nodes=batch_nodes, total_nodes=len(nodes)) - await asyncio.gather(*(node.send_prompt_async(objective=objective) for node in batch_nodes)) + tasks = [asyncio.create_task(node.send_prompt_async(objective=objective)) for node in batch_nodes] + try: + await asyncio.gather(*tasks) + except BaseException: + for task in tasks: + task.cancel() + await asyncio.gather(*tasks, return_exceptions=True) + raise yield batch_start, batch_nodes @@ -2198,6 +2212,7 @@ def _create_attack_node( attack_id=self.get_identifier(), attack_strategy_name=self.__class__.__name__, modality_router=self._modality_router, + record_objective_conversation=context._record_objective_target_invocation, use_score_as_feedback=self._attack_scoring_config.use_score_as_feedback, memory_labels=context.memory_labels, desired_response_prefix=self._configuration.desired_response_prefix, diff --git a/pyrit/executor/attack/single_turn/prompt_sending.py b/pyrit/executor/attack/single_turn/prompt_sending.py index c37889d7e9..6307b9adf3 100644 --- a/pyrit/executor/attack/single_turn/prompt_sending.py +++ b/pyrit/executor/attack/single_turn/prompt_sending.py @@ -317,6 +317,7 @@ async def _send_prompt_to_objective_target_async( objective_target_conversation_id=context.conversation_id, objective=context.params.objective, ): + context._record_objective_target_invocation(conversation_id=context.conversation_id) return await self._prompt_normalizer.send_prompt_async( message=message, target=self._objective_target, diff --git a/pyrit/prompt_target/common/prompt_target.py b/pyrit/prompt_target/common/prompt_target.py index 504c3e35ee..790abcbba4 100644 --- a/pyrit/prompt_target/common/prompt_target.py +++ b/pyrit/prompt_target/common/prompt_target.py @@ -360,6 +360,24 @@ def set_system_prompt( ).to_message(), ) + async def reset_conversation_async(self, *, conversation_id: str) -> None: + """ + Release any target-side state held for a conversation. + + The attack execution scope calls this for objective-target conversations + recorded at the common dispatch boundary. Targets that keep external state + keyed by conversation (a websocket connection, a browser page, an upstream + session) override this to close or discard it. Targets that are stateless + between calls need not override it. + + This is best-effort cleanup, so implementations should not raise for a + conversation id they do not recognize, and should be safe to call more + than once for the same id. + + Args: + conversation_id (str): The conversation id to release state for. + """ + def dispose_db_engine(self) -> None: """ Dispose database engine to release database connections and resources. diff --git a/pyrit/prompt_target/openai/openai_realtime_target.py b/pyrit/prompt_target/openai/openai_realtime_target.py index 54e7a6f57a..511f228b44 100644 --- a/pyrit/prompt_target/openai/openai_realtime_target.py +++ b/pyrit/prompt_target/openai/openai_realtime_target.py @@ -11,6 +11,7 @@ from openai import AsyncOpenAI from pyrit.common import forward_init_parameters +from pyrit.common.deprecation import print_deprecation_message from pyrit.exceptions import ( pyrit_target_retry, ) @@ -479,21 +480,45 @@ async def cleanup_target_async(self) -> None: logger.warning(f"Error closing realtime client: {e}") self._realtime_client = None - async def cleanup_conversation_async(self, conversation_id: str) -> None: + async def reset_conversation_async(self, *, conversation_id: str) -> None: """ Disconnects from the Realtime API for a specific conversation. + Closes the cached connection for ``conversation_id`` and drops it from + ``_existing_conversation``. Errors while closing are logged and + swallowed, and an unknown conversation id is a no-op, so this is safe + to call from attack lifecycle cleanup. + Args: conversation_id (str): The conversation ID to disconnect from. """ - connection = self._existing_conversation.get(conversation_id) - if connection: - try: - await connection.close() - logger.info(f"Disconnected from {self._endpoint} with conversation ID: {conversation_id}") - except Exception as e: - logger.warning(f"Error closing connection for {conversation_id}: {e}") - del self._existing_conversation[conversation_id] + connection = self._existing_conversation.pop(conversation_id, None) + if not connection: + return + + try: + await connection.close() + except Exception as error: # noqa: BLE001 - cleanup must not replace the attack outcome + logger.warning(f"Error closing connection for {conversation_id}: {error}") + return + + logger.info(f"Disconnected from {self._endpoint} with conversation ID: {conversation_id}") + + async def cleanup_conversation_async(self, conversation_id: str) -> None: + """ + Disconnect from the Realtime API for a specific conversation. + + Deprecated. Use ``reset_conversation_async`` instead. + + Args: + conversation_id (str): The conversation ID to disconnect from. + """ + print_deprecation_message( + old_item="RealtimeTarget.cleanup_conversation_async", + new_item="RealtimeTarget.reset_conversation_async", + removed_in="1.3.0", + ) + await self.reset_conversation_async(conversation_id=conversation_id) async def _connect_async(self, *, conversation_id: str) -> Any: """ diff --git a/pyrit/prompt_target/websocket_target.py b/pyrit/prompt_target/websocket_target.py index d8fb94a1b4..ee196d61bb 100644 --- a/pyrit/prompt_target/websocket_target.py +++ b/pyrit/prompt_target/websocket_target.py @@ -11,6 +11,7 @@ from websockets.asyncio.client import ClientConnection from websockets.protocol import State +from pyrit.common.deprecation import print_deprecation_message 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 @@ -182,32 +183,44 @@ async def _send_text_async(self, *, text: str, conversation_id: str) -> str: 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: + async def reset_conversation_async(self, *, conversation_id: str) -> None: """ Close and remove one conversation connection. + Called from attack lifecycle cleanup once a conversation is finished. An + unknown conversation id is a no-op, so this is safe to call more than once. + 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 - 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 + await websocket.close() + except Exception as error: # noqa: BLE001 - cleanup must not replace the attack outcome + logger.warning("Error closing WebSocket conversation %s: %s", conversation_id, error) + return logger.info("Disconnected WebSocket conversation: %s", conversation_id) + async def cleanup_conversation_async(self, conversation_id: str) -> None: + """ + Close and remove one conversation connection. + + Deprecated. Use ``reset_conversation_async`` instead. + + Args: + conversation_id (str): PyRIT conversation ID. + """ + print_deprecation_message( + old_item="WebsocketTarget.cleanup_conversation_async", + new_item="WebsocketTarget.reset_conversation_async", + removed_in="1.3.0", + ) + await self.reset_conversation_async(conversation_id=conversation_id) + async def cleanup_target_async(self) -> None: """ Close and remove all conversation connections. diff --git a/tests/unit/executor/attack/core/test_attack_strategy.py b/tests/unit/executor/attack/core/test_attack_strategy.py index 9cf0fbeb9d..2332b321e2 100644 --- a/tests/unit/executor/attack/core/test_attack_strategy.py +++ b/tests/unit/executor/attack/core/test_attack_strategy.py @@ -4,7 +4,7 @@ import asyncio import logging from dataclasses import replace -from unittest.mock import MagicMock, patch +from unittest.mock import AsyncMock, MagicMock, patch import pytest @@ -20,6 +20,7 @@ AttackContext, AttackStrategy, _DefaultAttackStrategyEventHandler, + _ObjectiveTargetConversationLifecycle, ) from pyrit.executor.attack.multi_turn.multi_turn_attack_strategy import ConversationSession, MultiTurnAttackContext from pyrit.executor.attack.multi_turn.tree_of_attacks import TAPAttackContext @@ -103,6 +104,70 @@ def event_handler(mock_logger): return _DefaultAttackStrategyEventHandler(logger=mock_logger) +async def test_objective_target_conversation_lifecycle_resets_unique_conversations() -> None: + target = MagicMock(spec=PromptTarget) + target.reset_conversation_async = AsyncMock() + lifecycle = _ObjectiveTargetConversationLifecycle( + objective_target=target, + logger=logging.getLogger(__name__), + ) + + async with lifecycle: + lifecycle.record_invocation(conversation_id="conversation-1") + lifecycle.record_invocation(conversation_id="conversation-1") + lifecycle.record_invocation(conversation_id="conversation-2") + + assert target.reset_conversation_async.await_count == 2 + reset_ids = {call.kwargs["conversation_id"] for call in target.reset_conversation_async.await_args_list} + assert reset_ids == {"conversation-1", "conversation-2"} + + +async def test_objective_target_cleanup_error_does_not_replace_attack_error() -> None: + target = MagicMock(spec=PromptTarget) + target.reset_conversation_async = AsyncMock(side_effect=RuntimeError("cleanup failed")) + mock_logger = MagicMock(spec=logging.Logger) + lifecycle = _ObjectiveTargetConversationLifecycle( + objective_target=target, + logger=mock_logger, + ) + + with pytest.raises(ValueError, match="attack failed"): + async with lifecycle: + lifecycle.record_invocation(conversation_id="conversation-id") + raise ValueError("attack failed") + + mock_logger.warning.assert_called_once() + + +async def test_objective_target_cleanup_propagates_cancellation() -> None: + target = MagicMock(spec=PromptTarget) + target.reset_conversation_async = AsyncMock(side_effect=asyncio.CancelledError()) + lifecycle = _ObjectiveTargetConversationLifecycle( + objective_target=target, + logger=logging.getLogger(__name__), + ) + + with pytest.raises(asyncio.CancelledError): + async with lifecycle: + lifecycle.record_invocation(conversation_id="conversation-id") + + +async def test_objective_target_cleanup_attempts_all_resets_under_cancellation() -> None: + target = MagicMock(spec=PromptTarget) + target.reset_conversation_async = AsyncMock(side_effect=asyncio.CancelledError()) + lifecycle = _ObjectiveTargetConversationLifecycle( + objective_target=target, + logger=logging.getLogger(__name__), + ) + + with pytest.raises(asyncio.CancelledError): + async with lifecycle: + lifecycle.record_invocation(conversation_id="conversation-1") + lifecycle.record_invocation(conversation_id="conversation-2") + + assert target.reset_conversation_async.await_count == 2 + + def test_next_message_override_can_clear_parameter_value_and_survive_copy(): """An explicit None override must not fall back to the immutable parameter after copying.""" diff --git a/tests/unit/executor/attack/multi_turn/test_adversarial_chat_schema_forwarding.py b/tests/unit/executor/attack/multi_turn/test_adversarial_chat_schema_forwarding.py index 910a3573b1..fc945e3c48 100644 --- a/tests/unit/executor/attack/multi_turn/test_adversarial_chat_schema_forwarding.py +++ b/tests/unit/executor/attack/multi_turn/test_adversarial_chat_schema_forwarding.py @@ -115,6 +115,7 @@ async def test_tap_forwards_schema_to_adversarial_target(patch_central_database) attack_id=attack.get_identifier(), attack_strategy_name="TreeOfAttacksWithPruningAttack", modality_router=_ModalityFeedbackRouter(adversarial_chat=adversarial, objective_target=objective), + record_objective_conversation=lambda *, conversation_id: None, ) await node._send_to_adversarial_chat_async(prompt_text="hello") diff --git a/tests/unit/executor/attack/multi_turn/test_prepended_history_normalization.py b/tests/unit/executor/attack/multi_turn/test_prepended_history_normalization.py index 067d3b4b6b..15d4015d4d 100644 --- a/tests/unit/executor/attack/multi_turn/test_prepended_history_normalization.py +++ b/tests/unit/executor/attack/multi_turn/test_prepended_history_normalization.py @@ -296,6 +296,7 @@ def _make_tap_node(*, target: PromptTarget) -> _TreeOfAttacksNode: adversarial_chat=adversarial_chat, objective_target=target, ), + record_objective_conversation=lambda *, conversation_id: None, ) diff --git a/tests/unit/executor/attack/multi_turn/test_red_teaming.py b/tests/unit/executor/attack/multi_turn/test_red_teaming.py index 8ba062cc05..4ce0d28865 100644 --- a/tests/unit/executor/attack/multi_turn/test_red_teaming.py +++ b/tests/unit/executor/attack/multi_turn/test_red_teaming.py @@ -21,6 +21,7 @@ ) from pyrit.executor.attack.component import ConversationManager, PrependedConversationConfig from pyrit.executor.attack.core.attack_config import DEFAULT_ADVERSARIAL_FIRST_MESSAGE +from pyrit.executor.attack.core.attack_strategy import _ObjectiveTargetConversationLifecycle from pyrit.memory import CentralMemory from pyrit.message_normalizer import MessageStringNormalizer from pyrit.models import ( @@ -997,10 +998,16 @@ async def test_second_turn_uses_configured_message_normalizer_without_rotation( ) basic_context.executed_turns = 1 - await attack._send_prompt_to_objective_target_async( - context=basic_context, - message=Message.from_prompt(prompt="Second request", role="user"), - ) + async with _ObjectiveTargetConversationLifecycle( + objective_target=objective_target, + logger=attack._logger, + ) as lifecycle: + basic_context._objective_target_conversation_lifecycle = lifecycle + await attack._send_prompt_to_objective_target_async( + context=basic_context, + message=Message.from_prompt(prompt="Second request", role="user"), + ) + basic_context._objective_target_conversation_lifecycle = None assert basic_context.session.conversation_id == old_conversation_id assert objective_target.prompt_sent == ["custom formatted request"] diff --git a/tests/unit/executor/attack/multi_turn/test_supports_multi_turn_attacks.py b/tests/unit/executor/attack/multi_turn/test_supports_multi_turn_attacks.py index 7cc2ce0625..81b279d855 100644 --- a/tests/unit/executor/attack/multi_turn/test_supports_multi_turn_attacks.py +++ b/tests/unit/executor/attack/multi_turn/test_supports_multi_turn_attacks.py @@ -491,6 +491,7 @@ def _make_tap_node(self, *, supports_multi_turn: bool): adversarial_chat=adversarial_chat, objective_target=target, ), + record_objective_conversation=lambda *, conversation_id: None, ) def test_single_turn_target_duplicates_logical_history_without_seed_boundary(self): @@ -877,6 +878,7 @@ def _make_tap_node(self, *, supports_multi_turn: bool): adversarial_chat=adversarial_chat, objective_target=target, ), + record_objective_conversation=lambda *, conversation_id: None, ) def test_branching_single_turn_target_preserves_system_across_depths(self): diff --git a/tests/unit/executor/attack/multi_turn/test_tree_of_attacks.py b/tests/unit/executor/attack/multi_turn/test_tree_of_attacks.py index 854dd032a5..f5275243e7 100644 --- a/tests/unit/executor/attack/multi_turn/test_tree_of_attacks.py +++ b/tests/unit/executor/attack/multi_turn/test_tree_of_attacks.py @@ -29,6 +29,7 @@ TAPAttackScoringConfig, _TAPAttackConfiguration, _TreeOfAttacksNode, + _TreeOfAttacksNodeExecutor, ) from pyrit.models import ( JSON_SCHEMA_METADATA_KEY, @@ -50,6 +51,40 @@ logger = logging.getLogger(__name__) +async def test_node_executor_cancels_and_awaits_siblings_after_failure() -> None: + sibling_started = asyncio.Event() + sibling_finished = asyncio.Event() + + async def fail_after_sibling_starts(*, objective: str) -> None: + await sibling_started.wait() + raise RuntimeError("node failed") + + async def block_until_cancelled(*, objective: str) -> None: + sibling_started.set() + try: + await asyncio.Event().wait() + finally: + sibling_finished.set() + + failed_node = MagicMock() + failed_node.send_prompt_async = AsyncMock(side_effect=fail_after_sibling_starts) + sibling_node = MagicMock() + sibling_node.send_prompt_async = AsyncMock(side_effect=block_until_cancelled) + executor = _TreeOfAttacksNodeExecutor( + batch_size=2, + logger=logger, + ) + + with pytest.raises(RuntimeError, match="node failed"): + async for _ in executor.execute_nodes_async( + nodes=[failed_node, sibling_node], + objective="objective", + ): + pass + + assert sibling_finished.is_set() + + # Mirrors the shipped ``adversarial_chat.yaml``: every key required, no extras allowed. Used to # exercise the strict validation TAP/PAIR now inherit by delegating to the shared parser. _STRICT_ADVERSARIAL_CHAT_SCHEMA: dict = { @@ -1087,6 +1122,7 @@ async def test_score_response_delegates_to_scorer_for_blocked(self, attack_build adversarial_chat=builder.adversarial_chat, objective_target=builder.objective_target, ), + record_objective_conversation=lambda *, conversation_id: None, desired_response_prefix="Sure, here is", prompt_normalizer=normalizer, ) @@ -1154,6 +1190,7 @@ async def test_score_response_delegates_to_scorer_for_unknown_error(self, attack adversarial_chat=builder.adversarial_chat, objective_target=builder.objective_target, ), + record_objective_conversation=lambda *, conversation_id: None, desired_response_prefix="Sure, here is", prompt_normalizer=normalizer, ) @@ -1563,6 +1600,7 @@ def node_components(self, attack_builder): "attack_id": {"id": "test_attack"}, "attack_strategy_name": "TreeOfAttacksWithPruningAttack", "modality_router": modality_router, + "record_objective_conversation": lambda *, conversation_id: None, "memory_labels": {"test": "label"}, "parent_id": None, "prompt_normalizer": prompt_normalizer, @@ -2995,6 +3033,7 @@ def node_components(self, attack_builder): "attack_id": {"id": "test_attack"}, "attack_strategy_name": "TreeOfAttacksWithPruningAttack", "modality_router": modality_router, + "record_objective_conversation": lambda *, conversation_id: None, "memory_labels": {}, "parent_id": None, "prompt_normalizer": prompt_normalizer, diff --git a/tests/unit/executor/attack/single_turn/test_prompt_sending.py b/tests/unit/executor/attack/single_turn/test_prompt_sending.py index 288a2dc9d5..21dc5481a5 100644 --- a/tests/unit/executor/attack/single_turn/test_prompt_sending.py +++ b/tests/unit/executor/attack/single_turn/test_prompt_sending.py @@ -1,6 +1,7 @@ # Copyright (c) Microsoft Corporation. # Licensed under the MIT license. +import asyncio import base64 import uuid from unittest.mock import AsyncMock, MagicMock, patch @@ -36,6 +37,107 @@ from pyrit.score import Scorer, TrueFalseScorer +@pytest.mark.usefixtures("patch_central_database") +async def test_execute_resets_invoked_objective_target_conversation() -> None: + target = MockPromptTarget() + target.reset_conversation_async = AsyncMock() # type: ignore[method-assign] + attack = PromptSendingAttack(objective_target=target) + + await attack.execute_async(objective="Test objective") + + target.reset_conversation_async.assert_awaited_once() + assert target.reset_conversation_async.await_args.kwargs["conversation_id"] + + +@pytest.mark.usefixtures("patch_central_database") +async def test_execute_resets_objective_target_conversation_after_send_failure() -> None: + target = MockPromptTarget() + target._send_prompt_to_target_async = AsyncMock(side_effect=RuntimeError("send failed")) # type: ignore[method-assign] + target.reset_conversation_async = AsyncMock() # type: ignore[method-assign] + attack = PromptSendingAttack(objective_target=target) + + with pytest.raises(Exception, match="Error sending prompt"): + await attack.execute_async(objective="Test objective") + + target.reset_conversation_async.assert_awaited_once() + + +@pytest.mark.usefixtures("patch_central_database") +async def test_cancelled_execute_resets_blocked_objective_target_conversation() -> None: + target = MockPromptTarget() + send_started = asyncio.Event() + wait_forever = asyncio.Event() + sent_conversation_id: str | None = None + + async def block_send(*, normalized_conversation: list[Message]) -> list[Message]: + nonlocal sent_conversation_id + sent_conversation_id = normalized_conversation[-1].get_piece().conversation_id + send_started.set() + await wait_forever.wait() + return [] + + target._send_prompt_to_target_async = block_send # type: ignore[method-assign] + target.reset_conversation_async = AsyncMock() # type: ignore[method-assign] + attack = PromptSendingAttack(objective_target=target) + task = asyncio.create_task(attack.execute_async(objective="Test objective")) + await send_started.wait() + + task.cancel() + with pytest.raises(asyncio.CancelledError): + await task + + target.reset_conversation_async.assert_awaited_once_with(conversation_id=sent_conversation_id) + + +@pytest.mark.usefixtures("patch_central_database") +async def test_concurrent_attacks_reset_only_their_own_conversations() -> None: + target = MockPromptTarget() + send_started = { + "first objective": asyncio.Event(), + "second objective": asyncio.Event(), + } + release_send = { + "first objective": asyncio.Event(), + "second objective": asyncio.Event(), + } + conversation_ids: dict[str, str] = {} + + async def controlled_send(*, normalized_conversation: list[Message]) -> list[Message]: + request = normalized_conversation[-1] + objective = request.get_value() + conversation_id = request.get_piece().conversation_id + conversation_ids[objective] = conversation_id + send_started[objective].set() + await release_send[objective].wait() + return [ + MessagePiece( + role="assistant", + original_value="response", + conversation_id=conversation_id, + ).to_message() + ] + + target._send_prompt_to_target_async = controlled_send # type: ignore[method-assign] + target.reset_conversation_async = AsyncMock() # type: ignore[method-assign] + first_attack = PromptSendingAttack(objective_target=target) + second_attack = PromptSendingAttack(objective_target=target) + first_task = asyncio.create_task(first_attack.execute_async(objective="first objective")) + second_task = asyncio.create_task(second_attack.execute_async(objective="second objective")) + await asyncio.gather(*(event.wait() for event in send_started.values())) + + release_send["first objective"].set() + await first_task + + target.reset_conversation_async.assert_awaited_once_with(conversation_id=conversation_ids["first objective"]) + assert not second_task.done() + + release_send["second objective"].set() + await second_task + + assert target.reset_conversation_async.await_count == 2 + target.reset_conversation_async.assert_any_await(conversation_id=conversation_ids["second objective"]) + + @pytest.fixture def mock_target(): """Create a mock prompt target for testing""" diff --git a/tests/unit/prompt_target/target/test_realtime_target.py b/tests/unit/prompt_target/target/test_realtime_target.py index a3a8164a74..b29355d554 100644 --- a/tests/unit/prompt_target/target/test_realtime_target.py +++ b/tests/unit/prompt_target/target/test_realtime_target.py @@ -1412,31 +1412,54 @@ async def test_send_prompt_audio_path_calls_send_audio_async(target, tmp_path): target.send_audio_async.assert_awaited_once() -async def test_cleanup_conversation_async_closes_and_removes(target): +async def test_reset_conversation_async_closes_and_removes(target): mock_connection = AsyncMock() target._existing_conversation["conv"] = mock_connection - await target.cleanup_conversation_async(conversation_id="conv") + await target.reset_conversation_async(conversation_id="conv") mock_connection.close.assert_awaited_once() assert "conv" not in target._existing_conversation -async def test_cleanup_conversation_async_swallows_close_error(target): +async def test_reset_conversation_async_swallows_close_error(target): mock_connection = AsyncMock() mock_connection.close.side_effect = RuntimeError("close failed") target._existing_conversation["conv"] = mock_connection # The error is swallowed and the conversation is still removed. - await target.cleanup_conversation_async(conversation_id="conv") + await target.reset_conversation_async(conversation_id="conv") assert "conv" not in target._existing_conversation -async def test_cleanup_conversation_async_unknown_id_is_noop(target): +async def test_reset_conversation_async_propagates_cancellation(target): + mock_connection = AsyncMock() + mock_connection.close.side_effect = asyncio.CancelledError + target._existing_conversation["conv"] = mock_connection + + with pytest.raises(asyncio.CancelledError): + await target.reset_conversation_async(conversation_id="conv") + + mock_connection.close.assert_awaited_once() + assert "conv" not in target._existing_conversation + + +async def test_cleanup_conversation_async_warns_and_delegates(target): + mock_connection = AsyncMock() + target._existing_conversation["conv"] = mock_connection + + with pytest.warns(DeprecationWarning, match="reset_conversation_async"): + await target.cleanup_conversation_async(conversation_id="conv") + + mock_connection.close.assert_awaited_once() + assert "conv" not in target._existing_conversation + + +async def test_reset_conversation_async_unknown_id_is_noop(target): target._existing_conversation["conv"] = AsyncMock() - await target.cleanup_conversation_async(conversation_id="missing") + await target.reset_conversation_async(conversation_id="missing") assert "conv" in target._existing_conversation diff --git a/tests/unit/prompt_target/target/test_websocket_target.py b/tests/unit/prompt_target/target/test_websocket_target.py index 60a13a0a35..3274c73280 100644 --- a/tests/unit/prompt_target/target/test_websocket_target.py +++ b/tests/unit/prompt_target/target/test_websocket_target.py @@ -3,7 +3,6 @@ import asyncio import json -import sys from collections.abc import Callable from unittest.mock import AsyncMock, patch @@ -684,77 +683,57 @@ async def wait_forever(*, websocket: ClientConnection) -> str: await target._initialize_connection_async(websocket=connection) -async def test_cleanup_conversation_async_removes_connection(websocket_target: WebsocketTarget) -> None: +async def test_reset_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") + await websocket_target.reset_conversation_async(conversation_id="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") +async def test_cleanup_conversation_async_warns_and_delegates(websocket_target: WebsocketTarget) -> None: + connection = AsyncMock(spec=ClientConnection) + websocket_target._existing_conversation["conversation"] = connection + + with pytest.warns(DeprecationWarning, match="cleanup_conversation_async"): + await websocket_target.cleanup_conversation_async("conversation") + + connection.close.assert_awaited_once() + assert websocket_target._existing_conversation == {} + + +async def test_reset_conversation_async_does_not_retain_unknown_lock(websocket_target: WebsocketTarget) -> None: + await websocket_target.reset_conversation_async(conversation_id="missing") assert "missing" not in websocket_target._conversation_locks -async def test_cleanup_conversation_async_cancellation_finishes_closing_connection( +async def test_reset_conversation_async_swallows_close_error( websocket_target: WebsocketTarget, ) -> None: connection = AsyncMock(spec=ClientConnection) + connection.close.side_effect = ConnectionError("close failed") 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 + # The error is swallowed and the conversation is still removed. + await websocket_target.reset_conversation_async(conversation_id="conversation") connection.close.assert_awaited_once() assert websocket_target._existing_conversation == {} -async def test_cleanup_conversation_async_cancellation_preserved_when_close_fails( +async def test_reset_conversation_async_propagates_cancellation( websocket_target: WebsocketTarget, ) -> None: connection = AsyncMock(spec=ClientConnection) + connection.close.side_effect = asyncio.CancelledError 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 + with pytest.raises(asyncio.CancelledError): + await websocket_target.reset_conversation_async(conversation_id="conversation") - # Python 3.10 replaces the task's CancelledError and drops its chained cause. - if sys.version_info >= (3, 11): - assert exc_info.value.__cause__ is close_error connection.close.assert_awaited_once() assert websocket_target._existing_conversation == {} diff --git a/tests/unit/prompt_target/test_text_target.py b/tests/unit/prompt_target/test_text_target.py index 5ba4b9520f..03c1b41714 100644 --- a/tests/unit/prompt_target/test_text_target.py +++ b/tests/unit/prompt_target/test_text_target.py @@ -94,3 +94,10 @@ async def test_cleanup_target_does_nothing(): target = TextTarget(text_stream=io.StringIO()) # Should not raise await target.cleanup_target_async() + + +@pytest.mark.usefixtures("patch_central_database") +async def test_reset_conversation_does_nothing_for_stateless_target(): + target = TextTarget(text_stream=io.StringIO()) + # A target that keeps no per-conversation state inherits the base no-op. + await target.reset_conversation_async(conversation_id="some-conversation-id")