From bd9c288b60e85ae5c0c3e866de9491c433d5b3a1 Mon Sep 17 00:00:00 2001 From: jsong468 Date: Wed, 8 Jul 2026 17:16:17 -0700 Subject: [PATCH 1/2] surface error piece content instead of generic 500 --- pyrit/backend/services/attack_service.py | 35 ++++++++-- tests/unit/backend/test_attack_service.py | 81 +++++++++++++++++++++++ 2 files changed, 109 insertions(+), 7 deletions(-) diff --git a/pyrit/backend/services/attack_service.py b/pyrit/backend/services/attack_service.py index 03b3a3e30f..5a8849de77 100644 --- a/pyrit/backend/services/attack_service.py +++ b/pyrit/backend/services/attack_service.py @@ -15,6 +15,7 @@ - AI-generated attacks may have multiple related conversations """ +import logging import mimetypes import uuid from collections.abc import Sequence @@ -68,6 +69,8 @@ ) from pyrit.prompt_normalizer import PromptConverterConfiguration, PromptNormalizer +logger = logging.getLogger(__name__) + class AttackService: """ @@ -623,13 +626,31 @@ async def add_message_async(self, *, attack_result_id: str, request: AddMessageR if request.send: assert target_registry_name is not None # validated above - await self._send_and_store_message_async( - conversation_id=msg_conversation_id, - target_registry_name=target_registry_name, - request=request, - sequence=sequence, - labels=attack_labels, # deprecated - ) + try: + await self._send_and_store_message_async( + conversation_id=msg_conversation_id, + target_registry_name=target_registry_name, + request=request, + sequence=sequence, + labels=attack_labels, # deprecated + ) + except Exception: + # PromptNormalizer persists a full error piece (response_error + + # traceback) to memory *before* re-raising. Surface that stored + # piece inline so the send (POST) response matches the + # conversation-reload (GET) view instead of collapsing to a + # generic 500. If no new error piece was stored (the failure + # happened before the send, e.g. target lookup), re-raise so the + # route still reports a real error. + prior_ids = {p.id for p in existing} + current_pieces = self._memory.get_message_pieces(conversation_id=msg_conversation_id) + if not any(p.id not in prior_ids and p.has_error() for p in current_pieces): + raise + logger.exception( + "Send failed for attack '%s' conversation '%s'; surfacing stored error piece.", + attack_result_id, + msg_conversation_id, + ) else: existing_metadata = self._memory._get_conversation(conversation_id=msg_conversation_id) await self._store_message_only_async( diff --git a/tests/unit/backend/test_attack_service.py b/tests/unit/backend/test_attack_service.py index 0afeb1b1b0..1139eea7eb 100644 --- a/tests/unit/backend/test_attack_service.py +++ b/tests/unit/backend/test_attack_service.py @@ -1104,6 +1104,87 @@ async def test_add_message_with_send_raises_when_target_not_found(self, attack_s with pytest.raises(ValueError, match="Target object .* not found"): await attack_service.add_message_async(attack_result_id="test-id", request=request) + async def test_add_message_surfaces_stored_error_piece_on_send_failure(self, attack_service, mock_memory) -> None: + """When the normalizer stores an error piece then raises, the send returns it inline (no raise).""" + ar = make_attack_result(conversation_id="test-id") + mock_memory.get_attack_results.return_value = [ar] + mock_memory.get_conversation_messages.return_value = [] + + # The PromptNormalizer persists a full error piece before re-raising; model + # that by flipping to return the stored error piece only after send fails. + error_piece = MessagePiece( + role="assistant", + original_value="Connection error.\nAPIConnectionError('Connection error.')\nTraceback...", + original_value_data_type="error", + converted_value="Connection error.\nAPIConnectionError('Connection error.')\nTraceback...", + converted_value_data_type="error", + conversation_id="test-id", + sequence=1, + response_error="processing", + ) + state = {"sent": False} + mock_memory.get_message_pieces.side_effect = lambda **_: [error_piece] if state["sent"] else [] + + async def _raise_after_store(**_): + state["sent"] = True + raise Exception("Error sending prompt with conversation ID: test-id") + + with ( + patch("pyrit.backend.services.attack_service.get_target_service") as mock_get_target_svc, + patch("pyrit.backend.services.attack_service.PromptNormalizer") as mock_normalizer_cls, + ): + mock_target_svc = MagicMock() + mock_target_svc.get_target_object.return_value = _make_matching_target_mock() + mock_get_target_svc.return_value = mock_target_svc + + mock_normalizer = MagicMock() + mock_normalizer.send_prompt_async = AsyncMock(side_effect=_raise_after_store) + mock_normalizer_cls.return_value = mock_normalizer + + request = AddMessageRequest( + pieces=[MessagePieceRequest(original_value="Hello")], + target_conversation_id="test-id", + send=True, + target_registry_name="test-target", + ) + + # Should NOT raise: the stored error piece is surfaced via the normal response. + result = await attack_service.add_message_async(attack_result_id="test-id", request=request) + + assert result.attack is not None + mock_normalizer.send_prompt_async.assert_called_once() + + async def test_add_message_reraises_when_send_fails_without_stored_error_piece( + self, attack_service, mock_memory + ) -> None: + """If the send fails but no error piece was stored, the exception propagates (real error).""" + ar = make_attack_result(conversation_id="test-id") + mock_memory.get_attack_results.return_value = [ar] + mock_memory.get_message_pieces.return_value = [] # no error piece ever stored + mock_memory.get_conversation_messages.return_value = [] + + with ( + patch("pyrit.backend.services.attack_service.get_target_service") as mock_get_target_svc, + patch("pyrit.backend.services.attack_service.PromptNormalizer") as mock_normalizer_cls, + ): + mock_target_svc = MagicMock() + mock_target_svc.get_target_object.return_value = _make_matching_target_mock() + mock_get_target_svc.return_value = mock_target_svc + + mock_normalizer = MagicMock() + mock_normalizer.send_prompt_async = AsyncMock(side_effect=RuntimeError("boom")) + mock_normalizer_cls.return_value = mock_normalizer + + request = AddMessageRequest( + pieces=[MessagePieceRequest(original_value="Hello")], + target_conversation_id="test-id", + send=True, + target_registry_name="test-target", + ) + + with pytest.raises(RuntimeError, match="boom"): + await attack_service.add_message_async(attack_result_id="test-id", request=request) + async def test_add_message_with_converter_ids_gets_converters(self, attack_service, mock_memory) -> None: """Test that add_message with converter_ids gets converters from service.""" ar = make_attack_result(conversation_id="test-id") From b1149e185515086b7d6b54ae95bbc9653b166efc Mon Sep 17 00:00:00 2001 From: jsong468 Date: Thu, 9 Jul 2026 12:12:13 -0700 Subject: [PATCH 2/2] pr feedback test buff --- tests/unit/backend/test_attack_service.py | 24 +++++++++++++++++------ 1 file changed, 18 insertions(+), 6 deletions(-) diff --git a/tests/unit/backend/test_attack_service.py b/tests/unit/backend/test_attack_service.py index 7601b3422f..41691af022 100644 --- a/tests/unit/backend/test_attack_service.py +++ b/tests/unit/backend/test_attack_service.py @@ -1106,18 +1106,18 @@ async def test_add_message_with_send_raises_when_target_not_found(self, attack_s await attack_service.add_message_async(attack_result_id="test-id", request=request) async def test_add_message_surfaces_stored_error_piece_on_send_failure(self, attack_service, mock_memory) -> None: - """When the normalizer stores an error piece then raises, the send returns it inline (no raise).""" + """When the normalizer stores an error piece then raises, the send returns that turn inline (no raise).""" ar = make_attack_result(conversation_id="test-id") mock_memory.get_attack_results.return_value = [ar] - mock_memory.get_conversation_messages.return_value = [] # The PromptNormalizer persists a full error piece before re-raising; model # that by flipping to return the stored error piece only after send fails. + traceback_text = "Connection error.\nAPIConnectionError('Connection error.')\nTraceback..." error_piece = MessagePiece( role="assistant", - original_value="Connection error.\nAPIConnectionError('Connection error.')\nTraceback...", + original_value=traceback_text, original_value_data_type="error", - converted_value="Connection error.\nAPIConnectionError('Connection error.')\nTraceback...", + converted_value=traceback_text, converted_value_data_type="error", conversation_id="test-id", sequence=1, @@ -1125,6 +1125,11 @@ async def test_add_message_surfaces_stored_error_piece_on_send_failure(self, att ) state = {"sent": False} mock_memory.get_message_pieces.side_effect = lambda **_: [error_piece] if state["sent"] else [] + # The conversation-messages read (used to build the response DTO) must include + # the stored error turn so we can assert it is surfaced to the caller. + mock_memory.get_conversation_messages.side_effect = lambda **_: ( + [Message(message_pieces=[error_piece])] if state["sent"] else [] + ) async def _raise_after_store(**_): state["sent"] = True @@ -1149,11 +1154,18 @@ async def _raise_after_store(**_): target_registry_name="test-target", ) - # Should NOT raise: the stored error piece is surfaced via the normal response. + # Should NOT raise: the stored error turn is surfaced via the normal response. result = await attack_service.add_message_async(attack_result_id="test-id", request=request) - assert result.attack is not None mock_normalizer.send_prompt_async.assert_called_once() + assert result.attack is not None + + # The error turn (response_error="processing" + traceback) must come back in the response, + # so the send-time view matches the conversation-reload view. + returned_pieces = [piece for message in result.messages.messages for piece in message.message_pieces] + error_views = [piece for piece in returned_pieces if piece.response_error == "processing"] + assert len(error_views) == 1 + assert "APIConnectionError" in error_views[0].converted_value async def test_add_message_reraises_when_send_fails_without_stored_error_piece( self, attack_service, mock_memory