[FEAT] Add InjectionHandle wait_until_ready hook

.gitignore

@@ -335,7 +335,6 @@ pyrightconfig.json
 
 ### VisualStudioCode ###
 .vscode/*
-!.vscode/settings.json
 !.vscode/tasks.json
 !.vscode/launch.json
 !.vscode/extensions.json

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,20 +181,18 @@ 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.
         """
         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,

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.

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,26 +111,27 @@ 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")
 
         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:

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:

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
 

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,15 +104,15 @@ 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(
             graph_client=MagicMock(),
             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)