diff --git a/rampart/attacks/_xpia.py b/rampart/attacks/_xpia.py index d7746365..245f33c7 100644 --- a/rampart/attacks/_xpia.py +++ b/rampart/attacks/_xpia.py @@ -160,13 +160,28 @@ 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. """ - for handle in self._handles: - await stack.enter_async_context(handle) + # 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: total = max of all wait times + # Concurrent readiness wait (indexing delays) async with asyncio.TaskGroup() as tg: for handle in self._handles: tg.create_task(handle.wait_until_ready_async()) diff --git a/tests/unit/attacks/test_xpia.py b/tests/unit/attacks/test_xpia.py index c932e327..96fbad66 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,49 @@ async def test_handle_activation_failure_async(self) -> None: assert result.status is SafetyStatus.ERROR assert "SharePoint 503" in result.summary + async def test_partial_activation_failure_still_cleans_up_siblings_async( + 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_async(self) -> None: adapter = AsyncMock() adapter.create_session_async.side_effect = InfrastructureError(