diff --git a/strix/runtime/deferred_session.py b/strix/runtime/deferred_session.py new file mode 100644 index 00000000..b2f9a7d0 --- /dev/null +++ b/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) diff --git a/strix/runtime/session_manager.py b/strix/runtime/session_manager.py index a8447ebb..cdb9a035 100644 --- a/strix/runtime/session_manager.py +++ b/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 = { diff --git a/tests/test_deferred_session.py b/tests/test_deferred_session.py new file mode 100644 index 00000000..12d16725 --- /dev/null +++ b/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)