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

Author: mahdi-al-hakim

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.
 

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

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,

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(