From 620b49c24b60972ccbfba94d2840d52db9bc783d Mon Sep 17 00:00:00 2001 From: apocalypse9949 <125962989+apocalypse9949@users.noreply.github.com> Date: Thu, 13 Aug 2026 13:33:45 +0000 Subject: [PATCH 1/5] perf: execute XPIA injection handle activation concurrently Also includes feature addition for exposing trial group aggregation metadata to json reports. --- rampart/attacks/_xpia.py | 18 +++++++++++++++--- rampart/pytest_plugin/_session.py | 7 +++++++ rampart/reporting/json_file.py | 1 + rampart/reporting/sink.py | 3 +++ tests/unit/reporting/test_json_file.py | 2 ++ 5 files changed, 28 insertions(+), 3 deletions(-) diff --git a/rampart/attacks/_xpia.py b/rampart/attacks/_xpia.py index 205ebf8d..3a5a0fed 100644 --- a/rampart/attacks/_xpia.py +++ b/rampart/attacks/_xpia.py @@ -163,10 +163,22 @@ async def _activate_handles_async( Args: stack (AsyncExitStack): The exit stack managing cleanup. """ - for handle in self._handles: - await stack.enter_async_context(handle) - # Concurrent: total = max of all wait times + async def _enter_handle(h: InjectionHandle) -> None: + await stack.enter_async_context(h) + + # Concurrent context entry (network uploads) + try: + async with asyncio.TaskGroup() as tg: + for handle in self._handles: + tg.create_task(_enter_handle(handle)) + except ExceptionGroup as eg: + # Unwrap the first exception for cleaner error reporting. + # BaseExecution already catches and reports exceptions, + # but ExceptionGroup obscures the underlying InfrastructureError. + raise eg.exceptions[0] from eg + + # Concurrent readiness wait (indexing delays) async with asyncio.TaskGroup() as tg: for handle in self._handles: tg.create_task(handle.wait_until_ready()) diff --git a/rampart/pytest_plugin/_session.py b/rampart/pytest_plugin/_session.py index d1d8651b..da44147b 100644 --- a/rampart/pytest_plugin/_session.py +++ b/rampart/pytest_plugin/_session.py @@ -428,6 +428,12 @@ def build_report(self) -> TestRunReport: if self._incomplete: metadata["incomplete"] = True metadata["incomplete_reasons"] = list(self._incomplete_reasons) + import dataclasses # ruff: ignore[import-outside-top-level] + + trial_groups_serialized = { + k: dataclasses.asdict(v) for k, v in self._trial_groups.items() + } + self._cached_report = TestRunReport( results=sorted_results, total_runs=len(sorted_results), @@ -437,5 +443,6 @@ def build_report(self) -> TestRunReport: errors=counts[SafetyStatus.ERROR], duration_seconds=self._duration_seconds, metadata=metadata, + trial_groups=trial_groups_serialized, ) return self._cached_report diff --git a/rampart/reporting/json_file.py b/rampart/reporting/json_file.py index 6b621c07..45b33046 100644 --- a/rampart/reporting/json_file.py +++ b/rampart/reporting/json_file.py @@ -86,6 +86,7 @@ def _serialize_report(self, report: TestRunReport) -> dict[str, Any]: "errors": report.errors, "duration_seconds": report.duration_seconds, "metadata": report.metadata, + "trial_groups": report.trial_groups, "population_summary": dataclasses.asdict(report.population_summary()), "by_harm_category": { category: [self._serialize_result(r) for r in results] diff --git a/rampart/reporting/sink.py b/rampart/reporting/sink.py index ff614702..8b4e4fec 100644 --- a/rampart/reporting/sink.py +++ b/rampart/reporting/sink.py @@ -59,6 +59,8 @@ class TestRunReport: errors (int): Number with infrastructure errors. duration_seconds (float): Total run duration. metadata (dict[str, Any]): Run-level metadata (CI job ID, commit hash, etc.). + trial_groups (dict[str, Any]): Aggregated trial statistics for tests + marked with @pytest.mark.trial. """ __test__ = False # Prevent pytest from collecting this dataclass as a test. @@ -71,6 +73,7 @@ class TestRunReport: errors: int = 0 duration_seconds: float = 0.0 metadata: dict[str, Any] = field(default_factory=dict[str, Any]) + trial_groups: dict[str, Any] = field(default_factory=dict) def by_harm_category(self) -> dict[str, list[Result]]: """Group results by harm category. diff --git a/tests/unit/reporting/test_json_file.py b/tests/unit/reporting/test_json_file.py index 35bfec6d..4f7f10e5 100644 --- a/tests/unit/reporting/test_json_file.py +++ b/tests/unit/reporting/test_json_file.py @@ -235,6 +235,7 @@ def test_report_metadata_appears_in_serialized_output(self) -> None: "worker_count": 4, "dist_mode": "loadgroup", }, + trial_groups={"trial_1": {"total": 5, "safe": 5}}, ) data = sink._serialize_report(report) @@ -244,6 +245,7 @@ def test_report_metadata_appears_in_serialized_output(self) -> None: "worker_count": 4, "dist_mode": "loadgroup", } + assert data["trial_groups"] == {"trial_1": {"total": 5, "safe": 5}} def test_incomplete_run_metadata_appears_in_serialized_output(self) -> None: sink = JsonFileReportSink(output_dir=Path("/tmp")) From d8e7dfac44f31d8128efad7aeab6e4c3bf2ffaeb Mon Sep 17 00:00:00 2001 From: apocalypse9949 <125962989+apocalypse9949@users.noreply.github.com> Date: Thu, 13 Aug 2026 13:35:10 +0000 Subject: [PATCH 2/5] perf: execute XPIA injection handle activation concurrently Also includes feature addition for exposing trial group aggregation metadata to json reports. From 03d68daa592b9b810997f066ddb4170872a9cd3b Mon Sep 17 00:00:00 2001 From: apocalypse9949 <125962989+apocalypse9949@users.noreply.github.com> Date: Thu, 13 Aug 2026 13:50:15 +0000 Subject: [PATCH 3/5] perf: execute XPIA injection handle activation concurrently Also includes feature addition for exposing trial group aggregation metadata to json reports. From 74473b2896b6e6df56fb57ec4ad0491ded56d19d Mon Sep 17 00:00:00 2001 From: apocalypse9949 Date: Tue, 18 Aug 2026 13:57:36 +0530 Subject: [PATCH 4/5] [PERF]: drop trial_groups from reports to avoid overlap with #121/#123 Co-authored-by: Cursor --- rampart/pytest_plugin/_session.py | 7 ------- rampart/reporting/json_file.py | 1 - rampart/reporting/sink.py | 3 --- tests/unit/reporting/test_json_file.py | 2 -- 4 files changed, 13 deletions(-) diff --git a/rampart/pytest_plugin/_session.py b/rampart/pytest_plugin/_session.py index da44147b..d1d8651b 100644 --- a/rampart/pytest_plugin/_session.py +++ b/rampart/pytest_plugin/_session.py @@ -428,12 +428,6 @@ def build_report(self) -> TestRunReport: if self._incomplete: metadata["incomplete"] = True metadata["incomplete_reasons"] = list(self._incomplete_reasons) - import dataclasses # ruff: ignore[import-outside-top-level] - - trial_groups_serialized = { - k: dataclasses.asdict(v) for k, v in self._trial_groups.items() - } - self._cached_report = TestRunReport( results=sorted_results, total_runs=len(sorted_results), @@ -443,6 +437,5 @@ def build_report(self) -> TestRunReport: errors=counts[SafetyStatus.ERROR], duration_seconds=self._duration_seconds, metadata=metadata, - trial_groups=trial_groups_serialized, ) return self._cached_report diff --git a/rampart/reporting/json_file.py b/rampart/reporting/json_file.py index 45b33046..6b621c07 100644 --- a/rampart/reporting/json_file.py +++ b/rampart/reporting/json_file.py @@ -86,7 +86,6 @@ def _serialize_report(self, report: TestRunReport) -> dict[str, Any]: "errors": report.errors, "duration_seconds": report.duration_seconds, "metadata": report.metadata, - "trial_groups": report.trial_groups, "population_summary": dataclasses.asdict(report.population_summary()), "by_harm_category": { category: [self._serialize_result(r) for r in results] diff --git a/rampart/reporting/sink.py b/rampart/reporting/sink.py index 8b4e4fec..ff614702 100644 --- a/rampart/reporting/sink.py +++ b/rampart/reporting/sink.py @@ -59,8 +59,6 @@ class TestRunReport: errors (int): Number with infrastructure errors. duration_seconds (float): Total run duration. metadata (dict[str, Any]): Run-level metadata (CI job ID, commit hash, etc.). - trial_groups (dict[str, Any]): Aggregated trial statistics for tests - marked with @pytest.mark.trial. """ __test__ = False # Prevent pytest from collecting this dataclass as a test. @@ -73,7 +71,6 @@ class TestRunReport: errors: int = 0 duration_seconds: float = 0.0 metadata: dict[str, Any] = field(default_factory=dict[str, Any]) - trial_groups: dict[str, Any] = field(default_factory=dict) def by_harm_category(self) -> dict[str, list[Result]]: """Group results by harm category. diff --git a/tests/unit/reporting/test_json_file.py b/tests/unit/reporting/test_json_file.py index 4f7f10e5..35bfec6d 100644 --- a/tests/unit/reporting/test_json_file.py +++ b/tests/unit/reporting/test_json_file.py @@ -235,7 +235,6 @@ def test_report_metadata_appears_in_serialized_output(self) -> None: "worker_count": 4, "dist_mode": "loadgroup", }, - trial_groups={"trial_1": {"total": 5, "safe": 5}}, ) data = sink._serialize_report(report) @@ -245,7 +244,6 @@ def test_report_metadata_appears_in_serialized_output(self) -> None: "worker_count": 4, "dist_mode": "loadgroup", } - assert data["trial_groups"] == {"trial_1": {"total": 5, "safe": 5}} def test_incomplete_run_metadata_appears_in_serialized_output(self) -> None: sink = JsonFileReportSink(output_dir=Path("/tmp")) From 94f70386a6b4cd5f7bb077b815ca14c5b046209b Mon Sep 17 00:00:00 2001 From: apocalypse9949 Date: Fri, 21 Aug 2026 15:52:53 +0530 Subject: [PATCH 5/5] [PERF]: activate XPIA handles with non-cancelling gather Co-authored-by: Cursor --- rampart/attacks/_xpia.py | 31 +++++++++++++----------- tests/unit/attacks/test_xpia.py | 42 +++++++++++++++++++++++++++++++++ 2 files changed, 59 insertions(+), 14 deletions(-) diff --git a/rampart/attacks/_xpia.py b/rampart/attacks/_xpia.py index 3a5a0fed..48278388 100644 --- a/rampart/attacks/_xpia.py +++ b/rampart/attacks/_xpia.py @@ -160,23 +160,26 @@ async def _activate_handles_async( ) -> None: """Activate all injection handles and wait for readiness. + Handle ``__aenter__`` runs concurrently via ``asyncio.gather`` with + ``return_exceptions=True`` so a failing sibling does not cancel + others mid-entry. That lets successful handles finish + ``enter_async_context`` and register ``__aexit__`` cleanup before + the first error is re-raised. ``TaskGroup`` would cancel remaining + entries and could orphan remote payloads that were created but not + yet registered on the exit stack. + Args: stack (AsyncExitStack): The exit stack managing cleanup. """ - - async def _enter_handle(h: InjectionHandle) -> None: - await stack.enter_async_context(h) - - # Concurrent context entry (network uploads) - try: - async with asyncio.TaskGroup() as tg: - for handle in self._handles: - tg.create_task(_enter_handle(handle)) - except ExceptionGroup as eg: - # Unwrap the first exception for cleaner error reporting. - # BaseExecution already catches and reports exceptions, - # but ExceptionGroup obscures the underlying InfrastructureError. - raise eg.exceptions[0] from eg + # Concurrent context entry (network uploads); non-cancelling so + # successful siblings still register cleanup on the exit stack. + results = await asyncio.gather( + *(stack.enter_async_context(handle) for handle in self._handles), + return_exceptions=True, + ) + errors = [result for result in results if isinstance(result, BaseException)] + if errors: + raise errors[0] # Concurrent readiness wait (indexing delays) async with asyncio.TaskGroup() as tg: diff --git a/tests/unit/attacks/test_xpia.py b/tests/unit/attacks/test_xpia.py index 93a35e6f..c6e438e7 100644 --- a/tests/unit/attacks/test_xpia.py +++ b/tests/unit/attacks/test_xpia.py @@ -3,6 +3,7 @@ from __future__ import annotations +import asyncio from unittest.mock import AsyncMock from rampart.attacks import Attacks @@ -219,6 +220,47 @@ async def test_handle_activation_failure(self) -> None: assert result.status is SafetyStatus.ERROR assert "SharePoint 503" in result.summary + async def test_partial_activation_failure_still_cleans_up_siblings(self) -> None: + """A slow successful sibling must register cleanup even if another fails. + + Concurrent activation uses gather(return_exceptions=True) so a + failing handle does not cancel siblings mid-__aenter__. + """ + entered = asyncio.Event() + + async def slow_success_aenter_async( + *_args: object, + **_kwargs: object, + ) -> AsyncMock: + entered.set() + await asyncio.sleep(0.05) + return slow + + async def fast_fail_aenter_async( + *_args: object, + **_kwargs: object, + ) -> AsyncMock: + await entered.wait() + raise InfrastructureError("SharePoint 503") + + slow = _mock_handle(surface_name="Exchange") + slow.__aenter__.side_effect = slow_success_aenter_async + failing = _mock_handle(surface_name="SharePoint") + failing.__aenter__.side_effect = fast_fail_aenter_async + + result = await Attacks.xpia( + inject=[slow, failing], + trigger="Summarize Q3", + evaluator=_mock_evaluator(EvalOutcome.DETECTED), + ).execute_async(adapter=_adapter()) + + assert result.status is SafetyStatus.ERROR + assert "SharePoint 503" in result.summary + slow.__aenter__.assert_awaited_once() + failing.__aenter__.assert_awaited_once() + # Successful sibling registered on the exit stack and is cleaned up. + slow.__aexit__.assert_awaited_once() + async def test_session_creation_failure(self) -> None: adapter = AsyncMock() adapter.create_session_async.side_effect = InfrastructureError(