FEAT: Add reset_conversation_async hook to PromptTarget

.github/instructions/targets.instructions.md

@@ -43,6 +43,12 @@ class MyTarget(PromptTarget):
 ``send_prompt_async`` (the public entry point) is ``@final`` and MUST NOT
 be overridden. Override ``_send_prompt_to_target_async`` instead.
 
+Targets that hold external state keyed by conversation (a websocket
+connection, browser page, or upstream session) SHOULD override
+``reset_conversation_async``. It must be safe to call more than once and for
+unknown conversation IDs. Whole-target cleanup stays in
+``cleanup_target_async``.
+
 ## Keyword-only ``__init__`` is enforced
 
 Every ``PromptTarget`` subclass MUST make all ``__init__`` parameters

pyrit/executor/attack/core/attack_strategy.py

@@ -3,6 +3,7 @@
 
 from __future__ import annotations
 
+import asyncio
 import dataclasses
 import logging  # noqa: TC003
 import time
@@ -47,6 +48,8 @@
 from pyrit.prompt_target.common.target_requirements import TargetRequirements
 
 if TYPE_CHECKING:
+    from types import TracebackType
+
     from pyrit.executor.attack.component.prepended_conversation_config import (
         PrependedConversationConfig,
     )
@@ -72,6 +75,61 @@ class _NextMessageOverrideState(Enum):
     UNSET = "unset"
 
 
+class _ObjectiveTargetConversationLifecycle:
+    """Track and release objective-target conversations for one attack execution."""
+
+    def __init__(
+        self,
+        *,
+        objective_target: PromptTarget,
+        logger: logging.Logger | logging.LoggerAdapter[logging.Logger],
+    ) -> None:
+        self._objective_target = objective_target
+        self._logger = logger
+        self._conversation_ids: set[str] = set()
+
+    async def __aenter__(self) -> _ObjectiveTargetConversationLifecycle:
+        """
+        Start tracking target invocations.
+
+        Returns:
+            _ObjectiveTargetConversationLifecycle: This lifecycle instance.
+        """
+        return self
+
+    async def __aexit__(
+        self,
+        exc_type: type[BaseException] | None,
+        exc_value: BaseException | None,
+        traceback: TracebackType | None,
+    ) -> None:
+        """Release each conversation invoked during the attack."""
+        pending_cancellation: asyncio.CancelledError | None = None
+        for conversation_id in self._conversation_ids:
+            try:
+                await self._objective_target.reset_conversation_async(conversation_id=conversation_id)
+            except asyncio.CancelledError as cancellation:
+                # Attempt every reset, then honor the first cancellation after the loop.
+                pending_cancellation = pending_cancellation or cancellation
+            except Exception as error:  # noqa: BLE001 - cleanup must not replace the attack outcome
+                self._logger.warning(
+                    "Failed to reset objective-target conversation %s: %s",
+                    conversation_id,
+                    error,
+                )
+        if pending_cancellation is not None:
+            raise pending_cancellation
+
+    def record_invocation(self, *, conversation_id: str) -> None:
+        """
+        Record one objective-target invocation.
+
+        Args:
+            conversation_id (str): The conversation ID used by the target.
+        """
+        self._conversation_ids.add(conversation_id)
+
+
 @dataclass
 class AttackContext(StrategyContext, ABC, Generic[AttackParamsT]):
     """
@@ -101,6 +159,12 @@ class AttackContext(StrategyContext, ABC, Generic[AttackParamsT]):
     _prepended_conversation_override: list[Message] | None = None
     _memory_labels_override: dict[str, str] | None = None
     _error_result_persistence_error: Exception | None = field(default=None, init=False, repr=False)
+    _objective_target_conversation_lifecycle: _ObjectiveTargetConversationLifecycle | None = field(
+        default=None,
+        init=False,
+        repr=False,
+        compare=False,
+    )
 
     # Per-execution prepended-history boundary and send lifecycle. Never persisted.
     prepended_history_send_context: PrependedHistorySendContext | None = field(
@@ -167,6 +231,21 @@ def next_message(self, value: Message | None) -> None:
         """Set the next message (for attacks that generate internally)."""
         self._next_message_override = value
 
+    def _record_objective_target_invocation(self, *, conversation_id: str) -> None:
+        """
+        Record an objective-target invocation for lifecycle cleanup.
+
+        Recording is a no-op when no objective-target scope is active, so attack
+        helpers remain callable outside a full execution.
+
+        Args:
+            conversation_id (str): The conversation ID used by the target.
+        """
+        lifecycle = self._objective_target_conversation_lifecycle
+        if lifecycle is None:
+            return
+        lifecycle.record_invocation(conversation_id=conversation_id)
+
 
 class _DefaultAttackStrategyEventHandler(StrategyEventHandler[AttackStrategyContextT, AttackStrategyResultT]):
     """
@@ -699,16 +778,25 @@ async def execute_with_context_async(self, *, context: AttackStrategyContextT) -
             ExceptionGroup: If attack execution and recording its error result both fail.
         """
         context._error_result_persistence_error = None
+        lifecycle = _ObjectiveTargetConversationLifecycle(
+            objective_target=self._objective_target,
+            logger=self._logger,
+        )
+        context._objective_target_conversation_lifecycle = lifecycle
         try:
-            result = await super().execute_with_context_async(context=context)
-        except Exception as attack_error:
-            persistence_error = context._error_result_persistence_error
-            if persistence_error is not None:
-                raise ExceptionGroup(
-                    "Attack execution and error result persistence failed",
-                    [attack_error, persistence_error],
-                ) from None
-            raise
+            async with lifecycle:
+                try:
+                    result = await super().execute_with_context_async(context=context)
+                except Exception as attack_error:
+                    persistence_error = context._error_result_persistence_error
+                    if persistence_error is not None:
+                        raise ExceptionGroup(
+                            "Attack execution and error result persistence failed",
+                            [attack_error, persistence_error],
+                        ) from None
+                    raise
+        finally:
+            context._objective_target_conversation_lifecycle = None
 
         self._default_event_handler._persist_result(result=result)
         return result

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,

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,

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,

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,

pyrit/executor/attack/multi_turn/tree_of_attacks.py

@@ -74,7 +74,7 @@
 from pyrit.score.true_false.true_false_inverter_scorer import TrueFalseInverterScorer
 
 if TYPE_CHECKING:
-    from collections.abc import AsyncIterator
+    from collections.abc import AsyncIterator, Callable
     from pathlib import Path
 
     from pyrit.models.literals import PromptDataType
@@ -377,6 +377,7 @@ def __init__(
         attack_id: ComponentIdentifier,
         attack_strategy_name: str,
         modality_router: _ModalityFeedbackRouter,
+        record_objective_conversation: Callable[..., None],
         use_score_as_feedback: bool = True,
         memory_labels: dict[str, str] | None = None,
         parent_id: str | None = None,
@@ -405,6 +406,8 @@ def __init__(
                 whether prior media should travel back to the adversarial chat or forward to
                 the objective target, and fills adversarial-placeholder pieces in seed
                 messages. Typically shared across all nodes of the same attack.
+            record_objective_conversation (Callable[..., None]): Records an objective-target
+                conversation ID for cleanup before each objective send.
             use_score_as_feedback (bool): Whether subsequent adversarial prompts include
                 the objective score. Defaults to True.
             memory_labels (dict[str, str] | None): Labels for memory storage.
@@ -432,6 +435,7 @@ def __init__(
         self._attack_strategy_name = attack_strategy_name
         self._memory_labels = memory_labels or {}
         self._modality_router = modality_router
+        self._record_objective_conversation = record_objective_conversation
         self._prepended_conversation_config = prepended_conversation_config or PrependedConversationConfig()
         self._use_score_as_feedback = use_score_as_feedback
 
@@ -682,6 +686,7 @@ async def _send_prompt_to_target_async(self, prompt: str) -> Message:
             objective_target_conversation_id=self.objective_target_conversation_id,
             objective=self._objective,
         ):
+            self._record_objective_conversation(conversation_id=self.objective_target_conversation_id)
             response = await self._prompt_normalizer.send_prompt_async(
                 message=message,
                 request_converter_configurations=self._request_converters,
@@ -761,6 +766,7 @@ async def _send_initial_prompt_to_target_async(self) -> Message:
             objective_target_conversation_id=self.objective_target_conversation_id,
             objective=self._objective,
         ):
+            self._record_objective_conversation(conversation_id=self.objective_target_conversation_id)
             response = await self._prompt_normalizer.send_prompt_async(
                 message=message,
                 request_converter_configurations=self._request_converters,
@@ -967,6 +973,7 @@ def duplicate(self) -> _TreeOfAttacksNode:
             attack_id=self._attack_id,
             attack_strategy_name=self._attack_strategy_name,
             modality_router=self._modality_router,
+            record_objective_conversation=self._record_objective_conversation,
             use_score_as_feedback=self._use_score_as_feedback,
             memory_labels=self._memory_labels,
             desired_response_prefix=self._desired_response_prefix,
@@ -1392,9 +1399,9 @@ async def execute_nodes_async(
         """
         Execute nodes in ordered batches and yield each completed batch.
 
-        Node instances own all branch-specific mutable state. This executor only
-        schedules their existing execution protocol, so failures and cancellation
-        retain ``asyncio.gather`` semantics.
+        Node instances own all branch-specific mutable state. If one node fails,
+        the executor cancels and awaits the other nodes before it propagates the
+        error.
 
         Args:
             nodes (list[_TreeOfAttacksNode]): Nodes to execute.
@@ -1408,7 +1415,14 @@ async def execute_nodes_async(
             batch_nodes = nodes[batch_start : batch_start + self._batch_size]
             self._log_batch_start(batch_start=batch_start, batch_nodes=batch_nodes, total_nodes=len(nodes))
 
-            await asyncio.gather(*(node.send_prompt_async(objective=objective) for node in batch_nodes))
+            tasks = [asyncio.create_task(node.send_prompt_async(objective=objective)) for node in batch_nodes]
+            try:
+                await asyncio.gather(*tasks)
+            except BaseException:
+                for task in tasks:
+                    task.cancel()
+                await asyncio.gather(*tasks, return_exceptions=True)
+                raise
 
             yield batch_start, batch_nodes
 
@@ -2198,6 +2212,7 @@ def _create_attack_node(
             attack_id=self.get_identifier(),
             attack_strategy_name=self.__class__.__name__,
             modality_router=self._modality_router,
+            record_objective_conversation=context._record_objective_target_invocation,
             use_score_as_feedback=self._attack_scoring_config.use_score_as_feedback,
             memory_labels=context.memory_labels,
             desired_response_prefix=self._configuration.desired_response_prefix,

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,

pyrit/prompt_target/common/prompt_target.py

@@ -360,6 +360,24 @@ def set_system_prompt(
             ).to_message(),
         )
 
+    async def reset_conversation_async(self, *, conversation_id: str) -> None:
+        """
+        Release any target-side state held for a conversation.
+
+        The attack execution scope calls this for objective-target conversations
+        recorded at the common dispatch boundary. Targets that keep external state
+        keyed by conversation (a websocket connection, a browser page, an upstream
+        session) override this to close or discard it. Targets that are stateless
+        between calls need not override it.
+
+        This is best-effort cleanup, so implementations should not raise for a
+        conversation id they do not recognize, and should be safe to call more
+        than once for the same id.
+
+        Args:
+            conversation_id (str): The conversation id to release state for.
+        """
+
     def dispose_db_engine(self) -> None:
         """
         Dispose database engine to release database connections and resources.

pyrit/prompt_target/openai/openai_realtime_target.py

@@ -11,6 +11,7 @@
 from openai import AsyncOpenAI
 
 from pyrit.common import forward_init_parameters
+from pyrit.common.deprecation import print_deprecation_message
 from pyrit.exceptions import (
     pyrit_target_retry,
 )
@@ -479,21 +480,45 @@ async def cleanup_target_async(self) -> None:
                 logger.warning(f"Error closing realtime client: {e}")
             self._realtime_client = None
 
-    async def cleanup_conversation_async(self, conversation_id: str) -> None:
+    async def reset_conversation_async(self, *, conversation_id: str) -> None:
         """
         Disconnects from the Realtime API for a specific conversation.
 
+        Closes the cached connection for ``conversation_id`` and drops it from
+        ``_existing_conversation``. Errors while closing are logged and
+        swallowed, and an unknown conversation id is a no-op, so this is safe
+        to call from attack lifecycle cleanup.
+
         Args:
             conversation_id (str): The conversation ID to disconnect from.
         """
-        connection = self._existing_conversation.get(conversation_id)
-        if connection:
-            try:
-                await connection.close()
-                logger.info(f"Disconnected from {self._endpoint} with conversation ID: {conversation_id}")
-            except Exception as e:
-                logger.warning(f"Error closing connection for {conversation_id}: {e}")
-            del self._existing_conversation[conversation_id]
+        connection = self._existing_conversation.pop(conversation_id, None)
+        if not connection:
+            return
+
+        try:
+            await connection.close()
+        except Exception as error:  # noqa: BLE001 - cleanup must not replace the attack outcome
+            logger.warning(f"Error closing connection for {conversation_id}: {error}")
+            return
+
+        logger.info(f"Disconnected from {self._endpoint} with conversation ID: {conversation_id}")
+
+    async def cleanup_conversation_async(self, conversation_id: str) -> None:
+        """
+        Disconnect from the Realtime API for a specific conversation.
+
+        Deprecated. Use ``reset_conversation_async`` instead.
+
+        Args:
+            conversation_id (str): The conversation ID to disconnect from.
+        """
+        print_deprecation_message(
+            old_item="RealtimeTarget.cleanup_conversation_async",
+            new_item="RealtimeTarget.reset_conversation_async",
+            removed_in="1.3.0",
+        )
+        await self.reset_conversation_async(conversation_id=conversation_id)
 
     async def _connect_async(self, *, conversation_id: str) -> Any:
         """