diff --git a/.gitignore b/.gitignore index fd410cfa..7e05e426 100644 --- a/.gitignore +++ b/.gitignore @@ -335,7 +335,6 @@ pyrightconfig.json ### VisualStudioCode ### .vscode/* -!.vscode/settings.json !.vscode/tasks.json !.vscode/launch.json !.vscode/extensions.json diff --git a/rampart/attacks/_xpia.py b/rampart/attacks/_xpia.py index e582bf6d..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. @@ -189,12 +189,10 @@ 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) + # 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, diff --git a/rampart/core/injection.py b/rampart/core/injection.py index d1866c45..85fceba7 100644 --- a/rampart/core/injection.py +++ b/rampart/core/injection.py @@ -5,10 +5,14 @@ Two protocols serving two audiences: Surface is what surface authors implement; InjectionHandle is what execution strategies consume. + +``sleep_until_ready`` is a helper function for surfaces that only need +a simple delay-based readiness wait. """ from __future__ import annotations +import asyncio from typing import TYPE_CHECKING, Protocol, Self, runtime_checkable if TYPE_CHECKING: @@ -29,11 +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 payload_id(self) -> str | None: """The injected payload's identifier, for reporting.""" @@ -44,6 +43,14 @@ 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 should raise `TimeoutError` if readiness + operations are long-running to prevent indefinite blocking. + """ + ... + async def __aenter__(self) -> Self: """Activate the injection (write payload to data source).""" ... @@ -58,6 +65,15 @@ async def __aexit__( ... +async def sleep_until_ready(delay: float) -> None: + """Sleep for `delay` seconds. Default readiness strategy for simple surfaces. + + Args: + delay: Seconds to sleep before the injection is considered ready. + """ + await asyncio.sleep(delay) + + @runtime_checkable class Surface(Protocol): """An injectable data source. diff --git a/rampart/surfaces/onedrive.py b/rampart/surfaces/onedrive.py index 447834fe..d2397439 100644 --- a/rampart/surfaces/onedrive.py +++ b/rampart/surfaces/onedrive.py @@ -13,6 +13,7 @@ from typing import TYPE_CHECKING, Self from rampart.core.errors import InfrastructureError +from rampart.core.injection import sleep_until_ready if TYPE_CHECKING: import types @@ -110,17 +111,18 @@ 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: msg = ( f"Binary payload format {payload.format.value} " - f"requires an artifact path." + "requires an artifact path." ) raise ValueError( msg, ) + content = payload.artifact.read_bytes() else: content = payload.content.encode("utf-8") @@ -128,8 +130,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, @@ -138,15 +140,15 @@ async def upload_async(self, *, payload: Payload) -> str: # 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: 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, @@ -156,7 +158,7 @@ async def upload_async(self, *, payload: Payload) -> str: 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, ) @@ -165,14 +167,14 @@ async def upload_async(self, *, payload: Payload) -> str: 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, ) @@ -184,11 +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 payload_id(self) -> str | None: """The injected payload's identifier.""" @@ -199,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 content 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_protocols.py b/tests/unit/core/test_protocols.py index d956a9ba..8a0bb69c 100644 --- a/tests/unit/core/test_protocols.py +++ b/tests/unit/core/test_protocols.py @@ -78,10 +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 payload_id(self) -> str | None: return "abc" @@ -90,6 +86,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) -> Self: return self @@ -107,10 +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 payload_id(self) -> str | None: return None @@ -119,6 +114,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) -> Self: return self diff --git a/tests/unit/surfaces/test_onedrive.py b/tests/unit/surfaces/test_onedrive.py index 92646554..12434ceb 100644 --- a/tests/unit/surfaces/test_onedrive.py +++ b/tests/unit/surfaces/test_onedrive.py @@ -6,7 +6,7 @@ from __future__ import annotations from typing import Any -from unittest.mock import AsyncMock, MagicMock, call +from unittest.mock import AsyncMock, MagicMock, call, patch import pytest @@ -93,9 +93,9 @@ 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 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: @@ -160,17 +160,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 - class TestOneDriveInjectionLifecycle: """Test the async context manager lifecycle (upload + delete).""" @@ -361,3 +350,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)