diff --git a/pyrit/executor/attack/multi_turn/tree_of_attacks.py b/pyrit/executor/attack/multi_turn/tree_of_attacks.py index a4baf452bf..5dccd697f9 100644 --- a/pyrit/executor/attack/multi_turn/tree_of_attacks.py +++ b/pyrit/executor/attack/multi_turn/tree_of_attacks.py @@ -686,6 +686,8 @@ async def _send_prompt_to_target_async(self, prompt: str) -> Message: Side Effects: - Sets self.last_response to the target's response text """ + # Check the turn before a single-turn target gets a fresh conversation, which is always empty. + turn_index = 0 if (await self._is_first_turn_async()) else 1 self._rotate_unseeded_single_turn_conversation() # Build the request message via the modality router so prior media (if any) @@ -693,7 +695,7 @@ async def _send_prompt_to_target_async(self, prompt: str) -> Message: message = self._modality_router.build_objective_input_message( text=prompt, last_response=self.last_response, - turn_index=0 if (await self._is_first_turn_async()) else 1, + turn_index=turn_index, ) # Send prompt with configured converters 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 737fe2e208..32eeed65e1 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 @@ -3337,8 +3337,9 @@ def _make_image_response(conversation_id: str) -> Message: ] ) - async def test_node_send_prompt_to_target_forwards_prev_image(self, node_components): - """Objective accepting {text, image_path} gets prev image attached on turn N.""" + @pytest.mark.parametrize("supports_multi_turn", [True, False]) + async def test_node_send_prompt_to_target_forwards_prev_image(self, node_components, supports_multi_turn): + """Objective accepting {text, image_path} gets prev image attached on turn N, also when it is single-turn.""" # Override the router to point at an image-capable objective. node_components["objective_target"].configuration.capabilities.input_modalities = frozenset( {frozenset({"text"}), frozenset({"text", "image_path"})} @@ -3346,6 +3347,9 @@ async def test_node_send_prompt_to_target_forwards_prev_image(self, node_compone node_components["objective_target"].configuration.capabilities.output_modalities = frozenset( {frozenset({"image_path"})} ) + node_components["objective_target"].configuration.includes.side_effect = lambda capability: ( + capability == CapabilityName.MULTI_TURN and supports_multi_turn + ) node_components["modality_router"] = _ModalityFeedbackRouter( adversarial_chat=node_components["adversarial_chat"], objective_target=node_components["objective_target"],