FIX: keep an audit record when an in-flight target call is cancelled

pyrit/executor/attack/core/attack_cancellation.py

@@ -0,0 +1,56 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT license.
+
+"""Typed signal for attacks cancelled after they started sending to the objective target."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, Any, ClassVar
+
+if TYPE_CHECKING:
+    from pyrit.models import AttackResult
+
+_DEFAULT_REASON = "Cancelled after a request was sent to the objective target; its response is unknown."
+
+
+@dataclass(frozen=True)
+class AttackCancellation:
+    """
+    A typed record that an attack was cancelled after it started sending to the objective target.
+
+    The attack carries this on ``AttackResult.metadata`` so the attempt keeps an audit trail.
+    Results carrying it use ``AttackOutcome.UNDETERMINED``: no verdict was reached, and what the
+    objective target did with an in-flight request is unknown. The row records an interrupted
+    attempt rather than a measurement, so scenario resume treats its objective as still pending.
+    """
+
+    reason: str = _DEFAULT_REASON
+
+    #: Key under which the signal travels in ``AttackResult.metadata``.
+    METADATA_KEY: ClassVar[str] = "attack_cancelled"
+
+    def to_metadata(self) -> dict[str, Any]:
+        """
+        Render the signal for ``AttackResult.metadata``.
+
+        Returns:
+            dict[str, Any]: The metadata fragment marking the result as cancelled.
+        """
+        return {self.METADATA_KEY: True}
+
+    @classmethod
+    def from_result(cls, *, result: AttackResult) -> AttackCancellation | None:
+        """
+        Read the signal back off a result.
+
+        Args:
+            result (AttackResult): The result to inspect.
+
+        Returns:
+            AttackCancellation | None: The recorded cancellation, or ``None`` when the result
+            carries no cancellation signal.
+        """
+        if result.metadata.get(cls.METADATA_KEY) is not True:
+            return None
+        return cls(reason=result.outcome_reason or _DEFAULT_REASON)

pyrit/executor/attack/core/attack_strategy.py

@@ -21,6 +21,7 @@
 from pyrit.exceptions.retry_collector import (
     get_retry_collector,
 )
+from pyrit.executor.attack.core.attack_cancellation import AttackCancellation
 from pyrit.executor.attack.core.attack_parameters import AttackParameters, AttackParamsT
 from pyrit.executor.attack.core.attack_scoring import prepare_attack_scoring
 from pyrit.executor.core import (
@@ -180,6 +181,11 @@ async def __aexit__(
         if pending_cancellation is not None:
             raise pending_cancellation
 
+    @property
+    def invoked(self) -> bool:
+        """Whether the attack has started sending to the objective target."""
+        return bool(self._conversation_ids)
+
     def record_invocation(self, *, conversation_id: str) -> None:
         """
         Record one objective-target invocation.
@@ -677,14 +683,8 @@ async def _on_error_async(
         collector = get_retry_collector()
         retry_events = collector.events if collector else []
 
-        # 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,
+            conversation_id=self._result_conversation_id(context=context),
             objective=context.objective,
             attack_result_id=context.attack_result_id or str(uuid.uuid4()),
             outcome=AttackOutcome.ERROR,
@@ -721,6 +721,70 @@ async def _on_error_async(
 
         self._logger.error(f"Attack failed with {type(error).__name__}: {error}")
 
+    @staticmethod
+    def _result_conversation_id(*, context: AttackStrategyContextT) -> str:
+        """
+        Return the conversation an interrupted attack's result should point to.
+
+        Args:
+            context (AttackStrategyContextT): The interrupted attack's context.
+
+        Returns:
+            str: The attack's active conversation ID, or a new ID if it has none.
+        """
+        # 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)
+        return str(conversation_id) if conversation_id else str(uuid.uuid4())
+
+    async def _record_cancellation_async(self, *, context: AttackStrategyContextT) -> None:
+        """
+        Persist an ``UNDETERMINED`` result for an attack cancelled after it started sending to the target.
+
+        The result carries ``AttackCancellation`` and no verdict. A failure to persist it is logged
+        rather than raised, so it never replaces the cancellation.
+
+        Args:
+            context (AttackStrategyContextT): The cancelled attack's context.
+        """
+        if not context._persist_attack_result:
+            return
+
+        cancellation = AttackCancellation()
+        collector = get_retry_collector()
+        retry_events = collector.events if collector else []
+        cancelled_result = AttackResult(
+            conversation_id=self._result_conversation_id(context=context),
+            objective=context.objective,
+            attack_result_id=context.attack_result_id or str(uuid.uuid4()),
+            outcome=AttackOutcome.UNDETERMINED,
+            outcome_reason=cancellation.reason,
+            atomic_attack_identifier=self._build_error_result_identifier(context=context),
+            labels=context.memory_labels,
+            related_conversations=context.related_conversations,
+            retry_events=retry_events,
+            total_retries=len(retry_events),
+            metadata={**context._error_result_metadata, **cancellation.to_metadata()},
+        )
+        if context.start_time:
+            cancelled_result.execution_time_ms = int((time.perf_counter() - context.start_time) * 1000)
+        self._apply_attribution(context=context, result=cancelled_result)
+        self._apply_targeted_harm_categories(context=context, result=cancelled_result)
+
+        try:
+            await self._save_result_async(result=cancelled_result)
+        except Exception:
+            self._logger.warning(
+                "Could not persist the result of cancelled attack %s.", cancelled_result.attack_result_id, exc_info=True
+            )
+            await self._confirm_result_persistence_async(
+                context=context, attack_result_id=cancelled_result.attack_result_id
+            )
+        else:
+            context._persisted_attack_result_id = cancelled_result.attack_result_id
+        self._logger.warning(f"Attack was cancelled. {cancellation.reason}")
+
 
 class AttackStrategy(Strategy[AttackStrategyContextT, AttackStrategyResultT], Identifiable, ABC):
     """
@@ -1055,6 +1119,9 @@ async def execute_with_context_async(self, *, context: AttackStrategyContextT) -
 
         Raises:
             ExceptionGroup: If attack execution and recording its error result both fail.
+            asyncio.CancelledError: If the attack is cancelled. An attack cancelled after it started
+                sending to the objective target first persists an ``UNDETERMINED`` result carrying
+                ``AttackCancellation``.
         """
         context._error_result_persistence_error = None
         context._error_result_metadata.clear()
@@ -1082,6 +1149,11 @@ async def execute_with_context_async(self, *, context: AttackStrategyContextT) -
                         _operation_admission_scope(admitted),
                     ):
                         result = await super().execute_with_context_async(context=context)
+                except asyncio.CancelledError:
+                    # Keep an audit record of an attempt the objective target may have received.
+                    if lifecycle.invoked:
+                        await self._default_event_handler._record_cancellation_async(context=context)
+                    raise
                 except Exception as attack_error:
                     persistence_error = context._error_result_persistence_error
                     if persistence_error is not None:

pyrit/prompt_normalizer/prompt_normalizer.py

@@ -125,6 +125,8 @@ async def send_prompt_async(
         Raises:
             Exception: If an error occurs during the request processing.
             ValueError: If the message pieces are not part of the same sequence.
+            asyncio.CancelledError: If the send is cancelled. A request the target may have
+                received is persisted first, without a response.
         """
         # Validates that the MessagePieces in the Message are part of the same sequence
         request_converter_configurations = request_converter_configurations or []
@@ -165,6 +167,12 @@ async def send_prompt_async(
                 send_context=send_context,
             )
             (await self.memory.add_message_to_memory_async(request=request))
+        except asyncio.CancelledError:
+            # Cancelled while the target may have been handling the request: keep the request as the
+            # record of the attempt, but no response, since none was received. Then propagate.
+            if not (send_context and send_context.target_invocation_count == target_invocation_count_before_send):
+                await self._persist_cancelled_request_async(request=request)
+            raise
         except EmptyResponseException as ex:
             if send_context and send_context.target_invocation_count == target_invocation_count_before_send:
                 cid = request.message_pieces[0].conversation_id if request.message_pieces else None
@@ -244,6 +252,21 @@ async def send_prompt_async(
         # Return the last response for backward compatibility
         return responses[-1]
 
+    async def _persist_cancelled_request_async(self, *, request: Message) -> None:
+        """
+        Persist the request of a send that was cancelled after reaching the target.
+
+        A failure to persist is logged rather than raised, so it never replaces the cancellation.
+
+        Args:
+            request (Message): The converted request that was sent.
+        """
+        try:
+            await self.memory.add_message_to_memory_async(request=request)
+        except Exception:
+            cid = request.message_pieces[0].conversation_id if request.message_pieces else None
+            logger.warning("Could not persist the request of a cancelled send in conversation %s.", cid, exc_info=True)
+
     async def send_prompt_batch_to_target_async(
         self,
         *,

pyrit/scenario/core/scenario.py

@@ -27,6 +27,7 @@
 from pyrit.exceptions import ScenarioPartialFailureException, ScenarioRunInProgressException
 from pyrit.exceptions.exceptions_helpers import summarize_exception
 from pyrit.executor.attack import AttackExecutor, AttackExecutorResult
+from pyrit.executor.attack.core.attack_cancellation import AttackCancellation
 from pyrit.executor.attack.core.attack_executor import ActiveAttackExecution, _observe_interrupted_executions
 from pyrit.executor.attack.core.attack_preparation import AttackPreparationFailure
 from pyrit.executor.attack.core.attack_strategy import AttackContext
@@ -1517,8 +1518,9 @@ async def _get_completed_objective_hashes_by_attack_async(self) -> dict[tuple[st
         )
         completed_by_attack: dict[tuple[str, str | None], set[str]] = {}
         for row in rows:
-            # ERROR rows hit infrastructure problems, and preparation failures never reached the
-            # objective target at all, so neither measured the objective and both stay pending.
+            # ERROR rows hit infrastructure problems, preparation failures never reached the
+            # objective target at all, and cancelled attempts were interrupted before a verdict,
+            # so none of them measured the objective and all stay pending.
             # Every other row did reach the target and recorded the best verdict available --
             # including UNDETERMINED when no scorer was configured or the scorer abstained. Those
             # are measured results of a deterministic configuration; retrying them would re-send
@@ -1529,6 +1531,8 @@ async def _get_completed_objective_hashes_by_attack_async(self) -> dict[tuple[st
                 continue
             if AttackPreparationFailure.from_result(result=row) is not None:
                 continue
+            if AttackCancellation.from_result(result=row) is not None:
+                continue
             name = row.attribution_data.get("parent_collection")
             eval_hash = row.attribution_data.get("parent_eval_hash")
             if not isinstance(name, str) or (eval_hash is not None and not isinstance(eval_hash, str)):

tests/unit/executor/attack/multi_turn/test_crescendo_resilience.py

@@ -21,6 +21,7 @@
     CrescendoAttackContext,
     CrescendoAttackResult,
 )
+from pyrit.executor.attack.core.attack_cancellation import AttackCancellation
 from pyrit.memory import CentralMemory
 from pyrit.models import AttackOutcome, ComponentIdentifier, ConversationType, Message, MessagePiece, Score
 from pyrit.prompt_normalizer import PromptNormalizer
@@ -1043,12 +1044,17 @@ async def test_cancellation_runs_teardown_preserves_completed_turns_and_clears_r
             "objective",
         ]
         pieces = await attack._memory.get_message_pieces_async(conversation_id=context.session.conversation_id)
-        assert len(pieces) == 8
+        # Four completed turns, plus the fifth request that was in flight, with no response.
+        assert len(pieces) == 9
+        assert pieces[-1].api_role == "user"
         assert [piece.original_value for piece in pieces if piece.api_role == "assistant"] == [
             "response-1",
             "response-2",
             "response-3",
             "response-4",
         ]
-        assert (await attack._memory.get_attack_results_async(objective=_OBJECTIVE)) == []
+        [cancelled_result] = await attack._memory.get_attack_results_async(objective=_OBJECTIVE)
+        assert cancelled_result.outcome == AttackOutcome.UNDETERMINED
+        assert cancelled_result.conversation_id == context.session.conversation_id
+        assert AttackCancellation.from_result(result=cancelled_result) is not None
         assert get_retry_collector() is None

tests/unit/executor/attack/single_turn/test_prompt_sending.py

@@ -5,6 +5,7 @@
 import base64
 import codecs
 import uuid
+from typing import Any
 from unittest.mock import AsyncMock, MagicMock, patch
 
 import pytest
@@ -34,6 +35,7 @@
     RTASystemPromptPaths,
     SingleTurnAttackContext,
 )
+from pyrit.executor.attack.core.attack_cancellation import AttackCancellation
 from pyrit.executor.attack.core.attack_preparation import (
     AttackPreparationFailure,
     AttackPreparationFailureKind,
@@ -1818,3 +1820,93 @@ async def test_retry_stores_unsuccessful_conversation_and_updates_id(
         # Verify that the AttackResult includes the related conversations
         assert result.related_conversations == basic_context.related_conversations
         assert len(result.related_conversations) == 2
+
+
+_CANCELLED_OBJECTIVE = "cancellation audit probe"
+
+
+async def _cancel_once_entered_async(*, attack: PromptSendingAttack, entered: asyncio.Event) -> None:
+    task = asyncio.create_task(attack.execute_async(objective=_CANCELLED_OBJECTIVE))
+    await entered.wait()
+    task.cancel()
+    with pytest.raises(asyncio.CancelledError):
+        await task
+
+
+def _blocking_target(entered: asyncio.Event) -> MockPromptTarget:
+    target = MockPromptTarget()
+
+    async def block_send(*, normalized_conversation: list[Message]) -> list[Message]:
+        entered.set()
+        await asyncio.Event().wait()
+        return []
+
+    target._send_prompt_to_target_async = block_send  # type: ignore[method-assign]
+    return target
+
+
+@pytest.mark.usefixtures("patch_central_database")
+async def test_cancelled_target_call_keeps_the_request_and_an_undetermined_result() -> None:
+    entered = asyncio.Event()
+    attack = PromptSendingAttack(objective_target=_blocking_target(entered))
+
+    await _cancel_once_entered_async(attack=attack, entered=entered)
+
+    memory = CentralMemory.get_memory_instance()
+    pieces = await memory.get_message_pieces_async()
+    assert [(piece.api_role, piece.original_value) for piece in pieces] == [("user", _CANCELLED_OBJECTIVE)]
+    [result] = await memory.get_attack_results_async(objective=_CANCELLED_OBJECTIVE)
+    assert result.outcome == AttackOutcome.UNDETERMINED
+    assert result.conversation_id == pieces[0].conversation_id
+    assert AttackCancellation.from_result(result=result) is not None
+    assert result.error_type is None
+
+
+@pytest.mark.usefixtures("patch_central_database")
+async def test_cancellation_before_the_target_call_records_nothing() -> None:
+    entered = asyncio.Event()
+    attack = PromptSendingAttack(objective_target=MockPromptTarget())
+
+    async def block_setup(*, context: Any) -> None:
+        entered.set()
+        await asyncio.Event().wait()
+
+    with patch.object(attack, "_setup_async", side_effect=block_setup):
+        await _cancel_once_entered_async(attack=attack, entered=entered)
+
+    memory = CentralMemory.get_memory_instance()
+    assert await memory.get_message_pieces_async() == []
+    assert await memory.get_attack_results_async(objective=_CANCELLED_OBJECTIVE) == []
+
+
+@pytest.mark.usefixtures("patch_central_database")
+async def test_cancellation_after_the_target_responds_keeps_the_exchange_and_an_undetermined_result() -> None:
+    entered = asyncio.Event()
+    attack = PromptSendingAttack(objective_target=MockPromptTarget())
+
+    async def block_scoring(**kwargs: Any) -> None:
+        entered.set()
+        await asyncio.Event().wait()
+
+    with patch.object(attack, "_evaluate_response_async", side_effect=block_scoring):
+        await _cancel_once_entered_async(attack=attack, entered=entered)
+
+    memory = CentralMemory.get_memory_instance()
+    pieces = await memory.get_message_pieces_async()
+    assert [piece.api_role for piece in pieces] == ["user", "assistant"]
+    [result] = await memory.get_attack_results_async(objective=_CANCELLED_OBJECTIVE)
+    assert result.outcome == AttackOutcome.UNDETERMINED
+    assert AttackCancellation.from_result(result=result) is not None
+
+
+@pytest.mark.usefixtures("patch_central_database")
+async def test_failure_to_persist_a_cancelled_result_does_not_mask_the_cancellation() -> None:
+    entered = asyncio.Event()
+    attack = PromptSendingAttack(objective_target=_blocking_target(entered))
+    memory = CentralMemory.get_memory_instance()
+
+    with patch.object(memory, "add_attack_results_to_memory_async", side_effect=RuntimeError("database unavailable")):
+        await _cancel_once_entered_async(attack=attack, entered=entered)
+
+    assert await memory.get_attack_results_async(objective=_CANCELLED_OBJECTIVE) == []
+    assert [piece.api_role for piece in await memory.get_message_pieces_async()] == ["user"]