diff --git a/pyrit/executor/attack/core/attack_strategy.py b/pyrit/executor/attack/core/attack_strategy.py index 6bab4492b7..11c58b60ad 100644 --- a/pyrit/executor/attack/core/attack_strategy.py +++ b/pyrit/executor/attack/core/attack_strategy.py @@ -355,8 +355,11 @@ async def _on_error_async( collector = get_retry_collector() retry_events = collector.events if collector else [] - # Build a conversation_id — use context's if available, otherwise generate one - conversation_id = getattr(context, "conversation_id", None) 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, diff --git a/pyrit/executor/attack/multi_turn/tree_of_attacks.py b/pyrit/executor/attack/multi_turn/tree_of_attacks.py index 1fe9a4e01c..ec1ab0f8e3 100644 --- a/pyrit/executor/attack/multi_turn/tree_of_attacks.py +++ b/pyrit/executor/attack/multi_turn/tree_of_attacks.py @@ -179,6 +179,15 @@ class TAPAttackContext(MultiTurnAttackContext[Any]): best_objective_score: Score | None = None best_adversarial_conversation_id: str | None = None + @property + def conversation_id(self) -> str | None: + """The best objective-target conversation, or the first active branch.""" + if self.best_conversation_id: + return self.best_conversation_id + if self.nodes: + return self.nodes[0].objective_target_conversation_id + return None + class TAPAttackResult(AttackResult): """ diff --git a/tests/unit/executor/attack/core/test_attack_strategy.py b/tests/unit/executor/attack/core/test_attack_strategy.py index 7ef9156c5e..6d588c447c 100644 --- a/tests/unit/executor/attack/core/test_attack_strategy.py +++ b/tests/unit/executor/attack/core/test_attack_strategy.py @@ -16,6 +16,8 @@ AttackStrategy, _DefaultAttackStrategyEventHandler, ) +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.core import StrategyEvent, StrategyEventData from pyrit.memory.central_memory import CentralMemory from pyrit.models import ( @@ -632,6 +634,49 @@ async def test_on_error_persists_result_to_memory(self, sample_attack_context, m assert stored_result.error_type == "ValueError" assert stored_result.execution_time_ms == 500 + async def test_on_error_uses_multi_turn_session_conversation_id(self, mock_memory): + """Test that multi-turn failures remain correlated with their active conversation.""" + context = MultiTurnAttackContext( + params=AttackParameters(objective="Test harmful objective"), + session=ConversationSession(conversation_id="active-conversation-id"), + ) + + with patch("pyrit.memory.central_memory.CentralMemory.get_memory_instance", return_value=mock_memory): + handler = _DefaultAttackStrategyEventHandler() + event_data = StrategyEventData( + event=StrategyEvent.ON_ERROR, + strategy_name="TestStrategy", + strategy_id="test-id", + context=context, + error=TimeoutError("target timed out"), + ) + await handler.on_event_async(event_data) + + stored_result = mock_memory.add_attack_results_to_memory.call_args.kwargs["attack_results"][0] + assert stored_result.conversation_id == "active-conversation-id" + + async def test_on_error_uses_tap_best_conversation_id(self, mock_memory): + """Test that TAP failures remain correlated with the best objective-target conversation.""" + context = TAPAttackContext( + params=AttackParameters(objective="Test harmful objective"), + session=ConversationSession(conversation_id="unused-session-id"), + best_conversation_id="best-conversation-id", + ) + + with patch("pyrit.memory.central_memory.CentralMemory.get_memory_instance", return_value=mock_memory): + handler = _DefaultAttackStrategyEventHandler() + event_data = StrategyEventData( + event=StrategyEvent.ON_ERROR, + strategy_name="TestStrategy", + strategy_id="test-id", + context=context, + error=TimeoutError("target timed out"), + ) + await handler.on_event_async(event_data) + + stored_result = mock_memory.add_attack_results_to_memory.call_args.kwargs["attack_results"][0] + assert stored_result.conversation_id == "best-conversation-id" + async def test_on_error_skips_when_no_error_or_context(self, mock_memory): """Test that error handler returns early when error or context is None""" with patch("pyrit.memory.central_memory.CentralMemory.get_memory_instance", return_value=mock_memory): diff --git a/tests/unit/executor/attack/multi_turn/test_pair.py b/tests/unit/executor/attack/multi_turn/test_pair.py index ee62a90dff..343dfd5800 100644 --- a/tests/unit/executor/attack/multi_turn/test_pair.py +++ b/tests/unit/executor/attack/multi_turn/test_pair.py @@ -18,6 +18,7 @@ from pyrit.executor.attack import ( AttackAdversarialConfig, + AttackParameters, AttackScoringConfig, PAIRAttack, TreeOfAttacksWithPruningAttack, @@ -148,6 +149,22 @@ def test_pair_uses_tap_context_type(self, objective_target, adversarial_config): ) assert attack._context_type is TAPAttackContext + def test_pair_context_uses_best_objective_target_conversation(self, objective_target, adversarial_config): + attack = PAIRAttack( + objective_target=objective_target, + attack_adversarial_config=adversarial_config, + ) + context = attack._context_type(params=AttackParameters(objective="Test objective")) + active_node = MagicMock() + active_node.objective_target_conversation_id = "pair-active-conversation" + context.nodes = [active_node] + + assert context.conversation_id == "pair-active-conversation" + + context.best_conversation_id = "pair-best-conversation" + + assert context.conversation_id == "pair-best-conversation" + def test_pair_validates_adversarial_target_capabilities(self, objective_target): """An adversarial target lacking native MULTI_TURN/SYSTEM_PROMPT must be rejected (inherited from TAP).""" bad_adversarial = MagicMock(spec=PromptTarget) 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..f4f962c224 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 @@ -433,6 +433,26 @@ def mock_prompt_loading(attack: TreeOfAttacksWithPruningAttack): attack._load_adversarial_prompts() +class TestTAPAttackContext: + """Tests for TAP conversation correlation.""" + + def test_conversation_id_uses_first_active_branch_before_best_is_selected(self, node_factory): + context = TestHelpers.create_basic_context() + context.nodes = [ + node_factory.create_node(NodeMockConfig(objective_target_conversation_id="first-branch")), + node_factory.create_node(NodeMockConfig(objective_target_conversation_id="second-branch")), + ] + + assert context.conversation_id == "first-branch" + + def test_conversation_id_prefers_best_branch(self, node_factory): + context = TestHelpers.create_basic_context() + context.nodes = [node_factory.create_node(NodeMockConfig(objective_target_conversation_id="active-branch"))] + context.best_conversation_id = "best-branch" + + assert context.conversation_id == "best-branch" + + @pytest.fixture def node_factory(): """Fixture providing the MockNodeFactory."""