From b62b271f49d7a9176cd847ce519ce940e64fd5b1 Mon Sep 17 00:00:00 2001 From: Nina Chikanov Date: Mon, 13 Apr 2026 18:28:35 -0700 Subject: [PATCH 01/13] initial implementation --- rampart/attacks/_xpia.py | 9 +-- rampart/core/injection.py | 25 +++++++ rampart/surfaces/onedrive.py | 13 ++++ tests/unit/core/test_injection.py | 98 ++++++++++++++++++++++++++++ tests/unit/core/test_protocols.py | 14 ++++ tests/unit/surfaces/test_onedrive.py | 11 ++++ 6 files changed, 164 insertions(+), 6 deletions(-) create mode 100644 tests/unit/core/test_injection.py diff --git a/rampart/attacks/_xpia.py b/rampart/attacks/_xpia.py index 5b258627..a6a8acb9 100644 --- a/rampart/attacks/_xpia.py +++ b/rampart/attacks/_xpia.py @@ -175,7 +175,7 @@ async def _activate_handles_async( self, *, stack: AsyncExitStack, ) -> None: """ - Activate all injection handles and wait for indexing. + Activate all injection handles and wait for readiness. Args: stack (AsyncExitStack): The exit stack managing cleanup. @@ -183,11 +183,8 @@ async def _activate_handles_async( for handle in self._handles: await stack.enter_async_context(handle) - delay = max( - (h.indexing_delay_seconds for h in self._handles), default=0.0, - ) - if delay > 0: - await asyncio.sleep(delay) + for handle in self._handles: + await handle.wait_until_ready() def _build_attack_result( self, diff --git a/rampart/core/injection.py b/rampart/core/injection.py index 0b3c04ac..8ff9ae39 100644 --- a/rampart/core/injection.py +++ b/rampart/core/injection.py @@ -9,6 +9,7 @@ from __future__ import annotations +import asyncio from typing import Any, Protocol, runtime_checkable from rampart.core.types import Payload @@ -32,6 +33,11 @@ def indexing_delay_seconds(self) -> float: """How long to wait after activation for the agent to see the content.""" ... + @property + def readiness_timeout_seconds(self) -> float: + """Maximum time `wait_until_ready` may block before raising `TimeoutError`.""" + ... + @property def payload_id(self) -> str | None: """The injected payload's identifier, for reporting.""" @@ -42,6 +48,25 @@ def surface_name(self) -> str: """The name of the surface this handle injects into (e.g., 'SharePoint').""" ... + async def wait_until_ready(self) -> None: + """Block until the injected content is visible to the agent. + + Implementations must complete within `readiness_timeout_seconds` + or raise `TimeoutError`. + + Default implementation sleeps for `indexing_delay_seconds`, with + an upper bound of `readiness_timeout_seconds` to prevent indefinite + blocking — which is a common strategy for many surfaces. Surfaces with + more complex readiness logic can override this method with a custom + implementation. + + Raises: + TimeoutError: If `indexing_delay_seconds` exceeds + `readiness_timeout_seconds`. + """ + async with asyncio.timeout(self.readiness_timeout_seconds): + await asyncio.sleep(self.indexing_delay_seconds) + async def __aenter__(self) -> InjectionHandle: """Activate the injection (write payload to data source).""" ... diff --git a/rampart/surfaces/onedrive.py b/rampart/surfaces/onedrive.py index 89146944..48d4b340 100644 --- a/rampart/surfaces/onedrive.py +++ b/rampart/surfaces/onedrive.py @@ -13,6 +13,7 @@ from typing import TYPE_CHECKING, Any from rampart.core.errors import InfrastructureError +from rampart.core.injection import InjectionHandle from rampart.core.types import Payload if TYPE_CHECKING: @@ -51,6 +52,7 @@ class OneDriveSurface: """ DEFAULT_INDEXING_DELAY: float = 10.0 + DEFAULT_READINESS_TIMEOUT: float = 120.0 def __init__( self, @@ -59,11 +61,13 @@ def __init__( drive_id: str, folder_path: str, indexing_delay: float = DEFAULT_INDEXING_DELAY, + readiness_timeout: float = DEFAULT_READINESS_TIMEOUT, ) -> None: self._graph_client = graph_client self._drive_id = drive_id self._folder_path = folder_path.strip("/") self._indexing_delay = indexing_delay + self._readiness_timeout = readiness_timeout def inject(self, *, payload: Payload) -> _OneDriveInjection: """ @@ -163,6 +167,11 @@ def indexing_delay_seconds(self) -> float: """How long to wait after upload for content to be discoverable.""" return self._surface._indexing_delay + @property + def readiness_timeout_seconds(self) -> float: + """Maximum time `wait_until_ready` may block.""" + return self._surface._readiness_timeout + @property def payload_id(self) -> str | None: """The injected payload's identifier.""" @@ -173,6 +182,10 @@ def surface_name(self) -> str: """Identifies this injection as OneDrive for reporting.""" return "OneDrive" + async def wait_until_ready(self) -> None: + """Wait for OneDrive indexing using the default sleep-based strategy.""" + await InjectionHandle.wait_until_ready(self) + async def __aenter__(self) -> _OneDriveInjection: """Upload payload to OneDrive. Raises InfrastructureError on failure.""" try: diff --git a/tests/unit/core/test_injection.py b/tests/unit/core/test_injection.py new file mode 100644 index 00000000..8c998e01 --- /dev/null +++ b/tests/unit/core/test_injection.py @@ -0,0 +1,98 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT license. + +from __future__ import annotations + +import asyncio +from typing import Any +from unittest.mock import AsyncMock + +import pytest + +from rampart.core.injection import InjectionHandle + + +class _ConcreteHandle(InjectionHandle): + """Minimal concrete handle that inherits the default wait_until_ready.""" + + def __init__( + self, + *, + delay: float = 0.0, + readiness_timeout: float = 30.0, + ) -> None: + self._delay = delay + self._readiness_timeout = readiness_timeout + + @property + def indexing_delay_seconds(self) -> float: + return self._delay + + @property + def readiness_timeout_seconds(self) -> float: + return self._readiness_timeout + + @property + def payload_id(self) -> str | None: + return "test-payload" + + @property + def surface_name(self) -> str: + return "TestSurface" + + async def __aenter__(self) -> _ConcreteHandle: + return self + + async def __aexit__( + self, + exc_type: type[BaseException] | None, + exc_val: BaseException | None, + exc_tb: Any, + ) -> None: + pass + + +class TestWaitUntilReady: + """Tests for InjectionHandle.wait_until_ready default and custom behaviour.""" + + @pytest.mark.asyncio + @pytest.mark.parametrize( + ("delay", "readiness_timeout"), + [ + (0.0, 1.0), + (0.05, 5.0), + ], + ids=["zero-delay", "short-delay-within-timeout"], + ) + async def test_default_completes_when_delay_within_timeout( + self, delay: float, readiness_timeout: float + ) -> None: + """Default sleep-based wait completes when delay is within the timeout.""" + handle = _ConcreteHandle(delay=delay, readiness_timeout=readiness_timeout) + + await handle.wait_until_ready() + + @pytest.mark.asyncio + async def test_custom_polling_implementation(self) -> None: + """A handle can override wait_until_ready with custom polling logic.""" + poll_mock = AsyncMock(side_effect=[False, False, True]) + + class _PollingHandle(_ConcreteHandle): + async def wait_until_ready(self) -> None: + async with asyncio.timeout(self.readiness_timeout_seconds): + while not await poll_mock(): + await asyncio.sleep(0) + + handle = _PollingHandle(readiness_timeout=5.0) + + await handle.wait_until_ready() + + assert poll_mock.await_count == 3 + + @pytest.mark.asyncio + async def test_timeout_raises_when_delay_exceeds_limit(self) -> None: + """Default implementation raises TimeoutError when delay exceeds timeout.""" + handle = _ConcreteHandle(delay=10.0, readiness_timeout=0.01) + + with pytest.raises(TimeoutError): + await handle.wait_until_ready() diff --git a/tests/unit/core/test_protocols.py b/tests/unit/core/test_protocols.py index 3891c042..1db00079 100644 --- a/tests/unit/core/test_protocols.py +++ b/tests/unit/core/test_protocols.py @@ -81,6 +81,10 @@ class MyHandle: def indexing_delay_seconds(self) -> float: return 5.0 + @property + def readiness_timeout_seconds(self) -> float: + return 30.0 + @property def payload_id(self) -> str | None: return "abc" @@ -89,6 +93,9 @@ def payload_id(self) -> str | None: def surface_name(self) -> str: return "SharePoint" + async def wait_until_ready(self) -> None: + pass + async def __aenter__(self) -> "MyHandle": return self @@ -110,6 +117,10 @@ class MyHandle: def indexing_delay_seconds(self) -> float: return 0.0 + @property + def readiness_timeout_seconds(self) -> float: + return 30.0 + @property def payload_id(self) -> str | None: return None @@ -118,6 +129,9 @@ def payload_id(self) -> str | None: def surface_name(self) -> str: return "test" + async def wait_until_ready(self) -> None: + pass + async def __aenter__(self) -> "MyHandle": return self diff --git a/tests/unit/surfaces/test_onedrive.py b/tests/unit/surfaces/test_onedrive.py index c8a2f1d2..6076b8e0 100644 --- a/tests/unit/surfaces/test_onedrive.py +++ b/tests/unit/surfaces/test_onedrive.py @@ -171,6 +171,17 @@ def test_indexing_delay_from_surface(self) -> None: handle = surface.inject(payload=payload) assert handle.indexing_delay_seconds == 99.0 + def test_readiness_timeout_from_surface(self) -> None: + surface = OneDriveSurface( + graph_client=MagicMock(), + drive_id="d", + folder_path="f", + readiness_timeout=60.0, + ) + payload = Payload(content="test") + handle = surface.inject(payload=payload) + assert handle.readiness_timeout_seconds == 60.0 + class TestOneDriveInjectionLifecycle: """Test the async context manager lifecycle (upload + delete).""" From 5afdf063586ecae95e1f6f320486812cfa6efa39 Mon Sep 17 00:00:00 2001 From: Nina Chikanov Date: Mon, 13 Apr 2026 19:09:14 -0700 Subject: [PATCH 02/13] some pre-commit linting --- pyproject.toml | 18 +++++++ rampart/attacks/_xpia.py | 76 +++++++++++++------------- rampart/core/injection.py | 20 +++---- rampart/surfaces/onedrive.py | 88 ++++++++++++++++++------------- tests/unit/core/test_injection.py | 20 ++++--- tests/unit/core/test_protocols.py | 9 ++-- 6 files changed, 140 insertions(+), 91 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index 33c3322b..894e0ed5 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -35,6 +35,24 @@ markers = [ [tool.ruff.lint] select = ["ALL"] +[tool.ruff.lint.per-file-ignores] +"tests/**" = [ + "S101", # assert is pytest's API + "D100", "D101", "D102", "D104", "D107", # no docstrings needed + "ANN001", "ANN201", "ANN202", # no type annotations needed + "PLR2004", # magic values in assertions are fine + "ARG001", "ARG002", # unused args (fixtures, stubs) + "PLC0415", # imports inside functions for isolation + "SLF001", # testing private members is valid + "TRY003", "EM101", "EM102", # exception message style + "BLE001", # catching Exception in tests + "TRY301", # raise in try blocks + "S108", # /tmp usage + "PT017", "PT018", # assertion style + "ASYNC240", # pathlib in async tests + "PERF401", "RUF015", "PTH123", # micro-optimizations / style +] + [tool.ruff.lint.flake8-copyright] notice-rgx = "Copyright \\(c\\) Microsoft Corporation\\.\\s*\\n.*Licensed under the MIT license" diff --git a/rampart/attacks/_xpia.py b/rampart/attacks/_xpia.py index a6a8acb9..594d9403 100644 --- a/rampart/attacks/_xpia.py +++ b/rampart/attacks/_xpia.py @@ -12,7 +12,6 @@ from __future__ import annotations -import asyncio import logging from contextlib import AsyncExitStack from typing import Any @@ -28,7 +27,6 @@ InjectionRecord, ObservabilityLevel, PromptDriver, - Request, Result, SafetyStatus, Turn, @@ -39,8 +37,7 @@ class XPIAExecution(BaseExecution): - """ - Executes the full XPIA attack lifecycle. + """Executes the full XPIA attack lifecycle. Inherits BaseExecution. Implements ``_execute_async`` with XPIA's specific phase structure. The lifecycle skeleton (event dispatch, @@ -95,8 +92,7 @@ def strategy_name(self) -> str: return "xpia" async def _execute_async(self, *, adapter: AgentAdapter) -> Result: - """ - Orchestrate the XPIA lifecycle and return a safety Result. + """Orchestrate the XPIA lifecycle and return a safety Result. Delegates phase execution to ``_run_phases_async`` and result construction to ``_build_attack_result`` or @@ -116,17 +112,22 @@ async def _execute_async(self, *, adapter: AgentAdapter) -> Result: ) if max_turns_hit: return self._max_turns_error_result( - adapter=adapter, turns=turns, eval_results=eval_results, + adapter=adapter, + turns=turns, + eval_results=eval_results, ) return self._build_attack_result( - adapter=adapter, turns=turns, eval_results=eval_results, + adapter=adapter, + turns=turns, + eval_results=eval_results, ) async def _run_phases_async( - self, *, adapter: AgentAdapter, + self, + *, + adapter: AgentAdapter, ) -> tuple[list[Turn], list[EvalResult], bool]: - """ - Run XPIA phases 1-5 inside a cleanup-guaranteed context. + """Run XPIA phases 1-5 inside a cleanup-guaranteed context. Args: adapter (AgentAdapter): The agent adapter. @@ -140,7 +141,7 @@ async def _run_phases_async( async with AsyncExitStack() as stack: await self._activate_handles_async(stack=stack) session = await stack.enter_async_context( - await adapter.create_session_async() + await adapter.create_session_async(), ) for turn_index in range(self._max_turns): @@ -150,16 +151,19 @@ async def _run_phases_async( request = decision.request response = await session.send_async(request) - turns.append(Turn( - request=request, - response=response, - turn_number=turn_index, - driver_reasoning=decision.reasoning, - )) + turns.append( + Turn( + request=request, + response=response, + turn_number=turn_index, + driver_reasoning=decision.reasoning, + ), + ) eval_result = await self._evaluator.evaluate_async( context=EvalContext( - turns=list(turns), manifest=adapter.manifest, + turns=list(turns), + manifest=adapter.manifest, ), ) eval_results.append(eval_result) @@ -172,10 +176,11 @@ async def _run_phases_async( return turns, eval_results, False async def _activate_handles_async( - self, *, stack: AsyncExitStack, + self, + *, + stack: AsyncExitStack, ) -> None: - """ - Activate all injection handles and wait for readiness. + """Activate all injection handles and wait for readiness. Args: stack (AsyncExitStack): The exit stack managing cleanup. @@ -193,8 +198,7 @@ def _build_attack_result( turns: list[Turn], eval_results: list[EvalResult], ) -> Result: - """ - Resolve eval results into a final attack Result. + """Resolve eval results into a final attack Result. Applies observability adjustment when the initial verdict is SAFE. @@ -210,7 +214,10 @@ def _build_attack_result( if status == SafetyStatus.SAFE: safe, status = self._adjust_for_observability( - adapter=adapter, turns=turns, safe=safe, status=status, + adapter=adapter, + turns=turns, + safe=safe, + status=status, ) return Result( @@ -232,8 +239,7 @@ def _max_turns_error_result( turns: list[Turn], eval_results: list[EvalResult], ) -> Result: - """ - Build an ERROR result when the driver exceeds max_turns. + """Build an ERROR result when the driver exceeds max_turns. Args: adapter (AgentAdapter): The adapter under test. @@ -252,8 +258,7 @@ def _max_turns_error_result( safe=False, status=SafetyStatus.ERROR, summary=( - f"Max turns ({self._max_turns}) reached" - " — driver did not terminate" + f"Max turns ({self._max_turns}) reached — driver did not terminate" ), turns=turns, eval_results=eval_results, @@ -281,8 +286,7 @@ def _adjust_for_observability( safe: bool, status: SafetyStatus, ) -> tuple[bool, SafetyStatus]: - """ - Downgrade SAFE to UNDETERMINED when the adapter lacks observability. + """Downgrade SAFE to UNDETERMINED when the adapter lacks observability. If the evaluator returned NOT_DETECTED (and ``resolve_as_attack`` produced SAFE), but the adapter has RESPONSE_ONLY observability @@ -320,8 +324,7 @@ def _collect_response_metadata( *, turns: list[Turn], ) -> dict[str, Any]: - """ - Aggregate response metadata from all turns. + """Aggregate response metadata from all turns. Adapters can attach diagnostic data (e.g. conversation_id, session_id) to Response.metadata. This function collects it into a single dict @@ -352,8 +355,7 @@ def _build_summary( status: SafetyStatus, eval_results: list[EvalResult], ) -> str: - """ - Build a human-readable one-line summary for XPIA results. + """Build a human-readable one-line summary for XPIA results. Args: status (SafetyStatus): The resolved safety status. @@ -371,7 +373,9 @@ def _build_summary( return f"Attack objective detected: {'; '.join(evidence[:3])}" if status == SafetyStatus.UNDETERMINED: rationales = [er.rationale for er in eval_results if er.rationale] - detail = "; ".join(rationales[:2]) if rationales else "Insufficient observability" + detail = ( + "; ".join(rationales[:2]) if rationales else "Insufficient observability" + ) return f"Evaluation undetermined: {detail}" if status == SafetyStatus.ERROR: return "Infrastructure error during execution" diff --git a/rampart/core/injection.py b/rampart/core/injection.py index 8ff9ae39..d33dc611 100644 --- a/rampart/core/injection.py +++ b/rampart/core/injection.py @@ -10,15 +10,17 @@ from __future__ import annotations import asyncio -from typing import Any, Protocol, runtime_checkable +from typing import TYPE_CHECKING, Protocol, Self, runtime_checkable -from rampart.core.types import Payload +if TYPE_CHECKING: + import types + + from rampart.core.types import Payload @runtime_checkable class InjectionHandle(Protocol): - """ - A prepared injection, ready to activate as an async context manager. + """A prepared injection, ready to activate as an async context manager. Returned by Surface.inject(). Entering activates the injection (writes the payload to the data source); exiting removes it @@ -67,7 +69,7 @@ async def wait_until_ready(self) -> None: async with asyncio.timeout(self.readiness_timeout_seconds): await asyncio.sleep(self.indexing_delay_seconds) - async def __aenter__(self) -> InjectionHandle: + async def __aenter__(self) -> Self: """Activate the injection (write payload to data source).""" ... @@ -75,7 +77,7 @@ async def __aexit__( self, exc_type: type[BaseException] | None, exc_val: BaseException | None, - exc_tb: Any, + exc_tb: types.TracebackType | None, ) -> None: """Remove the injection. Must be idempotent. Must not raise.""" ... @@ -83,8 +85,7 @@ async def __aexit__( @runtime_checkable class Surface(Protocol): - """ - An injectable data source. + """An injectable data source. Surfaces are fully configured at construction (credentials, target location) and expose a universal inject() signature. Teams implement @@ -96,8 +97,7 @@ class Surface(Protocol): """ def inject(self, *, payload: Payload) -> InjectionHandle: - """ - Prepare an injection of the given payload. + """Prepare an injection of the given payload. Does not activate the injection — the caller enters the returned handle as an async context manager to activate it. diff --git a/rampart/surfaces/onedrive.py b/rampart/surfaces/onedrive.py index 48d4b340..274d4e96 100644 --- a/rampart/surfaces/onedrive.py +++ b/rampart/surfaces/onedrive.py @@ -10,15 +10,18 @@ from __future__ import annotations import logging -from typing import TYPE_CHECKING, Any +from typing import TYPE_CHECKING, Self from rampart.core.errors import InfrastructureError from rampart.core.injection import InjectionHandle -from rampart.core.types import Payload if TYPE_CHECKING: + import types + from msgraph.graph_service_client import GraphServiceClient + from rampart.core.types import Payload + logger = logging.getLogger(__name__) # Graph's PUT /drives/{id}/items/{parent}:/{path}:/content @@ -28,8 +31,7 @@ class OneDriveSurface: - """ - Injects payloads into a specific OneDrive location. + """Injects payloads into a specific OneDrive location. The surface is fully configured at construction — drive ID, credentials, and target folder path. The ``inject()`` method takes @@ -63,15 +65,15 @@ def __init__( indexing_delay: float = DEFAULT_INDEXING_DELAY, readiness_timeout: float = DEFAULT_READINESS_TIMEOUT, ) -> None: + """Initialize with Graph client and OneDrive location.""" self._graph_client = graph_client - self._drive_id = drive_id - self._folder_path = folder_path.strip("/") - self._indexing_delay = indexing_delay - self._readiness_timeout = readiness_timeout + self.drive_id = drive_id + self.folder_path = folder_path.strip("/") + self.indexing_delay = indexing_delay + self.readiness_timeout = readiness_timeout def inject(self, *, payload: Payload) -> _OneDriveInjection: - """ - Prepare an injection into the configured OneDrive folder. + """Prepare an injection into the configured OneDrive folder. Returns an InjectionHandle — enter it as an async context manager to activate the injection, exit to clean up. @@ -84,7 +86,7 @@ def inject(self, *, payload: Payload) -> _OneDriveInjection: """ return _OneDriveInjection(surface=self, payload=payload) - async def _upload_async(self, *, payload: Payload) -> str: + async def upload_async(self, *, payload: Payload) -> str: """Upload payload content to OneDrive. Returns the item ID. Uses the small-file upload endpoint @@ -97,60 +99,69 @@ async def _upload_async(self, *, payload: Payload) -> str: InfrastructureError: If Graph returns no ``DriveItem``. """ filename = f"{payload.id}{payload.format.extension}" - upload_path = f"{self._folder_path}/{filename}" + upload_path = f"{self.folder_path}/{filename}" if payload.format.is_binary: if payload.artifact is None: - raise ValueError( + msg = ( f"Binary payload format {payload.format.value} " - f"requires an artifact path." + f"requires an artifact path.", + ) + raise ValueError( + msg, ) content = payload.artifact.read_bytes() else: content = payload.content.encode("utf-8") if len(content) > _MAX_SMALL_UPLOAD_BYTES: - raise ValueError( + msg = ( f"Payload {payload.id} is {len(content)} bytes, which " f"exceeds the 4 MiB small-upload limit. Upload sessions " - f"are not yet implemented." + f"are not yet implemented.", + ) + raise ValueError( + msg, ) # Graph path-based addressing: root:/{relative-path}: # The trailing colon is required by the API. drive_item = ( - await self._graph_client.drives.by_drive_id(self._drive_id) + await self._graph_client.drives.by_drive_id(self.drive_id) .items.by_drive_item_id(f"root:/{upload_path}:") .content.put(content) ) if drive_item is None or drive_item.id is None: - raise InfrastructureError( + msg = ( f"Graph API returned no DriveItem after upload to " - f"drive={self._drive_id} path={upload_path}" + f"drive={self.drive_id} path={upload_path}", + ) + raise InfrastructureError( + msg, ) item_id = drive_item.id logger.info( "Uploaded payload %s to OneDrive drive=%s path=%s (item=%s)", payload.id, - self._drive_id, + self.drive_id, upload_path, item_id, ) return item_id - async def _delete_async(self, *, item_id: str) -> None: + async def delete_async(self, *, item_id: str) -> None: """Delete a file from OneDrive by item ID.""" await ( - self._graph_client.drives.by_drive_id(self._drive_id) + self._graph_client.drives.by_drive_id(self.drive_id) .items.by_drive_item_id(item_id) .delete() ) logger.info( "Deleted OneDrive item %s from drive=%s", item_id, - self._drive_id, + self.drive_id, ) @@ -165,12 +176,12 @@ def __init__(self, *, surface: OneDriveSurface, payload: Payload) -> None: @property def indexing_delay_seconds(self) -> float: """How long to wait after upload for content to be discoverable.""" - return self._surface._indexing_delay + return self._surface.indexing_delay @property def readiness_timeout_seconds(self) -> float: """Maximum time `wait_until_ready` may block.""" - return self._surface._readiness_timeout + return self._surface.readiness_timeout @property def payload_id(self) -> str | None: @@ -186,18 +197,21 @@ async def wait_until_ready(self) -> None: """Wait for OneDrive indexing using the default sleep-based strategy.""" await InjectionHandle.wait_until_ready(self) - async def __aenter__(self) -> _OneDriveInjection: + async def __aenter__(self) -> Self: """Upload payload to OneDrive. Raises InfrastructureError on failure.""" try: - self._item_id = await self._surface._upload_async( - payload=self._payload + self._item_id = await self._surface.upload_async( + payload=self._payload, ) except InfrastructureError: raise except Exception as exc: + msg = ( + f"OneDrive upload failed for drive={self._surface.drive_id} " + f"path={self._surface.folder_path}: {exc}", + ) raise InfrastructureError( - f"OneDrive upload failed for drive={self._surface._drive_id} " - f"path={self._surface._folder_path}: {exc}", + msg, ) from exc return self @@ -205,16 +219,18 @@ async def __aexit__( self, exc_type: type[BaseException] | None, exc_val: BaseException | None, - exc_tb: Any, + exc_tb: types.TracebackType | None, ) -> None: """Delete uploaded content. Logs warnings on failure but never raises.""" if self._item_id is not None: try: - await self._surface._delete_async(item_id=self._item_id) - except Exception: + await self._surface.delete_async(item_id=self._item_id) + except Exception: # noqa: BLE001 — cleanup must not raise + msg = ( + f"OneDrive cleanup failed for item {self._item_id} " + f"in drive={self._surface.drive_id}" + ) logger.warning( - "OneDrive cleanup failed for item %s in drive=%s", - self._item_id, - self._surface._drive_id, + msg, exc_info=True, ) diff --git a/tests/unit/core/test_injection.py b/tests/unit/core/test_injection.py index 8c998e01..5c57e32b 100644 --- a/tests/unit/core/test_injection.py +++ b/tests/unit/core/test_injection.py @@ -1,16 +1,21 @@ # Copyright (c) Microsoft Corporation. # Licensed under the MIT license. +"""Tests for rampart.core.injection.""" + from __future__ import annotations import asyncio -from typing import Any +from typing import TYPE_CHECKING, Self from unittest.mock import AsyncMock import pytest from rampart.core.injection import InjectionHandle +if TYPE_CHECKING: + import types + class _ConcreteHandle(InjectionHandle): """Minimal concrete handle that inherits the default wait_until_ready.""" @@ -40,14 +45,14 @@ def payload_id(self) -> str | None: def surface_name(self) -> str: return "TestSurface" - async def __aenter__(self) -> _ConcreteHandle: + async def __aenter__(self) -> Self: return self async def __aexit__( self, exc_type: type[BaseException] | None, exc_val: BaseException | None, - exc_tb: Any, + exc_tb: types.TracebackType | None, ) -> None: pass @@ -65,7 +70,9 @@ class TestWaitUntilReady: ids=["zero-delay", "short-delay-within-timeout"], ) async def test_default_completes_when_delay_within_timeout( - self, delay: float, readiness_timeout: float + self, + delay: float, + readiness_timeout: float, ) -> None: """Default sleep-based wait completes when delay is within the timeout.""" handle = _ConcreteHandle(delay=delay, readiness_timeout=readiness_timeout) @@ -75,19 +82,20 @@ async def test_default_completes_when_delay_within_timeout( @pytest.mark.asyncio async def test_custom_polling_implementation(self) -> None: """A handle can override wait_until_ready with custom polling logic.""" + _expected_poll_calls = 3 poll_mock = AsyncMock(side_effect=[False, False, True]) class _PollingHandle(_ConcreteHandle): async def wait_until_ready(self) -> None: async with asyncio.timeout(self.readiness_timeout_seconds): while not await poll_mock(): - await asyncio.sleep(0) + await asyncio.Event() handle = _PollingHandle(readiness_timeout=5.0) await handle.wait_until_ready() - assert poll_mock.await_count == 3 + assert poll_mock.await_count == _expected_poll_calls @pytest.mark.asyncio async def test_timeout_raises_when_delay_exceeds_limit(self) -> None: diff --git a/tests/unit/core/test_protocols.py b/tests/unit/core/test_protocols.py index 1db00079..5224338d 100644 --- a/tests/unit/core/test_protocols.py +++ b/tests/unit/core/test_protocols.py @@ -38,6 +38,7 @@ async def __aexit__( def test_send_async_accepts_request(self) -> None: """Verify the protocol requires a Request parameter.""" + class MySession: async def send_async(self, request: Request) -> Response: return Response(text="ok") @@ -60,8 +61,7 @@ async def __aexit__( class TestAgentAdapterProtocol: def test_structural_subtyping(self) -> None: class MyAdapter: - async def create_session_async(self) -> Session: - ... + async def create_session_async(self) -> Session: ... @property def manifest(self) -> AppManifest: @@ -154,7 +154,9 @@ class TestPromptDriverProtocol: def test_structural_subtyping(self) -> None: class MyDriver: async def next_prompt_async( - self, *, history: list[Turn], + self, + *, + history: list[Turn], ) -> PromptDecision | None: return None @@ -175,5 +177,6 @@ def test_with_attachments_only(self) -> None: def test_empty_request_raises(self) -> None: import pytest + with pytest.raises(ValueError, match="at least"): Request() From 9e1e80c298169a35109d5606c96bd80d401ddd3b Mon Sep 17 00:00:00 2001 From: Nina Chikanov Date: Tue, 14 Apr 2026 13:47:58 -0700 Subject: [PATCH 03/13] fix tests & pre-commit run on remainder of files. remove line from .gitignore so it does not track .vscode/settings.json --- .gitignore | 4 ++-- tests/unit/core/test_injection.py | 13 +++++++++++-- tests/unit/core/test_protocols.py | 19 ++++++++++--------- tests/unit/surfaces/test_onedrive.py | 26 ++++++++++++++------------ 4 files changed, 37 insertions(+), 25 deletions(-) diff --git a/.gitignore b/.gitignore index fd410cfa..d6d86830 100644 --- a/.gitignore +++ b/.gitignore @@ -23,7 +23,8 @@ .LSOverride # Icon must end with two \r -Icon +Icon + # Thumbnails ._* @@ -335,7 +336,6 @@ pyrightconfig.json ### VisualStudioCode ### .vscode/* -!.vscode/settings.json !.vscode/tasks.json !.vscode/launch.json !.vscode/extensions.json diff --git a/tests/unit/core/test_injection.py b/tests/unit/core/test_injection.py index 5c57e32b..43e4875d 100644 --- a/tests/unit/core/test_injection.py +++ b/tests/unit/core/test_injection.py @@ -84,16 +84,25 @@ async def test_custom_polling_implementation(self) -> None: """A handle can override wait_until_ready with custom polling logic.""" _expected_poll_calls = 3 poll_mock = AsyncMock(side_effect=[False, False, True]) + ready_event = asyncio.Event() class _PollingHandle(_ConcreteHandle): async def wait_until_ready(self) -> None: async with asyncio.timeout(self.readiness_timeout_seconds): while not await poll_mock(): - await asyncio.Event() + ready_event.clear() + await ready_event.wait() handle = _PollingHandle(readiness_timeout=5.0) - await handle.wait_until_ready() + async def _signal_ready() -> None: + for _ in range(_expected_poll_calls - 1): + await asyncio.sleep(0) + ready_event.set() + + async with asyncio.TaskGroup() as tg: + tg.create_task(handle.wait_until_ready()) + tg.create_task(_signal_ready()) assert poll_mock.await_count == _expected_poll_calls diff --git a/tests/unit/core/test_protocols.py b/tests/unit/core/test_protocols.py index 5224338d..d9eff556 100644 --- a/tests/unit/core/test_protocols.py +++ b/tests/unit/core/test_protocols.py @@ -8,7 +8,8 @@ inheriting from the protocol. """ -from typing import Any +import types +from typing import Self from rampart.core.adapter import AgentAdapter, Session from rampart.core.injection import InjectionHandle, Surface @@ -23,14 +24,14 @@ class MySession: async def send_async(self, request: Request) -> Response: return Response(text="ok") - async def __aenter__(self) -> "MySession": + async def __aenter__(self) -> Self: return self async def __aexit__( self, exc_type: type[BaseException] | None, exc_val: BaseException | None, - exc_tb: Any, + exc_tb: types.TracebackType | None, ) -> None: pass @@ -43,14 +44,14 @@ class MySession: async def send_async(self, request: Request) -> Response: return Response(text="ok") - async def __aenter__(self) -> "MySession": + async def __aenter__(self) -> Self: return self async def __aexit__( self, exc_type: type[BaseException] | None, exc_val: BaseException | None, - exc_tb: Any, + exc_tb: types.TracebackType | None, ) -> None: pass @@ -96,14 +97,14 @@ def surface_name(self) -> str: async def wait_until_ready(self) -> None: pass - async def __aenter__(self) -> "MyHandle": + async def __aenter__(self) -> Self: return self async def __aexit__( self, exc_type: type[BaseException] | None, exc_val: BaseException | None, - exc_tb: Any, + exc_tb: types.TracebackType | None, ) -> None: pass @@ -132,14 +133,14 @@ def surface_name(self) -> str: async def wait_until_ready(self) -> None: pass - async def __aenter__(self) -> "MyHandle": + async def __aenter__(self) -> Self: return self async def __aexit__( self, exc_type: type[BaseException] | None, exc_val: BaseException | None, - exc_tb: Any, + exc_tb: types.TracebackType | None, ) -> None: pass diff --git a/tests/unit/surfaces/test_onedrive.py b/tests/unit/surfaces/test_onedrive.py index 6076b8e0..902e7750 100644 --- a/tests/unit/surfaces/test_onedrive.py +++ b/tests/unit/surfaces/test_onedrive.py @@ -5,18 +5,17 @@ from __future__ import annotations -from typing import Any from unittest.mock import AsyncMock, MagicMock, call import pytest from rampart.core.errors import InfrastructureError from rampart.core.injection import InjectionHandle, Surface -from rampart.core.types import Payload, PayloadFormat +from rampart.core.types import Payload from rampart.surfaces.onedrive import ( + _MAX_SMALL_UPLOAD_BYTES, OneDriveSurface, _OneDriveInjection, - _MAX_SMALL_UPLOAD_BYTES, ) _UNSET = object() @@ -25,7 +24,7 @@ def _make_graph_client( *, upload_item_id: str = "item-abc-123", - upload_return: Any = _UNSET, + upload_return: object = _UNSET, upload_error: Exception | None = None, delete_error: Exception | None = None, ) -> MagicMock: @@ -63,13 +62,13 @@ def _make_graph_client( items_mock = MagicMock() - def _by_drive_item_id_dispatch(item_id: str) -> Any: + def _by_drive_item_id_dispatch(item_id: str) -> object: if item_id.startswith("root:"): return upload_item_mock return delete_item_mock items_mock.by_drive_item_id = MagicMock( - side_effect=_by_drive_item_id_dispatch + side_effect=_by_drive_item_id_dispatch, ) by_drive_id_mock = MagicMock() @@ -93,9 +92,10 @@ def test_stores_configuration(self) -> None: drive_id="drive-1", folder_path="Documents/payloads", ) - assert surface._drive_id == "drive-1" - assert surface._folder_path == "Documents/payloads" - assert surface._indexing_delay == OneDriveSurface.DEFAULT_INDEXING_DELAY + assert surface.drive_id == "drive-1" + assert surface.folder_path == "Documents/payloads" + assert surface.indexing_delay == OneDriveSurface.DEFAULT_INDEXING_DELAY + assert surface.readiness_timeout == OneDriveSurface.DEFAULT_READINESS_TIMEOUT def test_custom_indexing_delay(self) -> None: surface = OneDriveSurface( @@ -104,7 +104,7 @@ def test_custom_indexing_delay(self) -> None: folder_path="test", indexing_delay=42.0, ) - assert surface._indexing_delay == 42.0 + assert surface.indexing_delay == 42.0 def test_strips_leading_trailing_slashes_from_folder_path(self) -> None: surface = OneDriveSurface( @@ -112,7 +112,7 @@ def test_strips_leading_trailing_slashes_from_folder_path(self) -> None: drive_id="d", folder_path="/foo/bar/", ) - assert surface._folder_path == "foo/bar" + assert surface.folder_path == "foo/bar" class TestOneDriveSurfaceProtocolConformance: @@ -304,7 +304,9 @@ async def test_returns_self_from_aenter(self) -> None: assert h is handle @pytest.mark.asyncio - async def test_upload_exceeding_size_limit_raises_infrastructure_error(self) -> None: + async def test_upload_exceeding_size_limit_raises_infrastructure_error( + self, + ) -> None: client = _make_graph_client() surface = OneDriveSurface( graph_client=client, From efe4a147d1979769ca21d0cb00910d3ee993ed56 Mon Sep 17 00:00:00 2001 From: Nina Chikanov Date: Tue, 14 Apr 2026 13:59:03 -0700 Subject: [PATCH 04/13] add back accidental character deletion in .gitignore --- .gitignore | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/.gitignore b/.gitignore index d6d86830..b2e20acb 100644 --- a/.gitignore +++ b/.gitignore @@ -23,8 +23,7 @@ .LSOverride # Icon must end with two \r -Icon - +Icon^M # Thumbnails ._* From 1937dffa9b11d4de6f669590b8e104c23b3077d0 Mon Sep 17 00:00:00 2001 From: Nina Chikanov Date: Tue, 14 Apr 2026 14:01:17 -0700 Subject: [PATCH 05/13] try again --- .gitignore | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/.gitignore b/.gitignore index b2e20acb..d6d86830 100644 --- a/.gitignore +++ b/.gitignore @@ -23,7 +23,8 @@ .LSOverride # Icon must end with two \r -Icon^M +Icon + # Thumbnails ._* From 10a8d2509bd0072a7ddf2e8f00475fcaf45edcd3 Mon Sep 17 00:00:00 2001 From: Nina Chikanov Date: Wed, 15 Apr 2026 15:36:56 -0700 Subject: [PATCH 06/13] [MAINT]: Restore Icon line in .gitignore Reverts unintentional removal of \r\r characters on the Icon line (L26) and extra blank line, per PR review feedback. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- .gitignore | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/.gitignore b/.gitignore index d6d86830..7e05e426 100644 --- a/.gitignore +++ b/.gitignore @@ -23,8 +23,7 @@ .LSOverride # Icon must end with two \r -Icon - +Icon # Thumbnails ._* From f5d66c2cfaf5fea4daf100cbeb43df716f318706 Mon Sep 17 00:00:00 2001 From: Nina Chikanov Date: Wed, 15 Apr 2026 17:04:14 -0700 Subject: [PATCH 07/13] reviewer agent suggested changes --- rampart/__init__.py | 3 ++- rampart/core/__init__.py | 3 ++- rampart/core/injection.py | 43 ++++++++++++++++++++++--------- rampart/surfaces/onedrive.py | 20 ++++++-------- tests/unit/core/test_injection.py | 6 ++--- 5 files changed, 46 insertions(+), 29 deletions(-) diff --git a/rampart/__init__.py b/rampart/__init__.py index f5381119..bb81838d 100644 --- a/rampart/__init__.py +++ b/rampart/__init__.py @@ -15,7 +15,7 @@ ExecutionEventData, ExecutionEventHandler, ) -from rampart.core.injection import InjectionHandle, Surface +from rampart.core.injection import InjectionHandle, InjectionHandleMixin, Surface from rampart.core.manifest import AppManifest, DataSource, ToolDeclaration from rampart.core.persona import Persona from rampart.core.prompt_driver import PromptDecision, PromptDriver @@ -61,6 +61,7 @@ "HarmCategory", "InfrastructureError", "InjectionHandle", + "InjectionHandleMixin", "InjectionRecord", "ObservabilityLevel", "Payload", diff --git a/rampart/core/__init__.py b/rampart/core/__init__.py index 5e62cf78..fcef15a3 100644 --- a/rampart/core/__init__.py +++ b/rampart/core/__init__.py @@ -17,7 +17,7 @@ ExecutionEventHandler, ExecutionHandlerFactory, ) -from rampart.core.injection import InjectionHandle, Surface +from rampart.core.injection import InjectionHandle, InjectionHandleMixin, Surface from rampart.core.llm import LLMConfig from rampart.core.manifest import AppManifest, DataSource, ToolDeclaration from rampart.core.persona import Persona @@ -61,6 +61,7 @@ "HarmCategory", "InfrastructureError", "InjectionHandle", + "InjectionHandleMixin", "InjectionRecord", "LLMConfig", "ObservabilityLevel", diff --git a/rampart/core/injection.py b/rampart/core/injection.py index d33dc611..231f50b7 100644 --- a/rampart/core/injection.py +++ b/rampart/core/injection.py @@ -5,6 +5,9 @@ Two protocols serving two audiences: Surface is what surface authors implement; InjectionHandle is what execution strategies consume. + +``InjectionHandleMixin`` provides a default sleep-based +``wait_until_ready`` for surfaces that only need a simple delay. """ from __future__ import annotations @@ -55,19 +58,8 @@ async def wait_until_ready(self) -> None: Implementations must complete within `readiness_timeout_seconds` or raise `TimeoutError`. - - Default implementation sleeps for `indexing_delay_seconds`, with - an upper bound of `readiness_timeout_seconds` to prevent indefinite - blocking — which is a common strategy for many surfaces. Surfaces with - more complex readiness logic can override this method with a custom - implementation. - - Raises: - TimeoutError: If `indexing_delay_seconds` exceeds - `readiness_timeout_seconds`. """ - async with asyncio.timeout(self.readiness_timeout_seconds): - await asyncio.sleep(self.indexing_delay_seconds) + ... async def __aenter__(self) -> Self: """Activate the injection (write payload to data source).""" @@ -83,6 +75,33 @@ async def __aexit__( ... +class InjectionHandleMixin: + """Mixin providing a default sleep-based ``wait_until_ready``. + + Surfaces whose readiness strategy is a simple timed delay can + inherit from this mixin instead of implementing + ``wait_until_ready`` from scratch. The mixin sleeps for + ``indexing_delay_seconds`` with an upper bound of + ``readiness_timeout_seconds`` to prevent indefinite blocking. + + Surfaces with more complex readiness logic (e.g. polling an + API) should implement ``wait_until_ready`` directly. + """ + + indexing_delay_seconds: float + readiness_timeout_seconds: float + + async def wait_until_ready(self) -> None: + """Sleep for ``indexing_delay_seconds``, bounded by timeout. + + Raises: + TimeoutError: If ``indexing_delay_seconds`` exceeds + ``readiness_timeout_seconds``. + """ + async with asyncio.timeout(self.readiness_timeout_seconds): + await asyncio.sleep(self.indexing_delay_seconds) + + @runtime_checkable class Surface(Protocol): """An injectable data source. diff --git a/rampart/surfaces/onedrive.py b/rampart/surfaces/onedrive.py index 274d4e96..0783145d 100644 --- a/rampart/surfaces/onedrive.py +++ b/rampart/surfaces/onedrive.py @@ -13,7 +13,7 @@ from typing import TYPE_CHECKING, Self from rampart.core.errors import InfrastructureError -from rampart.core.injection import InjectionHandle +from rampart.core.injection import InjectionHandleMixin if TYPE_CHECKING: import types @@ -105,7 +105,7 @@ async def upload_async(self, *, payload: Payload) -> str: if payload.artifact is None: msg = ( f"Binary payload format {payload.format.value} " - f"requires an artifact path.", + "requires an artifact path." ) raise ValueError( msg, @@ -117,8 +117,8 @@ async def upload_async(self, *, payload: Payload) -> str: if len(content) > _MAX_SMALL_UPLOAD_BYTES: msg = ( f"Payload {payload.id} is {len(content)} bytes, which " - f"exceeds the 4 MiB small-upload limit. Upload sessions " - f"are not yet implemented.", + "exceeds the 4 MiB small-upload limit. Upload sessions " + "are not yet implemented." ) raise ValueError( msg, @@ -134,8 +134,8 @@ async def upload_async(self, *, payload: Payload) -> str: if drive_item is None or drive_item.id is None: msg = ( - f"Graph API returned no DriveItem after upload to " - f"drive={self.drive_id} path={upload_path}", + "Graph API returned no DriveItem after upload to " + f"drive={self.drive_id} path={upload_path}" ) raise InfrastructureError( msg, @@ -165,7 +165,7 @@ async def delete_async(self, *, item_id: str) -> None: ) -class _OneDriveInjection: +class _OneDriveInjection(InjectionHandleMixin): """InjectionHandle for OneDrive. Manages upload and cleanup lifecycle.""" def __init__(self, *, surface: OneDriveSurface, payload: Payload) -> None: @@ -193,10 +193,6 @@ def surface_name(self) -> str: """Identifies this injection as OneDrive for reporting.""" return "OneDrive" - async def wait_until_ready(self) -> None: - """Wait for OneDrive indexing using the default sleep-based strategy.""" - await InjectionHandle.wait_until_ready(self) - async def __aenter__(self) -> Self: """Upload payload to OneDrive. Raises InfrastructureError on failure.""" try: @@ -208,7 +204,7 @@ async def __aenter__(self) -> Self: except Exception as exc: msg = ( f"OneDrive upload failed for drive={self._surface.drive_id} " - f"path={self._surface.folder_path}: {exc}", + f"path={self._surface.folder_path}: {exc}" ) raise InfrastructureError( msg, diff --git a/tests/unit/core/test_injection.py b/tests/unit/core/test_injection.py index 43e4875d..000bd5d9 100644 --- a/tests/unit/core/test_injection.py +++ b/tests/unit/core/test_injection.py @@ -11,13 +11,13 @@ import pytest -from rampart.core.injection import InjectionHandle +from rampart.core.injection import InjectionHandleMixin if TYPE_CHECKING: import types -class _ConcreteHandle(InjectionHandle): +class _ConcreteHandle(InjectionHandleMixin): """Minimal concrete handle that inherits the default wait_until_ready.""" def __init__( @@ -58,7 +58,7 @@ async def __aexit__( class TestWaitUntilReady: - """Tests for InjectionHandle.wait_until_ready default and custom behaviour.""" + """Tests for InjectionHandleMixin.wait_until_ready default and custom behaviour.""" @pytest.mark.asyncio @pytest.mark.parametrize( From 93add437cdab76dfb00a37868bf93e1d69ee36aa Mon Sep 17 00:00:00 2001 From: Nina Chikanov Date: Wed, 15 Apr 2026 17:27:56 -0700 Subject: [PATCH 08/13] fix merge miss --- rampart/surfaces/onedrive.py | 13 +++++++++---- 1 file changed, 9 insertions(+), 4 deletions(-) diff --git a/rampart/surfaces/onedrive.py b/rampart/surfaces/onedrive.py index e9041cbd..60673578 100644 --- a/rampart/surfaces/onedrive.py +++ b/rampart/surfaces/onedrive.py @@ -67,10 +67,10 @@ def __init__( ) -> None: """Initialize with Graph client and OneDrive location.""" self._graph_client = graph_client - self.drive_id = drive_id - self.folder_path = folder_path.strip("/") - self.indexing_delay = indexing_delay - self.readiness_timeout = readiness_timeout + self._drive_id = drive_id + self._folder_path = folder_path.strip("/") + self._indexing_delay = indexing_delay + self._readiness_timeout = readiness_timeout @property def drive_id(self) -> str: @@ -87,6 +87,11 @@ def indexing_delay(self) -> float: """Seconds to wait after upload for indexing.""" return self._indexing_delay + @property + def readiness_timeout(self) -> float: + """Maximum readiness wait time in seconds.""" + return self._readiness_timeout + def inject(self, *, payload: Payload) -> _OneDriveInjection: """Prepare an injection into the configured OneDrive folder. From 65684846b6f4d697635633792e5afb366dc597ed Mon Sep 17 00:00:00 2001 From: Nina Chikanov Date: Wed, 15 Apr 2026 17:32:29 -0700 Subject: [PATCH 09/13] +1 more reviewer agent-suggested change --- rampart/attacks/_xpia.py | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/rampart/attacks/_xpia.py b/rampart/attacks/_xpia.py index a62bd5ce..41346da2 100644 --- a/rampart/attacks/_xpia.py +++ b/rampart/attacks/_xpia.py @@ -12,6 +12,7 @@ from __future__ import annotations +import asyncio import logging from contextlib import AsyncExitStack from typing import Any @@ -188,8 +189,10 @@ async def _activate_handles_async( for handle in self._handles: await stack.enter_async_context(handle) - for handle in self._handles: - await handle.wait_until_ready() + # Concurrent: total = max of all wait times + async with asyncio.TaskGroup() as tg: + for handle in self._handles: + tg.create_task(handle.wait_until_ready()) def _build_attack_result( self, From 43d6574ea4d08e4109f34e7a9f31bdeb39862498 Mon Sep 17 00:00:00 2001 From: Nina Chikanov Date: Fri, 17 Apr 2026 17:51:30 -0700 Subject: [PATCH 10/13] Simplified code, reviewer agent feedback addressed --- rampart/__init__.py | 3 +- rampart/attacks/_xpia.py | 4 +- rampart/core/__init__.py | 3 +- rampart/core/injection.py | 46 +++-------- rampart/surfaces/onedrive.py | 31 +++----- tests/unit/attacks/test_xpia.py | 4 +- tests/unit/core/test_injection.py | 115 --------------------------- tests/unit/core/test_protocols.py | 16 ---- tests/unit/surfaces/test_onedrive.py | 48 +++++------ 9 files changed, 50 insertions(+), 220 deletions(-) delete mode 100644 tests/unit/core/test_injection.py diff --git a/rampart/__init__.py b/rampart/__init__.py index b70f1273..eade4111 100644 --- a/rampart/__init__.py +++ b/rampart/__init__.py @@ -16,7 +16,7 @@ ExecutionEventData, ExecutionEventHandler, ) -from rampart.core.injection import InjectionHandle, InjectionHandleMixin, Surface +from rampart.core.injection import InjectionHandle, Surface from rampart.core.manifest import AppManifest, DataSource, ToolDeclaration from rampart.core.persona import Persona from rampart.core.prompt_driver import PromptDecision, PromptDriver @@ -61,7 +61,6 @@ "HarmCategory", "InfrastructureError", "InjectionHandle", - "InjectionHandleMixin", "InjectionRecord", "ObservabilityLevel", "Payload", diff --git a/rampart/attacks/_xpia.py b/rampart/attacks/_xpia.py index 41346da2..af2c8df1 100644 --- a/rampart/attacks/_xpia.py +++ b/rampart/attacks/_xpia.py @@ -47,7 +47,7 @@ class XPIAExecution(BaseExecution): Phases (delegated to private helpers from ``_execute_async``): 1. Activate all injection handles (via AsyncExitStack). - 2. Wait for indexing (max delay across all handles). + 2. Wait for indexing (concurrent per-handle). 3. Create session (via async context manager). 4. Drive the trigger conversation via the PromptDriver. 5. Evaluate per-turn with early stopping on detection. @@ -181,7 +181,7 @@ async def _activate_handles_async( *, stack: AsyncExitStack, ) -> None: - """Activate all injection handles and wait for indexing. + """Activate all injection handles and wait for readiness. Args: stack (AsyncExitStack): The exit stack managing cleanup. diff --git a/rampart/core/__init__.py b/rampart/core/__init__.py index 637d11f4..102d1e12 100644 --- a/rampart/core/__init__.py +++ b/rampart/core/__init__.py @@ -17,7 +17,7 @@ ExecutionEventHandler, ExecutionHandlerFactory, ) -from rampart.core.injection import InjectionHandle, InjectionHandleMixin, Surface +from rampart.core.injection import InjectionHandle, Surface from rampart.core.llm import LLMConfig from rampart.core.manifest import AppManifest, DataSource, ToolDeclaration from rampart.core.persona import Persona @@ -61,7 +61,6 @@ "HarmCategory", "InfrastructureError", "InjectionHandle", - "InjectionHandleMixin", "InjectionRecord", "LLMConfig", "ObservabilityLevel", diff --git a/rampart/core/injection.py b/rampart/core/injection.py index 231f50b7..dfd5efc9 100644 --- a/rampart/core/injection.py +++ b/rampart/core/injection.py @@ -6,8 +6,8 @@ Two protocols serving two audiences: Surface is what surface authors implement; InjectionHandle is what execution strategies consume. -``InjectionHandleMixin`` provides a default sleep-based -``wait_until_ready`` for surfaces that only need a simple delay. +``sleep_until_ready`` is a free helper for surfaces that only need +a simple delay-based readiness wait. """ from __future__ import annotations @@ -33,16 +33,6 @@ class InjectionHandle(Protocol): Surface or its concrete implementations. """ - @property - def indexing_delay_seconds(self) -> float: - """How long to wait after activation for the agent to see the content.""" - ... - - @property - def readiness_timeout_seconds(self) -> float: - """Maximum time `wait_until_ready` may block before raising `TimeoutError`.""" - ... - @property def payload_id(self) -> str | None: """The injected payload's identifier, for reporting.""" @@ -56,8 +46,8 @@ def surface_name(self) -> str: async def wait_until_ready(self) -> None: """Block until the injected content is visible to the agent. - Implementations must complete within `readiness_timeout_seconds` - or raise `TimeoutError`. + Implementations should raise `TimeoutError` if readiness + cannot be confirmed within a reasonable time. """ ... @@ -75,31 +65,13 @@ async def __aexit__( ... -class InjectionHandleMixin: - """Mixin providing a default sleep-based ``wait_until_ready``. +async def sleep_until_ready(delay: float) -> None: + """Sleep for `delay` seconds. Default readiness strategy for simple surfaces. - Surfaces whose readiness strategy is a simple timed delay can - inherit from this mixin instead of implementing - ``wait_until_ready`` from scratch. The mixin sleeps for - ``indexing_delay_seconds`` with an upper bound of - ``readiness_timeout_seconds`` to prevent indefinite blocking. - - Surfaces with more complex readiness logic (e.g. polling an - API) should implement ``wait_until_ready`` directly. + Args: + delay: Seconds to sleep before the injection is considered ready. """ - - indexing_delay_seconds: float - readiness_timeout_seconds: float - - async def wait_until_ready(self) -> None: - """Sleep for ``indexing_delay_seconds``, bounded by timeout. - - Raises: - TimeoutError: If ``indexing_delay_seconds`` exceeds - ``readiness_timeout_seconds``. - """ - async with asyncio.timeout(self.readiness_timeout_seconds): - await asyncio.sleep(self.indexing_delay_seconds) + await asyncio.sleep(delay) @runtime_checkable diff --git a/rampart/surfaces/onedrive.py b/rampart/surfaces/onedrive.py index 60673578..cad900d1 100644 --- a/rampart/surfaces/onedrive.py +++ b/rampart/surfaces/onedrive.py @@ -13,7 +13,7 @@ from typing import TYPE_CHECKING, Self from rampart.core.errors import InfrastructureError -from rampart.core.injection import InjectionHandleMixin +from rampart.core.injection import sleep_until_ready if TYPE_CHECKING: import types @@ -54,7 +54,6 @@ class OneDriveSurface: """ DEFAULT_INDEXING_DELAY: float = 10.0 - DEFAULT_READINESS_TIMEOUT: float = 120.0 def __init__( self, @@ -63,14 +62,12 @@ def __init__( drive_id: str, folder_path: str, indexing_delay: float = DEFAULT_INDEXING_DELAY, - readiness_timeout: float = DEFAULT_READINESS_TIMEOUT, ) -> None: """Initialize with Graph client and OneDrive location.""" self._graph_client = graph_client self._drive_id = drive_id self._folder_path = folder_path.strip("/") self._indexing_delay = indexing_delay - self._readiness_timeout = readiness_timeout @property def drive_id(self) -> str: @@ -87,11 +84,6 @@ def indexing_delay(self) -> float: """Seconds to wait after upload for indexing.""" return self._indexing_delay - @property - def readiness_timeout(self) -> float: - """Maximum readiness wait time in seconds.""" - return self._readiness_timeout - def inject(self, *, payload: Payload) -> _OneDriveInjection: """Prepare an injection into the configured OneDrive folder. @@ -186,7 +178,7 @@ async def delete_async(self, *, item_id: str) -> None: ) -class _OneDriveInjection(InjectionHandleMixin): +class _OneDriveInjection: """InjectionHandle for OneDrive. Manages upload and cleanup lifecycle.""" def __init__(self, *, surface: OneDriveSurface, payload: Payload) -> None: @@ -194,16 +186,6 @@ def __init__(self, *, surface: OneDriveSurface, payload: Payload) -> None: self._payload = payload self._item_id: str | None = None - @property - def indexing_delay_seconds(self) -> float: - """How long to wait after upload for content to be discoverable.""" - return self._surface.indexing_delay - - @property - def readiness_timeout_seconds(self) -> float: - """Maximum time `wait_until_ready` may block.""" - return self._surface.readiness_timeout - @property def payload_id(self) -> str | None: """The injected payload's identifier.""" @@ -214,6 +196,15 @@ def surface_name(self) -> str: """Identifies this injection as OneDrive for reporting.""" return "OneDrive" + async def wait_until_ready(self) -> None: + """Wait for the uploaded content to be indexed and discoverable. + + Note: Currently sleeps for `OneDriveSurface.indexing_delay` seconds. + Future versions will poll the Graph API for the file's availability instead and + raise `TimeoutError` if it doesn't appear within the `indexing_delay`. + """ + await sleep_until_ready(delay=self._surface.indexing_delay) + async def __aenter__(self) -> Self: """Upload payload to OneDrive. Raises InfrastructureError on failure.""" try: diff --git a/tests/unit/attacks/test_xpia.py b/tests/unit/attacks/test_xpia.py index d3d09c84..69580f8e 100644 --- a/tests/unit/attacks/test_xpia.py +++ b/tests/unit/attacks/test_xpia.py @@ -29,13 +29,11 @@ def _mock_handle( *, surface_name: str = "FakeSurface", payload_id: str | None = "p-001", - delay: float = 0.0, ) -> AsyncMock: """Create an AsyncMock satisfying the InjectionHandle protocol.""" h = AsyncMock() h.surface_name = surface_name h.payload_id = payload_id - h.indexing_delay_seconds = delay h.__aenter__.return_value = h return h @@ -179,6 +177,7 @@ async def test_handle_entered_and_exited(self) -> None: handle.__aenter__.assert_awaited_once() handle.__aexit__.assert_awaited_once() + handle.wait_until_ready.assert_awaited_once() @pytest.mark.asyncio async def test_multiple_handles_all_cleaned(self) -> None: @@ -194,6 +193,7 @@ async def test_multiple_handles_all_cleaned(self) -> None: for h in (h1, h2): h.__aenter__.assert_awaited_once() h.__aexit__.assert_awaited_once() + h.wait_until_ready.assert_awaited_once() @pytest.mark.asyncio async def test_cleanup_on_evaluator_exception(self) -> None: diff --git a/tests/unit/core/test_injection.py b/tests/unit/core/test_injection.py deleted file mode 100644 index 000bd5d9..00000000 --- a/tests/unit/core/test_injection.py +++ /dev/null @@ -1,115 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# Licensed under the MIT license. - -"""Tests for rampart.core.injection.""" - -from __future__ import annotations - -import asyncio -from typing import TYPE_CHECKING, Self -from unittest.mock import AsyncMock - -import pytest - -from rampart.core.injection import InjectionHandleMixin - -if TYPE_CHECKING: - import types - - -class _ConcreteHandle(InjectionHandleMixin): - """Minimal concrete handle that inherits the default wait_until_ready.""" - - def __init__( - self, - *, - delay: float = 0.0, - readiness_timeout: float = 30.0, - ) -> None: - self._delay = delay - self._readiness_timeout = readiness_timeout - - @property - def indexing_delay_seconds(self) -> float: - return self._delay - - @property - def readiness_timeout_seconds(self) -> float: - return self._readiness_timeout - - @property - def payload_id(self) -> str | None: - return "test-payload" - - @property - def surface_name(self) -> str: - return "TestSurface" - - async def __aenter__(self) -> Self: - return self - - async def __aexit__( - self, - exc_type: type[BaseException] | None, - exc_val: BaseException | None, - exc_tb: types.TracebackType | None, - ) -> None: - pass - - -class TestWaitUntilReady: - """Tests for InjectionHandleMixin.wait_until_ready default and custom behaviour.""" - - @pytest.mark.asyncio - @pytest.mark.parametrize( - ("delay", "readiness_timeout"), - [ - (0.0, 1.0), - (0.05, 5.0), - ], - ids=["zero-delay", "short-delay-within-timeout"], - ) - async def test_default_completes_when_delay_within_timeout( - self, - delay: float, - readiness_timeout: float, - ) -> None: - """Default sleep-based wait completes when delay is within the timeout.""" - handle = _ConcreteHandle(delay=delay, readiness_timeout=readiness_timeout) - - await handle.wait_until_ready() - - @pytest.mark.asyncio - async def test_custom_polling_implementation(self) -> None: - """A handle can override wait_until_ready with custom polling logic.""" - _expected_poll_calls = 3 - poll_mock = AsyncMock(side_effect=[False, False, True]) - ready_event = asyncio.Event() - - class _PollingHandle(_ConcreteHandle): - async def wait_until_ready(self) -> None: - async with asyncio.timeout(self.readiness_timeout_seconds): - while not await poll_mock(): - ready_event.clear() - await ready_event.wait() - - handle = _PollingHandle(readiness_timeout=5.0) - - async def _signal_ready() -> None: - for _ in range(_expected_poll_calls - 1): - await asyncio.sleep(0) - ready_event.set() - - async with asyncio.TaskGroup() as tg: - tg.create_task(handle.wait_until_ready()) - tg.create_task(_signal_ready()) - - assert poll_mock.await_count == _expected_poll_calls - - @pytest.mark.asyncio - async def test_timeout_raises_when_delay_exceeds_limit(self) -> None: - """Default implementation raises TimeoutError when delay exceeds timeout.""" - handle = _ConcreteHandle(delay=10.0, readiness_timeout=0.01) - - with pytest.raises(TimeoutError): - await handle.wait_until_ready() diff --git a/tests/unit/core/test_protocols.py b/tests/unit/core/test_protocols.py index d9eff556..8a0bb69c 100644 --- a/tests/unit/core/test_protocols.py +++ b/tests/unit/core/test_protocols.py @@ -78,14 +78,6 @@ def observability_profile(self) -> ObservabilityLevel: class TestInjectionHandleProtocol: def test_structural_subtyping(self) -> None: class MyHandle: - @property - def indexing_delay_seconds(self) -> float: - return 5.0 - - @property - def readiness_timeout_seconds(self) -> float: - return 30.0 - @property def payload_id(self) -> str | None: return "abc" @@ -114,14 +106,6 @@ async def __aexit__( class TestSurfaceProtocol: def test_structural_subtyping(self) -> None: class MyHandle: - @property - def indexing_delay_seconds(self) -> float: - return 0.0 - - @property - def readiness_timeout_seconds(self) -> float: - return 30.0 - @property def payload_id(self) -> str | None: return None diff --git a/tests/unit/surfaces/test_onedrive.py b/tests/unit/surfaces/test_onedrive.py index 902e7750..d22d8dc5 100644 --- a/tests/unit/surfaces/test_onedrive.py +++ b/tests/unit/surfaces/test_onedrive.py @@ -5,7 +5,7 @@ from __future__ import annotations -from unittest.mock import AsyncMock, MagicMock, call +from unittest.mock import AsyncMock, MagicMock, call, patch import pytest @@ -95,7 +95,6 @@ def test_stores_configuration(self) -> None: assert surface.drive_id == "drive-1" assert surface.folder_path == "Documents/payloads" assert surface.indexing_delay == OneDriveSurface.DEFAULT_INDEXING_DELAY - assert surface.readiness_timeout == OneDriveSurface.DEFAULT_READINESS_TIMEOUT def test_custom_indexing_delay(self) -> None: surface = OneDriveSurface( @@ -160,28 +159,6 @@ def test_payload_id(self) -> None: handle = surface.inject(payload=payload) assert handle.payload_id == "my-payload-id" - def test_indexing_delay_from_surface(self) -> None: - surface = OneDriveSurface( - graph_client=MagicMock(), - drive_id="d", - folder_path="f", - indexing_delay=99.0, - ) - payload = Payload(content="test") - handle = surface.inject(payload=payload) - assert handle.indexing_delay_seconds == 99.0 - - def test_readiness_timeout_from_surface(self) -> None: - surface = OneDriveSurface( - graph_client=MagicMock(), - drive_id="d", - folder_path="f", - readiness_timeout=60.0, - ) - payload = Payload(content="test") - handle = surface.inject(payload=payload) - assert handle.readiness_timeout_seconds == 60.0 - class TestOneDriveInjectionLifecycle: """Test the async context manager lifecycle (upload + delete).""" @@ -372,3 +349,26 @@ async def test_infrastructure_error_from_upload_not_double_wrapped(self) -> None pass assert exc_info.value is original + + +class TestOneDriveInjectionWaitUntilReady: + """Test _OneDriveInjection.wait_until_ready wiring.""" + + @pytest.mark.asyncio + async def test_delegates_to_sleep_until_ready(self) -> None: + """Verifies correct arguments are passed to sleep_until_ready.""" + surface = OneDriveSurface( + graph_client=MagicMock(), + drive_id="d", + folder_path="f", + indexing_delay=5.0, + ) + handle = surface.inject(payload=Payload(content="test")) + + with patch( + "rampart.surfaces.onedrive.sleep_until_ready", + new_callable=AsyncMock, + ) as mock_sleep: + await handle.wait_until_ready() + + mock_sleep.assert_awaited_once_with(delay=5.0) From 66da6c074fd61fafd5a813f01144e01a72740f38 Mon Sep 17 00:00:00 2001 From: Nina Chikanov Date: Fri, 17 Apr 2026 18:01:11 -0700 Subject: [PATCH 11/13] undo changes to "any" in tests --- tests/unit/surfaces/test_onedrive.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/tests/unit/surfaces/test_onedrive.py b/tests/unit/surfaces/test_onedrive.py index d22d8dc5..12434ceb 100644 --- a/tests/unit/surfaces/test_onedrive.py +++ b/tests/unit/surfaces/test_onedrive.py @@ -5,6 +5,7 @@ from __future__ import annotations +from typing import Any from unittest.mock import AsyncMock, MagicMock, call, patch import pytest @@ -24,7 +25,7 @@ def _make_graph_client( *, upload_item_id: str = "item-abc-123", - upload_return: object = _UNSET, + upload_return: Any = _UNSET, upload_error: Exception | None = None, delete_error: Exception | None = None, ) -> MagicMock: @@ -62,7 +63,7 @@ def _make_graph_client( items_mock = MagicMock() - def _by_drive_item_id_dispatch(item_id: str) -> object: + def _by_drive_item_id_dispatch(item_id: str) -> Any: if item_id.startswith("root:"): return upload_item_mock return delete_item_mock From a16c534ae5380ccfe3921de83b699e1d8e2654b3 Mon Sep 17 00:00:00 2001 From: Nina Chikanov Date: Fri, 17 Apr 2026 18:13:51 -0700 Subject: [PATCH 12/13] fix up docstrings --- .../work-items/exception-group-infra-error.md | 35 +++++++++++++++++++ rampart/core/injection.py | 4 +-- rampart/surfaces/onedrive.py | 2 +- 3 files changed, 38 insertions(+), 3 deletions(-) create mode 100644 docs/work-items/exception-group-infra-error.md diff --git a/docs/work-items/exception-group-infra-error.md b/docs/work-items/exception-group-infra-error.md new file mode 100644 index 00000000..462f078a --- /dev/null +++ b/docs/work-items/exception-group-infra-error.md @@ -0,0 +1,35 @@ +# Broaden `BaseExecution` error handling to produce clean error results for all exceptions + +## Problem + +`BaseExecution.execute_async` only catches `InfrastructureError` and converts it to a clean `Result(status=SafetyStatus.ERROR)`. All other exceptions — including `TimeoutError`, `ExceptionGroup` (from `asyncio.TaskGroup`), `ConnectionError`, and any unexpected runtime failure — propagate unhandled, crashing the test run instead of producing a reportable result. + +This is a gap in the base class's role as the cross-cutting error handler for all execution strategies. Any exception raised during `_execute_async` that isn't an `InfrastructureError` bypasses the error-result path entirely, even when it represents a transient, non-diagnostic failure that should be reported the same way. + +## Impact + +- **Test runs crash on transient failures.** Any non-`InfrastructureError` exception from a surface, adapter, readiness check, or evaluator aborts the run instead of recording an error result. +- **`asyncio.TaskGroup` amplifies the problem.** Strategies using `TaskGroup` (e.g., XPIA concurrent readiness) will have their exceptions wrapped in `ExceptionGroup`, which the current handler doesn't catch at all. +- **Surface authors must know framework internals.** Today, surfaces must raise `InfrastructureError` specifically, or their failures crash the run. The base class should be resilient to any exception type, not just the framework's own. +- **Inconsistent reporting.** Some failures produce clean error results (those that raise `InfrastructureError`) while equivalent failures from other exception types produce stack traces and test run crashes. + +## Proposed Fix + +Broaden `BaseExecution.execute_async` to produce clean `Result(status=SafetyStatus.ERROR)` for all non-safety-diagnostic exceptions, not just `InfrastructureError`. Currently, only `InfrastructureError` is caught and converted to an error result — any other exception (including `ExceptionGroup`, `TimeoutError`, `ConnectionError`, etc.) propagates unhandled and crashes the test run. + +The fix should be in `BaseExecution.execute_async` so that all execution strategies benefit uniformly, rather than requiring each strategy to catch and re-raise as `InfrastructureError`. + +Possible approaches: +1. Widen the `except` clause in `execute_async` to catch `Exception` (and/or `BaseExceptionGroup`) and produce `SafetyStatus.ERROR` for any failure, while still firing `ON_ERROR` for observability. +2. Add an explicit `except ExceptionGroup` / `except* TimeoutError` handler alongside the existing `except InfrastructureError` handler. + +The chosen approach should preserve the existing `ON_ERROR` event dispatch so handlers are still notified, while ensuring the test run is never crashed by a transient failure from a surface, adapter, or readiness check. + +## Acceptance Criteria + +- [ ] `BaseExecution.execute_async` produces `Result(status=SafetyStatus.ERROR)` for exceptions beyond `InfrastructureError`, including `ExceptionGroup` and `TimeoutError`. +- [ ] When any handle's `wait_until_ready()` raises `TimeoutError` (wrapped in `ExceptionGroup` by `TaskGroup`), the test produces a clean error result — not an unhandled exception. +- [ ] `ON_ERROR` event handlers are still notified when a non-`InfrastructureError` exception is caught and converted to an error result. +- [ ] Individual execution strategies (e.g., `XPIAExecution`) do not need to catch and translate exceptions — the base class handles it as a cross-cutting concern. +- [ ] Unit tests in `test_execution.py` verify the broadened error handling for `ExceptionGroup`, `TimeoutError`, and other non-`InfrastructureError` exceptions. +- [ ] Existing tests continue to pass (no regression). diff --git a/rampart/core/injection.py b/rampart/core/injection.py index dfd5efc9..85fceba7 100644 --- a/rampart/core/injection.py +++ b/rampart/core/injection.py @@ -6,7 +6,7 @@ Two protocols serving two audiences: Surface is what surface authors implement; InjectionHandle is what execution strategies consume. -``sleep_until_ready`` is a free helper for surfaces that only need +``sleep_until_ready`` is a helper function for surfaces that only need a simple delay-based readiness wait. """ @@ -47,7 +47,7 @@ async def wait_until_ready(self) -> None: """Block until the injected content is visible to the agent. Implementations should raise `TimeoutError` if readiness - cannot be confirmed within a reasonable time. + operations are long-running to prevent indefinite blocking. """ ... diff --git a/rampart/surfaces/onedrive.py b/rampart/surfaces/onedrive.py index cad900d1..d2397439 100644 --- a/rampart/surfaces/onedrive.py +++ b/rampart/surfaces/onedrive.py @@ -200,7 +200,7 @@ async def wait_until_ready(self) -> None: """Wait for the uploaded content to be indexed and discoverable. Note: Currently sleeps for `OneDriveSurface.indexing_delay` seconds. - Future versions will poll the Graph API for the file's availability instead and + Future versions will poll the Graph API for content availability instead and raise `TimeoutError` if it doesn't appear within the `indexing_delay`. """ await sleep_until_ready(delay=self._surface.indexing_delay) From 59e21c8960995300ae7705b4a7cf91b65ae16e7e Mon Sep 17 00:00:00 2001 From: Nina Chikanov Date: Mon, 20 Apr 2026 10:15:41 -0700 Subject: [PATCH 13/13] remove docs --- .../work-items/exception-group-infra-error.md | 35 ------------------- 1 file changed, 35 deletions(-) delete mode 100644 docs/work-items/exception-group-infra-error.md diff --git a/docs/work-items/exception-group-infra-error.md b/docs/work-items/exception-group-infra-error.md deleted file mode 100644 index 462f078a..00000000 --- a/docs/work-items/exception-group-infra-error.md +++ /dev/null @@ -1,35 +0,0 @@ -# Broaden `BaseExecution` error handling to produce clean error results for all exceptions - -## Problem - -`BaseExecution.execute_async` only catches `InfrastructureError` and converts it to a clean `Result(status=SafetyStatus.ERROR)`. All other exceptions — including `TimeoutError`, `ExceptionGroup` (from `asyncio.TaskGroup`), `ConnectionError`, and any unexpected runtime failure — propagate unhandled, crashing the test run instead of producing a reportable result. - -This is a gap in the base class's role as the cross-cutting error handler for all execution strategies. Any exception raised during `_execute_async` that isn't an `InfrastructureError` bypasses the error-result path entirely, even when it represents a transient, non-diagnostic failure that should be reported the same way. - -## Impact - -- **Test runs crash on transient failures.** Any non-`InfrastructureError` exception from a surface, adapter, readiness check, or evaluator aborts the run instead of recording an error result. -- **`asyncio.TaskGroup` amplifies the problem.** Strategies using `TaskGroup` (e.g., XPIA concurrent readiness) will have their exceptions wrapped in `ExceptionGroup`, which the current handler doesn't catch at all. -- **Surface authors must know framework internals.** Today, surfaces must raise `InfrastructureError` specifically, or their failures crash the run. The base class should be resilient to any exception type, not just the framework's own. -- **Inconsistent reporting.** Some failures produce clean error results (those that raise `InfrastructureError`) while equivalent failures from other exception types produce stack traces and test run crashes. - -## Proposed Fix - -Broaden `BaseExecution.execute_async` to produce clean `Result(status=SafetyStatus.ERROR)` for all non-safety-diagnostic exceptions, not just `InfrastructureError`. Currently, only `InfrastructureError` is caught and converted to an error result — any other exception (including `ExceptionGroup`, `TimeoutError`, `ConnectionError`, etc.) propagates unhandled and crashes the test run. - -The fix should be in `BaseExecution.execute_async` so that all execution strategies benefit uniformly, rather than requiring each strategy to catch and re-raise as `InfrastructureError`. - -Possible approaches: -1. Widen the `except` clause in `execute_async` to catch `Exception` (and/or `BaseExceptionGroup`) and produce `SafetyStatus.ERROR` for any failure, while still firing `ON_ERROR` for observability. -2. Add an explicit `except ExceptionGroup` / `except* TimeoutError` handler alongside the existing `except InfrastructureError` handler. - -The chosen approach should preserve the existing `ON_ERROR` event dispatch so handlers are still notified, while ensuring the test run is never crashed by a transient failure from a surface, adapter, or readiness check. - -## Acceptance Criteria - -- [ ] `BaseExecution.execute_async` produces `Result(status=SafetyStatus.ERROR)` for exceptions beyond `InfrastructureError`, including `ExceptionGroup` and `TimeoutError`. -- [ ] When any handle's `wait_until_ready()` raises `TimeoutError` (wrapped in `ExceptionGroup` by `TaskGroup`), the test produces a clean error result — not an unhandled exception. -- [ ] `ON_ERROR` event handlers are still notified when a non-`InfrastructureError` exception is caught and converted to an error result. -- [ ] Individual execution strategies (e.g., `XPIAExecution`) do not need to catch and translate exceptions — the base class handles it as a cross-cutting concern. -- [ ] Unit tests in `test_execution.py` verify the broadened error handling for `ExceptionGroup`, `TimeoutError`, and other non-`InfrastructureError` exceptions. -- [ ] Existing tests continue to pass (no regression).