diff --git a/strix/runtime/backends.py b/strix/runtime/backends.py index 0cb2218fd..763c40282 100644 --- a/strix/runtime/backends.py +++ b/strix/runtime/backends.py @@ -57,6 +57,8 @@ _BACKENDS: dict[str, SandboxBackend] = { _BIND_MOUNT_BACKENDS: set[str] = {"docker"} +_PTY_BACKENDS: set[str] = {"docker"} + def get_backend(name: str) -> SandboxBackend: """Return the backend factory for ``name`` or raise. @@ -81,6 +83,7 @@ def register_backend( backend: SandboxBackend, *, supports_bind_mounts: bool = False, + supports_pty: bool = False, ) -> None: """Register a custom backend under ``name``. @@ -89,18 +92,34 @@ def register_backend( an existing name overwrites the prior entry. ``supports_bind_mounts`` defaults to False: a remote runtime cannot see the caller's filesystem, so it is handed local sources as manifest entries to upload instead. + ``supports_pty`` declares whether the backend's sessions support PTY exec; + deferred startup reports it before the session exists so the shell tools + pick the same mode they would after bring-up. """ _BACKENDS[name] = backend if supports_bind_mounts: _BIND_MOUNT_BACKENDS.add(name) else: _BIND_MOUNT_BACKENDS.discard(name) - logger.info("Registered sandbox backend: %s (bind mounts: %s)", name, supports_bind_mounts) + if supports_pty: + _PTY_BACKENDS.add(name) + else: + _PTY_BACKENDS.discard(name) + logger.info( + "Registered sandbox backend: %s (bind mounts: %s, pty: %s)", + name, + supports_bind_mounts, + supports_pty, + ) def backend_supports_bind_mounts(name: str) -> bool: return name in _BIND_MOUNT_BACKENDS +def backend_supports_pty(name: str) -> bool: + return name in _PTY_BACKENDS + + def supported_backends() -> list[str]: return sorted(_BACKENDS) diff --git a/strix/runtime/deferred_session.py b/strix/runtime/deferred_session.py index b2f9a7d02..5ce6a2db9 100644 --- a/strix/runtime/deferred_session.py +++ b/strix/runtime/deferred_session.py @@ -37,8 +37,10 @@ class DeferredSandboxSession(SandboxSession): *, manifest: Manifest, exposed_ports: tuple[int, ...] = (), + supports_pty: bool = False, ) -> None: self._bring_up = bring_up + self._pending_supports_pty = supports_pty self._pending_state: SandboxSessionState = _PendingSessionState( snapshot=NoopSnapshot(id="pending"), manifest=manifest, @@ -126,7 +128,9 @@ class DeferredSandboxSession(SandboxSession): 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 + if self._resolved is None: + return self._pending_supports_pty + return self._resolved.supports_pty() async def running(self) -> bool: if self._resolved is None: diff --git a/strix/runtime/session_manager.py b/strix/runtime/session_manager.py index d92b1fdbd..86020baa0 100644 --- a/strix/runtime/session_manager.py +++ b/strix/runtime/session_manager.py @@ -18,7 +18,11 @@ from agents.sandbox.entries import BaseEntry, LocalDir from agents.sandbox.manifest import Environment, Manifest from strix.config import load_settings -from strix.runtime.backends import backend_supports_bind_mounts, get_backend +from strix.runtime.backends import ( + backend_supports_bind_mounts, + backend_supports_pty, + 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 @@ -378,6 +382,7 @@ async def create_or_reuse( bring_up, manifest=manifest, exposed_ports=(_CONTAINER_CAIDO_PORT,), + supports_pty=backend_supports_pty(backend_name), ) async def _caido_when_up() -> Any: diff --git a/tests/test_deferred_session.py b/tests/test_deferred_session.py index 774426fb0..f8f46019e 100644 --- a/tests/test_deferred_session.py +++ b/tests/test_deferred_session.py @@ -2,6 +2,7 @@ from __future__ import annotations import asyncio import io +import tempfile from dataclasses import dataclass from pathlib import Path from typing import TYPE_CHECKING, Any @@ -10,10 +11,18 @@ import pytest from agents.run_config import SandboxRunConfig from agents.sandbox import SandboxAgent from agents.sandbox.capabilities import Filesystem, Shell +from agents.sandbox.capabilities.tools.shell_tool import ExecCommandArgs, ExecCommandTool +from agents.sandbox.manifest import Manifest from agents.sandbox.runtime_session_manager import SandboxRuntimeSessionManager +from agents.sandbox.sandboxes.unix_local import UnixLocalSandboxClient from strix.runtime import session_manager -from strix.runtime.backends import _BACKENDS, _BIND_MOUNT_BACKENDS, register_backend +from strix.runtime.backends import ( + _BACKENDS, + _BIND_MOUNT_BACKENDS, + _PTY_BACKENDS, + register_backend, +) from strix.runtime.deferred_session import DeferredSandboxClient, DeferredSandboxSession @@ -119,7 +128,7 @@ class _Backend: 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) + register_backend(name, fake, supports_bind_mounts=True, supports_pty="pty" in request.node.name) monkeypatch.setattr(session_manager, "bootstrap_caido", _no_caido) monkeypatch.setattr( session_manager, "load_settings", lambda: _Settings(runtime=_RuntimeSettings(backend=name)) @@ -127,6 +136,7 @@ def backend(monkeypatch: pytest.MonkeyPatch, request: pytest.FixtureRequest) -> yield fake _BACKENDS.pop(name, None) _BIND_MOUNT_BACKENDS.discard(name) + _PTY_BACKENDS.discard(name) async def _settle() -> None: @@ -314,3 +324,64 @@ async def test_sdk_session_prep_does_not_wait_for_the_bring_up(backend: _Backend assert not session.pending finally: await session_manager.cleanup(scan_id) + + +@pytest.mark.asyncio +@pytest.mark.usefixtures("backend") +async def test_pending_session_reports_the_backends_pty_support() -> None: + scan_id = "deferred-pty" + try: + bundle = await _create(scan_id) + session = bundle["session"] + assert session.pending + assert session.supports_pty() + assert [tool.name for tool in Shell(session=session).tools()] == [ + "exec_command", + "write_stdin", + ] + finally: + await session_manager.cleanup(scan_id) + + +@pytest.mark.asyncio +@pytest.mark.usefixtures("backend") +async def test_pending_session_without_declared_support_stays_one_shot() -> None: + scan_id = "deferred-one-shot" + try: + bundle = await _create(scan_id) + session = bundle["session"] + assert session.pending + assert not session.supports_pty() + assert [tool.name for tool in Shell(session=session).tools()] == ["exec_command"] + finally: + await session_manager.cleanup(scan_id) + + +@pytest.mark.asyncio +async def test_first_exec_while_pending_keeps_the_pty_session_id() -> None: + with tempfile.TemporaryDirectory() as directory: + manifest = Manifest(root=directory) + client = UnixLocalSandboxClient() + actual = await client.create(manifest=manifest) + await actual.start() + gate = asyncio.Event() + + async def _bring_up() -> tuple[Any, Any]: + await gate.wait() + return client, actual + + deferred = DeferredSandboxSession( + asyncio.create_task(_bring_up()), manifest=manifest, supports_pty=True + ) + args = ExecCommandArgs(cmd="sleep 2; printf finished", yield_time_ms=5, login=False) + try: + first = asyncio.create_task(ExecCommandTool(session=deferred).run(args)) + await asyncio.sleep(0) + gate.set() + output = await asyncio.wait_for(first, timeout=5) + assert "Process running with session ID" in output + assert "timed out" not in output + finally: + await actual.pty_terminate_all() + await actual.aclose() + await client.delete(actual)