FEAT: Add reset_conversation_async hook to PromptTarget

.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

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`.

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] = []
 

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,

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

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.

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:

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.

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

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.

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.

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.

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.