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 3a7943e45d..9c21c6db0e 100644 --- a/doc/code/targets/0_prompt_targets.md +++ b/doc/code/targets/0_prompt_targets.md @@ -19,6 +19,20 @@ async def send_prompt_async(self, *, message: Message) -> Message: A `Message` object is a normalized object with all the information a target will need to send a prompt, including a way to get a history for that prompt (in the cases that also needs to be sent). This is discussed in more depth [here](../memory/3_memory_data_types.md). +## 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 6bab4492b7..3edf939d17 100644 --- a/pyrit/executor/attack/core/attack_strategy.py +++ b/pyrit/executor/attack/core/attack_strategy.py @@ -32,6 +32,7 @@ AttackResult, ComponentIdentifier, ConversationReference, + ConversationType, ConverterIdentifier, Identifiable, Message, @@ -614,6 +615,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}") + @overload async def execute_async( self, diff --git a/pyrit/executor/attack/multi_turn/chunked_request.py b/pyrit/executor/attack/multi_turn/chunked_request.py index 6db7537722..f7c4ca2765 100644 --- a/pyrit/executor/attack/multi_turn/chunked_request.py +++ b/pyrit/executor/attack/multi_turn/chunked_request.py @@ -378,11 +378,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 584065aa87..9b0c14dd17 100644 --- a/pyrit/executor/attack/multi_turn/crescendo.py +++ b/pyrit/executor/attack/multi_turn/crescendo.py @@ -470,15 +470,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 d431ec20db..a475d92a7a 100644 --- a/pyrit/executor/attack/multi_turn/multi_prompt_sending.py +++ b/pyrit/executor/attack/multi_turn/multi_prompt_sending.py @@ -335,10 +335,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 402eb1303b..03f7a153ae 100644 --- a/pyrit/executor/attack/multi_turn/red_teaming.py +++ b/pyrit/executor/attack/multi_turn/red_teaming.py @@ -369,10 +369,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 1fe9a4e01c..99ebb2bafa 100644 --- a/pyrit/executor/attack/multi_turn/tree_of_attacks.py +++ b/pyrit/executor/attack/multi_turn/tree_of_attacks.py @@ -1668,22 +1668,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 508b72d924..cb62564961 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 717e269d71..1233d5f646 100644 --- a/pyrit/executor/attack/streaming/barge_in.py +++ b/pyrit/executor/attack/streaming/barge_in.py @@ -137,10 +137,6 @@ async def _setup_async(self, *, context: BargeInAttackContext[Any]) -> None: request_converters=self._request_converters, ) - 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 aaf918f4cf..48fb31a017 100644 --- a/pyrit/prompt_target/common/prompt_target.py +++ b/pyrit/prompt_target/common/prompt_target.py @@ -361,6 +361,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 3894ddc623..6cd2dfb0a4 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, ) @@ -475,10 +476,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. """ @@ -491,6 +497,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 f4f8701245..7702fbbdb0 100644 --- a/tests/unit/executor/attack/core/test_attack_strategy.py +++ b/tests/unit/executor/attack/core/test_attack_strategy.py @@ -15,12 +15,16 @@ AttackStrategy, _DefaultAttackStrategyEventHandler, ) +from pyrit.executor.attack.multi_turn.multi_turn_attack_strategy import MultiTurnAttackContext +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, ) @@ -321,6 +325,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 e6a4467089..ed74b98485 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 @@ -688,9 +688,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 4603152590..2a49fdf730 100644 --- a/tests/unit/executor/attack/multi_turn/test_red_teaming.py +++ b/tests/unit/executor/attack/multi_turn/test_red_teaming.py @@ -1415,14 +1415,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) @@ -1432,9 +1432,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 bca5535b0f..e23056a9dc 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 @@ -3013,3 +3013,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 5d4fda2209..494b41d8e8 100644 --- a/tests/unit/executor/attack/single_turn/test_prompt_sending.py +++ b/tests/unit/executor/attack/single_turn/test_prompt_sending.py @@ -1004,12 +1004,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 6a4959e76d..5cd18be1c5 100644 --- a/tests/unit/prompt_target/target/test_realtime_target.py +++ b/tests/unit/prompt_target/target/test_realtime_target.py @@ -927,31 +927,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")