From f9664de0fb9c16b81ccd3d8c189f415edc11db5f Mon Sep 17 00:00:00 2001 From: Mahdi Alhakim Date: Mon, 3 Aug 2026 22:29:03 +0300 Subject: [PATCH 1/7] FEAT: Add reset_conversation_async hook to PromptTarget Targets that hold external state keyed by conversation had no standard way to release it when an attack finished. RealtimeTarget grew its own cleanup_conversation_async, and the issue reporter had to monkey-patch _teardown_async to call it. PromptTarget gains a no-op reset_conversation_async(*, conversation_id), and AttackStrategy._teardown_async now hands it every objective-target conversation the run used. That is the live conversation plus the ones recorded as PRUNED, since a PromptSendingAttack retry, a Crescendo backtrack, and the single-turn rotation in multi-turn attacks all leave earlier conversations behind. TAP keys conversations per tree node instead, so it overrides the lookup. The reset runs in the lifecycle finally block, so a target that raises is logged rather than replacing whatever error the attack was reporting. RealtimeTarget now implements the hook, and cleanup_conversation_async delegates to it with a deprecation warning. cleanup_target_async is left alone since closing the whole target is a different concern. Towards #1247 --- .github/instructions/targets.instructions.md | 26 ++++++ doc/code/targets/0_prompt_targets.md | 14 +++ .../attack/compound/sequential_attack.py | 3 - pyrit/executor/attack/core/attack_strategy.py | 61 +++++++++++++ .../attack/multi_turn/chunked_request.py | 8 -- pyrit/executor/attack/multi_turn/crescendo.py | 9 -- .../attack/multi_turn/multi_prompt_sending.py | 4 - .../executor/attack/multi_turn/red_teaming.py | 4 - .../attack/multi_turn/tree_of_attacks.py | 30 ++++--- .../attack/single_turn/prompt_sending.py | 4 - pyrit/executor/attack/streaming/barge_in.py | 4 - pyrit/prompt_target/common/prompt_target.py | 18 ++++ .../openai/openai_realtime_target.py | 24 +++++- .../attack/core/test_attack_strategy.py | 86 +++++++++++++++++++ .../multi_turn/test_multi_prompt_sending.py | 8 +- .../attack/multi_turn/test_red_teaming.py | 10 ++- .../attack/multi_turn/test_tree_of_attacks.py | 50 +++++++++++ .../attack/single_turn/test_prompt_sending.py | 6 +- .../target/test_realtime_target.py | 23 +++-- tests/unit/prompt_target/test_text_target.py | 7 ++ 20 files changed, 335 insertions(+), 64 deletions(-) diff --git a/.github/instructions/targets.instructions.md b/.github/instructions/targets.instructions.md index 7d21935b4e..d8b81ab1e6 100644 --- a/.github/instructions/targets.instructions.md +++ b/.github/instructions/targets.instructions.md @@ -43,6 +43,32 @@ class MyTarget(PromptTarget): ``send_prompt_async`` (the public entry point) is ``@final`` and MUST NOT be overridden. Override ``_send_prompt_to_target_async`` instead. +## Releasing per-conversation state + +Attacks call ``reset_conversation_async(*, conversation_id)`` from +``_teardown_async`` when they are done with a conversation id. The base +implementation is a no-op, so a target that keeps no state between calls +does not need to do anything. + +Targets that hold external state keyed by conversation (a websocket +connection, a browser page, an upstream session) SHOULD override it to +release that state: + +```python +async def reset_conversation_async(self, *, conversation_id: str) -> None: + connection = self._connections.pop(conversation_id, None) + if connection: + await connection.close() +``` + +It is best-effort cleanup, so an implementation SHOULD NOT raise for an +unknown conversation id and SHOULD be safe to call more than once for the +same id. The attack logs and swallows anything that does raise, so a +failure here never replaces the error the attack was reporting. + +Closing the whole target rather than one conversation is a different +concern and stays in ``cleanup_target_async``. + ## Keyword-only ``__init__`` is enforced Every ``PromptTarget`` subclass MUST make all ``__init__`` parameters diff --git a/doc/code/targets/0_prompt_targets.md b/doc/code/targets/0_prompt_targets.md index 85de4c9f50..de744503b2 100644 --- a/doc/code/targets/0_prompt_targets.md +++ b/doc/code/targets/0_prompt_targets.md @@ -57,6 +57,20 @@ Target-side rate limiting remains independent and continues to use `limit_reques `_send_prompt_to_target_async(*, normalized_conversation: list[Message]) -> list[Message]` instead of overriding `send_prompt_async`. +## Releasing per-conversation state + +Some targets hold state for a conversation outside of PyRIT's memory: an open websocket, a browser page, a session on the far side of an HTTP API. When an attack is finished with a conversation, it calls + +``` +async def reset_conversation_async(self, *, conversation_id: str) -> None: +``` + +on the objective target for every conversation the run used. That includes conversations the attack abandoned partway through, such as a `PromptSendingAttack` retry or a `CrescendoAttack` backtrack. + +The base implementation does nothing, so a target that keeps no state between calls does not need to override it. `RealtimeTarget` overrides it to close the websocket it caches per conversation. If you write a target that holds something similar, override it and release that state there. Do not raise for a conversation id you do not recognize, since the attack calls this while it is tearing down and treats it as best effort. + +Closing the target as a whole, rather than one conversation, is separate and is not part of this hook. + ## Chat-style targets vs general targets A `PromptTarget` is a generic place to send a prompt. With PyRIT, the idea is that it will eventually be consumed by an AI application, but that doesn't have to be immediate. For example, you could have a SharePoint target. Everything you send a prompt to is a `PromptTarget`. Many attacks work generically with any `PromptTarget` including `RedTeamingAttack` and `PromptSendingAttack`. diff --git a/pyrit/executor/attack/compound/sequential_attack.py b/pyrit/executor/attack/compound/sequential_attack.py index cf212378fa..82eab47355 100644 --- a/pyrit/executor/attack/compound/sequential_attack.py +++ b/pyrit/executor/attack/compound/sequential_attack.py @@ -242,9 +242,6 @@ def _validate_context(self, *, context: AttackContext[AttackParameters]) -> None async def _setup_async(self, *, context: AttackContext[AttackParameters]) -> None: """No-op: per-child-attack setup is owned by each inner strategy's executor.""" - async def _teardown_async(self, *, context: AttackContext[AttackParameters]) -> None: - """No-op: per-child-attack teardown is owned by each inner strategy's executor.""" - async def _perform_async(self, *, context: AttackContext[AttackParameters]) -> SequentialAttackResult: results: list[AttackResult] = [] diff --git a/pyrit/executor/attack/core/attack_strategy.py b/pyrit/executor/attack/core/attack_strategy.py index c1e260fc12..f88456729e 100644 --- a/pyrit/executor/attack/core/attack_strategy.py +++ b/pyrit/executor/attack/core/attack_strategy.py @@ -37,6 +37,7 @@ AttackResult, ComponentIdentifier, ConversationReference, + ConversationType, ConverterIdentifier, Identifiable, Message, @@ -685,6 +686,66 @@ def get_request_converters(self) -> list[Any]: """ return self._request_converters + def _get_objective_conversation_ids(self, *, context: AttackStrategyContextT) -> list[str]: + """ + Collect every objective-target conversation id this run used. + + The live conversation sits directly on single-turn contexts and on + ``session`` for multi-turn ones. A run can also leave earlier + conversations behind: a retry in ``PromptSendingAttack``, a Crescendo + backtrack, or the rotation multi-turn attacks do for single-turn + targets all mint a fresh id and record the old one as ``PRUNED``. + Those still hold target-side state, so they are collected too. + + Attacks that key conversations somewhere else should override this. + + Args: + context (AttackStrategyContextT): The context for the attack. + + Returns: + list[str]: Conversation ids to release, in no particular order and + without duplicates. + """ + ids: list[str] = [] + + live = getattr(context, "conversation_id", None) or getattr( + getattr(context, "session", None), "conversation_id", None + ) + if live: + ids.append(live) + + ids.extend( + ref.conversation_id + for ref in context.related_conversations + if ref.conversation_type == ConversationType.PRUNED + ) + return list(dict.fromkeys(ids)) + + async def _teardown_async(self, *, context: AttackStrategyContextT) -> None: + """ + Release the objective target's state for the run's conversations. + + Hands each conversation id to ``PromptTarget.reset_conversation_async`` + so targets holding external state keyed by conversation (a websocket + connection, a browser page) can close it. The base target + implementation is a no-op, so this is inert for stateless targets. + + This runs in the ``finally`` of the execution lifecycle, so a target + that raises here is logged rather than allowed to replace whatever + error the attack was already reporting. + + Subclasses that need their own teardown should override this and call + ``await super()._teardown_async(context=context)``. + + Args: + context (AttackStrategyContextT): The context for the attack. + """ + for conversation_id in self._get_objective_conversation_ids(context=context): + try: + await self._objective_target.reset_conversation_async(conversation_id=conversation_id) + except Exception as e: # noqa: BLE001 - teardown runs in a finally; never mask the attack's own error + self._logger.warning(f"Error resetting conversation {conversation_id} on the objective target: {e}") + async def execute_with_context_async(self, *, context: AttackStrategyContextT) -> AttackStrategyResultT: """ Execute an attack and persist its completed result after teardown. diff --git a/pyrit/executor/attack/multi_turn/chunked_request.py b/pyrit/executor/attack/multi_turn/chunked_request.py index b90eca4d09..9e3b0266c6 100644 --- a/pyrit/executor/attack/multi_turn/chunked_request.py +++ b/pyrit/executor/attack/multi_turn/chunked_request.py @@ -387,11 +387,3 @@ async def _score_combined_value_async( ): scores = await self._objective_scorer.score_text_async(text=combined_value, objective=objective) return scores[0] if scores else None - - async def _teardown_async(self, *, context: ChunkedRequestAttackContext) -> None: - """ - Teardown the attack by cleaning up conversation context. - - Args: - context (ChunkedRequestAttackContext): The attack context containing conversation session. - """ diff --git a/pyrit/executor/attack/multi_turn/crescendo.py b/pyrit/executor/attack/multi_turn/crescendo.py index b9fb98a677..c9bd9a5fc5 100644 --- a/pyrit/executor/attack/multi_turn/crescendo.py +++ b/pyrit/executor/attack/multi_turn/crescendo.py @@ -462,15 +462,6 @@ async def _perform_async(self, *, context: CrescendoAttackContext) -> CrescendoA result.backtrack_count = context.backtrack_count return result - async def _teardown_async(self, *, context: CrescendoAttackContext) -> None: - """ - Clean up after attack execution. - - Args: - context (CrescendoAttackContext): The attack context. - """ - # Nothing to be done here, no-op - def _build_adversarial_manager(self, *, context: CrescendoAttackContext) -> _AdversarialConversationManager: """ Build the adversarial-conversation manager that owns Crescendo's adversarial-chat turn. diff --git a/pyrit/executor/attack/multi_turn/multi_prompt_sending.py b/pyrit/executor/attack/multi_turn/multi_prompt_sending.py index eb065fb6d0..318259194b 100644 --- a/pyrit/executor/attack/multi_turn/multi_prompt_sending.py +++ b/pyrit/executor/attack/multi_turn/multi_prompt_sending.py @@ -340,10 +340,6 @@ def _determine_attack_outcome( # At least one prompt was filtered or failed to get a response return AttackOutcome.FAILURE, "At least one prompt was filtered or failed to get a response" - async def _teardown_async(self, *, context: MultiTurnAttackContext[Any]) -> None: - """Clean up after attack execution.""" - # Nothing to be done here, no-op - async def _send_prompt_to_objective_target_async( self, *, current_message: Message, context: MultiTurnAttackContext[Any] ) -> Message | None: diff --git a/pyrit/executor/attack/multi_turn/red_teaming.py b/pyrit/executor/attack/multi_turn/red_teaming.py index 44893bb755..624e86a2be 100644 --- a/pyrit/executor/attack/multi_turn/red_teaming.py +++ b/pyrit/executor/attack/multi_turn/red_teaming.py @@ -377,10 +377,6 @@ async def _perform_async(self, *, context: MultiTurnAttackContext[Any]) -> Attac labels=context.memory_labels, ) - async def _teardown_async(self, *, context: MultiTurnAttackContext[Any]) -> None: - """Clean up after attack execution.""" - # Nothing to be done here, no-op - def _build_adversarial_manager(self, *, context: MultiTurnAttackContext[Any]) -> _AdversarialConversationManager: """ Build the adversarial conversation manager for this execution. diff --git a/pyrit/executor/attack/multi_turn/tree_of_attacks.py b/pyrit/executor/attack/multi_turn/tree_of_attacks.py index 7e86386c4c..004b029575 100644 --- a/pyrit/executor/attack/multi_turn/tree_of_attacks.py +++ b/pyrit/executor/attack/multi_turn/tree_of_attacks.py @@ -1874,22 +1874,30 @@ async def _perform_async(self, *, context: TAPAttackContext) -> TAPAttackResult: return self._create_failure_result(context) - async def _teardown_async(self, *, context: TAPAttackContext) -> None: + def _get_objective_conversation_ids(self, *, context: TAPAttackContext) -> list[str]: """ - Clean up after attack execution. + Collect the objective-target conversations across the whole tree. - This method is called automatically after attack execution completes, - regardless of success or failure. It provides an opportunity to clean - up resources, close connections, or perform other finalization tasks. - - Currently, the TAP attack does not require any specific cleanup operations - as all resources are managed by the parent components. + TAP keeps one objective conversation per node rather than a single one + on ``session``, so the surviving nodes and the best conversation are + collected alongside the pruned ones the base class already finds. Args: - context (TAPAttackContext): The attack context containing the final - state after execution. + context (TAPAttackContext): The attack context. + + Returns: + list[str]: Conversation ids to release, without duplicates. """ - # No specific teardown needed for TAP attack + ids = [node.objective_target_conversation_id for node in context.nodes] + if context.best_conversation_id: + ids.append(context.best_conversation_id) + + ids.extend( + ref.conversation_id + for ref in context.related_conversations + if ref.conversation_type == ConversationType.PRUNED + ) + return list(dict.fromkeys(ids)) async def _prepare_nodes_for_iteration_async(self, context: TAPAttackContext) -> None: """ diff --git a/pyrit/executor/attack/single_turn/prompt_sending.py b/pyrit/executor/attack/single_turn/prompt_sending.py index c37889d7e9..5db843d030 100644 --- a/pyrit/executor/attack/single_turn/prompt_sending.py +++ b/pyrit/executor/attack/single_turn/prompt_sending.py @@ -272,10 +272,6 @@ def _determine_attack_outcome( # No response at all (all attempts filtered/failed) return AttackOutcome.FAILURE, "All attempts were filtered or failed to get a response" - async def _teardown_async(self, *, context: SingleTurnAttackContext[Any]) -> None: - """Clean up after attack execution.""" - # Nothing to be done here, no-op - def _get_message(self, context: SingleTurnAttackContext[Any]) -> Message: """ Prepare the message for the attack. diff --git a/pyrit/executor/attack/streaming/barge_in.py b/pyrit/executor/attack/streaming/barge_in.py index 947b2c3221..616bc2dfad 100644 --- a/pyrit/executor/attack/streaming/barge_in.py +++ b/pyrit/executor/attack/streaming/barge_in.py @@ -154,10 +154,6 @@ async def _setup_async(self, *, context: BargeInAttackContext[Any]) -> None: ] context.prepended_history_send_context = None - async def _teardown_async(self, *, context: BargeInAttackContext[Any]) -> None: - """No-op teardown — connection / dispatcher are closed inside the session's ``run_async``.""" - return - async def _perform_async(self, *, context: BargeInAttackContext[Any]) -> AttackResult: """ Drive the realtime streaming session and collect per-turn assistant messages. diff --git a/pyrit/prompt_target/common/prompt_target.py b/pyrit/prompt_target/common/prompt_target.py index 504c3e35ee..e90b17c457 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. + + Attacks call this from ``_teardown_async`` once they are done with a + conversation id. 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..ac1886a74b 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,10 +480,15 @@ 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 teardown. + Args: conversation_id (str): The conversation ID to disconnect from. """ @@ -495,6 +501,22 @@ async def cleanup_conversation_async(self, conversation_id: str) -> None: logger.warning(f"Error closing connection for {conversation_id}: {e}") del self._existing_conversation[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: """ Open a fresh Realtime API websocket connection and return the connection handle. diff --git a/tests/unit/executor/attack/core/test_attack_strategy.py b/tests/unit/executor/attack/core/test_attack_strategy.py index 9cf0fbeb9d..12a165e782 100644 --- a/tests/unit/executor/attack/core/test_attack_strategy.py +++ b/tests/unit/executor/attack/core/test_attack_strategy.py @@ -23,12 +23,15 @@ ) from pyrit.executor.attack.multi_turn.multi_turn_attack_strategy import ConversationSession, MultiTurnAttackContext from pyrit.executor.attack.multi_turn.tree_of_attacks import TAPAttackContext +from pyrit.executor.attack.single_turn.single_turn_attack_strategy import SingleTurnAttackContext from pyrit.executor.core import StrategyEvent, StrategyEventData from pyrit.memory.central_memory import CentralMemory from pyrit.models import ( AttackOutcome, AttackResult, ComponentIdentifier, + ConversationReference, + ConversationType, Message, SeedPrompt, ) @@ -329,6 +332,89 @@ async def test_execute_async_allows_optional_parameters_as_none(self, mock_attac assert result is not None +@pytest.mark.usefixtures("patch_central_database") +class TestAttackStrategyTeardown: + """Tests for the objective target conversation reset in _teardown_async""" + + def _strategy(self, target): + class TeardownStrategy(AttackStrategy): + def __init__(self, **kwargs): + super().__init__(context_type=AttackContext, logger=logging.getLogger(), **kwargs) + + def _validate_context(self, *, context): + pass + + async def _setup_async(self, *, context): + pass + + async def _perform_async(self, *, context): + raise NotImplementedError + + return TeardownStrategy(objective_target=target) + + def _target(self): + target = MagicMock(spec=PromptTarget) + target.get_identifier.return_value = _mock_target_id() + return target + + async def test_teardown_resets_single_turn_conversation(self): + target = self._target() + context = SingleTurnAttackContext(params=AttackParameters(objective="o")) + + await self._strategy(target)._teardown_async(context=context) + + target.reset_conversation_async.assert_awaited_once_with(conversation_id=context.conversation_id) + + async def test_teardown_resets_multi_turn_session_conversation(self): + target = self._target() + context = MultiTurnAttackContext(params=AttackParameters(objective="o")) + + await self._strategy(target)._teardown_async(context=context) + + target.reset_conversation_async.assert_awaited_once_with(conversation_id=context.session.conversation_id) + + async def test_teardown_also_resets_pruned_conversations(self): + target = self._target() + context = SingleTurnAttackContext(params=AttackParameters(objective="o")) + context.related_conversations.add( + ConversationReference(conversation_id="pruned-1", conversation_type=ConversationType.PRUNED) + ) + + await self._strategy(target)._teardown_async(context=context) + + reset_ids = {call.kwargs["conversation_id"] for call in target.reset_conversation_async.await_args_list} + assert reset_ids == {context.conversation_id, "pruned-1"} + + async def test_teardown_ignores_non_pruned_related_conversations(self): + target = self._target() + context = SingleTurnAttackContext(params=AttackParameters(objective="o")) + context.related_conversations.add( + ConversationReference(conversation_id="adv-1", conversation_type=ConversationType.ADVERSARIAL) + ) + + await self._strategy(target)._teardown_async(context=context) + + target.reset_conversation_async.assert_awaited_once_with(conversation_id=context.conversation_id) + + async def test_teardown_skips_reset_without_conversation_id(self, sample_attack_context): + target = self._target() + + # The base AttackContext carries neither a conversation_id nor a session. + await self._strategy(target)._teardown_async(context=sample_attack_context) + + target.reset_conversation_async.assert_not_awaited() + + async def test_teardown_swallows_target_errors(self): + target = self._target() + target.reset_conversation_async.side_effect = RuntimeError("connection already closed") + context = SingleTurnAttackContext(params=AttackParameters(objective="o")) + + # Teardown runs in a finally block, so it must not replace the attack's own error. + await self._strategy(target)._teardown_async(context=context) + + target.reset_conversation_async.assert_awaited_once() + + @pytest.mark.usefixtures("patch_central_database") class TestDefaultAttackStrategyEventHandler: """Tests for the default attack strategy event handler""" diff --git a/tests/unit/executor/attack/multi_turn/test_multi_prompt_sending.py b/tests/unit/executor/attack/multi_turn/test_multi_prompt_sending.py index 40dee2a434..54d2025a01 100644 --- a/tests/unit/executor/attack/multi_turn/test_multi_prompt_sending.py +++ b/tests/unit/executor/attack/multi_turn/test_multi_prompt_sending.py @@ -727,9 +727,11 @@ def test_attack_has_same_identifier_for_same_config(self, mock_target): assert attack1.get_identifier().hash == attack2.get_identifier().hash assert attack1.get_identifier().class_name == "MultiPromptSendingAttack" - async def test_teardown_async_is_noop(self, mock_target, basic_context): + async def test_teardown_async_resets_target_conversation(self, mock_target, basic_context): attack = MultiPromptSendingAttack(objective_target=mock_target) - # Should complete without error await attack._teardown_async(context=basic_context) - # No assertions needed - we just want to ensure it runs without exceptions + + mock_target.reset_conversation_async.assert_awaited_once_with( + conversation_id=basic_context.session.conversation_id + ) 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..406b0d365f 100644 --- a/tests/unit/executor/attack/multi_turn/test_red_teaming.py +++ b/tests/unit/executor/attack/multi_turn/test_red_teaming.py @@ -1509,14 +1509,14 @@ async def test_execute_with_context_async_successful( assert result.outcome == AttackOutcome.SUCCESS assert result.objective == basic_context.objective - async def test_teardown_async_is_noop( + async def test_teardown_async_resets_target_conversation( self, mock_objective_target: MagicMock, mock_objective_scorer: MagicMock, mock_adversarial_chat: MagicMock, basic_context: MultiTurnAttackContext, ): - """Test that teardown completes without errors.""" + """Test that teardown releases the objective target's conversation.""" adversarial_config = AttackAdversarialConfig(target=mock_adversarial_chat) scoring_config = AttackScoringConfig(objective_scorer=mock_objective_scorer) @@ -1526,9 +1526,11 @@ async def test_teardown_async_is_noop( attack_scoring_config=scoring_config, ) - # Should complete without error await attack._teardown_async(context=basic_context) - # No assertions needed - we just want to ensure it runs without exceptions + + mock_objective_target.reset_conversation_async.assert_awaited_once_with( + conversation_id=basic_context.session.conversation_id + ) @pytest.mark.usefixtures("patch_central_database") 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..11786097f6 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 @@ -3278,3 +3278,53 @@ def test_inline_system_prompt_string_resolved_and_in_identity(self): ) assert attack._adversarial_chat_system_seed_prompt.value == "tap persona {{ desired_prefix }}" assert attack.get_identifier().params["adversarial_system_prompt"] == "tap persona {{ desired_prefix }}" + + +@pytest.mark.usefixtures("patch_central_database") +class TestTAPConversationReset: + """TAP keeps one objective conversation per node, not a single one on session.""" + + def _context_with_nodes(self, *node_ids: str) -> TAPAttackContext: + context = TAPAttackContext(params=AttackParameters(objective="Test objective")) + for node_id in node_ids: + node = MagicMock(spec=_TreeOfAttacksNode) + node.objective_target_conversation_id = node_id + context.nodes.append(node) + return context + + def test_collects_surviving_node_conversations(self, basic_attack): + context = self._context_with_nodes("node-a", "node-b") + + ids = basic_attack._get_objective_conversation_ids(context=context) + + assert set(ids) == {"node-a", "node-b"} + + def test_collects_best_and_pruned_conversations(self, basic_attack): + context = self._context_with_nodes("node-a") + context.best_conversation_id = "best-1" + context.related_conversations.add( + ConversationReference(conversation_id="pruned-1", conversation_type=ConversationType.PRUNED) + ) + + ids = basic_attack._get_objective_conversation_ids(context=context) + + assert set(ids) == {"node-a", "best-1", "pruned-1"} + + def test_does_not_use_the_unused_session_conversation_id(self, basic_attack): + context = self._context_with_nodes("node-a") + + ids = basic_attack._get_objective_conversation_ids(context=context) + + # TAP never sends anything on session.conversation_id. + assert context.session.conversation_id not in ids + + def test_returns_no_duplicates(self, basic_attack): + context = self._context_with_nodes("node-a") + context.best_conversation_id = "node-a" + context.related_conversations.add( + ConversationReference(conversation_id="node-a", conversation_type=ConversationType.PRUNED) + ) + + ids = basic_attack._get_objective_conversation_ids(context=context) + + assert ids == ["node-a"] 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..62671376c3 100644 --- a/tests/unit/executor/attack/single_turn/test_prompt_sending.py +++ b/tests/unit/executor/attack/single_turn/test_prompt_sending.py @@ -1105,12 +1105,12 @@ async def test_execute_async_execution_error_still_calls_teardown(self, mock_tar attack._perform_async.assert_called_once_with(context=basic_context) attack._teardown_async.assert_called_once_with(context=basic_context) - async def test_teardown_async_is_noop(self, mock_target, basic_context): + async def test_teardown_async_resets_target_conversation(self, mock_target, basic_context): attack = PromptSendingAttack(objective_target=mock_target) - # Should complete without error await attack._teardown_async(context=basic_context) - # No assertions needed - we just want to ensure it runs without raising + + mock_target.reset_conversation_async.assert_awaited_once_with(conversation_id=basic_context.conversation_id) async def test_execute_async_with_parameters(self, mock_target, sample_response): """Test execute_async creates context using factory method and executes attack""" diff --git a/tests/unit/prompt_target/target/test_realtime_target.py b/tests/unit/prompt_target/target/test_realtime_target.py index a3a8164a74..8212aa65d7 100644 --- a/tests/unit/prompt_target/target/test_realtime_target.py +++ b/tests/unit/prompt_target/target/test_realtime_target.py @@ -1412,31 +1412,42 @@ 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_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/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") From 48d4b26d2aa4a1fda950edc6f0ec0057785cd922 Mon Sep 17 00:00:00 2001 From: Mahdi Alhakim Date: Fri, 28 Aug 2026 01:38:17 +0300 Subject: [PATCH 2/7] FIX: Record every objective conversation TAP opens, as it opens it TAP kept objective-target conversations that nothing could name afterwards, so nothing could release the target-side state behind them. Two ways they went missing. Against a single-turn objective target a node mints a fresh conversation id every turn and dropped the one it replaced; a node holds no reference to the context, so the old id went nowhere. And the branches still standing when the run ended were never recorded, so only the winner survived as result.conversation_id. Measured on a real tree at depth 3: a single-turn target served 3 conversations with 1 still reachable, and a multi-turn target served 6 with the result reporting 5. PAIRAttack subclasses TAP and overrides none of this, so it had both gaps too. A node now reports each conversation as its send returns, which is the moment the target starts holding state for it, and the attack records it straight onto the context. Recording it there rather than while building the result is what makes it survive a run that raises or is cancelled, which are the runs most likely to leave a connection open: raising in the second iteration of a depth 3 width 3 tree used to lose 2 of 3 conversations. The winning branch is taken back out when the result is built, since it becomes result.conversation_id and would otherwise be reported twice. Reporting on send rather than on rotation also means a branched node does not record the conversation it was cloned from but never sent on. Single-turn with branching_factor 2 serves 10 conversations and records exactly those 10. Towards #1247 --- .../attack/multi_turn/tree_of_attacks.py | 82 ++++-- ...test_adversarial_chat_schema_forwarding.py | 1 + .../test_prepended_history_normalization.py | 1 + .../test_supports_multi_turn_attacks.py | 2 + .../attack/multi_turn/test_tree_of_attacks.py | 238 ++++++++++++++++-- 5 files changed, 294 insertions(+), 30 deletions(-) diff --git a/pyrit/executor/attack/multi_turn/tree_of_attacks.py b/pyrit/executor/attack/multi_turn/tree_of_attacks.py index 004b029575..2e8993385d 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, + report_objective_conversation: Callable[[str], None], use_score_as_feedback: bool = True, memory_labels: dict[str, str] | None = None, parent_id: str | None = None, @@ -405,6 +406,10 @@ 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. + report_objective_conversation (Callable[[str], None]): Called with each + objective-target conversation id this node sends on, as soon as the send + returns. The attack records them so a conversation stays nameable even + if the run ends before a result is built. 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. @@ -449,6 +454,9 @@ def __init__( # Conversation tracking self.objective_target_conversation_id = str(uuid.uuid4()) self.adversarial_chat_conversation_id = str(uuid.uuid4()) + # Reports every objective-target conversation this node actually sends on, so + # the attack can record it while the run is still going. + self._report_objective_conversation = report_objective_conversation # Execution results (populated after send_prompt_async) self.completed = False @@ -695,6 +703,11 @@ async def _send_prompt_to_target_async(self, prompt: str) -> Message: send_context=self._prepended_history_send_context, ) + # Report before returning. From here the target holds state for this + # conversation, and a single-turn rotation replaces the id on the next turn + # without telling anything else. + self._report_objective_conversation(self.objective_target_conversation_id) + # Store the full response so subsequent turns can forward media when supported. self.last_response = response logger.debug(f"Node {self.node_id}: Received response from target") @@ -774,6 +787,11 @@ async def _send_initial_prompt_to_target_async(self) -> Message: send_context=self._prepended_history_send_context, ) + # Report before returning. From here the target holds state for this + # conversation, and a single-turn rotation replaces the id on the next turn + # without telling anything else. + self._report_objective_conversation(self.objective_target_conversation_id) + # Store the full response so subsequent turns can forward media when supported. self.last_response = response logger.debug(f"Node {self.node_id}: Received response from target") @@ -967,6 +985,7 @@ def duplicate(self) -> _TreeOfAttacksNode: attack_id=self._attack_id, attack_strategy_name=self._attack_strategy_name, modality_router=self._modality_router, + report_objective_conversation=self._report_objective_conversation, use_score_as_feedback=self._use_score_as_feedback, memory_labels=self._memory_labels, desired_response_prefix=self._desired_response_prefix, @@ -1874,30 +1893,54 @@ async def _perform_async(self, *, context: TAPAttackContext) -> TAPAttackResult: return self._create_failure_result(context) - def _get_objective_conversation_ids(self, *, context: TAPAttackContext) -> list[str]: + def _make_objective_conversation_recorder(self, *, context: TAPAttackContext) -> Callable[[str], None]: """ - Collect the objective-target conversations across the whole tree. + Build the callback nodes use to report an objective-target conversation. - TAP keeps one objective conversation per node rather than a single one - on ``session``, so the surviving nodes and the best conversation are - collected alongside the pruned ones the base class already finds. + TAP keeps one conversation per node and, against a single-turn target, mints + a fresh id every turn. Recording them only when the result is built would + lose every conversation opened by a run that raises or is cancelled, which + are the runs most likely to leave a connection open. Recording as each send + returns means a conversation is nameable from the moment the target holds + state for it. Args: - context (TAPAttackContext): The attack context. + context (TAPAttackContext): The attack context to record onto. Returns: - list[str]: Conversation ids to release, without duplicates. + Callable[[str], None]: Recorder for one objective-target conversation id. """ - ids = [node.objective_target_conversation_id for node in context.nodes] - if context.best_conversation_id: - ids.append(context.best_conversation_id) - ids.extend( - ref.conversation_id - for ref in context.related_conversations - if ref.conversation_type == ConversationType.PRUNED + def record(conversation_id: str) -> None: + context.related_conversations.add( + ConversationReference( + conversation_id=conversation_id, + conversation_type=ConversationType.PRUNED, + ) + ) + + return record + + def _release_best_conversation(self, context: TAPAttackContext) -> None: + """ + Stop reporting the winning branch as pruned. + + Every conversation is recorded while the run is in flight, before there is + any way to know which branch will win. The one that does becomes + ``result.conversation_id``, so leaving it in ``related_conversations`` would + report it twice. + + Args: + context (TAPAttackContext): The attack context. + """ + if not context.best_conversation_id: + return + context.related_conversations.discard( + ConversationReference( + conversation_id=context.best_conversation_id, + conversation_type=ConversationType.PRUNED, + ) ) - return list(dict.fromkeys(ids)) async def _prepare_nodes_for_iteration_async(self, context: TAPAttackContext) -> None: """ @@ -2206,6 +2249,7 @@ def _create_attack_node( attack_id=self.get_identifier(), attack_strategy_name=self.__class__.__name__, modality_router=self._modality_router, + report_objective_conversation=self._make_objective_conversation_recorder(context=context), 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, @@ -2392,6 +2436,10 @@ def _create_attack_result( from the top node, calculates tree statistics, and populates all TAP-specific metadata fields. + Drops the winning branch from ``context.related_conversations`` before + copying it onto the result, since which branch wins is only known once the + run is over. Both endings come through here, so that happens exactly once. + Args: context (TAPAttackContext): The attack context containing the final state after execution, including best conversation ID, score, and tree visualization. @@ -2403,6 +2451,8 @@ def _create_attack_result( about the attack execution, including conversation ID, objective, outcome, outcome reason, executed turns, last response, last score, and additional metadata. """ + self._release_best_conversation(context) + last_response = self._get_result_response( conversation_id=context.best_conversation_id, score=context.best_objective_score, 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..808c71f5d6 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), + report_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..83818e521b 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, ), + report_objective_conversation=lambda conversation_id: None, ) 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..7c50add347 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, ), + report_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, ), + report_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 11786097f6..0469b3a3bd 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 @@ -30,6 +30,7 @@ _TAPAttackConfiguration, _TreeOfAttacksNode, ) +from pyrit.memory.central_memory import CentralMemory from pyrit.models import ( JSON_SCHEMA_METADATA_KEY, AttackOutcome, @@ -1087,6 +1088,7 @@ async def test_score_response_delegates_to_scorer_for_blocked(self, attack_build adversarial_chat=builder.adversarial_chat, objective_target=builder.objective_target, ), + report_objective_conversation=lambda conversation_id: None, desired_response_prefix="Sure, here is", prompt_normalizer=normalizer, ) @@ -1154,6 +1156,7 @@ async def test_score_response_delegates_to_scorer_for_unknown_error(self, attack adversarial_chat=builder.adversarial_chat, objective_target=builder.objective_target, ), + report_objective_conversation=lambda conversation_id: None, desired_response_prefix="Sure, here is", prompt_normalizer=normalizer, ) @@ -1563,6 +1566,7 @@ def node_components(self, attack_builder): "attack_id": {"id": "test_attack"}, "attack_strategy_name": "TreeOfAttacksWithPruningAttack", "modality_router": modality_router, + "report_objective_conversation": lambda conversation_id: None, "memory_labels": {"test": "label"}, "parent_id": None, "prompt_normalizer": prompt_normalizer, @@ -2995,6 +2999,7 @@ def node_components(self, attack_builder): "attack_id": {"id": "test_attack"}, "attack_strategy_name": "TreeOfAttacksWithPruningAttack", "modality_router": modality_router, + "report_objective_conversation": lambda conversation_id: None, "memory_labels": {}, "parent_id": None, "prompt_normalizer": prompt_normalizer, @@ -3282,7 +3287,13 @@ def test_inline_system_prompt_string_resolved_and_in_identity(self): @pytest.mark.usefixtures("patch_central_database") class TestTAPConversationReset: - """TAP keeps one objective conversation per node, not a single one on session.""" + """Every objective-target conversation TAP opens has to stay reachable. + + TAP keeps one conversation per node and rotates it per turn against a + single-turn target. Anything it abandons without recording is a conversation + no caller can name afterwards, so the base teardown cannot release it and the + result under-reports what the run opened. + """ def _context_with_nodes(self, *node_ids: str) -> TAPAttackContext: context = TAPAttackContext(params=AttackParameters(objective="Test objective")) @@ -3292,32 +3303,65 @@ def _context_with_nodes(self, *node_ids: str) -> TAPAttackContext: context.nodes.append(node) return context - def test_collects_surviving_node_conversations(self, basic_attack): - context = self._context_with_nodes("node-a", "node-b") - - ids = basic_attack._get_objective_conversation_ids(context=context) - - assert set(ids) == {"node-a", "node-b"} - - def test_collects_best_and_pruned_conversations(self, basic_attack): + def test_the_base_lookup_finds_the_best_conversation(self, basic_attack): context = self._context_with_nodes("node-a") context.best_conversation_id = "best-1" - context.related_conversations.add( - ConversationReference(conversation_id="pruned-1", conversation_type=ConversationType.PRUNED) - ) ids = basic_attack._get_objective_conversation_ids(context=context) - assert set(ids) == {"node-a", "best-1", "pruned-1"} + assert "best-1" in ids - def test_does_not_use_the_unused_session_conversation_id(self, basic_attack): + def test_the_base_lookup_does_not_use_the_unused_session_conversation_id(self, basic_attack): context = self._context_with_nodes("node-a") + context.best_conversation_id = "best-1" ids = basic_attack._get_objective_conversation_ids(context=context) # TAP never sends anything on session.conversation_id. assert context.session.conversation_id not in ids + def test_a_conversation_is_recorded_as_soon_as_it_is_sent_on(self, basic_attack): + context = self._context_with_nodes() + record = basic_attack._make_objective_conversation_recorder(context=context) + + record("turn-1") + record("turn-2") + + # Recorded while the run is still going, so a run that raises or is + # cancelled can still name them. related_conversations is a set, so the + # order the lookup returns them in is not meaningful. + assert set(basic_attack._get_objective_conversation_ids(context=context)) == {"turn-1", "turn-2"} + + def test_recording_the_same_conversation_twice_records_it_once(self, basic_attack): + context = self._context_with_nodes() + record = basic_attack._make_objective_conversation_recorder(context=context) + + record("turn-1") + record("turn-1") + + assert {ref.conversation_id for ref in context.related_conversations} == {"turn-1"} + + def test_the_winning_branch_stops_being_reported_as_pruned(self, basic_attack): + context = self._context_with_nodes() + record = basic_attack._make_objective_conversation_recorder(context=context) + record("node-a") + record("node-b") + context.best_conversation_id = "node-a" + + basic_attack._release_best_conversation(context) + + # It is still released, through the live conversation rather than the pruned list. + assert {ref.conversation_id for ref in context.related_conversations} == {"node-b"} + assert set(basic_attack._get_objective_conversation_ids(context=context)) == {"node-a", "node-b"} + + def test_releasing_the_best_branch_without_one_is_a_noop(self, basic_attack): + context = self._context_with_nodes() + basic_attack._make_objective_conversation_recorder(context=context)("node-a") + + basic_attack._release_best_conversation(context) + + assert {ref.conversation_id for ref in context.related_conversations} == {"node-a"} + def test_returns_no_duplicates(self, basic_attack): context = self._context_with_nodes("node-a") context.best_conversation_id = "node-a" @@ -3328,3 +3372,169 @@ def test_returns_no_duplicates(self, basic_attack): ids = basic_attack._get_objective_conversation_ids(context=context) assert ids == ["node-a"] + + +@pytest.mark.usefixtures("patch_central_database") +class TestTAPConversationsAreAllReachable: + """End to end, with real nodes: nothing the objective target served goes missing. + + ``TestTAPConversationReset`` covers the pieces in isolation. This drives the + real ``_TreeOfAttacksNode`` so the wiring is covered too, which is where a + conversation actually goes missing: the per-turn rotation against a + single-turn target, and the branch a run walks away from. + """ + + def _run_and_collect(self, *, attack_builder, supports_multi_turn, depth, width, branching): + """Run TAP and return (ids the target served, ids the run can still name).""" + served: list[str] = [] + + attack = ( + attack_builder.with_supports_multi_turn(supports_multi_turn) + .with_default_mocks() + .with_tree_params(tree_depth=depth, tree_width=width, branching_factor=branching) + .build() + ) + objective_target = attack._objective_target + memory = CentralMemory.get_memory_instance() + + async def record_and_reply(**kwargs): + conversation_id = kwargs.get("conversation_id") + reply = Message( + message_pieces=[ + MessagePiece( + role="assistant", + original_value="response", + converted_value="response", + conversation_id=conversation_id, + ) + ] + ) + if kwargs.get("target") is not objective_target: + return reply + served.append(conversation_id) + request = Message( + message_pieces=[ + MessagePiece( + role="user", + original_value="request", + converted_value="request", + conversation_id=conversation_id, + ) + ] + ) + for message in (request, reply): + for piece in message.message_pieces: + piece.not_in_memory = False + memory.add_message_to_memory(request=message) + return reply + + normalizer = MagicMock(spec=PromptNormalizer) + normalizer.send_prompt_async = AsyncMock(side_effect=record_and_reply) + attack._prompt_normalizer = normalizer + attack._node_executor._prompt_normalizer = normalizer + + return attack, served + + async def _execute(self, attack, context): + async def score(node_self, *, response, objective): + node_self.objective_score = MagicMock( + spec=Score, get_value=MagicMock(return_value=0.1), score_metadata=None + ) + + with patch.object( + _TreeOfAttacksNode, "_generate_adversarial_prompt_async", new_callable=AsyncMock, return_value="prompt" + ): + with patch.object(_TreeOfAttacksNode, "_score_response_async", new=score): + await attack._setup_async(context=context) + return await attack._perform_async(context=context) + + @pytest.mark.parametrize( + "supports_multi_turn, depth, width, branching", + [ + pytest.param(False, 3, 1, 1, id="single_turn_rotates_per_turn"), + pytest.param(True, 3, 2, 2, id="multi_turn_branches_per_node"), + ], + ) + async def test_every_conversation_the_target_served_stays_reachable( + self, attack_builder, supports_multi_turn, depth, width, branching + ): + attack, served = self._run_and_collect( + attack_builder=attack_builder, + supports_multi_turn=supports_multi_turn, + depth=depth, + width=width, + branching=branching, + ) + context = TAPAttackContext(params=AttackParameters(objective="Test objective")) + + result = await self._execute(attack, context) + + assert served, "the run has to have sent something for this to mean anything" + # Nothing the target served may be left without a name: the teardown reset + # and the result readers both work from these two. + assert set(served) <= set(attack._get_objective_conversation_ids(context=context)) + assert set(served) <= result.get_active_conversation_ids() + + async def test_a_conversation_that_was_never_used_is_not_recorded(self, attack_builder): + attack, served = self._run_and_collect( + attack_builder=attack_builder, supports_multi_turn=False, depth=3, width=1, branching=1 + ) + context = TAPAttackContext(params=AttackParameters(objective="Test objective")) + + result = await self._execute(attack, context) + + # A node is constructed with a conversation id it rotates away from before + # its first send. That conversation has no messages, so recording it would + # put an empty conversation in front of the user. Checked on an unbranched + # tree, where every conversation that exists is one the target served; + # branching also duplicates conversations, which are real but never sent. + assert result.get_active_conversation_ids() == set(served) + + async def test_the_context_and_the_result_agree(self, attack_builder): + attack, _ = self._run_and_collect( + attack_builder=attack_builder, supports_multi_turn=True, depth=3, width=2, branching=2 + ) + context = TAPAttackContext(params=AttackParameters(objective="Test objective")) + + result = await self._execute(attack, context) + + assert set(attack._get_objective_conversation_ids(context=context)) == result.get_active_conversation_ids() + + async def test_the_winning_conversation_is_not_also_reported_as_pruned(self, attack_builder): + attack, served = self._run_and_collect( + attack_builder=attack_builder, supports_multi_turn=True, depth=3, width=2, branching=2 + ) + context = TAPAttackContext(params=AttackParameters(objective="Test objective")) + + result = await self._execute(attack, context) + + assert result.conversation_id, "the run has to have picked a best branch" + # Every conversation is recorded while the run is in flight, before there is + # any way to know which branch wins. The winner has to come back out. + assert result.conversation_id not in result.get_pruned_conversation_ids() + assert result.conversation_id in result.get_active_conversation_ids() + + async def test_a_run_that_raises_still_names_everything_it_served(self, attack_builder): + """The path that matters most, because a run that blew up is the one holding connections.""" + attack, served = self._run_and_collect( + attack_builder=attack_builder, supports_multi_turn=True, depth=3, width=3, branching=1 + ) + context = TAPAttackContext(params=AttackParameters(objective="Test objective")) + + iterations = {"count": 0} + prepare = type(attack)._prepare_nodes_for_iteration_async + + async def fail_on_the_second_iteration(self, context): + iterations["count"] += 1 + if iterations["count"] >= 2: + raise ValueError("blew up mid-run") + await prepare(self, context=context) + + with patch.object(type(attack), "_prepare_nodes_for_iteration_async", new=fail_on_the_second_iteration): + with pytest.raises(ValueError): + await self._execute(attack, context) + + assert served, "the run has to have sent something for this to mean anything" + # No result is built on this path, so anything recorded only at result time + # would be lost, and teardown is the only hook that still runs. + assert set(served) <= set(attack._get_objective_conversation_ids(context=context)) From b29f81c36559ca4b0dcd6f90fc1ade2af630110c Mon Sep 17 00:00:00 2001 From: Mahdi Alhakim Date: Fri, 28 Aug 2026 01:39:40 +0300 Subject: [PATCH 3/7] MAINT: Resolve the objective conversation in one place Two copies of the same lookup had grown up next to each other. #2322 taught the error-result builder to read context.conversation_id and fall back to context.session.conversation_id; _get_objective_conversation_ids was doing the same walk a few hundred lines away. They are one resolver now, so the conversation an error result is filed under and the conversations teardown releases cannot drift apart. That commit also gave TAPAttackContext a conversation_id property, which is what the TAP override of _get_objective_conversation_ids existed to work around. With TAP's branch bookkeeping fixed the base lookup covers it, so the override is gone and no attack overrides the lookup. What is left is AttackResult.get_active_conversation_ids() read off the context rather than the result, with a test pinning the two equal. The context is what teardown has: execute_with_context_async re-raises rather than returning, so a run that fails or is cancelled produces no result for anything downstream to read, and those are the runs that leave connections open. Tests cover all three endings. An attack whose context keeps the live conversation somewhere else should expose it as a conversation_id property, the way TAPAttackContext reports the best branch, rather than overriding the lookup and putting the answer back in two places. Written down in the docstring and in the target instructions, alongside what the pass covers: the objective target only, since adversarial, scorer and converter targets have their own lifetimes. Towards #1247 --- .github/instructions/targets.instructions.md | 9 +++ doc/code/targets/0_prompt_targets.md | 4 +- pyrit/executor/attack/core/attack_strategy.py | 68 ++++++++++++++----- .../attack/core/test_attack_strategy.py | 63 +++++++++++++++++ 4 files changed, 125 insertions(+), 19 deletions(-) diff --git a/.github/instructions/targets.instructions.md b/.github/instructions/targets.instructions.md index d8b81ab1e6..75da569f9f 100644 --- a/.github/instructions/targets.instructions.md +++ b/.github/instructions/targets.instructions.md @@ -50,6 +50,15 @@ Attacks call ``reset_conversation_async(*, conversation_id)`` from implementation is a no-op, so a target that keeps no state between calls does not need to do anything. +The attack decides when a conversation is over and says so; the target only +releases what it holds. Only the **objective** target is reset. Adversarial, +scorer and converter targets have their own lifetimes and are out of scope. + +An attack whose context keeps the live conversation somewhere other than +``conversation_id`` or ``session.conversation_id`` should expose it as a +``conversation_id`` property, the way ``TAPAttackContext`` reports the best +branch. That is the same property the error-result builder reads. + Targets that hold external state keyed by conversation (a websocket connection, a browser page, an upstream session) SHOULD override it to release that state: diff --git a/doc/code/targets/0_prompt_targets.md b/doc/code/targets/0_prompt_targets.md index de744503b2..67481005d4 100644 --- a/doc/code/targets/0_prompt_targets.md +++ b/doc/code/targets/0_prompt_targets.md @@ -69,7 +69,9 @@ on the objective target for every conversation the run used. That includes conve The base implementation does nothing, so a target that keeps no state between calls does not need to override it. `RealtimeTarget` overrides it to close the websocket it caches per conversation. If you write a target that holds something similar, override it and release that state there. Do not raise for a conversation id you do not recognize, since the attack calls this while it is tearing down and treats it as best effort. -Closing the target as a whole, rather than one conversation, is separate and is not part of this hook. +The reset runs from the attack's teardown, which is in the `finally` of the execution lifecycle, so it covers runs that succeed, runs that raise and runs that are cancelled. + +Two things are deliberately out of scope. **Only the objective target is reset.** An attack can also drive an adversarial chat target, a scorer target and converter targets; those have their own lifetimes and are not released here, which is why the adversarial conversations an attack records are skipped. And **closing the target as a whole** is a different lifetime from releasing one conversation, so it stays where it is rather than moving into this hook. ## Chat-style targets vs general targets diff --git a/pyrit/executor/attack/core/attack_strategy.py b/pyrit/executor/attack/core/attack_strategy.py index f88456729e..ac4def2f00 100644 --- a/pyrit/executor/attack/core/attack_strategy.py +++ b/pyrit/executor/attack/core/attack_strategy.py @@ -169,6 +169,29 @@ def next_message(self, value: Message | None) -> None: self._next_message_override = value +def _resolve_live_conversation_id(*, context: AttackContext[Any]) -> str | None: + """ + Return the objective-target conversation the run is currently using. + + A context that declares ``conversation_id`` answers for itself, including when + the answer is that it has none: ``TAPAttackContext`` reports the best branch + and has nothing to report before the first one is chosen. Only a context + without the attribute falls through to its conversation session, which is + where multi-turn contexts keep it. + + Args: + context (AttackContext[Any]): The context for the attack. + + Returns: + str | None: The conversation id, or ``None`` when the context has none. + """ + if hasattr(context, "conversation_id"): + candidate = getattr(context, "conversation_id", None) + else: + candidate = getattr(getattr(context, "session", None), "conversation_id", None) + return candidate if isinstance(candidate, str) and candidate else None + + class _DefaultAttackStrategyEventHandler(StrategyEventHandler[AttackStrategyContextT, AttackStrategyResultT]): """ Default event handler for attack strategies. @@ -385,11 +408,7 @@ async def _on_error_async( collector = get_retry_collector() retry_events = collector.events if collector else [] - # Multi-turn contexts keep the active ID on their conversation session. - conversation_id = getattr(context, "conversation_id", None) - if not conversation_id: - conversation_id = getattr(getattr(context, "session", None), "conversation_id", None) - conversation_id = conversation_id or str(uuid.uuid4()) + conversation_id = _resolve_live_conversation_id(context=context) or str(uuid.uuid4()) error_result = AttackResult( conversation_id=conversation_id, @@ -690,14 +709,25 @@ def _get_objective_conversation_ids(self, *, context: AttackStrategyContextT) -> """ Collect every objective-target conversation id this run used. - The live conversation sits directly on single-turn contexts and on - ``session`` for multi-turn ones. A run can also leave earlier - conversations behind: a retry in ``PromptSendingAttack``, a Crescendo - backtrack, or the rotation multi-turn attacks do for single-turn - targets all mint a fresh id and record the old one as ``PRUNED``. - Those still hold target-side state, so they are collected too. + This is ``AttackResult.get_active_conversation_ids()`` read off the + context instead of the result: the live conversation plus the ones + recorded as ``PRUNED``. A run leaves conversations behind whenever it + mints a fresh id mid-run, which a ``PromptSendingAttack`` retry, a + Crescendo backtrack, the single-turn rotation in multi-turn attacks and + TAP branching all do. Those still hold target-side state. + + The two are pinned equal by test. Reading the context rather than the + result is what lets teardown run on the paths where no result exists, + which is every failed and every cancelled run. + + Adversarial, scorer and converter conversations belong to other targets + and are deliberately not included; ``get_active_conversation_ids()`` + excludes them for the same reason. - Attacks that key conversations somewhere else should override this. + An attack that keeps its live conversation somewhere else should expose + it as ``conversation_id`` on its context, the way ``TAPAttackContext`` + reports the best branch, rather than overriding this. That keeps one + lookup, and it is the same property the error-result builder reads. Args: context (AttackStrategyContextT): The context for the attack. @@ -708,9 +738,7 @@ def _get_objective_conversation_ids(self, *, context: AttackStrategyContextT) -> """ ids: list[str] = [] - live = getattr(context, "conversation_id", None) or getattr( - getattr(context, "session", None), "conversation_id", None - ) + live = _resolve_live_conversation_id(context=context) if live: ids.append(live) @@ -730,9 +758,13 @@ async def _teardown_async(self, *, context: AttackStrategyContextT) -> None: connection, a browser page) can close it. The base target implementation is a no-op, so this is inert for stateless targets. - This runs in the ``finally`` of the execution lifecycle, so a target - that raises here is logged rather than allowed to replace whatever - error the attack was already reporting. + This pass covers the objective target only. Adversarial, scorer and + converter targets have their own lifetimes and are not released here. + + This runs in the ``finally`` of the execution lifecycle, so it covers + runs that succeed, runs that raise and runs that are cancelled, and a + target that raises here is logged rather than allowed to replace + whatever error the attack was already reporting. Subclasses that need their own teardown should override this and call ``await super()._teardown_async(context=context)``. diff --git a/tests/unit/executor/attack/core/test_attack_strategy.py b/tests/unit/executor/attack/core/test_attack_strategy.py index 12a165e782..bcec0795fb 100644 --- a/tests/unit/executor/attack/core/test_attack_strategy.py +++ b/tests/unit/executor/attack/core/test_attack_strategy.py @@ -23,6 +23,7 @@ ) from pyrit.executor.attack.multi_turn.multi_turn_attack_strategy import ConversationSession, MultiTurnAttackContext from pyrit.executor.attack.multi_turn.tree_of_attacks import TAPAttackContext +from pyrit.executor.attack.single_turn.prompt_sending import PromptSendingAttack from pyrit.executor.attack.single_turn.single_turn_attack_strategy import SingleTurnAttackContext from pyrit.executor.core import StrategyEvent, StrategyEventData from pyrit.memory.central_memory import CentralMemory @@ -41,6 +42,7 @@ ) from pyrit.models.retry_event import RetryEvent from pyrit.prompt_target import PromptTarget +from tests.unit.mocks import MockPromptTarget def _mock_target_id(name: str = "MockTarget") -> ComponentIdentifier: @@ -415,6 +417,67 @@ async def test_teardown_swallows_target_errors(self): target.reset_conversation_async.assert_awaited_once() +class _RecordingTarget(MockPromptTarget): + """Objective target that records every conversation it is asked to release.""" + + def __init__(self, *, failure: Exception | None = None, block: asyncio.Event | None = None) -> None: + super().__init__() + self.reset_calls: list[str] = [] + self.started = asyncio.Event() + self._failure = failure + self._block = block + + async def _send_prompt_to_target_async(self, *, normalized_conversation: list[Message]) -> list[Message]: + self.started.set() + if self._failure: + raise self._failure + if self._block: + await self._block.wait() + return await super()._send_prompt_to_target_async(normalized_conversation=normalized_conversation) + + async def reset_conversation_async(self, *, conversation_id: str) -> None: + self.reset_calls.append(conversation_id) + + +@pytest.mark.usefixtures("patch_central_database") +class TestObjectiveConversationRelease: + """The reset has to reach every way a run can end, not only the ones that return a result.""" + + async def test_the_context_lookup_matches_the_result_contract(self): + target = _RecordingTarget() + attack = PromptSendingAttack(objective_target=target) + + result = await attack.execute_async(objective="o") + + # _get_objective_conversation_ids is AttackResult.get_active_conversation_ids + # read off the context. Pinned equal here so the two cannot drift apart. + assert set(target.reset_calls) == result.get_active_conversation_ids() + + async def test_a_failed_run_still_releases_its_conversation(self): + target = _RecordingTarget(failure=RuntimeError("target exploded mid-run")) + attack = PromptSendingAttack(objective_target=target) + + with pytest.raises(Exception): # noqa: B017 - the wrapper type is not the point + await attack.execute_async(objective="o") + + # execute_async re-raises rather than returning, so a caller holding only the + # return value has no conversation id to release. + assert len(target.reset_calls) == 1 + + async def test_a_cancelled_run_still_releases_its_conversation(self): + block = asyncio.Event() + target = _RecordingTarget(block=block) + attack = PromptSendingAttack(objective_target=target) + + task = asyncio.create_task(attack.execute_async(objective="o")) + await asyncio.wait_for(target.started.wait(), timeout=10) + task.cancel() + with pytest.raises(asyncio.CancelledError): + await task + + assert len(target.reset_calls) == 1 + + @pytest.mark.usefixtures("patch_central_database") class TestDefaultAttackStrategyEventHandler: """Tests for the default attack strategy event handler""" From 548157e17e280a3f42423a86b089ecb499eb7791 Mon Sep 17 00:00:00 2001 From: Mahdi Alhakim Date: Fri, 28 Aug 2026 01:39:40 +0300 Subject: [PATCH 4/7] FEAT: Implement the reset hook on WebsocketTarget too WebsocketTarget has the same shape as RealtimeTarget: a PromptTarget that caches one connection per conversation id in _existing_conversation, with its own cleanup_conversation_async that nothing calls. Leaving it out would have deprecated that method on one target while the identical method stayed live on the other, and a scenario sharing a WebsocketTarget would keep accumulating connections for exactly the reason this PR exists. Its close path is more careful than RealtimeTarget's, holding a per-conversation lock and shielding the close so a cancellation still finishes it, so the body is unchanged and only the name and the keyword-only signature move. The four tests covering it move with it, and cleanup_conversation_async keeps working through the same deprecation shim. Towards #1247 --- doc/code/targets/0_prompt_targets.md | 2 +- pyrit/prompt_target/websocket_target.py | 22 ++++++++++++++- .../target/test_websocket_target.py | 27 +++++++++++++------ 3 files changed, 41 insertions(+), 10 deletions(-) diff --git a/doc/code/targets/0_prompt_targets.md b/doc/code/targets/0_prompt_targets.md index 67481005d4..80893834f6 100644 --- a/doc/code/targets/0_prompt_targets.md +++ b/doc/code/targets/0_prompt_targets.md @@ -67,7 +67,7 @@ async def reset_conversation_async(self, *, conversation_id: str) -> None: on the objective target for every conversation the run used. That includes conversations the attack abandoned partway through, such as a `PromptSendingAttack` retry or a `CrescendoAttack` backtrack. -The base implementation does nothing, so a target that keeps no state between calls does not need to override it. `RealtimeTarget` overrides it to close the websocket it caches per conversation. If you write a target that holds something similar, override it and release that state there. Do not raise for a conversation id you do not recognize, since the attack calls this while it is tearing down and treats it as best effort. +The base implementation does nothing, so a target that keeps no state between calls does not need to override it. `RealtimeTarget` and `WebsocketTarget` override it to close the websocket each caches per conversation. If you write a target that holds something similar, override it and release that state there. Do not raise for a conversation id you do not recognize, since the attack calls this while it is tearing down and treats it as best effort. The reset runs from the attack's teardown, which is in the `finally` of the execution lifecycle, so it covers runs that succeed, runs that raise and runs that are cancelled. diff --git a/pyrit/prompt_target/websocket_target.py b/pyrit/prompt_target/websocket_target.py index d8fb94a1b4..b49188a753 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,10 +183,13 @@ 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 teardown 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. @@ -208,6 +212,22 @@ async def cleanup_conversation_async(self, conversation_id: str) -> None: raise 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/prompt_target/target/test_websocket_target.py b/tests/unit/prompt_target/target/test_websocket_target.py index 60a13a0a35..d59273d8e0 100644 --- a/tests/unit/prompt_target/target/test_websocket_target.py +++ b/tests/unit/prompt_target/target/test_websocket_target.py @@ -684,23 +684,34 @@ 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_cancellation_finishes_closing_connection( websocket_target: WebsocketTarget, ) -> None: connection = AsyncMock(spec=ClientConnection) @@ -713,7 +724,7 @@ async def close_connection() -> None: await finish_close.wait() connection.close.side_effect = close_connection - cleanup_task = asyncio.create_task(websocket_target.cleanup_conversation_async("conversation")) + cleanup_task = asyncio.create_task(websocket_target.reset_conversation_async(conversation_id="conversation")) await close_started.wait() cleanup_task.cancel() @@ -728,7 +739,7 @@ async def close_connection() -> None: assert websocket_target._existing_conversation == {} -async def test_cleanup_conversation_async_cancellation_preserved_when_close_fails( +async def test_reset_conversation_async_cancellation_preserved_when_close_fails( websocket_target: WebsocketTarget, ) -> None: connection = AsyncMock(spec=ClientConnection) @@ -743,7 +754,7 @@ async def close_connection() -> None: raise close_error connection.close.side_effect = close_connection - cleanup_task = asyncio.create_task(websocket_target.cleanup_conversation_async("conversation")) + cleanup_task = asyncio.create_task(websocket_target.reset_conversation_async(conversation_id="conversation")) await close_started.wait() cleanup_task.cancel() From dbbc8face866c1dcee9748140c73789bd3e35f19 Mon Sep 17 00:00:00 2001 From: Mahdi Alhakim Date: Fri, 28 Aug 2026 04:34:25 +0300 Subject: [PATCH 5/7] FIX: Release the leading TAP branch as soon as the lead changes Two defects, both found by diffing this branch against main rather than by a test failing. TAP records every branch as PRUNED while the run is in flight, and the leading one was taken back out only while building the result. A run that raises never builds one, so its own conversation stayed in both places: it was the error result's conversation_id and a pruned entry at the same time. attack_service.list_attacks adds the main conversation's message count to the pruned ones and sums a list rather than a set, so that run's messages were counted twice in the backend, and the markdown and pretty printers rendered the conversation twice. The release now happens wherever the lead is recomputed, which is the last step of every iteration, so the invariant holds at every instant instead of only once a result exists. A branch that led and then lost it is an abandoned branch again, so it goes back. The fallback that picks a conversation when no node completed was setting the lead without releasing it; it goes through the same path now. That makes the release at result-build time unreachable, since the lead is always recomputed last, so it is gone rather than left as dead code. Second, _resolve_live_conversation_id had grown a hasattr dispatch that changed what the error-result builder does for a TAPAttackContext with no nodes and no best branch: main falls through to session.conversation_id, this returned None and the caller minted a fresh uuid. Neither id names anything real, but that is #2322's code and this PR was not asked to change it. It is back to main's exact lookup, verified by computing both over all six concrete context types: zero divergences. Towards #1247 --- doc/code/targets/0_prompt_targets.md | 2 +- pyrit/executor/attack/core/attack_strategy.py | 32 ++++++---- .../attack/multi_turn/tree_of_attacks.py | 30 +++++---- .../attack/multi_turn/test_tree_of_attacks.py | 61 +++++++++++++++++-- 4 files changed, 96 insertions(+), 29 deletions(-) diff --git a/doc/code/targets/0_prompt_targets.md b/doc/code/targets/0_prompt_targets.md index 80893834f6..28d9c349f8 100644 --- a/doc/code/targets/0_prompt_targets.md +++ b/doc/code/targets/0_prompt_targets.md @@ -69,7 +69,7 @@ on the objective target for every conversation the run used. That includes conve The base implementation does nothing, so a target that keeps no state between calls does not need to override it. `RealtimeTarget` and `WebsocketTarget` override it to close the websocket each caches per conversation. If you write a target that holds something similar, override it and release that state there. Do not raise for a conversation id you do not recognize, since the attack calls this while it is tearing down and treats it as best effort. -The reset runs from the attack's teardown, which is in the `finally` of the execution lifecycle, so it covers runs that succeed, runs that raise and runs that are cancelled. +The reset runs from the attack's teardown, which is in the `finally` of the execution lifecycle, so it covers runs that succeed, runs that raise and runs that are cancelled. An error from your implementation is logged and swallowed, but a `CancelledError` is not: cancelling a run while it is releasing stops the release, and whatever is left is `cleanup_target_async`'s job. Two things are deliberately out of scope. **Only the objective target is reset.** An attack can also drive an adversarial chat target, a scorer target and converter targets; those have their own lifetimes and are not released here, which is why the adversarial conversations an attack records are skipped. And **closing the target as a whole** is a different lifetime from releasing one conversation, so it stays where it is rather than moving into this hook. diff --git a/pyrit/executor/attack/core/attack_strategy.py b/pyrit/executor/attack/core/attack_strategy.py index ac4def2f00..a583450ab2 100644 --- a/pyrit/executor/attack/core/attack_strategy.py +++ b/pyrit/executor/attack/core/attack_strategy.py @@ -173,22 +173,24 @@ def _resolve_live_conversation_id(*, context: AttackContext[Any]) -> str | None: """ Return the objective-target conversation the run is currently using. - A context that declares ``conversation_id`` answers for itself, including when - the answer is that it has none: ``TAPAttackContext`` reports the best branch - and has nothing to report before the first one is chosen. Only a context - without the attribute falls through to its conversation session, which is - where multi-turn contexts keep it. + Single-turn contexts expose it directly and multi-turn contexts keep it on + their conversation session. ``TAPAttackContext`` overrides + ``conversation_id`` to report the best branch, so the first lookup covers it. + + This is the lookup #2322 gave the error-result builder, moved here so that + builder and the teardown reset resolve a run's conversation the same way + rather than walking the context twice. Args: context (AttackContext[Any]): The context for the attack. Returns: - str | None: The conversation id, or ``None`` when the context has none. + str | None: The conversation id, or ``None`` when the context exposes + neither layout. """ - if hasattr(context, "conversation_id"): - candidate = getattr(context, "conversation_id", None) - else: - candidate = getattr(getattr(context, "session", None), "conversation_id", None) + candidate = getattr(context, "conversation_id", None) or getattr( + getattr(context, "session", None), "conversation_id", None + ) return candidate if isinstance(candidate, str) and candidate else None @@ -762,9 +764,13 @@ async def _teardown_async(self, *, context: AttackStrategyContextT) -> None: converter targets have their own lifetimes and are not released here. This runs in the ``finally`` of the execution lifecycle, so it covers - runs that succeed, runs that raise and runs that are cancelled, and a - target that raises here is logged rather than allowed to replace - whatever error the attack was already reporting. + runs that succeed, runs that raise and runs that are cancelled. An + ``Exception`` from a target is logged rather than allowed to replace + whatever error the attack was already reporting. Cancellation is not + caught: if the run is cancelled while this is releasing, it propagates + and the conversations after it are left to ``cleanup_target_async``, + because swallowing a ``CancelledError`` to finish a cleanup loop is + worse than not finishing it. Subclasses that need their own teardown should override this and call ``await super()._teardown_async(context=context)``. diff --git a/pyrit/executor/attack/multi_turn/tree_of_attacks.py b/pyrit/executor/attack/multi_turn/tree_of_attacks.py index 2e8993385d..b09e4b7ef8 100644 --- a/pyrit/executor/attack/multi_turn/tree_of_attacks.py +++ b/pyrit/executor/attack/multi_turn/tree_of_attacks.py @@ -1921,18 +1921,29 @@ def record(conversation_id: str) -> None: return record - def _release_best_conversation(self, context: TAPAttackContext) -> None: + def _release_best_conversation(self, context: TAPAttackContext, *, previous_best: str | None = None) -> None: """ - Stop reporting the winning branch as pruned. + Stop reporting the current best branch as pruned. Every conversation is recorded while the run is in flight, before there is - any way to know which branch will win. The one that does becomes - ``result.conversation_id``, so leaving it in ``related_conversations`` would - report it twice. + any way to know which branch will lead. Whichever one does becomes + ``result.conversation_id``, so leaving it in ``related_conversations`` + would report it twice: the backend adds the main conversation's message + count to the pruned ones, and the report printers list it in both places. + + Called whenever the lead is recomputed, which is the last step of every + iteration, so the invariant holds at every instant rather than only once a + result exists. A run that raises never builds a result and would otherwise + report its own conversation twice. A branch that led and then lost it is an + abandoned branch again, so it goes back. Args: context (TAPAttackContext): The attack context. + previous_best (str | None): The branch that was leading before, if the + lead just changed. """ + if previous_best and previous_best != context.best_conversation_id: + self._make_objective_conversation_recorder(context=context)(previous_best) if not context.best_conversation_id: return context.related_conversations.discard( @@ -2193,6 +2204,7 @@ def _update_best_performing_node(self, context: TAPAttackContext) -> None: # but we ensure it is sorted to avoid making any assumptions # about the order of nodes in context.nodes. completed_nodes = self._get_completed_nodes_sorted_by_score(context.nodes) + previous_best = context.best_conversation_id if completed_nodes: best_node = completed_nodes[0] @@ -2210,6 +2222,8 @@ def _update_best_performing_node(self, context: TAPAttackContext) -> None: context.best_adversarial_conversation_id = node.adversarial_chat_conversation_id break + self._release_best_conversation(context, previous_best=previous_best) + def _create_attack_node( self, *, @@ -2436,10 +2450,6 @@ def _create_attack_result( from the top node, calculates tree statistics, and populates all TAP-specific metadata fields. - Drops the winning branch from ``context.related_conversations`` before - copying it onto the result, since which branch wins is only known once the - run is over. Both endings come through here, so that happens exactly once. - Args: context (TAPAttackContext): The attack context containing the final state after execution, including best conversation ID, score, and tree visualization. @@ -2451,8 +2461,6 @@ def _create_attack_result( about the attack execution, including conversation ID, objective, outcome, outcome reason, executed turns, last response, last score, and additional metadata. """ - self._release_best_conversation(context) - last_response = self._get_result_response( conversation_id=context.best_conversation_id, score=context.best_objective_score, 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 0469b3a3bd..acc2787911 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 @@ -1620,6 +1620,19 @@ def test_node_duplicate_creates_child(self, node_components): assert child_node.parent_id == parent_node.node_id assert child_node.completed is False + def test_node_duplicate_keeps_reporting_its_conversations(self, node_components): + """A child sends on its own conversation, so it needs the same recorder as its parent.""" + reported: list[str] = [] + components = {**node_components, "report_objective_conversation": reported.append} + parent_node = _TreeOfAttacksNode(**components) + + with patch.object(parent_node._memory, "duplicate_conversation", return_value="new_conv_id"): + child_node = parent_node.duplicate() + + child_node._report_objective_conversation(child_node.objective_target_conversation_id) + + assert reported == ["new_conv_id"] + def _node_with_schema(self, node_components, schema): """Build a real node whose adversarial system prompt advertises ``schema``. @@ -3321,16 +3334,16 @@ def test_the_base_lookup_does_not_use_the_unused_session_conversation_id(self, b assert context.session.conversation_id not in ids def test_a_conversation_is_recorded_as_soon_as_it_is_sent_on(self, basic_attack): - context = self._context_with_nodes() + context = self._context_with_nodes("node-a") record = basic_attack._make_objective_conversation_recorder(context=context) record("turn-1") record("turn-2") # Recorded while the run is still going, so a run that raises or is - # cancelled can still name them. related_conversations is a set, so the - # order the lookup returns them in is not meaningful. - assert set(basic_attack._get_objective_conversation_ids(context=context)) == {"turn-1", "turn-2"} + # cancelled can still name them. + ids = set(basic_attack._get_objective_conversation_ids(context=context)) + assert {"turn-1", "turn-2"} <= ids def test_recording_the_same_conversation_twice_records_it_once(self, basic_attack): context = self._context_with_nodes() @@ -3500,6 +3513,20 @@ async def test_the_context_and_the_result_agree(self, attack_builder): assert set(attack._get_objective_conversation_ids(context=context)) == result.get_active_conversation_ids() + async def test_a_branched_node_reports_its_own_conversations(self, attack_builder): + """A child gets its own conversation from duplicate(), and must report on it too.""" + attack, served = self._run_and_collect( + attack_builder=attack_builder, supports_multi_turn=True, depth=3, width=1, branching=2 + ) + context = TAPAttackContext(params=AttackParameters(objective="Test objective")) + + result = await self._execute(attack, context) + + # width=1 keeps one node per level, so every conversation beyond the first + # belongs to a branch, and nothing else records those for us. + assert len(set(served)) > 1, "branching has to have produced more than one conversation" + assert set(served) <= result.get_active_conversation_ids() + async def test_the_winning_conversation_is_not_also_reported_as_pruned(self, attack_builder): attack, served = self._run_and_collect( attack_builder=attack_builder, supports_multi_turn=True, depth=3, width=2, branching=2 @@ -3514,6 +3541,32 @@ async def test_the_winning_conversation_is_not_also_reported_as_pruned(self, att assert result.conversation_id not in result.get_pruned_conversation_ids() assert result.conversation_id in result.get_active_conversation_ids() + async def test_a_run_that_raises_does_not_report_its_own_conversation_as_pruned(self, attack_builder): + """The backend adds the main conversation's messages to the pruned ones, so it cannot be in both.""" + attack, served = self._run_and_collect( + attack_builder=attack_builder, supports_multi_turn=True, depth=3, width=2, branching=1 + ) + context = TAPAttackContext(params=AttackParameters(objective="Test objective")) + + iterations = {"count": 0} + prepare = type(attack)._prepare_nodes_for_iteration_async + + async def fail_on_the_second_iteration(self, context): + iterations["count"] += 1 + if iterations["count"] >= 2: + raise ValueError("blew up mid-run") + await prepare(self, context=context) + + with patch.object(type(attack), "_prepare_nodes_for_iteration_async", new=fail_on_the_second_iteration): + with pytest.raises(ValueError): + await self._execute(attack, context) + + # No result is built on this path, so the invariant has to already hold on + # the context the error result is assembled from. + assert context.best_conversation_id, "a branch has to have taken the lead" + pruned = {ref.conversation_id for ref in context.related_conversations} + assert context.best_conversation_id not in pruned + async def test_a_run_that_raises_still_names_everything_it_served(self, attack_builder): """The path that matters most, because a run that blew up is the one holding connections.""" attack, served = self._run_and_collect( From 284f3b63f565f7b96f00bbf34f7b279c65da32b0 Mon Sep 17 00:00:00 2001 From: Richard Lundeen Date: Fri, 28 Aug 2026 11:24:21 -0700 Subject: [PATCH 6/7] Reset objective-target conversations at the attack boundary Add an idempotent PromptTarget.reset_conversation_async hook and an execution-scoped lifecycle in AttackStrategy that records real objective-target invocations through a generic callback and releases each unique conversation after the attack completes. TAP now cancels and awaits sibling node sends before cleanup. RealtimeTarget and WebsocketTarget implement the reset hook; related_conversations stays reporting data only. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Copilot-Session: 1dc87b33-0580-45f9-a9cc-0dc9603bd2a2 --- .github/instructions/targets.instructions.md | 37 +- doc/code/targets/0_prompt_targets.md | 16 - .../attack/compound/sequential_attack.py | 3 + pyrit/executor/attack/core/attack_strategy.py | 202 +++++----- .../attack/multi_turn/chunked_request.py | 9 + pyrit/executor/attack/multi_turn/crescendo.py | 10 + .../attack/multi_turn/multi_prompt_sending.py | 5 + .../executor/attack/multi_turn/red_teaming.py | 5 + .../attack/multi_turn/tree_of_attacks.py | 112 ++---- .../attack/single_turn/prompt_sending.py | 5 + pyrit/executor/attack/streaming/barge_in.py | 4 + pyrit/prompt_normalizer/prompt_normalizer.py | 6 +- pyrit/prompt_target/common/prompt_target.py | 17 +- .../common/target_send_context.py | 13 + .../openai/openai_realtime_target.py | 29 +- pyrit/prompt_target/websocket_target.py | 4 +- .../attack/core/test_attack_strategy.py | 200 +++------- ...test_adversarial_chat_schema_forwarding.py | 2 +- .../multi_turn/test_multi_prompt_sending.py | 8 +- .../test_prepended_history_normalization.py | 2 +- .../attack/multi_turn/test_red_teaming.py | 25 +- .../test_supports_multi_turn_attacks.py | 4 +- .../attack/multi_turn/test_tree_of_attacks.py | 352 ++---------------- .../attack/single_turn/test_prompt_sending.py | 108 +++++- .../test_prompt_normalizer.py | 3 + .../target/test_prompt_target.py | 42 ++- .../target/test_realtime_target.py | 26 ++ 27 files changed, 506 insertions(+), 743 deletions(-) diff --git a/.github/instructions/targets.instructions.md b/.github/instructions/targets.instructions.md index 75da569f9f..36b985e736 100644 --- a/.github/instructions/targets.instructions.md +++ b/.github/instructions/targets.instructions.md @@ -43,40 +43,11 @@ class MyTarget(PromptTarget): ``send_prompt_async`` (the public entry point) is ``@final`` and MUST NOT be overridden. Override ``_send_prompt_to_target_async`` instead. -## Releasing per-conversation state - -Attacks call ``reset_conversation_async(*, conversation_id)`` from -``_teardown_async`` when they are done with a conversation id. The base -implementation is a no-op, so a target that keeps no state between calls -does not need to do anything. - -The attack decides when a conversation is over and says so; the target only -releases what it holds. Only the **objective** target is reset. Adversarial, -scorer and converter targets have their own lifetimes and are out of scope. - -An attack whose context keeps the live conversation somewhere other than -``conversation_id`` or ``session.conversation_id`` should expose it as a -``conversation_id`` property, the way ``TAPAttackContext`` reports the best -branch. That is the same property the error-result builder reads. - Targets that hold external state keyed by conversation (a websocket -connection, a browser page, an upstream session) SHOULD override it to -release that state: - -```python -async def reset_conversation_async(self, *, conversation_id: str) -> None: - connection = self._connections.pop(conversation_id, None) - if connection: - await connection.close() -``` - -It is best-effort cleanup, so an implementation SHOULD NOT raise for an -unknown conversation id and SHOULD be safe to call more than once for the -same id. The attack logs and swallows anything that does raise, so a -failure here never replaces the error the attack was reporting. - -Closing the whole target rather than one conversation is a different -concern and stays in ``cleanup_target_async``. +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 diff --git a/doc/code/targets/0_prompt_targets.md b/doc/code/targets/0_prompt_targets.md index 28d9c349f8..85de4c9f50 100644 --- a/doc/code/targets/0_prompt_targets.md +++ b/doc/code/targets/0_prompt_targets.md @@ -57,22 +57,6 @@ Target-side rate limiting remains independent and continues to use `limit_reques `_send_prompt_to_target_async(*, normalized_conversation: list[Message]) -> list[Message]` instead of overriding `send_prompt_async`. -## Releasing per-conversation state - -Some targets hold state for a conversation outside of PyRIT's memory: an open websocket, a browser page, a session on the far side of an HTTP API. When an attack is finished with a conversation, it calls - -``` -async def reset_conversation_async(self, *, conversation_id: str) -> None: -``` - -on the objective target for every conversation the run used. That includes conversations the attack abandoned partway through, such as a `PromptSendingAttack` retry or a `CrescendoAttack` backtrack. - -The base implementation does nothing, so a target that keeps no state between calls does not need to override it. `RealtimeTarget` and `WebsocketTarget` override it to close the websocket each caches per conversation. If you write a target that holds something similar, override it and release that state there. Do not raise for a conversation id you do not recognize, since the attack calls this while it is tearing down and treats it as best effort. - -The reset runs from the attack's teardown, which is in the `finally` of the execution lifecycle, so it covers runs that succeed, runs that raise and runs that are cancelled. An error from your implementation is logged and swallowed, but a `CancelledError` is not: cancelling a run while it is releasing stops the release, and whatever is left is `cleanup_target_async`'s job. - -Two things are deliberately out of scope. **Only the objective target is reset.** An attack can also drive an adversarial chat target, a scorer target and converter targets; those have their own lifetimes and are not released here, which is why the adversarial conversations an attack records are skipped. And **closing the target as a whole** is a different lifetime from releasing one conversation, so it stays where it is rather than moving into this hook. - ## Chat-style targets vs general targets A `PromptTarget` is a generic place to send a prompt. With PyRIT, the idea is that it will eventually be consumed by an AI application, but that doesn't have to be immediate. For example, you could have a SharePoint target. Everything you send a prompt to is a `PromptTarget`. Many attacks work generically with any `PromptTarget` including `RedTeamingAttack` and `PromptSendingAttack`. diff --git a/pyrit/executor/attack/compound/sequential_attack.py b/pyrit/executor/attack/compound/sequential_attack.py index 82eab47355..cf212378fa 100644 --- a/pyrit/executor/attack/compound/sequential_attack.py +++ b/pyrit/executor/attack/compound/sequential_attack.py @@ -242,6 +242,9 @@ def _validate_context(self, *, context: AttackContext[AttackParameters]) -> None async def _setup_async(self, *, context: AttackContext[AttackParameters]) -> None: """No-op: per-child-attack setup is owned by each inner strategy's executor.""" + async def _teardown_async(self, *, context: AttackContext[AttackParameters]) -> None: + """No-op: per-child-attack teardown is owned by each inner strategy's executor.""" + async def _perform_async(self, *, context: AttackContext[AttackParameters]) -> SequentialAttackResult: results: list[AttackResult] = [] diff --git a/pyrit/executor/attack/core/attack_strategy.py b/pyrit/executor/attack/core/attack_strategy.py index a583450ab2..6345f6b355 100644 --- a/pyrit/executor/attack/core/attack_strategy.py +++ b/pyrit/executor/attack/core/attack_strategy.py @@ -37,7 +37,6 @@ AttackResult, ComponentIdentifier, ConversationReference, - ConversationType, ConverterIdentifier, Identifiable, Message, @@ -48,6 +47,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, ) @@ -73,6 +74,55 @@ 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.""" + for conversation_id in self._conversation_ids: + try: + await self._objective_target.reset_conversation_async(conversation_id=conversation_id) + 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, + ) + + 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]): """ @@ -102,6 +152,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( @@ -168,30 +224,20 @@ 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. -def _resolve_live_conversation_id(*, context: AttackContext[Any]) -> str | None: - """ - Return the objective-target conversation the run is currently using. - - Single-turn contexts expose it directly and multi-turn contexts keep it on - their conversation session. ``TAPAttackContext`` overrides - ``conversation_id`` to report the best branch, so the first lookup covers it. - - This is the lookup #2322 gave the error-result builder, moved here so that - builder and the teardown reset resolve a run's conversation the same way - rather than walking the context twice. - - Args: - context (AttackContext[Any]): The context for the attack. + Args: + conversation_id (str): The conversation ID used by the target. - Returns: - str | None: The conversation id, or ``None`` when the context exposes - neither layout. - """ - candidate = getattr(context, "conversation_id", None) or getattr( - getattr(context, "session", None), "conversation_id", None - ) - return candidate if isinstance(candidate, str) and candidate else None + Raises: + RuntimeError: If called outside this context's attack execution. + """ + lifecycle = self._objective_target_conversation_lifecycle + if lifecycle is None: + raise RuntimeError("Objective-target invocation occurred outside the attack lifecycle.") + lifecycle.record_invocation(conversation_id=conversation_id) class _DefaultAttackStrategyEventHandler(StrategyEventHandler[AttackStrategyContextT, AttackStrategyResultT]): @@ -410,7 +456,11 @@ async def _on_error_async( collector = get_retry_collector() retry_events = collector.events if collector else [] - conversation_id = _resolve_live_conversation_id(context=context) or str(uuid.uuid4()) + # Multi-turn contexts keep the active ID on their conversation session. + conversation_id = getattr(context, "conversation_id", None) + if not conversation_id: + conversation_id = getattr(getattr(context, "session", None), "conversation_id", None) + conversation_id = conversation_id or str(uuid.uuid4()) error_result = AttackResult( conversation_id=conversation_id, @@ -707,83 +757,6 @@ def get_request_converters(self) -> list[Any]: """ return self._request_converters - def _get_objective_conversation_ids(self, *, context: AttackStrategyContextT) -> list[str]: - """ - Collect every objective-target conversation id this run used. - - This is ``AttackResult.get_active_conversation_ids()`` read off the - context instead of the result: the live conversation plus the ones - recorded as ``PRUNED``. A run leaves conversations behind whenever it - mints a fresh id mid-run, which a ``PromptSendingAttack`` retry, a - Crescendo backtrack, the single-turn rotation in multi-turn attacks and - TAP branching all do. Those still hold target-side state. - - The two are pinned equal by test. Reading the context rather than the - result is what lets teardown run on the paths where no result exists, - which is every failed and every cancelled run. - - Adversarial, scorer and converter conversations belong to other targets - and are deliberately not included; ``get_active_conversation_ids()`` - excludes them for the same reason. - - An attack that keeps its live conversation somewhere else should expose - it as ``conversation_id`` on its context, the way ``TAPAttackContext`` - reports the best branch, rather than overriding this. That keeps one - lookup, and it is the same property the error-result builder reads. - - Args: - context (AttackStrategyContextT): The context for the attack. - - Returns: - list[str]: Conversation ids to release, in no particular order and - without duplicates. - """ - ids: list[str] = [] - - live = _resolve_live_conversation_id(context=context) - if live: - ids.append(live) - - ids.extend( - ref.conversation_id - for ref in context.related_conversations - if ref.conversation_type == ConversationType.PRUNED - ) - return list(dict.fromkeys(ids)) - - async def _teardown_async(self, *, context: AttackStrategyContextT) -> None: - """ - Release the objective target's state for the run's conversations. - - Hands each conversation id to ``PromptTarget.reset_conversation_async`` - so targets holding external state keyed by conversation (a websocket - connection, a browser page) can close it. The base target - implementation is a no-op, so this is inert for stateless targets. - - This pass covers the objective target only. Adversarial, scorer and - converter targets have their own lifetimes and are not released here. - - This runs in the ``finally`` of the execution lifecycle, so it covers - runs that succeed, runs that raise and runs that are cancelled. An - ``Exception`` from a target is logged rather than allowed to replace - whatever error the attack was already reporting. Cancellation is not - caught: if the run is cancelled while this is releasing, it propagates - and the conversations after it are left to ``cleanup_target_async``, - because swallowing a ``CancelledError`` to finish a cleanup loop is - worse than not finishing it. - - Subclasses that need their own teardown should override this and call - ``await super()._teardown_async(context=context)``. - - Args: - context (AttackStrategyContextT): The context for the attack. - """ - for conversation_id in self._get_objective_conversation_ids(context=context): - try: - await self._objective_target.reset_conversation_async(conversation_id=conversation_id) - except Exception as e: # noqa: BLE001 - teardown runs in a finally; never mask the attack's own error - self._logger.warning(f"Error resetting conversation {conversation_id} on the objective target: {e}") - async def execute_with_context_async(self, *, context: AttackStrategyContextT) -> AttackStrategyResultT: """ Execute an attack and persist its completed result after teardown. @@ -798,16 +771,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 9e3b0266c6..51b7cf9352 100644 --- a/pyrit/executor/attack/multi_turn/chunked_request.py +++ b/pyrit/executor/attack/multi_turn/chunked_request.py @@ -297,6 +297,7 @@ async def _perform_async(self, *, context: ChunkedRequestAttackContext) -> Attac prepended_history_send_context=context.prepended_history_send_context, ), send_context=context.prepended_history_send_context, + target_invocation_callback=context._record_objective_target_invocation, ) # Store the response @@ -387,3 +388,11 @@ async def _score_combined_value_async( ): scores = await self._objective_scorer.score_text_async(text=combined_value, objective=objective) return scores[0] if scores else None + + async def _teardown_async(self, *, context: ChunkedRequestAttackContext) -> None: + """ + Teardown the attack by cleaning up conversation context. + + Args: + context (ChunkedRequestAttackContext): The attack context containing conversation session. + """ diff --git a/pyrit/executor/attack/multi_turn/crescendo.py b/pyrit/executor/attack/multi_turn/crescendo.py index c9bd9a5fc5..55901671b7 100644 --- a/pyrit/executor/attack/multi_turn/crescendo.py +++ b/pyrit/executor/attack/multi_turn/crescendo.py @@ -462,6 +462,15 @@ async def _perform_async(self, *, context: CrescendoAttackContext) -> CrescendoA result.backtrack_count = context.backtrack_count return result + async def _teardown_async(self, *, context: CrescendoAttackContext) -> None: + """ + Clean up after attack execution. + + Args: + context (CrescendoAttackContext): The attack context. + """ + # Nothing to be done here, no-op + def _build_adversarial_manager(self, *, context: CrescendoAttackContext) -> _AdversarialConversationManager: """ Build the adversarial-conversation manager that owns Crescendo's adversarial-chat turn. @@ -636,6 +645,7 @@ async def _send_prompt_to_objective_target_async( prepended_history_send_context=context.prepended_history_send_context, ), send_context=context.prepended_history_send_context, + target_invocation_callback=context._record_objective_target_invocation, ) if not response: diff --git a/pyrit/executor/attack/multi_turn/multi_prompt_sending.py b/pyrit/executor/attack/multi_turn/multi_prompt_sending.py index 318259194b..c6776fd64f 100644 --- a/pyrit/executor/attack/multi_turn/multi_prompt_sending.py +++ b/pyrit/executor/attack/multi_turn/multi_prompt_sending.py @@ -340,6 +340,10 @@ def _determine_attack_outcome( # At least one prompt was filtered or failed to get a response return AttackOutcome.FAILURE, "At least one prompt was filtered or failed to get a response" + async def _teardown_async(self, *, context: MultiTurnAttackContext[Any]) -> None: + """Clean up after attack execution.""" + # Nothing to be done here, no-op + async def _send_prompt_to_objective_target_async( self, *, current_message: Message, context: MultiTurnAttackContext[Any] ) -> Message | None: @@ -371,6 +375,7 @@ async def _send_prompt_to_objective_target_async( prepended_history_send_context=context.prepended_history_send_context, ), send_context=context.prepended_history_send_context, + target_invocation_callback=context._record_objective_target_invocation, ) async def _evaluate_response_async(self, *, response: Message, objective: str) -> Score | None: diff --git a/pyrit/executor/attack/multi_turn/red_teaming.py b/pyrit/executor/attack/multi_turn/red_teaming.py index 624e86a2be..b651d40ec5 100644 --- a/pyrit/executor/attack/multi_turn/red_teaming.py +++ b/pyrit/executor/attack/multi_turn/red_teaming.py @@ -377,6 +377,10 @@ async def _perform_async(self, *, context: MultiTurnAttackContext[Any]) -> Attac labels=context.memory_labels, ) + async def _teardown_async(self, *, context: MultiTurnAttackContext[Any]) -> None: + """Clean up after attack execution.""" + # Nothing to be done here, no-op + def _build_adversarial_manager(self, *, context: MultiTurnAttackContext[Any]) -> _AdversarialConversationManager: """ Build the adversarial conversation manager for this execution. @@ -494,6 +498,7 @@ async def _send_prompt_to_objective_target_async( prepended_history_send_context=context.prepended_history_send_context, ), send_context=context.prepended_history_send_context, + target_invocation_callback=context._record_objective_target_invocation, ) if response is None: diff --git a/pyrit/executor/attack/multi_turn/tree_of_attacks.py b/pyrit/executor/attack/multi_turn/tree_of_attacks.py index b09e4b7ef8..7c6beaa5eb 100644 --- a/pyrit/executor/attack/multi_turn/tree_of_attacks.py +++ b/pyrit/executor/attack/multi_turn/tree_of_attacks.py @@ -74,10 +74,11 @@ from pyrit.score.true_false.true_false_inverter_scorer import TrueFalseInverterScorer if TYPE_CHECKING: - from collections.abc import AsyncIterator, Callable + from collections.abc import AsyncIterator from pathlib import Path from pyrit.models.literals import PromptDataType + from pyrit.prompt_target.common.target_send_context import TargetInvocationCallback logger = logging.getLogger(__name__) @@ -377,7 +378,7 @@ def __init__( attack_id: ComponentIdentifier, attack_strategy_name: str, modality_router: _ModalityFeedbackRouter, - report_objective_conversation: Callable[[str], None], + target_invocation_callback: TargetInvocationCallback, use_score_as_feedback: bool = True, memory_labels: dict[str, str] | None = None, parent_id: str | None = None, @@ -406,10 +407,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. - report_objective_conversation (Callable[[str], None]): Called with each - objective-target conversation id this node sends on, as soon as the send - returns. The attack records them so a conversation stays nameable even - if the run ends before a result is built. + target_invocation_callback (TargetInvocationCallback): Callback for objective-target + provider invocations. 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. @@ -437,6 +436,7 @@ def __init__( self._attack_strategy_name = attack_strategy_name self._memory_labels = memory_labels or {} self._modality_router = modality_router + self._target_invocation_callback = target_invocation_callback self._prepended_conversation_config = prepended_conversation_config or PrependedConversationConfig() self._use_score_as_feedback = use_score_as_feedback @@ -454,9 +454,6 @@ def __init__( # Conversation tracking self.objective_target_conversation_id = str(uuid.uuid4()) self.adversarial_chat_conversation_id = str(uuid.uuid4()) - # Reports every objective-target conversation this node actually sends on, so - # the attack can record it while the run is still going. - self._report_objective_conversation = report_objective_conversation # Execution results (populated after send_prompt_async) self.completed = False @@ -701,13 +698,9 @@ async def _send_prompt_to_target_async(self, prompt: str) -> Message: prepended_history_send_context=self._prepended_history_send_context, ), send_context=self._prepended_history_send_context, + target_invocation_callback=self._target_invocation_callback, ) - # Report before returning. From here the target holds state for this - # conversation, and a single-turn rotation replaces the id on the next turn - # without telling anything else. - self._report_objective_conversation(self.objective_target_conversation_id) - # Store the full response so subsequent turns can forward media when supported. self.last_response = response logger.debug(f"Node {self.node_id}: Received response from target") @@ -785,13 +778,9 @@ async def _send_initial_prompt_to_target_async(self) -> Message: prepended_history_send_context=self._prepended_history_send_context, ), send_context=self._prepended_history_send_context, + target_invocation_callback=self._target_invocation_callback, ) - # Report before returning. From here the target holds state for this - # conversation, and a single-turn rotation replaces the id on the next turn - # without telling anything else. - self._report_objective_conversation(self.objective_target_conversation_id) - # Store the full response so subsequent turns can forward media when supported. self.last_response = response logger.debug(f"Node {self.node_id}: Received response from target") @@ -985,7 +974,7 @@ def duplicate(self) -> _TreeOfAttacksNode: attack_id=self._attack_id, attack_strategy_name=self._attack_strategy_name, modality_router=self._modality_router, - report_objective_conversation=self._report_objective_conversation, + target_invocation_callback=self._target_invocation_callback, use_score_as_feedback=self._use_score_as_feedback, memory_labels=self._memory_labels, desired_response_prefix=self._desired_response_prefix, @@ -1411,9 +1400,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. @@ -1427,7 +1416,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 @@ -1893,65 +1889,22 @@ async def _perform_async(self, *, context: TAPAttackContext) -> TAPAttackResult: return self._create_failure_result(context) - def _make_objective_conversation_recorder(self, *, context: TAPAttackContext) -> Callable[[str], None]: - """ - Build the callback nodes use to report an objective-target conversation. - - TAP keeps one conversation per node and, against a single-turn target, mints - a fresh id every turn. Recording them only when the result is built would - lose every conversation opened by a run that raises or is cancelled, which - are the runs most likely to leave a connection open. Recording as each send - returns means a conversation is nameable from the moment the target holds - state for it. - - Args: - context (TAPAttackContext): The attack context to record onto. - - Returns: - Callable[[str], None]: Recorder for one objective-target conversation id. - """ - - def record(conversation_id: str) -> None: - context.related_conversations.add( - ConversationReference( - conversation_id=conversation_id, - conversation_type=ConversationType.PRUNED, - ) - ) - - return record - - def _release_best_conversation(self, context: TAPAttackContext, *, previous_best: str | None = None) -> None: + async def _teardown_async(self, *, context: TAPAttackContext) -> None: """ - Stop reporting the current best branch as pruned. + Clean up after attack execution. - Every conversation is recorded while the run is in flight, before there is - any way to know which branch will lead. Whichever one does becomes - ``result.conversation_id``, so leaving it in ``related_conversations`` - would report it twice: the backend adds the main conversation's message - count to the pruned ones, and the report printers list it in both places. + This method is called automatically after attack execution completes, + regardless of success or failure. It provides an opportunity to clean + up resources, close connections, or perform other finalization tasks. - Called whenever the lead is recomputed, which is the last step of every - iteration, so the invariant holds at every instant rather than only once a - result exists. A run that raises never builds a result and would otherwise - report its own conversation twice. A branch that led and then lost it is an - abandoned branch again, so it goes back. + Currently, the TAP attack does not require any specific cleanup operations + as all resources are managed by the parent components. Args: - context (TAPAttackContext): The attack context. - previous_best (str | None): The branch that was leading before, if the - lead just changed. + context (TAPAttackContext): The attack context containing the final + state after execution. """ - if previous_best and previous_best != context.best_conversation_id: - self._make_objective_conversation_recorder(context=context)(previous_best) - if not context.best_conversation_id: - return - context.related_conversations.discard( - ConversationReference( - conversation_id=context.best_conversation_id, - conversation_type=ConversationType.PRUNED, - ) - ) + # No specific teardown needed for TAP attack async def _prepare_nodes_for_iteration_async(self, context: TAPAttackContext) -> None: """ @@ -2204,7 +2157,6 @@ def _update_best_performing_node(self, context: TAPAttackContext) -> None: # but we ensure it is sorted to avoid making any assumptions # about the order of nodes in context.nodes. completed_nodes = self._get_completed_nodes_sorted_by_score(context.nodes) - previous_best = context.best_conversation_id if completed_nodes: best_node = completed_nodes[0] @@ -2222,8 +2174,6 @@ def _update_best_performing_node(self, context: TAPAttackContext) -> None: context.best_adversarial_conversation_id = node.adversarial_chat_conversation_id break - self._release_best_conversation(context, previous_best=previous_best) - def _create_attack_node( self, *, @@ -2263,7 +2213,7 @@ def _create_attack_node( attack_id=self.get_identifier(), attack_strategy_name=self.__class__.__name__, modality_router=self._modality_router, - report_objective_conversation=self._make_objective_conversation_recorder(context=context), + target_invocation_callback=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 5db843d030..8f3f770c55 100644 --- a/pyrit/executor/attack/single_turn/prompt_sending.py +++ b/pyrit/executor/attack/single_turn/prompt_sending.py @@ -272,6 +272,10 @@ def _determine_attack_outcome( # No response at all (all attempts filtered/failed) return AttackOutcome.FAILURE, "All attempts were filtered or failed to get a response" + async def _teardown_async(self, *, context: SingleTurnAttackContext[Any]) -> None: + """Clean up after attack execution.""" + # Nothing to be done here, no-op + def _get_message(self, context: SingleTurnAttackContext[Any]) -> Message: """ Prepare the message for the attack. @@ -323,6 +327,7 @@ async def _send_prompt_to_objective_target_async( prepended_history_send_context=context.prepended_history_send_context, ), send_context=context.prepended_history_send_context, + target_invocation_callback=context._record_objective_target_invocation, ) async def _evaluate_response_async( diff --git a/pyrit/executor/attack/streaming/barge_in.py b/pyrit/executor/attack/streaming/barge_in.py index 616bc2dfad..947b2c3221 100644 --- a/pyrit/executor/attack/streaming/barge_in.py +++ b/pyrit/executor/attack/streaming/barge_in.py @@ -154,6 +154,10 @@ async def _setup_async(self, *, context: BargeInAttackContext[Any]) -> None: ] context.prepended_history_send_context = None + async def _teardown_async(self, *, context: BargeInAttackContext[Any]) -> None: + """No-op teardown — connection / dispatcher are closed inside the session's ``run_async``.""" + return + async def _perform_async(self, *, context: BargeInAttackContext[Any]) -> AttackResult: """ Drive the realtime streaming session and collect per-turn assistant messages. diff --git a/pyrit/prompt_normalizer/prompt_normalizer.py b/pyrit/prompt_normalizer/prompt_normalizer.py index 20f990f77d..e8e393560c 100644 --- a/pyrit/prompt_normalizer/prompt_normalizer.py +++ b/pyrit/prompt_normalizer/prompt_normalizer.py @@ -31,7 +31,7 @@ from pyrit.prompt_normalizer import ConverterConfiguration, NormalizerRequest from pyrit.prompt_target import CapabilityName, PromptTarget from pyrit.prompt_target.batch_helper import batch_task_async -from pyrit.prompt_target.common.target_send_context import TargetSendContext +from pyrit.prompt_target.common.target_send_context import TargetInvocationCallback, TargetSendContext logger = logging.getLogger(__name__) @@ -76,6 +76,7 @@ async def send_prompt_async( response_converter_configurations: list[ConverterConfiguration] | None = None, normalizer_overrides: Mapping[CapabilityName, MessageListNormalizer[Message]] | None = None, send_context: TargetSendContext | None = None, + target_invocation_callback: TargetInvocationCallback | None = None, ) -> Message: """ Send a single request to a target. @@ -91,6 +92,8 @@ async def send_prompt_async( normalizer_overrides: Optional per-send target normalizer overrides. send_context: Optional internal coordination contract for caller-owned history selection and send lifecycle state. + target_invocation_callback: Optional callback invoked immediately before + target-specific execution. Returns: Message: The response received from the target. @@ -130,6 +133,7 @@ async def send_prompt_async( message=request, normalizer_overrides=normalizer_overrides, send_context=send_context, + target_invocation_callback=target_invocation_callback, ) self.memory.add_message_to_memory(request=request) except EmptyResponseException as ex: diff --git a/pyrit/prompt_target/common/prompt_target.py b/pyrit/prompt_target/common/prompt_target.py index e90b17c457..3aa1a4d024 100644 --- a/pyrit/prompt_target/common/prompt_target.py +++ b/pyrit/prompt_target/common/prompt_target.py @@ -24,7 +24,7 @@ ) from pyrit.prompt_target.common.target_configuration import TargetConfiguration from pyrit.prompt_target.common.target_history import filter_non_replayable_messages -from pyrit.prompt_target.common.target_send_context import TargetSendContext +from pyrit.prompt_target.common.target_send_context import TargetInvocationCallback, TargetSendContext logger = logging.getLogger(__name__) @@ -142,6 +142,7 @@ async def send_prompt_async( message: Message, normalizer_overrides: Mapping[CapabilityName, MessageListNormalizer[Message]] | None = None, send_context: TargetSendContext | None = None, + target_invocation_callback: TargetInvocationCallback | None = None, ) -> list[Message]: """ Validate, normalize, and send a prompt to the target. @@ -161,6 +162,8 @@ async def send_prompt_async( normalizer_overrides: Optional per-send target normalizer overrides. send_context: Optional internal coordination contract for caller-owned history selection and send lifecycle state. + target_invocation_callback: Optional callback invoked immediately before + target-specific execution. Returns: list[Message]: Response messages from the target. @@ -185,6 +188,8 @@ async def send_prompt_async( if not normalized_conversation: raise ValueError("Normalization pipeline returned an empty conversation. Cannot send an empty request.") self._validate_request(normalized_conversation=normalized_conversation) + if target_invocation_callback: + target_invocation_callback(conversation_id=conversation_id) if send_context: send_context.mark_target_invoked() response = await self._send_prompt_to_target_async(normalized_conversation=normalized_conversation) @@ -364,11 +369,11 @@ async def reset_conversation_async(self, *, conversation_id: str) -> None: """ Release any target-side state held for a conversation. - Attacks call this from ``_teardown_async`` once they are done with a - conversation id. 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. + 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 diff --git a/pyrit/prompt_target/common/target_send_context.py b/pyrit/prompt_target/common/target_send_context.py index 2c0e4f3e4d..8dbc705f64 100644 --- a/pyrit/prompt_target/common/target_send_context.py +++ b/pyrit/prompt_target/common/target_send_context.py @@ -9,6 +9,19 @@ from pyrit.models import Message +class TargetInvocationCallback(Protocol): + """Callback invoked immediately before target-specific execution.""" + + def __call__(self, *, conversation_id: str) -> None: + """ + Record one target invocation. + + Args: + conversation_id (str): The conversation ID for the invocation. + """ + ... + + class TargetSendContext(Protocol): """Internal contract coordinating one target send with caller-owned state.""" diff --git a/pyrit/prompt_target/openai/openai_realtime_target.py b/pyrit/prompt_target/openai/openai_realtime_target.py index ac1886a74b..93a913000c 100644 --- a/pyrit/prompt_target/openai/openai_realtime_target.py +++ b/pyrit/prompt_target/openai/openai_realtime_target.py @@ -487,19 +487,32 @@ async def reset_conversation_async(self, *, conversation_id: str) -> None: 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 teardown. + to call from attack lifecycle cleanup. Args: conversation_id (str): The conversation ID to disconnect from. + + Raises: + asyncio.CancelledError: If cleanup is cancelled after the connection has finished closing. """ - connection = self._existing_conversation.get(conversation_id) - if connection: + connection = self._existing_conversation.pop(conversation_id, None) + if not connection: + return + + close_future = asyncio.ensure_future(connection.close()) + try: + await asyncio.shield(close_future) + except asyncio.CancelledError as cancellation_error: 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] + await close_future + except BaseException as close_error: + raise cancellation_error from close_error + raise + except Exception as error: + 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: """ diff --git a/pyrit/prompt_target/websocket_target.py b/pyrit/prompt_target/websocket_target.py index b49188a753..8190abe808 100644 --- a/pyrit/prompt_target/websocket_target.py +++ b/pyrit/prompt_target/websocket_target.py @@ -187,8 +187,8 @@ async def reset_conversation_async(self, *, conversation_id: str) -> None: """ Close and remove one conversation connection. - Called from attack teardown once a conversation is finished. An unknown - conversation id is a no-op, so this is safe to call more than once. + 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. diff --git a/tests/unit/executor/attack/core/test_attack_strategy.py b/tests/unit/executor/attack/core/test_attack_strategy.py index bcec0795fb..03111a0d34 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,19 +20,16 @@ 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 -from pyrit.executor.attack.single_turn.prompt_sending import PromptSendingAttack -from pyrit.executor.attack.single_turn.single_turn_attack_strategy import SingleTurnAttackContext from pyrit.executor.core import StrategyEvent, StrategyEventData from pyrit.memory.central_memory import CentralMemory from pyrit.models import ( AttackOutcome, AttackResult, ComponentIdentifier, - ConversationReference, - ConversationType, Message, SeedPrompt, ) @@ -42,7 +39,6 @@ ) from pyrit.models.retry_event import RetryEvent from pyrit.prompt_target import PromptTarget -from tests.unit.mocks import MockPromptTarget def _mock_target_id(name: str = "MockTarget") -> ComponentIdentifier: @@ -108,6 +104,54 @@ 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") + + 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.""" @@ -334,150 +378,6 @@ async def test_execute_async_allows_optional_parameters_as_none(self, mock_attac assert result is not None -@pytest.mark.usefixtures("patch_central_database") -class TestAttackStrategyTeardown: - """Tests for the objective target conversation reset in _teardown_async""" - - def _strategy(self, target): - class TeardownStrategy(AttackStrategy): - def __init__(self, **kwargs): - super().__init__(context_type=AttackContext, logger=logging.getLogger(), **kwargs) - - def _validate_context(self, *, context): - pass - - async def _setup_async(self, *, context): - pass - - async def _perform_async(self, *, context): - raise NotImplementedError - - return TeardownStrategy(objective_target=target) - - def _target(self): - target = MagicMock(spec=PromptTarget) - target.get_identifier.return_value = _mock_target_id() - return target - - async def test_teardown_resets_single_turn_conversation(self): - target = self._target() - context = SingleTurnAttackContext(params=AttackParameters(objective="o")) - - await self._strategy(target)._teardown_async(context=context) - - target.reset_conversation_async.assert_awaited_once_with(conversation_id=context.conversation_id) - - async def test_teardown_resets_multi_turn_session_conversation(self): - target = self._target() - context = MultiTurnAttackContext(params=AttackParameters(objective="o")) - - await self._strategy(target)._teardown_async(context=context) - - target.reset_conversation_async.assert_awaited_once_with(conversation_id=context.session.conversation_id) - - async def test_teardown_also_resets_pruned_conversations(self): - target = self._target() - context = SingleTurnAttackContext(params=AttackParameters(objective="o")) - context.related_conversations.add( - ConversationReference(conversation_id="pruned-1", conversation_type=ConversationType.PRUNED) - ) - - await self._strategy(target)._teardown_async(context=context) - - reset_ids = {call.kwargs["conversation_id"] for call in target.reset_conversation_async.await_args_list} - assert reset_ids == {context.conversation_id, "pruned-1"} - - async def test_teardown_ignores_non_pruned_related_conversations(self): - target = self._target() - context = SingleTurnAttackContext(params=AttackParameters(objective="o")) - context.related_conversations.add( - ConversationReference(conversation_id="adv-1", conversation_type=ConversationType.ADVERSARIAL) - ) - - await self._strategy(target)._teardown_async(context=context) - - target.reset_conversation_async.assert_awaited_once_with(conversation_id=context.conversation_id) - - async def test_teardown_skips_reset_without_conversation_id(self, sample_attack_context): - target = self._target() - - # The base AttackContext carries neither a conversation_id nor a session. - await self._strategy(target)._teardown_async(context=sample_attack_context) - - target.reset_conversation_async.assert_not_awaited() - - async def test_teardown_swallows_target_errors(self): - target = self._target() - target.reset_conversation_async.side_effect = RuntimeError("connection already closed") - context = SingleTurnAttackContext(params=AttackParameters(objective="o")) - - # Teardown runs in a finally block, so it must not replace the attack's own error. - await self._strategy(target)._teardown_async(context=context) - - target.reset_conversation_async.assert_awaited_once() - - -class _RecordingTarget(MockPromptTarget): - """Objective target that records every conversation it is asked to release.""" - - def __init__(self, *, failure: Exception | None = None, block: asyncio.Event | None = None) -> None: - super().__init__() - self.reset_calls: list[str] = [] - self.started = asyncio.Event() - self._failure = failure - self._block = block - - async def _send_prompt_to_target_async(self, *, normalized_conversation: list[Message]) -> list[Message]: - self.started.set() - if self._failure: - raise self._failure - if self._block: - await self._block.wait() - return await super()._send_prompt_to_target_async(normalized_conversation=normalized_conversation) - - async def reset_conversation_async(self, *, conversation_id: str) -> None: - self.reset_calls.append(conversation_id) - - -@pytest.mark.usefixtures("patch_central_database") -class TestObjectiveConversationRelease: - """The reset has to reach every way a run can end, not only the ones that return a result.""" - - async def test_the_context_lookup_matches_the_result_contract(self): - target = _RecordingTarget() - attack = PromptSendingAttack(objective_target=target) - - result = await attack.execute_async(objective="o") - - # _get_objective_conversation_ids is AttackResult.get_active_conversation_ids - # read off the context. Pinned equal here so the two cannot drift apart. - assert set(target.reset_calls) == result.get_active_conversation_ids() - - async def test_a_failed_run_still_releases_its_conversation(self): - target = _RecordingTarget(failure=RuntimeError("target exploded mid-run")) - attack = PromptSendingAttack(objective_target=target) - - with pytest.raises(Exception): # noqa: B017 - the wrapper type is not the point - await attack.execute_async(objective="o") - - # execute_async re-raises rather than returning, so a caller holding only the - # return value has no conversation id to release. - assert len(target.reset_calls) == 1 - - async def test_a_cancelled_run_still_releases_its_conversation(self): - block = asyncio.Event() - target = _RecordingTarget(block=block) - attack = PromptSendingAttack(objective_target=target) - - task = asyncio.create_task(attack.execute_async(objective="o")) - await asyncio.wait_for(target.started.wait(), timeout=10) - task.cancel() - with pytest.raises(asyncio.CancelledError): - await task - - assert len(target.reset_calls) == 1 - - @pytest.mark.usefixtures("patch_central_database") class TestDefaultAttackStrategyEventHandler: """Tests for the default attack strategy event handler""" 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 808c71f5d6..50c1e8a3fe 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,7 +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), - report_objective_conversation=lambda conversation_id: None, + target_invocation_callback=lambda *, conversation_id: None, ) await node._send_to_adversarial_chat_async(prompt_text="hello") diff --git a/tests/unit/executor/attack/multi_turn/test_multi_prompt_sending.py b/tests/unit/executor/attack/multi_turn/test_multi_prompt_sending.py index 54d2025a01..40dee2a434 100644 --- a/tests/unit/executor/attack/multi_turn/test_multi_prompt_sending.py +++ b/tests/unit/executor/attack/multi_turn/test_multi_prompt_sending.py @@ -727,11 +727,9 @@ def test_attack_has_same_identifier_for_same_config(self, mock_target): assert attack1.get_identifier().hash == attack2.get_identifier().hash assert attack1.get_identifier().class_name == "MultiPromptSendingAttack" - async def test_teardown_async_resets_target_conversation(self, mock_target, basic_context): + async def test_teardown_async_is_noop(self, mock_target, basic_context): attack = MultiPromptSendingAttack(objective_target=mock_target) + # Should complete without error await attack._teardown_async(context=basic_context) - - mock_target.reset_conversation_async.assert_awaited_once_with( - conversation_id=basic_context.session.conversation_id - ) + # No assertions needed - we just want to ensure it runs without exceptions 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 83818e521b..ffeeacbefb 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,7 +296,7 @@ def _make_tap_node(*, target: PromptTarget) -> _TreeOfAttacksNode: adversarial_chat=adversarial_chat, objective_target=target, ), - report_objective_conversation=lambda conversation_id: None, + target_invocation_callback=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 406b0d365f..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"] @@ -1509,14 +1516,14 @@ async def test_execute_with_context_async_successful( assert result.outcome == AttackOutcome.SUCCESS assert result.objective == basic_context.objective - async def test_teardown_async_resets_target_conversation( + async def test_teardown_async_is_noop( self, mock_objective_target: MagicMock, mock_objective_scorer: MagicMock, mock_adversarial_chat: MagicMock, basic_context: MultiTurnAttackContext, ): - """Test that teardown releases the objective target's conversation.""" + """Test that teardown completes without errors.""" adversarial_config = AttackAdversarialConfig(target=mock_adversarial_chat) scoring_config = AttackScoringConfig(objective_scorer=mock_objective_scorer) @@ -1526,11 +1533,9 @@ async def test_teardown_async_resets_target_conversation( attack_scoring_config=scoring_config, ) + # Should complete without error await attack._teardown_async(context=basic_context) - - mock_objective_target.reset_conversation_async.assert_awaited_once_with( - conversation_id=basic_context.session.conversation_id - ) + # No assertions needed - we just want to ensure it runs without exceptions @pytest.mark.usefixtures("patch_central_database") 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 7c50add347..6c25b97e14 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,7 +491,7 @@ def _make_tap_node(self, *, supports_multi_turn: bool): adversarial_chat=adversarial_chat, objective_target=target, ), - report_objective_conversation=lambda conversation_id: None, + target_invocation_callback=lambda *, conversation_id: None, ) def test_single_turn_target_duplicates_logical_history_without_seed_boundary(self): @@ -878,7 +878,7 @@ def _make_tap_node(self, *, supports_multi_turn: bool): adversarial_chat=adversarial_chat, objective_target=target, ), - report_objective_conversation=lambda conversation_id: None, + target_invocation_callback=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 acc2787911..ad14df6329 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,8 +29,8 @@ TAPAttackScoringConfig, _TAPAttackConfiguration, _TreeOfAttacksNode, + _TreeOfAttacksNodeExecutor, ) -from pyrit.memory.central_memory import CentralMemory from pyrit.models import ( JSON_SCHEMA_METADATA_KEY, AttackOutcome, @@ -51,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 = { @@ -1088,7 +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, ), - report_objective_conversation=lambda conversation_id: None, + target_invocation_callback=lambda *, conversation_id: None, desired_response_prefix="Sure, here is", prompt_normalizer=normalizer, ) @@ -1156,7 +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, ), - report_objective_conversation=lambda conversation_id: None, + target_invocation_callback=lambda *, conversation_id: None, desired_response_prefix="Sure, here is", prompt_normalizer=normalizer, ) @@ -1566,7 +1600,7 @@ def node_components(self, attack_builder): "attack_id": {"id": "test_attack"}, "attack_strategy_name": "TreeOfAttacksWithPruningAttack", "modality_router": modality_router, - "report_objective_conversation": lambda conversation_id: None, + "target_invocation_callback": lambda *, conversation_id: None, "memory_labels": {"test": "label"}, "parent_id": None, "prompt_normalizer": prompt_normalizer, @@ -1620,19 +1654,6 @@ def test_node_duplicate_creates_child(self, node_components): assert child_node.parent_id == parent_node.node_id assert child_node.completed is False - def test_node_duplicate_keeps_reporting_its_conversations(self, node_components): - """A child sends on its own conversation, so it needs the same recorder as its parent.""" - reported: list[str] = [] - components = {**node_components, "report_objective_conversation": reported.append} - parent_node = _TreeOfAttacksNode(**components) - - with patch.object(parent_node._memory, "duplicate_conversation", return_value="new_conv_id"): - child_node = parent_node.duplicate() - - child_node._report_objective_conversation(child_node.objective_target_conversation_id) - - assert reported == ["new_conv_id"] - def _node_with_schema(self, node_components, schema): """Build a real node whose adversarial system prompt advertises ``schema``. @@ -3012,7 +3033,7 @@ def node_components(self, attack_builder): "attack_id": {"id": "test_attack"}, "attack_strategy_name": "TreeOfAttacksWithPruningAttack", "modality_router": modality_router, - "report_objective_conversation": lambda conversation_id: None, + "target_invocation_callback": lambda *, conversation_id: None, "memory_labels": {}, "parent_id": None, "prompt_normalizer": prompt_normalizer, @@ -3296,298 +3317,3 @@ def test_inline_system_prompt_string_resolved_and_in_identity(self): ) assert attack._adversarial_chat_system_seed_prompt.value == "tap persona {{ desired_prefix }}" assert attack.get_identifier().params["adversarial_system_prompt"] == "tap persona {{ desired_prefix }}" - - -@pytest.mark.usefixtures("patch_central_database") -class TestTAPConversationReset: - """Every objective-target conversation TAP opens has to stay reachable. - - TAP keeps one conversation per node and rotates it per turn against a - single-turn target. Anything it abandons without recording is a conversation - no caller can name afterwards, so the base teardown cannot release it and the - result under-reports what the run opened. - """ - - def _context_with_nodes(self, *node_ids: str) -> TAPAttackContext: - context = TAPAttackContext(params=AttackParameters(objective="Test objective")) - for node_id in node_ids: - node = MagicMock(spec=_TreeOfAttacksNode) - node.objective_target_conversation_id = node_id - context.nodes.append(node) - return context - - def test_the_base_lookup_finds_the_best_conversation(self, basic_attack): - context = self._context_with_nodes("node-a") - context.best_conversation_id = "best-1" - - ids = basic_attack._get_objective_conversation_ids(context=context) - - assert "best-1" in ids - - def test_the_base_lookup_does_not_use_the_unused_session_conversation_id(self, basic_attack): - context = self._context_with_nodes("node-a") - context.best_conversation_id = "best-1" - - ids = basic_attack._get_objective_conversation_ids(context=context) - - # TAP never sends anything on session.conversation_id. - assert context.session.conversation_id not in ids - - def test_a_conversation_is_recorded_as_soon_as_it_is_sent_on(self, basic_attack): - context = self._context_with_nodes("node-a") - record = basic_attack._make_objective_conversation_recorder(context=context) - - record("turn-1") - record("turn-2") - - # Recorded while the run is still going, so a run that raises or is - # cancelled can still name them. - ids = set(basic_attack._get_objective_conversation_ids(context=context)) - assert {"turn-1", "turn-2"} <= ids - - def test_recording_the_same_conversation_twice_records_it_once(self, basic_attack): - context = self._context_with_nodes() - record = basic_attack._make_objective_conversation_recorder(context=context) - - record("turn-1") - record("turn-1") - - assert {ref.conversation_id for ref in context.related_conversations} == {"turn-1"} - - def test_the_winning_branch_stops_being_reported_as_pruned(self, basic_attack): - context = self._context_with_nodes() - record = basic_attack._make_objective_conversation_recorder(context=context) - record("node-a") - record("node-b") - context.best_conversation_id = "node-a" - - basic_attack._release_best_conversation(context) - - # It is still released, through the live conversation rather than the pruned list. - assert {ref.conversation_id for ref in context.related_conversations} == {"node-b"} - assert set(basic_attack._get_objective_conversation_ids(context=context)) == {"node-a", "node-b"} - - def test_releasing_the_best_branch_without_one_is_a_noop(self, basic_attack): - context = self._context_with_nodes() - basic_attack._make_objective_conversation_recorder(context=context)("node-a") - - basic_attack._release_best_conversation(context) - - assert {ref.conversation_id for ref in context.related_conversations} == {"node-a"} - - def test_returns_no_duplicates(self, basic_attack): - context = self._context_with_nodes("node-a") - context.best_conversation_id = "node-a" - context.related_conversations.add( - ConversationReference(conversation_id="node-a", conversation_type=ConversationType.PRUNED) - ) - - ids = basic_attack._get_objective_conversation_ids(context=context) - - assert ids == ["node-a"] - - -@pytest.mark.usefixtures("patch_central_database") -class TestTAPConversationsAreAllReachable: - """End to end, with real nodes: nothing the objective target served goes missing. - - ``TestTAPConversationReset`` covers the pieces in isolation. This drives the - real ``_TreeOfAttacksNode`` so the wiring is covered too, which is where a - conversation actually goes missing: the per-turn rotation against a - single-turn target, and the branch a run walks away from. - """ - - def _run_and_collect(self, *, attack_builder, supports_multi_turn, depth, width, branching): - """Run TAP and return (ids the target served, ids the run can still name).""" - served: list[str] = [] - - attack = ( - attack_builder.with_supports_multi_turn(supports_multi_turn) - .with_default_mocks() - .with_tree_params(tree_depth=depth, tree_width=width, branching_factor=branching) - .build() - ) - objective_target = attack._objective_target - memory = CentralMemory.get_memory_instance() - - async def record_and_reply(**kwargs): - conversation_id = kwargs.get("conversation_id") - reply = Message( - message_pieces=[ - MessagePiece( - role="assistant", - original_value="response", - converted_value="response", - conversation_id=conversation_id, - ) - ] - ) - if kwargs.get("target") is not objective_target: - return reply - served.append(conversation_id) - request = Message( - message_pieces=[ - MessagePiece( - role="user", - original_value="request", - converted_value="request", - conversation_id=conversation_id, - ) - ] - ) - for message in (request, reply): - for piece in message.message_pieces: - piece.not_in_memory = False - memory.add_message_to_memory(request=message) - return reply - - normalizer = MagicMock(spec=PromptNormalizer) - normalizer.send_prompt_async = AsyncMock(side_effect=record_and_reply) - attack._prompt_normalizer = normalizer - attack._node_executor._prompt_normalizer = normalizer - - return attack, served - - async def _execute(self, attack, context): - async def score(node_self, *, response, objective): - node_self.objective_score = MagicMock( - spec=Score, get_value=MagicMock(return_value=0.1), score_metadata=None - ) - - with patch.object( - _TreeOfAttacksNode, "_generate_adversarial_prompt_async", new_callable=AsyncMock, return_value="prompt" - ): - with patch.object(_TreeOfAttacksNode, "_score_response_async", new=score): - await attack._setup_async(context=context) - return await attack._perform_async(context=context) - - @pytest.mark.parametrize( - "supports_multi_turn, depth, width, branching", - [ - pytest.param(False, 3, 1, 1, id="single_turn_rotates_per_turn"), - pytest.param(True, 3, 2, 2, id="multi_turn_branches_per_node"), - ], - ) - async def test_every_conversation_the_target_served_stays_reachable( - self, attack_builder, supports_multi_turn, depth, width, branching - ): - attack, served = self._run_and_collect( - attack_builder=attack_builder, - supports_multi_turn=supports_multi_turn, - depth=depth, - width=width, - branching=branching, - ) - context = TAPAttackContext(params=AttackParameters(objective="Test objective")) - - result = await self._execute(attack, context) - - assert served, "the run has to have sent something for this to mean anything" - # Nothing the target served may be left without a name: the teardown reset - # and the result readers both work from these two. - assert set(served) <= set(attack._get_objective_conversation_ids(context=context)) - assert set(served) <= result.get_active_conversation_ids() - - async def test_a_conversation_that_was_never_used_is_not_recorded(self, attack_builder): - attack, served = self._run_and_collect( - attack_builder=attack_builder, supports_multi_turn=False, depth=3, width=1, branching=1 - ) - context = TAPAttackContext(params=AttackParameters(objective="Test objective")) - - result = await self._execute(attack, context) - - # A node is constructed with a conversation id it rotates away from before - # its first send. That conversation has no messages, so recording it would - # put an empty conversation in front of the user. Checked on an unbranched - # tree, where every conversation that exists is one the target served; - # branching also duplicates conversations, which are real but never sent. - assert result.get_active_conversation_ids() == set(served) - - async def test_the_context_and_the_result_agree(self, attack_builder): - attack, _ = self._run_and_collect( - attack_builder=attack_builder, supports_multi_turn=True, depth=3, width=2, branching=2 - ) - context = TAPAttackContext(params=AttackParameters(objective="Test objective")) - - result = await self._execute(attack, context) - - assert set(attack._get_objective_conversation_ids(context=context)) == result.get_active_conversation_ids() - - async def test_a_branched_node_reports_its_own_conversations(self, attack_builder): - """A child gets its own conversation from duplicate(), and must report on it too.""" - attack, served = self._run_and_collect( - attack_builder=attack_builder, supports_multi_turn=True, depth=3, width=1, branching=2 - ) - context = TAPAttackContext(params=AttackParameters(objective="Test objective")) - - result = await self._execute(attack, context) - - # width=1 keeps one node per level, so every conversation beyond the first - # belongs to a branch, and nothing else records those for us. - assert len(set(served)) > 1, "branching has to have produced more than one conversation" - assert set(served) <= result.get_active_conversation_ids() - - async def test_the_winning_conversation_is_not_also_reported_as_pruned(self, attack_builder): - attack, served = self._run_and_collect( - attack_builder=attack_builder, supports_multi_turn=True, depth=3, width=2, branching=2 - ) - context = TAPAttackContext(params=AttackParameters(objective="Test objective")) - - result = await self._execute(attack, context) - - assert result.conversation_id, "the run has to have picked a best branch" - # Every conversation is recorded while the run is in flight, before there is - # any way to know which branch wins. The winner has to come back out. - assert result.conversation_id not in result.get_pruned_conversation_ids() - assert result.conversation_id in result.get_active_conversation_ids() - - async def test_a_run_that_raises_does_not_report_its_own_conversation_as_pruned(self, attack_builder): - """The backend adds the main conversation's messages to the pruned ones, so it cannot be in both.""" - attack, served = self._run_and_collect( - attack_builder=attack_builder, supports_multi_turn=True, depth=3, width=2, branching=1 - ) - context = TAPAttackContext(params=AttackParameters(objective="Test objective")) - - iterations = {"count": 0} - prepare = type(attack)._prepare_nodes_for_iteration_async - - async def fail_on_the_second_iteration(self, context): - iterations["count"] += 1 - if iterations["count"] >= 2: - raise ValueError("blew up mid-run") - await prepare(self, context=context) - - with patch.object(type(attack), "_prepare_nodes_for_iteration_async", new=fail_on_the_second_iteration): - with pytest.raises(ValueError): - await self._execute(attack, context) - - # No result is built on this path, so the invariant has to already hold on - # the context the error result is assembled from. - assert context.best_conversation_id, "a branch has to have taken the lead" - pruned = {ref.conversation_id for ref in context.related_conversations} - assert context.best_conversation_id not in pruned - - async def test_a_run_that_raises_still_names_everything_it_served(self, attack_builder): - """The path that matters most, because a run that blew up is the one holding connections.""" - attack, served = self._run_and_collect( - attack_builder=attack_builder, supports_multi_turn=True, depth=3, width=3, branching=1 - ) - context = TAPAttackContext(params=AttackParameters(objective="Test objective")) - - iterations = {"count": 0} - prepare = type(attack)._prepare_nodes_for_iteration_async - - async def fail_on_the_second_iteration(self, context): - iterations["count"] += 1 - if iterations["count"] >= 2: - raise ValueError("blew up mid-run") - await prepare(self, context=context) - - with patch.object(type(attack), "_prepare_nodes_for_iteration_async", new=fail_on_the_second_iteration): - with pytest.raises(ValueError): - await self._execute(attack, context) - - assert served, "the run has to have sent something for this to mean anything" - # No result is built on this path, so anything recorded only at result time - # would be lost, and teardown is the only hook that still runs. - assert set(served) <= set(attack._get_objective_conversation_ids(context=context)) 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 62671376c3..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""" @@ -1105,12 +1207,12 @@ async def test_execute_async_execution_error_still_calls_teardown(self, mock_tar attack._perform_async.assert_called_once_with(context=basic_context) attack._teardown_async.assert_called_once_with(context=basic_context) - async def test_teardown_async_resets_target_conversation(self, mock_target, basic_context): + async def test_teardown_async_is_noop(self, mock_target, basic_context): attack = PromptSendingAttack(objective_target=mock_target) + # Should complete without error await attack._teardown_async(context=basic_context) - - mock_target.reset_conversation_async.assert_awaited_once_with(conversation_id=basic_context.conversation_id) + # No assertions needed - we just want to ensure it runs without raising async def test_execute_async_with_parameters(self, mock_target, sample_response): """Test execute_async creates context using factory method and executes attack""" diff --git a/tests/unit/prompt_normalizer/test_prompt_normalizer.py b/tests/unit/prompt_normalizer/test_prompt_normalizer.py index 90dfdb4969..7ba085bab9 100644 --- a/tests/unit/prompt_normalizer/test_prompt_normalizer.py +++ b/tests/unit/prompt_normalizer/test_prompt_normalizer.py @@ -145,6 +145,7 @@ async def test_send_prompt_async_forwards_normalizer_overrides_and_context(mock_ seed_message_ids=(uuid4(),), replay_seed_each_send=False, ) + target_invocation_callback = MagicMock() await normalizer.send_prompt_async( message=Message.from_prompt(prompt="first request", role="user"), @@ -152,11 +153,13 @@ async def test_send_prompt_async_forwards_normalizer_overrides_and_context(mock_ conversation_id=conversation_id, normalizer_overrides=normalizer_overrides, send_context=target_context, + target_invocation_callback=target_invocation_callback, ) call = prompt_target.send_prompt_async.await_args assert call.kwargs["normalizer_overrides"] == normalizer_overrides assert call.kwargs["send_context"] is target_context + assert call.kwargs["target_invocation_callback"] is target_invocation_callback async def test_send_prompt_async_conversion_failure_does_not_call_target(mock_memory_instance): diff --git a/tests/unit/prompt_target/target/test_prompt_target.py b/tests/unit/prompt_target/target/test_prompt_target.py index 453870e382..19b975db89 100644 --- a/tests/unit/prompt_target/target/test_prompt_target.py +++ b/tests/unit/prompt_target/target/test_prompt_target.py @@ -8,7 +8,7 @@ import pytest from openai.types.chat import ChatCompletion -from unit.mocks import get_sample_conversations, openai_chat_response_json_dict +from unit.mocks import MockPromptTarget, get_sample_conversations, openai_chat_response_json_dict from pyrit.executor.attack.core.attack_strategy import AttackStrategy from pyrit.memory.memory_interface import MemoryInterface @@ -56,6 +56,46 @@ def mock_attack_strategy(): return strategy +@pytest.mark.usefixtures("patch_central_database") +async def test_target_invocation_callback_runs_immediately_before_target_send() -> None: + target = MockPromptTarget() + events: list[str] = [] + + def record_invocation(*, conversation_id: str) -> None: + assert conversation_id == "conversation-id" + events.append("callback") + + async def send_to_target(*, normalized_conversation: list[Message]) -> list[Message]: + events.append("send") + return [] + + target._send_prompt_to_target_async = send_to_target # type: ignore[method-assign] + message = Message.from_prompt(prompt="request", role="user") + message.message_pieces[0].conversation_id = "conversation-id" + + await target.send_prompt_async( + message=message, + target_invocation_callback=record_invocation, + ) + + assert events == ["callback", "send"] + + +@pytest.mark.usefixtures("patch_central_database") +async def test_target_invocation_callback_does_not_run_when_validation_fails() -> None: + target = MockPromptTarget() + target._validate_request = MagicMock(side_effect=ValueError("invalid request")) # type: ignore[method-assign] + target_invocation_callback = MagicMock() + + with pytest.raises(ValueError, match="invalid request"): + await target.send_prompt_async( + message=Message.from_prompt(prompt="request", role="user"), + target_invocation_callback=target_invocation_callback, + ) + + target_invocation_callback.assert_not_called() + + def test_set_system_prompt(azure_openai_target: OpenAIChatTarget, mock_attack_strategy: AttackStrategy): azure_openai_target.set_system_prompt( system_prompt="system prompt", diff --git a/tests/unit/prompt_target/target/test_realtime_target.py b/tests/unit/prompt_target/target/test_realtime_target.py index 8212aa65d7..33b52764c0 100644 --- a/tests/unit/prompt_target/target/test_realtime_target.py +++ b/tests/unit/prompt_target/target/test_realtime_target.py @@ -1433,6 +1433,32 @@ async def test_reset_conversation_async_swallows_close_error(target): assert "conv" not in target._existing_conversation +async def test_reset_conversation_async_cancellation_finishes_closing_connection(target): + mock_connection = AsyncMock() + target._existing_conversation["conv"] = mock_connection + close_started = asyncio.Event() + finish_close = asyncio.Event() + + async def close_connection(): + close_started.set() + await finish_close.wait() + + mock_connection.close.side_effect = close_connection + cleanup_task = asyncio.create_task(target.reset_conversation_async(conversation_id="conv")) + 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 + + 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 From 00b114059c46269e6147a96f66eb12f03f528569 Mon Sep 17 00:00:00 2001 From: Richard Lundeen Date: Fri, 28 Aug 2026 11:51:09 -0700 Subject: [PATCH 7/7] Simplify objective-target conversation reset Record objective-target conversations directly at each attack send site instead of threading a callback through PromptNormalizer and PromptTarget. Remove the TargetInvocationCallback protocol so the target-dispatch surface no longer carries lifecycle concerns. Reduce RealtimeTarget and WebsocketTarget reset_conversation_async to plain try/except that logs close errors and lets cancellation propagate, dropping the asyncio.shield guard. Attempt every reset in the lifecycle __aexit__ even under cancellation, then re-raise the first stashed CancelledError. Recording outside an active scope is now a no-op. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Copilot-Session: 1dc87b33-0580-45f9-a9cc-0dc9603bd2a2 --- pyrit/executor/attack/core/attack_strategy.py | 15 ++++-- .../attack/multi_turn/chunked_request.py | 2 +- pyrit/executor/attack/multi_turn/crescendo.py | 2 +- .../attack/multi_turn/multi_prompt_sending.py | 2 +- .../executor/attack/multi_turn/red_teaming.py | 2 +- .../attack/multi_turn/tree_of_attacks.py | 19 ++++---- .../attack/single_turn/prompt_sending.py | 2 +- pyrit/prompt_normalizer/prompt_normalizer.py | 6 +-- pyrit/prompt_target/common/prompt_target.py | 7 +-- .../common/target_send_context.py | 13 ----- .../openai/openai_realtime_target.py | 14 +----- pyrit/prompt_target/websocket_target.py | 15 ++---- .../attack/core/test_attack_strategy.py | 16 +++++++ ...test_adversarial_chat_schema_forwarding.py | 2 +- .../test_prepended_history_normalization.py | 2 +- .../test_supports_multi_turn_attacks.py | 4 +- .../attack/multi_turn/test_tree_of_attacks.py | 8 ++-- .../test_prompt_normalizer.py | 3 -- .../target/test_prompt_target.py | 42 +--------------- .../target/test_realtime_target.py | 20 ++------ .../target/test_websocket_target.py | 48 ++++--------------- 21 files changed, 69 insertions(+), 175 deletions(-) diff --git a/pyrit/executor/attack/core/attack_strategy.py b/pyrit/executor/attack/core/attack_strategy.py index 6345f6b355..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 @@ -103,15 +104,21 @@ async def __aexit__( 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: """ @@ -228,15 +235,15 @@ 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. - - Raises: - RuntimeError: If called outside this context's attack execution. """ lifecycle = self._objective_target_conversation_lifecycle if lifecycle is None: - raise RuntimeError("Objective-target invocation occurred outside the attack lifecycle.") + return lifecycle.record_invocation(conversation_id=conversation_id) diff --git a/pyrit/executor/attack/multi_turn/chunked_request.py b/pyrit/executor/attack/multi_turn/chunked_request.py index 51b7cf9352..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, @@ -297,7 +298,6 @@ async def _perform_async(self, *, context: ChunkedRequestAttackContext) -> Attac prepended_history_send_context=context.prepended_history_send_context, ), send_context=context.prepended_history_send_context, - target_invocation_callback=context._record_objective_target_invocation, ) # Store the response diff --git a/pyrit/executor/attack/multi_turn/crescendo.py b/pyrit/executor/attack/multi_turn/crescendo.py index 55901671b7..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, @@ -645,7 +646,6 @@ async def _send_prompt_to_objective_target_async( prepended_history_send_context=context.prepended_history_send_context, ), send_context=context.prepended_history_send_context, - target_invocation_callback=context._record_objective_target_invocation, ) if not response: diff --git a/pyrit/executor/attack/multi_turn/multi_prompt_sending.py b/pyrit/executor/attack/multi_turn/multi_prompt_sending.py index c6776fd64f..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, @@ -375,7 +376,6 @@ async def _send_prompt_to_objective_target_async( prepended_history_send_context=context.prepended_history_send_context, ), send_context=context.prepended_history_send_context, - target_invocation_callback=context._record_objective_target_invocation, ) async def _evaluate_response_async(self, *, response: Message, objective: str) -> Score | None: diff --git a/pyrit/executor/attack/multi_turn/red_teaming.py b/pyrit/executor/attack/multi_turn/red_teaming.py index b651d40ec5..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, @@ -498,7 +499,6 @@ async def _send_prompt_to_objective_target_async( prepended_history_send_context=context.prepended_history_send_context, ), send_context=context.prepended_history_send_context, - target_invocation_callback=context._record_objective_target_invocation, ) if response is None: diff --git a/pyrit/executor/attack/multi_turn/tree_of_attacks.py b/pyrit/executor/attack/multi_turn/tree_of_attacks.py index 7c6beaa5eb..26ed278507 100644 --- a/pyrit/executor/attack/multi_turn/tree_of_attacks.py +++ b/pyrit/executor/attack/multi_turn/tree_of_attacks.py @@ -74,11 +74,10 @@ 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 - from pyrit.prompt_target.common.target_send_context import TargetInvocationCallback logger = logging.getLogger(__name__) @@ -378,7 +377,7 @@ def __init__( attack_id: ComponentIdentifier, attack_strategy_name: str, modality_router: _ModalityFeedbackRouter, - target_invocation_callback: TargetInvocationCallback, + record_objective_conversation: Callable[..., None], use_score_as_feedback: bool = True, memory_labels: dict[str, str] | None = None, parent_id: str | None = None, @@ -407,8 +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. - target_invocation_callback (TargetInvocationCallback): Callback for objective-target - provider invocations. + 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. @@ -436,7 +435,7 @@ def __init__( self._attack_strategy_name = attack_strategy_name self._memory_labels = memory_labels or {} self._modality_router = modality_router - self._target_invocation_callback = target_invocation_callback + 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 @@ -687,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, @@ -698,7 +698,6 @@ async def _send_prompt_to_target_async(self, prompt: str) -> Message: prepended_history_send_context=self._prepended_history_send_context, ), send_context=self._prepended_history_send_context, - target_invocation_callback=self._target_invocation_callback, ) # Store the full response so subsequent turns can forward media when supported. @@ -767,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, @@ -778,7 +778,6 @@ async def _send_initial_prompt_to_target_async(self) -> Message: prepended_history_send_context=self._prepended_history_send_context, ), send_context=self._prepended_history_send_context, - target_invocation_callback=self._target_invocation_callback, ) # Store the full response so subsequent turns can forward media when supported. @@ -974,7 +973,7 @@ def duplicate(self) -> _TreeOfAttacksNode: attack_id=self._attack_id, attack_strategy_name=self._attack_strategy_name, modality_router=self._modality_router, - target_invocation_callback=self._target_invocation_callback, + 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, @@ -2213,7 +2212,7 @@ def _create_attack_node( attack_id=self.get_identifier(), attack_strategy_name=self.__class__.__name__, modality_router=self._modality_router, - target_invocation_callback=context._record_objective_target_invocation, + 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 8f3f770c55..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, @@ -327,7 +328,6 @@ async def _send_prompt_to_objective_target_async( prepended_history_send_context=context.prepended_history_send_context, ), send_context=context.prepended_history_send_context, - target_invocation_callback=context._record_objective_target_invocation, ) async def _evaluate_response_async( diff --git a/pyrit/prompt_normalizer/prompt_normalizer.py b/pyrit/prompt_normalizer/prompt_normalizer.py index e8e393560c..20f990f77d 100644 --- a/pyrit/prompt_normalizer/prompt_normalizer.py +++ b/pyrit/prompt_normalizer/prompt_normalizer.py @@ -31,7 +31,7 @@ from pyrit.prompt_normalizer import ConverterConfiguration, NormalizerRequest from pyrit.prompt_target import CapabilityName, PromptTarget from pyrit.prompt_target.batch_helper import batch_task_async -from pyrit.prompt_target.common.target_send_context import TargetInvocationCallback, TargetSendContext +from pyrit.prompt_target.common.target_send_context import TargetSendContext logger = logging.getLogger(__name__) @@ -76,7 +76,6 @@ async def send_prompt_async( response_converter_configurations: list[ConverterConfiguration] | None = None, normalizer_overrides: Mapping[CapabilityName, MessageListNormalizer[Message]] | None = None, send_context: TargetSendContext | None = None, - target_invocation_callback: TargetInvocationCallback | None = None, ) -> Message: """ Send a single request to a target. @@ -92,8 +91,6 @@ async def send_prompt_async( normalizer_overrides: Optional per-send target normalizer overrides. send_context: Optional internal coordination contract for caller-owned history selection and send lifecycle state. - target_invocation_callback: Optional callback invoked immediately before - target-specific execution. Returns: Message: The response received from the target. @@ -133,7 +130,6 @@ async def send_prompt_async( message=request, normalizer_overrides=normalizer_overrides, send_context=send_context, - target_invocation_callback=target_invocation_callback, ) self.memory.add_message_to_memory(request=request) except EmptyResponseException as ex: diff --git a/pyrit/prompt_target/common/prompt_target.py b/pyrit/prompt_target/common/prompt_target.py index 3aa1a4d024..790abcbba4 100644 --- a/pyrit/prompt_target/common/prompt_target.py +++ b/pyrit/prompt_target/common/prompt_target.py @@ -24,7 +24,7 @@ ) from pyrit.prompt_target.common.target_configuration import TargetConfiguration from pyrit.prompt_target.common.target_history import filter_non_replayable_messages -from pyrit.prompt_target.common.target_send_context import TargetInvocationCallback, TargetSendContext +from pyrit.prompt_target.common.target_send_context import TargetSendContext logger = logging.getLogger(__name__) @@ -142,7 +142,6 @@ async def send_prompt_async( message: Message, normalizer_overrides: Mapping[CapabilityName, MessageListNormalizer[Message]] | None = None, send_context: TargetSendContext | None = None, - target_invocation_callback: TargetInvocationCallback | None = None, ) -> list[Message]: """ Validate, normalize, and send a prompt to the target. @@ -162,8 +161,6 @@ async def send_prompt_async( normalizer_overrides: Optional per-send target normalizer overrides. send_context: Optional internal coordination contract for caller-owned history selection and send lifecycle state. - target_invocation_callback: Optional callback invoked immediately before - target-specific execution. Returns: list[Message]: Response messages from the target. @@ -188,8 +185,6 @@ async def send_prompt_async( if not normalized_conversation: raise ValueError("Normalization pipeline returned an empty conversation. Cannot send an empty request.") self._validate_request(normalized_conversation=normalized_conversation) - if target_invocation_callback: - target_invocation_callback(conversation_id=conversation_id) if send_context: send_context.mark_target_invoked() response = await self._send_prompt_to_target_async(normalized_conversation=normalized_conversation) diff --git a/pyrit/prompt_target/common/target_send_context.py b/pyrit/prompt_target/common/target_send_context.py index 8dbc705f64..2c0e4f3e4d 100644 --- a/pyrit/prompt_target/common/target_send_context.py +++ b/pyrit/prompt_target/common/target_send_context.py @@ -9,19 +9,6 @@ from pyrit.models import Message -class TargetInvocationCallback(Protocol): - """Callback invoked immediately before target-specific execution.""" - - def __call__(self, *, conversation_id: str) -> None: - """ - Record one target invocation. - - Args: - conversation_id (str): The conversation ID for the invocation. - """ - ... - - class TargetSendContext(Protocol): """Internal contract coordinating one target send with caller-owned state.""" diff --git a/pyrit/prompt_target/openai/openai_realtime_target.py b/pyrit/prompt_target/openai/openai_realtime_target.py index 93a913000c..511f228b44 100644 --- a/pyrit/prompt_target/openai/openai_realtime_target.py +++ b/pyrit/prompt_target/openai/openai_realtime_target.py @@ -491,24 +491,14 @@ async def reset_conversation_async(self, *, conversation_id: str) -> None: Args: conversation_id (str): The conversation ID to disconnect from. - - Raises: - asyncio.CancelledError: If cleanup is cancelled after the connection has finished closing. """ connection = self._existing_conversation.pop(conversation_id, None) if not connection: return - close_future = asyncio.ensure_future(connection.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 - except Exception as error: + 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 diff --git a/pyrit/prompt_target/websocket_target.py b/pyrit/prompt_target/websocket_target.py index 8190abe808..ee196d61bb 100644 --- a/pyrit/prompt_target/websocket_target.py +++ b/pyrit/prompt_target/websocket_target.py @@ -192,24 +192,17 @@ async def reset_conversation_async(self, *, conversation_id: str) -> None: Args: conversation_id (str): PyRIT conversation ID. - - Raises: - asyncio.CancelledError: If cleanup is cancelled after the connection has finished closing. """ conversation_lock = self._conversation_locks.setdefault(conversation_id, asyncio.Lock()) async with conversation_lock: websocket = self._existing_conversation.pop(conversation_id, None) if websocket is None: return - 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: diff --git a/tests/unit/executor/attack/core/test_attack_strategy.py b/tests/unit/executor/attack/core/test_attack_strategy.py index 03111a0d34..2332b321e2 100644 --- a/tests/unit/executor/attack/core/test_attack_strategy.py +++ b/tests/unit/executor/attack/core/test_attack_strategy.py @@ -152,6 +152,22 @@ async def test_objective_target_cleanup_propagates_cancellation() -> None: 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 50c1e8a3fe..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,7 +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), - target_invocation_callback=lambda *, conversation_id: None, + 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 ffeeacbefb..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,7 +296,7 @@ def _make_tap_node(*, target: PromptTarget) -> _TreeOfAttacksNode: adversarial_chat=adversarial_chat, objective_target=target, ), - target_invocation_callback=lambda *, conversation_id: None, + record_objective_conversation=lambda *, conversation_id: None, ) 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 6c25b97e14..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,7 +491,7 @@ def _make_tap_node(self, *, supports_multi_turn: bool): adversarial_chat=adversarial_chat, objective_target=target, ), - target_invocation_callback=lambda *, conversation_id: None, + record_objective_conversation=lambda *, conversation_id: None, ) def test_single_turn_target_duplicates_logical_history_without_seed_boundary(self): @@ -878,7 +878,7 @@ def _make_tap_node(self, *, supports_multi_turn: bool): adversarial_chat=adversarial_chat, objective_target=target, ), - target_invocation_callback=lambda *, conversation_id: None, + 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 ad14df6329..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 @@ -1122,7 +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, ), - target_invocation_callback=lambda *, conversation_id: None, + record_objective_conversation=lambda *, conversation_id: None, desired_response_prefix="Sure, here is", prompt_normalizer=normalizer, ) @@ -1190,7 +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, ), - target_invocation_callback=lambda *, conversation_id: None, + record_objective_conversation=lambda *, conversation_id: None, desired_response_prefix="Sure, here is", prompt_normalizer=normalizer, ) @@ -1600,7 +1600,7 @@ def node_components(self, attack_builder): "attack_id": {"id": "test_attack"}, "attack_strategy_name": "TreeOfAttacksWithPruningAttack", "modality_router": modality_router, - "target_invocation_callback": lambda *, conversation_id: None, + "record_objective_conversation": lambda *, conversation_id: None, "memory_labels": {"test": "label"}, "parent_id": None, "prompt_normalizer": prompt_normalizer, @@ -3033,7 +3033,7 @@ def node_components(self, attack_builder): "attack_id": {"id": "test_attack"}, "attack_strategy_name": "TreeOfAttacksWithPruningAttack", "modality_router": modality_router, - "target_invocation_callback": lambda *, conversation_id: None, + "record_objective_conversation": lambda *, conversation_id: None, "memory_labels": {}, "parent_id": None, "prompt_normalizer": prompt_normalizer, diff --git a/tests/unit/prompt_normalizer/test_prompt_normalizer.py b/tests/unit/prompt_normalizer/test_prompt_normalizer.py index 7ba085bab9..90dfdb4969 100644 --- a/tests/unit/prompt_normalizer/test_prompt_normalizer.py +++ b/tests/unit/prompt_normalizer/test_prompt_normalizer.py @@ -145,7 +145,6 @@ async def test_send_prompt_async_forwards_normalizer_overrides_and_context(mock_ seed_message_ids=(uuid4(),), replay_seed_each_send=False, ) - target_invocation_callback = MagicMock() await normalizer.send_prompt_async( message=Message.from_prompt(prompt="first request", role="user"), @@ -153,13 +152,11 @@ async def test_send_prompt_async_forwards_normalizer_overrides_and_context(mock_ conversation_id=conversation_id, normalizer_overrides=normalizer_overrides, send_context=target_context, - target_invocation_callback=target_invocation_callback, ) call = prompt_target.send_prompt_async.await_args assert call.kwargs["normalizer_overrides"] == normalizer_overrides assert call.kwargs["send_context"] is target_context - assert call.kwargs["target_invocation_callback"] is target_invocation_callback async def test_send_prompt_async_conversion_failure_does_not_call_target(mock_memory_instance): diff --git a/tests/unit/prompt_target/target/test_prompt_target.py b/tests/unit/prompt_target/target/test_prompt_target.py index 19b975db89..453870e382 100644 --- a/tests/unit/prompt_target/target/test_prompt_target.py +++ b/tests/unit/prompt_target/target/test_prompt_target.py @@ -8,7 +8,7 @@ import pytest from openai.types.chat import ChatCompletion -from unit.mocks import MockPromptTarget, get_sample_conversations, openai_chat_response_json_dict +from unit.mocks import get_sample_conversations, openai_chat_response_json_dict from pyrit.executor.attack.core.attack_strategy import AttackStrategy from pyrit.memory.memory_interface import MemoryInterface @@ -56,46 +56,6 @@ def mock_attack_strategy(): return strategy -@pytest.mark.usefixtures("patch_central_database") -async def test_target_invocation_callback_runs_immediately_before_target_send() -> None: - target = MockPromptTarget() - events: list[str] = [] - - def record_invocation(*, conversation_id: str) -> None: - assert conversation_id == "conversation-id" - events.append("callback") - - async def send_to_target(*, normalized_conversation: list[Message]) -> list[Message]: - events.append("send") - return [] - - target._send_prompt_to_target_async = send_to_target # type: ignore[method-assign] - message = Message.from_prompt(prompt="request", role="user") - message.message_pieces[0].conversation_id = "conversation-id" - - await target.send_prompt_async( - message=message, - target_invocation_callback=record_invocation, - ) - - assert events == ["callback", "send"] - - -@pytest.mark.usefixtures("patch_central_database") -async def test_target_invocation_callback_does_not_run_when_validation_fails() -> None: - target = MockPromptTarget() - target._validate_request = MagicMock(side_effect=ValueError("invalid request")) # type: ignore[method-assign] - target_invocation_callback = MagicMock() - - with pytest.raises(ValueError, match="invalid request"): - await target.send_prompt_async( - message=Message.from_prompt(prompt="request", role="user"), - target_invocation_callback=target_invocation_callback, - ) - - target_invocation_callback.assert_not_called() - - def test_set_system_prompt(azure_openai_target: OpenAIChatTarget, mock_attack_strategy: AttackStrategy): azure_openai_target.set_system_prompt( system_prompt="system prompt", diff --git a/tests/unit/prompt_target/target/test_realtime_target.py b/tests/unit/prompt_target/target/test_realtime_target.py index 33b52764c0..b29355d554 100644 --- a/tests/unit/prompt_target/target/test_realtime_target.py +++ b/tests/unit/prompt_target/target/test_realtime_target.py @@ -1433,27 +1433,13 @@ async def test_reset_conversation_async_swallows_close_error(target): assert "conv" not in target._existing_conversation -async def test_reset_conversation_async_cancellation_finishes_closing_connection(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 - close_started = asyncio.Event() - finish_close = asyncio.Event() - - async def close_connection(): - close_started.set() - await finish_close.wait() - - mock_connection.close.side_effect = close_connection - cleanup_task = asyncio.create_task(target.reset_conversation_async(conversation_id="conv")) - 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 + await target.reset_conversation_async(conversation_id="conv") mock_connection.close.assert_awaited_once() assert "conv" not 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 d59273d8e0..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 @@ -711,61 +710,30 @@ async def test_reset_conversation_async_does_not_retain_unknown_lock(websocket_t assert "missing" not in websocket_target._conversation_locks -async def test_reset_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.reset_conversation_async(conversation_id="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_reset_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.reset_conversation_async(conversation_id="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 == {}