From 33c033b501d5074648c2603dd18b409b4af8e510 Mon Sep 17 00:00:00 2001
From: "Slai.life" <294879785+huong8373tt-beep@users.noreply.github.com>
Date: Sat, 17 Oct 2026 06:38:18 +0800
Subject: [PATCH] MAINT Add atomic attack context to scenario errors

---
 pyrit/scenario/core/scenario.py           |  3 +++
 tests/unit/scenario/core/test_scenario.py | 16 +++++++++++++---
 2 files changed, 16 insertions(+), 3 deletions(-)

diff --git a/pyrit/scenario/core/scenario.py b/pyrit/scenario/core/scenario.py
index 7b67c5a83d..c5bb563c43 100644
--- a/pyrit/scenario/core/scenario.py
+++ b/pyrit/scenario/core/scenario.py
@@ -1995,6 +1995,9 @@ async def worker_async() -> None:
                     if result.has_incomplete:
                         stop_event.set()
                 except Exception as exc:
+                    exc.add_note(
+                        f"Failed atomic attack {atomic_attack.atomic_attack_name!r} in scenario {self._name!r}."
+                    )
                     outcomes.append(exc)
                     stop_event.set()
                 except BaseException:
diff --git a/tests/unit/scenario/core/test_scenario.py b/tests/unit/scenario/core/test_scenario.py
index d1cd757cca..2bac1ebf45 100644
--- a/tests/unit/scenario/core/test_scenario.py
+++ b/tests/unit/scenario/core/test_scenario.py
@@ -2493,6 +2493,11 @@ async def _run(*args, **kwargs):
         messages = sorted(str(e) for e in exc_info.value.exceptions)
         assert messages == ["a boom", "b boom", "c boom"]
         assert all(isinstance(e, RuntimeError) for e in exc_info.value.exceptions)
+        assert sorted(error.__notes__ for error in exc_info.value.exceptions) == [
+            ["Failed atomic attack 'attack_run_1' in scenario 'Test Scenario'."],
+            ["Failed atomic attack 'attack_run_2' in scenario 'Test Scenario'."],
+            ["Failed atomic attack 'attack_run_3' in scenario 'Test Scenario'."],
+        ]
 
     async def test_single_failure_is_raised_directly_not_wrapped(
         self, mock_atomic_attacks, sample_attack_results, mock_objective_target
@@ -2503,8 +2508,10 @@ async def test_single_failure_is_raised_directly_not_wrapped(
                 [sample_attack_results[i]], atomic_attack=mock_atomic_attacks[i]
             )
 
+        failure = RuntimeError("solo boom")
+
         async def bad_run(*a, **k):
-            raise RuntimeError("solo boom")
+            raise failure
 
         mock_atomic_attacks[1].run_async = AsyncMock(side_effect=bad_run)
 
@@ -2521,10 +2528,13 @@ async def bad_run(*a, **k):
         )
         await scenario.initialize_async()
 
-        # Bare RuntimeError, not ExceptionGroup.
-        with pytest.raises(RuntimeError, match="solo boom"):
+        # Bare RuntimeError, not ExceptionGroup, with the original exception identity.
+        with pytest.raises(RuntimeError, match="solo boom") as exc_info:
             await scenario.run_async()
 
+        assert exc_info.value is failure
+        assert exc_info.value.__notes__ == ["Failed atomic attack 'attack_run_2' in scenario 'Test Scenario'."]
+
     async def test_max_concurrency_one_serializes_via_single_worker(
         self, mock_atomic_attacks, sample_attack_results, mock_objective_target
     ):