feat(runtime): defer the sandbox bring-up behind a DeferredSandboxSession when create_or_reuse is called with defer=True

strix/runtime/deferred_session.py

@@ -0,0 +1,255 @@
+from __future__ import annotations
+
+import asyncio
+import contextlib
+import logging
+from pathlib import Path
+from typing import TYPE_CHECKING, Any
+
+from agents.sandbox.session.manager import Instrumentation
+from agents.sandbox.session.sandbox_session import SandboxSession
+from agents.sandbox.session.sandbox_session_state import SandboxSessionState
+from agents.sandbox.snapshot import NoopSnapshot
+
+
+if TYPE_CHECKING:
+    import io
+    from collections.abc import Sequence
+
+    from agents.sandbox.entries import BaseEntry
+    from agents.sandbox.manifest import Manifest
+    from agents.sandbox.materialization import MaterializedFile
+    from agents.sandbox.session.base_sandbox_session import BaseSandboxSession
+    from agents.sandbox.types import ExecResult, ExposedPortEndpoint, User
+
+
+logger = logging.getLogger(__name__)
+
+
+class _PendingSessionState(SandboxSessionState):
+    type: str = "pending"
+
+
+class DeferredSandboxSession(SandboxSession):
+    def __init__(
+        self,
+        bring_up: asyncio.Task[tuple[Any, BaseSandboxSession]],
+        *,
+        manifest: Manifest,
+        exposed_ports: tuple[int, ...] = (),
+    ) -> None:
+        self._bring_up = bring_up
+        self._pending_state: SandboxSessionState = _PendingSessionState(
+            snapshot=NoopSnapshot(id="pending"),
+            manifest=manifest,
+            exposed_ports=exposed_ports,
+        )
+        self._pending_limits: Any = None
+        self._pending_archive_limits: Any = None
+        self._resolved: BaseSandboxSession | None = None
+        self._client: Any = None
+        self._instrumentation: Instrumentation = Instrumentation()
+        self._seq = 0
+
+    @property
+    def pending(self) -> bool:
+        return self._resolved is None
+
+    @property
+    def bring_up(self) -> asyncio.Task[tuple[Any, BaseSandboxSession]]:
+        return self._bring_up
+
+    async def resolve(self) -> BaseSandboxSession:
+        if self._resolved is not None:
+            return self._resolved
+        client, session = await asyncio.shield(self._bring_up)
+        self._adopt(client, session)
+        return session
+
+    def _adopt(self, client: Any, session: BaseSandboxSession) -> None:
+        if self._resolved is not None:
+            return
+        self._client = client
+        self._resolved = session
+        self._inner = session
+        if self._pending_limits is not None:
+            session._set_concurrency_limits(self._pending_limits)
+        if self._pending_archive_limits is not None:
+            session._set_archive_limits(self._pending_archive_limits)
+        if self._pending_state.manifest != session.state.manifest:
+            session.state = session.state.model_copy(
+                update={"manifest": self._pending_state.manifest}
+            )
+
+    @property
+    def client(self) -> Any:
+        return self._client
+
+    @property
+    def state(self) -> SandboxSessionState:
+        if self._resolved is None:
+            return self._pending_state
+        return self._resolved.state
+
+    @state.setter
+    def state(self, value: SandboxSessionState) -> None:
+        if self._resolved is None:
+            self._pending_state = value
+        else:
+            self._resolved.state = value
+
+    def _set_concurrency_limits(self, limits: Any) -> None:
+        if self._resolved is None:
+            self._pending_limits = limits
+        else:
+            self._resolved._set_concurrency_limits(limits)
+
+    def _set_archive_limits(self, limits: Any) -> None:
+        if self._resolved is None:
+            self._pending_archive_limits = limits
+        else:
+            self._resolved._set_archive_limits(limits)
+
+    def _manifest_base_dir(self) -> Path:
+        if self._resolved is None:
+            return Path.cwd()
+        return self._resolved._manifest_base_dir()
+
+    def normalize_path(self, path: Path | str, *, for_write: bool = False) -> Path:
+        if self._resolved is None:
+            return Path(path)
+        return self._resolved.normalize_path(path, for_write=for_write)
+
+    def register_persist_workspace_skip_path(self, path: Path | str) -> Path:
+        if self._resolved is None:
+            return Path(path)
+        return self._resolved.register_persist_workspace_skip_path(path)
+
+    def supports_pty(self) -> bool:
+        return self._resolved.supports_pty() if self._resolved is not None else False
+
+    async def running(self) -> bool:
+        if self._resolved is None:
+            if not self._bring_up.done():
+                return True
+            return not self._bring_up.cancelled() and self._bring_up.exception() is None
+        return await self._resolved.running()
+
+    async def start(self) -> None:
+        if self._resolved is None:
+            return
+        await self._resolved.start()
+
+    async def stop(self) -> None:
+        if self._resolved is None:
+            return
+        await self._resolved.stop()
+
+    async def shutdown(self) -> None:
+        if self._resolved is None:
+            return
+        await self._resolved.shutdown()
+
+    async def aclose(self) -> None:
+        if self._resolved is None:
+            return
+        await self._resolved.aclose()
+
+    async def _validate_manifest_application(self, *, only_ephemeral: bool = False) -> None:
+        session = await self.resolve()
+        await session._validate_manifest_application(only_ephemeral=only_ephemeral)
+
+    async def apply_manifest(self, *, only_ephemeral: bool = False) -> Any:
+        session = await self.resolve()
+        return await session.apply_manifest(only_ephemeral=only_ephemeral)
+
+    async def _apply_entry_batch(
+        self,
+        entries: Sequence[tuple[Path, BaseEntry]],
+        *,
+        base_dir: Path,
+    ) -> list[MaterializedFile]:
+        del base_dir
+        session = await self.resolve()
+        return await session._apply_entry_batch(entries, base_dir=session._manifest_base_dir())
+
+    async def exec(self, *command: str | Path, **kwargs: Any) -> ExecResult:
+        await self.resolve()
+        return await super().exec(*command, **kwargs)
+
+    async def ls(self, *args: Any, **kwargs: Any) -> Any:
+        await self.resolve()
+        return await super().ls(*args, **kwargs)
+
+    async def rm(self, *args: Any, **kwargs: Any) -> Any:
+        await self.resolve()
+        return await super().rm(*args, **kwargs)
+
+    async def mkdir(self, *args: Any, **kwargs: Any) -> Any:
+        await self.resolve()
+        return await super().mkdir(*args, **kwargs)
+
+    async def read(self, path: Path, *, user: str | User | None = None) -> io.IOBase:
+        await self.resolve()
+        return await super().read(path, user=user)
+
+    async def write(self, *args: Any, **kwargs: Any) -> Any:
+        await self.resolve()
+        return await super().write(*args, **kwargs)
+
+    async def resolve_exposed_port(self, port: int) -> ExposedPortEndpoint:
+        await self.resolve()
+        return await super().resolve_exposed_port(port)
+
+    async def _resolve_exposed_port(self, port: int) -> ExposedPortEndpoint:
+        session = await self.resolve()
+        return await session._resolve_exposed_port(port)
+
+    async def pty_exec_start(self, *args: Any, **kwargs: Any) -> Any:
+        await self.resolve()
+        return await super().pty_exec_start(*args, **kwargs)
+
+    async def pty_write_stdin(self, *args: Any, **kwargs: Any) -> Any:
+        await self.resolve()
+        return await super().pty_write_stdin(*args, **kwargs)
+
+    async def pty_terminate_all(self) -> None:
+        if self._resolved is None:
+            return
+        await super().pty_terminate_all()
+
+    async def persist_workspace(self) -> io.IOBase:
+        await self.resolve()
+        return await super().persist_workspace()
+
+    async def hydrate_workspace(self, data: io.IOBase) -> None:
+        await self.resolve()
+        await super().hydrate_workspace(data)
+
+    async def teardown(self) -> None:
+        if self._resolved is None:
+            self._bring_up.cancel()
+            with contextlib.suppress(BaseException):
+                client, session = await self._bring_up
+                self._adopt(client, session)
+        if self._resolved is None:
+            return
+        await self._client.delete(self._resolved)
+
+
+class DeferredSandboxClient:
+    def __init__(self, session: DeferredSandboxSession) -> None:
+        self._session = session
+
+    def __getattr__(self, name: str) -> Any:
+        client = self._session.client
+        if client is None:
+            raise AttributeError(f"sandbox client unavailable before bring-up finishes ({name})")
+        return getattr(client, name)
+
+    async def delete(self, session: Any) -> None:
+        if isinstance(session, DeferredSandboxSession):
+            await session.teardown()
+            return
+        real = await self._session.resolve()
+        await self._session.client.delete(real if session is self._session else session)

strix/runtime/session_manager.py

@@ -19,6 +19,7 @@ from strix.config import load_settings
 from strix.runtime.backends import backend_supports_bind_mounts, get_backend
 from strix.runtime.caido_bootstrap import bootstrap_caido
 from strix.runtime.caido_handle import CaidoBootstrapHandle
+from strix.runtime.deferred_session import DeferredSandboxClient, DeferredSandboxSession
 
 
 if TYPE_CHECKING:
@@ -262,6 +263,7 @@ async def create_or_reuse(
     local_sources: list[dict[str, Any]],
     extra_files: list[dict[str, Any]] | None = None,
     status_sink: StatusSink | None = None,
+    defer: bool = False,
 ) -> dict[str, Any]:
     """Return the existing session bundle for ``scan_id`` or create a new one.
 
@@ -323,41 +325,61 @@ async def create_or_reuse(
         backend_name,
         image,
     )
-    report("Starting sandbox container")
-    client, session = await backend(
-        image=image,
-        manifest=manifest,
-        exposed_ports=(_CONTAINER_CAIDO_PORT,),
-        bind_mounts=bind_mounts,
-    )
 
-    if extra_file_archive is not None:
-        report("Placing workspace files")
-        try:
-            await stage_extra_files(session, extra_file_archive)
-        except BaseException:
-            await _discard_session(client, session)
-            raise
-
-    report("Setting up the proxy")
-    caido_endpoint = await session.resolve_exposed_port(_CONTAINER_CAIDO_PORT)
-    scheme = "https" if caido_endpoint.tls else "http"
-    host_caido_url = f"{scheme}://{caido_endpoint.host}:{caido_endpoint.port}"
-    logger.debug("Caido host endpoint resolved: %s", host_caido_url)
-
-    # The Caido login + project setup polls the guest for a couple of seconds
-    # and nothing needs the client before the first proxy tool call, so it
-    # runs concurrently with the rest of scan start; consumers resolve the
-    # handle at first use (see CaidoBootstrapHandle).
-    caido_client = CaidoBootstrapHandle(
-        asyncio.create_task(
-            bootstrap_caido(
-                session,
-                host_url=host_caido_url,
-                container_url=container_caido_url,
-            ),
-            name=f"caido-bootstrap-{scan_id}",
+    async def _bring_up() -> tuple[Any, Any]:
+        report("Starting sandbox container")
+        client, session = await backend(
+            image=image,
+            manifest=manifest,
+            exposed_ports=(_CONTAINER_CAIDO_PORT,),
+            bind_mounts=bind_mounts,
         )
+        if extra_file_archive is not None:
+            report("Placing workspace files")
+            try:
+                await stage_extra_files(session, extra_file_archive)
+            except BaseException:
+                await _discard_session(client, session)
+                raise
+        return client, session
+
+    async def _caido(session: Any) -> Any:
+        report("Setting up the proxy")
+        caido_endpoint = await session.resolve_exposed_port(_CONTAINER_CAIDO_PORT)
+        scheme = "https" if caido_endpoint.tls else "http"
+        host_caido_url = f"{scheme}://{caido_endpoint.host}:{caido_endpoint.port}"
+        logger.debug("Caido host endpoint resolved: %s", host_caido_url)
+        return await bootstrap_caido(
+            session,
+            host_url=host_caido_url,
+            container_url=container_caido_url,
+        )
+
+    if defer:
+        bring_up = asyncio.create_task(_bring_up(), name=f"sandbox-bring-up-{scan_id}")
+        deferred = DeferredSandboxSession(
+            bring_up,
+            manifest=manifest,
+            exposed_ports=(_CONTAINER_CAIDO_PORT,),
+        )
+
+        async def _caido_when_up() -> Any:
+            return await _caido(await deferred.resolve())
+
+        bundle = {
+            "client": DeferredSandboxClient(deferred),
+            "session": deferred,
+            "caido_client": CaidoBootstrapHandle(
+                asyncio.create_task(_caido_when_up(), name=f"caido-bootstrap-{scan_id}")
+            ),
+        }
+        _SESSION_CACHE[scan_id] = bundle
+        logger.info("Sandbox session for scan %s deferred and cached", scan_id)
+        return bundle
+
+    client, session = await _bring_up()
+    caido_client = CaidoBootstrapHandle(
+        asyncio.create_task(_caido(session), name=f"caido-bootstrap-{scan_id}")
     )
 
     bundle = {

tests/test_deferred_session.py

@@ -0,0 +1,265 @@
+from __future__ import annotations
+
+import asyncio
+import io
+from dataclasses import dataclass
+from pathlib import Path
+from typing import TYPE_CHECKING, Any
+
+import pytest
+from agents.run_config import SandboxRunConfig
+from agents.sandbox import SandboxAgent
+from agents.sandbox.capabilities import Filesystem, Shell
+from agents.sandbox.runtime_session_manager import SandboxRuntimeSessionManager
+
+from strix.runtime import session_manager
+from strix.runtime.backends import _BACKENDS, _BIND_MOUNT_BACKENDS, register_backend
+from strix.runtime.deferred_session import DeferredSandboxClient, DeferredSandboxSession
+
+
+if TYPE_CHECKING:
+    from collections.abc import Iterator
+
+
+@dataclass
+class _RuntimeSettings:
+    backend: str
+
+
+@dataclass
+class _Settings:
+    runtime: _RuntimeSettings
+
+
+@dataclass
+class _Endpoint:
+    host: str = "127.0.0.1"
+    port: int = 8080
+    tls: bool = False
+
+
+@dataclass
+class _ExecResult:
+    exit_code: int = 0
+    stdout: bytes = b""
+    stderr: bytes = b""
+
+    def ok(self) -> bool:
+        return self.exit_code == 0
+
+
+class _State:
+    def __init__(self) -> None:
+        self.manifest: Any = None
+        self.session_id = "00000000-0000-4000-8000-000000000001"
+
+    def model_copy(self, *, update: dict[str, Any]) -> _State:
+        copy = _State()
+        copy.manifest = update.get("manifest", self.manifest)
+        return copy
+
+
+class _Session:
+    def __init__(self) -> None:
+        self.state = _State()
+        self.execs: list[tuple[str, ...]] = []
+        self.writes: list[tuple[Path, bytes]] = []
+        self.limits: list[Any] = []
+
+    def _set_concurrency_limits(self, limits: Any) -> None:
+        self.limits.append(limits)
+
+    def _set_archive_limits(self, limits: Any) -> None:
+        self.limits.append(limits)
+
+    async def running(self) -> bool:
+        return True
+
+    async def resolve_exposed_port(self, _port: int) -> _Endpoint:
+        return _Endpoint()
+
+    async def exec(self, *argv: str, timeout: float | None = None, **_kw: Any) -> _ExecResult:
+        del timeout
+        self.execs.append(argv)
+        return _ExecResult()
+
+    async def write(self, path: Path, data: io.IOBase, **_kw: Any) -> None:
+        self.writes.append((path, data.read()))
+
+
+class _Client:
+    def __init__(self) -> None:
+        self.deleted: list[Any] = []
+
+    async def delete(self, session: Any) -> None:
+        self.deleted.append(session)
+
+
+async def _no_caido(*_args: Any, **_kwargs: Any) -> None:
+    return None
+
+
+class _Backend:
+    def __init__(self, *, fail: bool = False) -> None:
+        self.gate = asyncio.Event()
+        self.calls = 0
+        self.fail = fail
+        self.client = _Client()
+        self.session = _Session()
+
+    async def __call__(self, **_kwargs: Any) -> tuple[Any, Any]:
+        self.calls += 1
+        await self.gate.wait()
+        if self.fail:
+            raise RuntimeError("bring-up exploded")
+        return self.client, self.session
+
+
+@pytest.fixture
+def backend(monkeypatch: pytest.MonkeyPatch, request: pytest.FixtureRequest) -> Iterator[_Backend]:
+    name = f"deferred-{request.node.name}"
+    fake = _Backend(fail="fail" in request.node.name)
+    register_backend(name, fake, supports_bind_mounts=True)
+    monkeypatch.setattr(session_manager, "bootstrap_caido", _no_caido)
+    monkeypatch.setattr(
+        session_manager, "load_settings", lambda: _Settings(runtime=_RuntimeSettings(backend=name))
+    )
+    yield fake
+    _BACKENDS.pop(name, None)
+    _BIND_MOUNT_BACKENDS.discard(name)
+
+
+async def _create(scan_id: str, **kwargs: Any) -> dict[str, Any]:
+    return await session_manager.create_or_reuse(
+        scan_id, image="img", local_sources=[], defer=True, **kwargs
+    )
+
+
+@pytest.mark.asyncio
+async def test_deferred_bundle_returns_before_the_backend_finishes(backend: _Backend) -> None:
+    scan_id = "deferred-returns"
+    try:
+        bundle = await asyncio.wait_for(_create(scan_id), timeout=1)
+        session = bundle["session"]
+        assert isinstance(session, DeferredSandboxSession)
+        assert isinstance(bundle["client"], DeferredSandboxClient)
+        assert session.pending
+        await asyncio.sleep(0)
+        assert backend.calls == 1
+        assert await session.running()
+        assert session.state.manifest.environment.value["PYTHONUNBUFFERED"] == "1"
+        again = await session_manager.create_or_reuse(scan_id, image="img", local_sources=[])
+        assert again is bundle
+        assert backend.calls == 1
+    finally:
+        await session_manager.cleanup(scan_id)
+
+
+@pytest.mark.asyncio
+async def test_first_io_resolves_the_sandbox_then_delegates(backend: _Backend) -> None:
+    scan_id = "deferred-io"
+    try:
+        bundle = await _create(scan_id)
+        session = bundle["session"]
+        session._set_concurrency_limits("limits")
+        exec_task = asyncio.create_task(session.exec("echo", "hi"))
+        await asyncio.sleep(0.05)
+        assert not exec_task.done()
+        backend.gate.set()
+        result = await asyncio.wait_for(exec_task, timeout=1)
+        assert result.ok()
+        assert backend.session.execs == [("echo", "hi")]
+        assert not session.pending
+        assert backend.session.limits == ["limits"]
+        assert backend.session.state.manifest is session.state.manifest
+        await session.write(Path("/workspace/x"), io.BytesIO(b"data"))
+        assert backend.session.writes == [(Path("/workspace/x"), b"data")]
+    finally:
+        await session_manager.cleanup(scan_id)
+    assert backend.client.deleted == [backend.session]
+
+
+@pytest.mark.asyncio
+async def test_failed_bring_up_surfaces_on_first_io(backend: _Backend) -> None:
+    scan_id = "deferred-fail"
+    try:
+        bundle = await _create(scan_id)
+        session = bundle["session"]
+        backend.gate.set()
+        with pytest.raises(RuntimeError, match="bring-up exploded"):
+            await asyncio.wait_for(session.exec("true"), timeout=1)
+        assert not await session.running()
+        with pytest.raises(RuntimeError, match="bring-up exploded"):
+            await bundle["caido_client"].get()
+    finally:
+        await session_manager.cleanup(scan_id)
+    assert backend.client.deleted == []
+
+
+@pytest.mark.asyncio
+async def test_cleanup_while_pending_cancels_the_bring_up(backend: _Backend) -> None:
+    scan_id = "deferred-cancel"
+    bundle = await _create(scan_id)
+    session = bundle["session"]
+    await asyncio.wait_for(session_manager.cleanup(scan_id), timeout=1)
+    assert session.bring_up.cancelled()
+    assert backend.client.deleted == []
+    assert session_manager._SESSION_CACHE.get(scan_id) is None
+
+
+@pytest.mark.asyncio
+async def test_extra_files_are_staged_inside_the_deferred_bring_up(backend: _Backend) -> None:
+    scan_id = "deferred-extra"
+    try:
+        bundle = await _create(
+            scan_id,
+            extra_files=[{"workspace_path": "/workspace/.strix/notes.md", "content": "hi"}],
+        )
+        assert backend.session.writes == []
+        backend.gate.set()
+        await asyncio.wait_for(bundle["session"].resolve(), timeout=1)
+        assert len(backend.session.writes) == 1
+        assert backend.session.execs
+    finally:
+        await session_manager.cleanup(scan_id)
+
+
+@pytest.mark.asyncio
+async def test_default_path_still_blocks_until_the_sandbox_is_up(backend: _Backend) -> None:
+    scan_id = "deferred-off"
+    backend.gate.set()
+    try:
+        bundle = await session_manager.create_or_reuse(scan_id, image="img", local_sources=[])
+        assert bundle["session"] is backend.session
+        assert bundle["client"] is backend.client
+    finally:
+        await session_manager.cleanup(scan_id)
+    assert backend.client.deleted == [backend.session]
+
+
+@pytest.mark.asyncio
+async def test_sdk_session_prep_does_not_wait_for_the_bring_up(backend: _Backend) -> None:
+    scan_id = "deferred-sdk-prep"
+    try:
+        bundle = await _create(scan_id)
+        agent = SandboxAgent(
+            name="root", instructions="x", tools=[], capabilities=[Filesystem(), Shell()]
+        )
+        manager = SandboxRuntimeSessionManager(
+            starting_agent=agent,
+            sandbox_config=SandboxRunConfig(client=bundle["client"], session=bundle["session"]),
+            run_state=None,
+        )
+        session = await asyncio.wait_for(
+            manager.ensure_session(
+                agent=agent, capabilities=list(agent.capabilities), is_resumed_state=False
+            ),
+            timeout=1,
+        )
+        assert session is bundle["session"]
+        assert session.pending
+        backend.gate.set()
+        await asyncio.wait_for(session.resolve(), timeout=1)
+        assert not session.pending
+    finally:
+        await session_manager.cleanup(scan_id)