From 71d798344cf8a0f795592c30a0ed24a951bd1b2d Mon Sep 17 00:00:00 2001 From: pranav-afk Date: Wed, 15 Jul 2026 12:34:05 +0530 Subject: [PATCH] perf(runtime): import local sources via single tar put_archive Replace SDK per-file LocalDir copy with temp-file tar + put_archive after session start. Includes greptile fixes: disk spill, cleanup on failure, walk onerror, image-aware ownership. --- strix/runtime/backends.py | 15 +- strix/runtime/docker_client.py | 5 +- strix/runtime/session_manager.py | 354 +++++++++++++++++++++++++---- tests/test_import_local_sources.py | 207 +++++++++++++++++ tests/test_session_entries.py | 215 +++++++++++++++--- 5 files changed, 707 insertions(+), 89 deletions(-) create mode 100644 tests/test_import_local_sources.py diff --git a/strix/runtime/backends.py b/strix/runtime/backends.py index b73fdbad3..f203de15c 100644 --- a/strix/runtime/backends.py +++ b/strix/runtime/backends.py @@ -155,12 +155,15 @@ async def _docker_backend( ``docker`` lazily so deployments that target a non-Docker backend don't need the docker-py library installed. - ``session.start()`` is what materializes the manifest entries - (LocalDir copies and manifest-declared volume/FUSE mounts) into the - running container — the SDK's ``client.create()`` only builds the inner - session object without applying the manifest. ``async with session:`` - would call it too, but Strix manages session lifetime explicitly via - ``client.delete()`` so we trigger ``start()`` ourselves. + ``session.start()`` materializes any manifest-declared volume/FUSE + mounts into the running container. Local source trees are no longer + carried as SDK ``LocalDir`` entries (that path is per-file and hangs + on large repos); the session manager imports them after ``start()`` + via a single tar ``put_archive``. The SDK's ``client.create()`` only + builds the inner session object without applying the manifest. + ``async with session:`` would call ``start()`` too, but Strix manages + session lifetime explicitly via ``client.delete()`` so we trigger + ``start()`` ourselves. ``bind_mounts`` are host directories (e.g. large repos passed via ``--mount``) bind-mounted read-only; unlike manifest entries they are diff --git a/strix/runtime/docker_client.py b/strix/runtime/docker_client.py index 2321fa50d..7710353fa 100644 --- a/strix/runtime/docker_client.py +++ b/strix/runtime/docker_client.py @@ -117,8 +117,9 @@ async def _create_container( extra_hosts = create_kwargs.setdefault("extra_hosts", {}) extra_hosts["host.docker.internal"] = "host-gateway" - # Strix injection: host bind mounts (e.g. large repos passed via --mount) - # that bypass the SDK's file-by-file LocalDir copy. + # Strix injection: host bind mounts (e.g. large repos passed via --mount). + # Non-mount local sources are imported after start via tar put_archive + # (session_manager._import_local_sources), not via SDK LocalDir. bind_mounts = getattr(self, "strix_bind_mounts", ()) if bind_mounts: mounts = create_kwargs.setdefault("mounts", []) diff --git a/strix/runtime/session_manager.py b/strix/runtime/session_manager.py index 5aa714238..137a2ddc2 100644 --- a/strix/runtime/session_manager.py +++ b/strix/runtime/session_manager.py @@ -2,18 +2,20 @@ from __future__ import annotations +import asyncio +import contextlib import logging -import shutil +import os +import tarfile +import tempfile from pathlib import Path from typing import Any -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 get_backend from strix.runtime.caido_bootstrap import bootstrap_caido -from strix.runtime.local_dir_staging import stage_symlink_safe_dir logger = logging.getLogger(__name__) @@ -25,31 +27,46 @@ _SESSION_CACHE: dict[str, dict[str, Any]] = {} -# Manifest root inside the container; entry keys hang off this path. +# Workspace root inside the container; sources land at ``/``. _WORKSPACE_ROOT = "/workspace" +# Preferred non-root account on the default Strix sandbox image. +_DEFAULT_SANDBOX_USER = "pentester" -def build_session_entries( + +def _is_safe_workspace_subdir(ws_subdir: str) -> bool: + """Reject subdirs that would escape ``/workspace`` when joined. + + ``ws_subdir`` becomes both a tar member prefix and a ``chown`` target, so a + value containing ``..`` (or an absolute path) could write outside the + intended ``/workspace/`` tree. Callers pass values with surrounding + slashes already stripped. + """ + if not ws_subdir: + return False + return not any(part in ("..", "") for part in ws_subdir.split("/")) + + +def split_local_sources( local_sources: list[dict[str, Any]], -) -> tuple[dict[str | Path, BaseEntry], list[dict[str, Any]], list[Path]]: - """Split local sources into copied manifest entries and host bind mounts. +) -> tuple[list[dict[str, str]], list[dict[str, Any]]]: + """Split local sources into tar-copied entries and host bind mounts. Sources flagged ``mount`` are bind-mounted read-only at - ``/workspace/`` (not added to the manifest, so the SDK - does not stream them in file-by-file). Every other source becomes a - ``LocalDir`` entry copied into the container as before. Trees containing - symlinks (which the SDK's ``LocalDir`` walker refuses outright) are first - staged into a symlink-safe temp copy; those temp dirs are returned so the - caller can remove them once the upload completes. + ``/workspace/`` (applied by Docker at container-create + time). Every other source is copied into the container after start via a + single tar ``put_archive`` (see :func:`_import_local_sources`). """ - entries: dict[str | Path, BaseEntry] = {} + copied: list[dict[str, str]] = [] bind_mounts: list[dict[str, Any]] = [] - staged_dirs: list[Path] = [] for src in local_sources: - ws_subdir = src.get("workspace_subdir") or "" + ws_subdir = (src.get("workspace_subdir") or "").strip("/") host_path = src.get("source_path") or "" if not ws_subdir or not host_path: continue + if not _is_safe_workspace_subdir(ws_subdir): + logger.warning("Skipping local source with unsafe workspace_subdir: %r", ws_subdir) + continue resolved = Path(host_path).expanduser().resolve() if src.get("mount"): bind_mounts.append( @@ -60,11 +77,241 @@ def build_session_entries( } ) else: - upload_path, staged = stage_symlink_safe_dir(resolved) - if staged is not None: - staged_dirs.append(staged) - entries[ws_subdir] = LocalDir(src=upload_path) - return entries, bind_mounts, staged_dirs + copied.append({"source_path": str(resolved), "workspace_subdir": ws_subdir}) + return copied, bind_mounts + + +def _walk_onerror(err: OSError) -> None: + """Fail the import when a directory cannot be listed. + + Default ``os.walk`` swallows listdir errors and silently omits subtrees, + which would leave the agent with an incomplete ``/workspace`` and no signal + that files are missing. + """ + raise err + + +def _build_source_tar(src_root: Path, arc_prefix: str) -> tuple[Path, int, int]: + """Pack ``src_root`` into a temporary tar file rooted at ``arc_prefix``. + + Returns ``(tar_path, added, skipped)``. The caller must delete ``tar_path``. + Spills to disk (not an in-memory ``BytesIO``) so multi-GB trees do not OOM + the host process before Docker receives the archive. + + Regular files and directories — including dotfiles such as ``.git`` and + empty dirs — are packed as-is. Symlinks are skipped (and counted) rather + than followed. Unreadable directories raise via :func:`_walk_onerror`. + """ + tmp = tempfile.NamedTemporaryFile( # noqa: SIM115 - closed explicitly below + prefix="strix-src-", + suffix=".tar", + delete=False, + ) + tar_path = Path(tmp.name) + added = 0 + skipped = 0 + try: + with tarfile.open(fileobj=tmp, mode="w") as tar: + for dirpath, dirnames, filenames in os.walk( + src_root, + followlinks=False, + onerror=_walk_onerror, + ): + kept_dirs: list[str] = [] + for name in dirnames: + if Path(dirpath, name).is_symlink(): + skipped += 1 + continue + kept_dirs.append(name) + dirnames[:] = kept_dirs + + dir_abs = Path(dirpath) + rel = dir_abs.relative_to(src_root).as_posix() + dir_arcname = arc_prefix if rel == "." else f"{arc_prefix}/{rel}" + tar.add(str(dir_abs), arcname=dir_arcname, recursive=False) + + for name in filenames: + full = dir_abs / name + if full.is_symlink() or not full.is_file(): + skipped += 1 + continue + arcname = f"{arc_prefix}/{full.relative_to(src_root).as_posix()}" + tar.add(str(full), arcname=arcname, recursive=False) + added += 1 + tmp.close() + except Exception: + tmp.close() + with contextlib.suppress(OSError): + tar_path.unlink(missing_ok=True) + raise + return tar_path, added, skipped + + +def _container_of(session: Any) -> Any: + """Reach the underlying docker-py ``Container`` from an SDK session. + + The SDK wraps the backend session in an outer object; the docker backend + exposes the real container as ``_inner._container``. Pinned to + openai-agents==0.14.6 — re-check these private attrs on SDK bumps. + """ + inner = getattr(session, "_inner", session) + container = getattr(inner, "_container", None) + if container is None: + raise RuntimeError("could not locate docker container on sandbox session") + return container + + +def _run_checked(container: Any, cmd: list[str]) -> None: + """Run ``cmd`` in the container as root and raise on non-zero exit.""" + result = container.exec_run(cmd, user="root") + if result.exit_code: + output = result.output + detail = output.decode("utf-8", "replace").strip() if isinstance(output, bytes) else output + raise RuntimeError(f"container command {cmd!r} failed (exit {result.exit_code}): {detail}") + + +def _decode_exec_output(output: Any) -> str: + if isinstance(output, bytes): + return output.decode("utf-8", "replace").strip() + return str(output or "").strip() + + +def _detect_runtime_owner(container: Any) -> str | None: + """Return ``user:group`` for workspace ownership, or None if unknown. + + Prefers the image ``Config.User`` (works for custom images), then probes + for the default Strix sandbox account. Never assumes a fixed username is + always present. + """ + try: + attrs = getattr(container, "attrs", None) + if not isinstance(attrs, dict) and hasattr(container, "reload"): + with contextlib.suppress(Exception): + container.reload() + attrs = getattr(container, "attrs", None) + cfg_user = str(((attrs or {}).get("Config") or {}).get("User") or "").strip() + if cfg_user and cfg_user not in {"0", "root", "0:0"}: + if ":" in cfg_user: + return cfg_user + return f"{cfg_user}:{cfg_user}" + except Exception: # noqa: BLE001 + logger.debug("Could not read container Config.User", exc_info=True) + + # Probe preferred default only when it actually exists in this image. + probe = container.exec_run( + [ + "sh", + "-c", + f"id -u {_DEFAULT_SANDBOX_USER} >/dev/null 2>&1 " + f"&& printf '%s' {_DEFAULT_SANDBOX_USER}", + ], + user="root", + ) + if probe.exit_code == 0 and _decode_exec_output(probe.output) == _DEFAULT_SANDBOX_USER: + return f"{_DEFAULT_SANDBOX_USER}:{_DEFAULT_SANDBOX_USER}" + return None + + +def _fix_workspace_ownership(container: Any, path: str) -> None: + """Best-effort chown of imported sources to the container runtime user. + + ``put_archive`` extracts as root with host UIDs. On the default image we + reassign to ``pentester``; on custom images we use ``Config.User`` when + set. Missing users or chown failures log a warning instead of aborting + import — the container is still usable (often as root). + """ + owner = _detect_runtime_owner(container) + if owner is None: + logger.warning( + "No non-root container user detected; leaving %s root-owned " + "(set USER in the sandbox image or create a runtime account)", + path, + ) + return + result = container.exec_run(["chown", "-R", owner, path], user="root") + if result.exit_code: + detail = _decode_exec_output(result.output) + logger.warning( + "chown %s to %s failed (exit %s): %s — continuing; tools may lack write access", + path, + owner, + result.exit_code, + detail, + ) + + +def _put_source_archive(container: Any, tar_path: Path, subdir: str) -> None: + """Stream a tar file into the container and fix ownership.""" + _run_checked(container, ["mkdir", "-p", _WORKSPACE_ROOT]) + with tar_path.open("rb") as tar_stream: + if not container.put_archive(_WORKSPACE_ROOT, tar_stream): + raise RuntimeError(f"put_archive failed for {_WORKSPACE_ROOT}/{subdir}") + _fix_workspace_ownership(container, f"{_WORKSPACE_ROOT}/{subdir}") + + +async def _import_local_sources( + session: Any, + copied_sources: list[dict[str, str]], +) -> None: + """Copy each host source tree into the container via a single tar import. + + Replaces the SDK's per-file ``LocalDir`` materialization, which issues + several ``docker exec`` calls per file. On large trees that serializes into + thousands of exec round-trips against a non-thread-safe docker-py client, + causing multi-minute hangs (the "stuck on loading" symptom) and occasional + ``ExecTransportError``. ``put_archive`` lands the whole tree in one shot. + + Archives are written to a temporary file (not fully buffered in RAM) so + large repos do not OOM the host process. + """ + if not copied_sources: + return + + container = _container_of(session) + loop = asyncio.get_running_loop() + + for src in copied_sources: + ws_subdir = src["workspace_subdir"] + src_root = Path(src["source_path"]) + if not src_root.is_dir(): + logger.warning("Skipping non-directory local source: %s", src_root) + continue + + tar_path, added, skipped = await loop.run_in_executor( + None, + _build_source_tar, + src_root, + ws_subdir, + ) + logger.info( + "Importing %d files into %s/%s (skipped %d symlink entries)", + added, + _WORKSPACE_ROOT, + ws_subdir, + skipped, + ) + try: + await loop.run_in_executor( + None, + _put_source_archive, + container, + tar_path, + ws_subdir, + ) + finally: + with contextlib.suppress(OSError): + tar_path.unlink(missing_ok=True) + + +async def _teardown_session(client: Any, session: Any) -> None: + """Best-effort delete of a session that never made it into the cache.""" + try: + await client.delete(session) + except Exception: + logger.exception( + "Failed to tear down sandbox session after setup error; " + "container may need manual reaping", + ) async def create_or_reuse( @@ -76,16 +323,22 @@ async def create_or_reuse( """Return the existing session bundle for ``scan_id`` or create a new one. Each ``local_sources`` entry exposes its host ``source_path`` at - ``/workspace/`` inside the container — copied in, or - bind-mounted read-only when the entry is flagged ``mount``. + ``/workspace/`` inside the container — copied in via a + single tar ``put_archive`` after start, or bind-mounted read-only when the + entry is flagged ``mount``. """ cached = _SESSION_CACHE.get(scan_id) if cached is not None: logger.info("Reusing existing sandbox session for scan %s", scan_id) return cached - entries, bind_mounts, staged_dirs = build_session_entries(local_sources) + copied_sources, bind_mounts = split_local_sources(local_sources) + # Copied source trees are imported after start() via a single tar + # ``put_archive`` (see ``_import_local_sources``) instead of the SDK's + # per-file ``LocalDir`` exec loop, which hangs on large trees. So the + # manifest itself carries no source entries. + # # Caido runs as an in-container sidecar; HTTP(S) traffic from any # process started via ``session.exec`` (the SDK's Shell tool, etc.) # picks up these env vars automatically. ``NO_PROXY`` keeps the @@ -93,7 +346,7 @@ async def create_or_reuse( # through Caido. container_caido_url = f"http://127.0.0.1:{_CONTAINER_CAIDO_PORT}" manifest = Manifest( - entries=entries, + entries={}, environment=Environment( value={ "PYTHONUNBUFFERED": "1", @@ -110,32 +363,43 @@ async def create_or_reuse( backend = get_backend(backend_name) logger.info( - "Creating sandbox session for scan %s (backend=%s, image=%s)", + "Creating sandbox session for scan %s (backend=%s, image=%s, copy_sources=%d, mounts=%d)", scan_id, backend_name, image, + len(copied_sources), + len(bind_mounts), ) + client, session = await backend( + image=image, + manifest=manifest, + exposed_ports=(_CONTAINER_CAIDO_PORT,), + bind_mounts=bind_mounts, + ) + + # Import / Caido bootstrap can fail after the container is already + # running. Tear it down before re-raising so we never leak an uncached + # container that ``cleanup(scan_id)`` cannot see. try: - client, session = await backend( - image=image, - manifest=manifest, - exposed_ports=(_CONTAINER_CAIDO_PORT,), - bind_mounts=bind_mounts, + await _import_local_sources(session, copied_sources) + + 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) + + caido_client = await bootstrap_caido( + session, + host_url=host_caido_url, + container_url=container_caido_url, ) - finally: - for staged in staged_dirs: - shutil.rmtree(staged, ignore_errors=True) - - 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) - - caido_client = await bootstrap_caido( - session, - host_url=host_caido_url, - container_url=container_caido_url, - ) + except Exception: + logger.exception( + "Sandbox setup failed for scan %s after container start; tearing down", + scan_id, + ) + await _teardown_session(client, session) + raise bundle = { "client": client, diff --git a/tests/test_import_local_sources.py b/tests/test_import_local_sources.py new file mode 100644 index 000000000..82e42175c --- /dev/null +++ b/tests/test_import_local_sources.py @@ -0,0 +1,207 @@ +"""Unit tests for tar-based local source import (no live Docker required).""" + +from __future__ import annotations + +from pathlib import Path +from types import SimpleNamespace +from typing import Any +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest + +from strix.runtime import session_manager + + +class _FakeExecResult: + def __init__(self, exit_code: int = 0, output: bytes = b"") -> None: + self.exit_code = exit_code + self.output = output + + +class _FakeContainer: + def __init__( + self, + *, + put_ok: bool = True, + chown_exit: int = 0, + config_user: str = "pentester", + pentester_exists: bool = True, + ) -> None: + self.put_ok = put_ok + self.chown_exit = chown_exit + self.pentester_exists = pentester_exists + self.attrs = {"Config": {"User": config_user}} + self.put_calls: list[tuple[str, bytes]] = [] + self.exec_calls: list[list[str]] = [] + + def put_archive(self, path: str, data: Any) -> bool: + payload = data.read() if hasattr(data, "read") else data + self.put_calls.append((path, payload)) + return self.put_ok + + def exec_run(self, cmd: list[str], user: str = "root") -> _FakeExecResult: # noqa: ARG002 + self.exec_calls.append(list(cmd)) + if cmd and cmd[0] == "chown": + return _FakeExecResult(self.chown_exit, b"chown failed" if self.chown_exit else b"") + # Probe: id -u pentester ... + joined = " ".join(cmd) + if "id -u pentester" in joined: + if self.pentester_exists: + return _FakeExecResult(0, b"pentester") + return _FakeExecResult(1, b"") + return _FakeExecResult(0, b"") + + +def _session_with(container: Any) -> SimpleNamespace: + return SimpleNamespace(_inner=SimpleNamespace(_container=container)) + + +@pytest.mark.asyncio +async def test_import_local_sources_put_archive_and_chown(tmp_path: Path) -> None: + (tmp_path / "a.txt").write_text("hello", encoding="utf-8") + container = _FakeContainer() + session = _session_with(container) + + await session_manager._import_local_sources( + session, + [{"source_path": str(tmp_path), "workspace_subdir": "repo"}], + ) + + assert len(container.put_calls) == 1 + path, payload = container.put_calls[0] + assert path == "/workspace" + assert payload # non-empty tar + assert any(cmd[0] == "mkdir" for cmd in container.exec_calls) + chown_cmds = [cmd for cmd in container.exec_calls if cmd[0] == "chown"] + assert chown_cmds == [ + ["chown", "-R", "pentester:pentester", "/workspace/repo"], + ] + + +@pytest.mark.asyncio +async def test_import_local_sources_raises_on_put_failure(tmp_path: Path) -> None: + (tmp_path / "a.txt").write_text("hello", encoding="utf-8") + container = _FakeContainer(put_ok=False) + session = _session_with(container) + + with pytest.raises(RuntimeError, match="put_archive failed"): + await session_manager._import_local_sources( + session, + [{"source_path": str(tmp_path), "workspace_subdir": "repo"}], + ) + + +@pytest.mark.asyncio +async def test_import_continues_when_chown_fails(tmp_path: Path) -> None: + """Custom images / missing users must not abort import after put_archive.""" + (tmp_path / "a.txt").write_text("hello", encoding="utf-8") + container = _FakeContainer(chown_exit=1) + session = _session_with(container) + + await session_manager._import_local_sources( + session, + [{"source_path": str(tmp_path), "workspace_subdir": "repo"}], + ) + + assert len(container.put_calls) == 1 + assert any(cmd[0] == "chown" for cmd in container.exec_calls) + + +@pytest.mark.asyncio +async def test_import_uses_config_user_for_chown(tmp_path: Path) -> None: + (tmp_path / "a.txt").write_text("hello", encoding="utf-8") + container = _FakeContainer(config_user="agent", pentester_exists=False) + session = _session_with(container) + + await session_manager._import_local_sources( + session, + [{"source_path": str(tmp_path), "workspace_subdir": "repo"}], + ) + + chown_cmds = [cmd for cmd in container.exec_calls if cmd[0] == "chown"] + assert chown_cmds == [["chown", "-R", "agent:agent", "/workspace/repo"]] + + +@pytest.mark.asyncio +async def test_import_skips_chown_when_no_owner(tmp_path: Path) -> None: + (tmp_path / "a.txt").write_text("hello", encoding="utf-8") + container = _FakeContainer(config_user="", pentester_exists=False) + session = _session_with(container) + + await session_manager._import_local_sources( + session, + [{"source_path": str(tmp_path), "workspace_subdir": "repo"}], + ) + + assert len(container.put_calls) == 1 + assert not any(cmd[0] == "chown" for cmd in container.exec_calls) + + +@pytest.mark.asyncio +async def test_import_skips_non_directory(tmp_path: Path) -> None: + file_path = tmp_path / "file.txt" + file_path.write_text("x", encoding="utf-8") + container = _FakeContainer() + session = _session_with(container) + + await session_manager._import_local_sources( + session, + [{"source_path": str(file_path), "workspace_subdir": "repo"}], + ) + + assert container.put_calls == [] + + +@pytest.mark.asyncio +async def test_create_or_reuse_tears_down_on_import_failure() -> None: + client = MagicMock() + client.delete = AsyncMock() + session = MagicMock() + + async def _backend(**_kwargs: Any) -> tuple[Any, Any]: + return client, session + + with ( + patch("strix.runtime.session_manager.load_settings") as settings, + patch("strix.runtime.session_manager.get_backend", return_value=_backend), + patch( + "strix.runtime.session_manager._import_local_sources", + new_callable=AsyncMock, + side_effect=RuntimeError("put failed"), + ), + ): + settings.return_value.runtime.backend = "docker" + with pytest.raises(RuntimeError, match="put failed"): + await session_manager.create_or_reuse( + "scan-leak-test", + image="test-image", + local_sources=[ + {"source_path": str(Path.cwd() / "missing-src"), "workspace_subdir": "repo"} + ], + ) + + client.delete.assert_awaited_once_with(session) + assert "scan-leak-test" not in session_manager._SESSION_CACHE + + +def test_container_of_missing_raises() -> None: + with pytest.raises(RuntimeError, match="could not locate docker container"): + session_manager._container_of(SimpleNamespace()) + + +def test_container_of_reads_inner() -> None: + container = MagicMock(name="container") + session = _session_with(container) + assert session_manager._container_of(session) is container + + +def test_build_source_tar_uses_temp_file_not_bytes(tmp_path: Path) -> None: + (tmp_path / "a.txt").write_text("hello", encoding="utf-8") + tar_path, added, skipped = session_manager._build_source_tar(tmp_path, "repo") + try: + assert tar_path.is_file() + assert added == 1 + assert skipped == 0 + assert tar_path.stat().st_size > 0 + finally: + tar_path.unlink(missing_ok=True) diff --git a/tests/test_session_entries.py b/tests/test_session_entries.py index e08c59aa1..a8fbea473 100644 --- a/tests/test_session_entries.py +++ b/tests/test_session_entries.py @@ -1,37 +1,71 @@ -"""Tests for build_session_entries: splitting copied vs bind-mounted sources.""" +"""Tests for local-source handling in session_manager. + +Covers splitting copied vs bind-mounted sources (``split_local_sources``) and +the on-disk tar builder (``_build_source_tar``) that replaces the SDK's +per-file ``LocalDir`` copy. +""" from __future__ import annotations +import io +import os +import tarfile +from contextlib import contextmanager from typing import TYPE_CHECKING, Any +from unittest.mock import patch -from agents.sandbox.entries import LocalDir +import pytest -from strix.runtime.session_manager import build_session_entries +from strix.runtime.session_manager import _build_source_tar, split_local_sources if TYPE_CHECKING: + from collections.abc import Iterator from pathlib import Path +_HAS_SYMLINK = hasattr(os, "symlink") + + def _source(subdir: str, path: str, *, mount: bool = False) -> dict[str, Any]: return {"source_path": path, "workspace_subdir": subdir, "mount": mount} -def test_copied_source_becomes_localdir_entry(tmp_path: Path) -> None: - entries, bind_mounts, staged_dirs = build_session_entries([_source("repo", str(tmp_path))]) +@contextmanager +def _built_tar(src_root: Path, arc_prefix: str) -> Iterator[tuple[bytes, int, int]]: + tar_path, added, skipped = _build_source_tar(src_root, arc_prefix) + try: + yield tar_path.read_bytes(), added, skipped + finally: + tar_path.unlink(missing_ok=True) + + +def _tar_names(tar_bytes: bytes) -> set[str]: + with tarfile.open(fileobj=io.BytesIO(tar_bytes), mode="r") as tar: + return set(tar.getnames()) + + +def _tar_file_names(tar_bytes: bytes) -> set[str]: + with tarfile.open(fileobj=io.BytesIO(tar_bytes), mode="r") as tar: + return {m.name for m in tar.getmembers() if m.isfile()} + + +def _tar_dir_names(tar_bytes: bytes) -> set[str]: + with tarfile.open(fileobj=io.BytesIO(tar_bytes), mode="r") as tar: + return {m.name for m in tar.getmembers() if m.isdir()} + + +def test_copied_source_is_returned_for_import(tmp_path: Path) -> None: + copied, bind_mounts = split_local_sources([_source("repo", str(tmp_path))]) assert bind_mounts == [] - assert staged_dirs == [] - assert isinstance(entries["repo"], LocalDir) - assert entries["repo"].src == tmp_path.resolve() + assert copied == [{"source_path": str(tmp_path.resolve()), "workspace_subdir": "repo"}] def test_mounted_source_becomes_bind_mount(tmp_path: Path) -> None: - entries, bind_mounts, _staged = build_session_entries( - [_source("repo", str(tmp_path), mount=True)] - ) + copied, bind_mounts = split_local_sources([_source("repo", str(tmp_path), mount=True)]) - assert entries == {} + assert copied == [] assert bind_mounts == [ { "source": str(tmp_path.resolve()), @@ -42,46 +76,155 @@ def test_mounted_source_becomes_bind_mount(tmp_path: Path) -> None: def test_mixed_sources_split_correctly(tmp_path: Path) -> None: - copied = tmp_path / "copied" - mounted = tmp_path / "mounted" - copied.mkdir() - mounted.mkdir() + copied_dir = tmp_path / "copied" + mounted_dir = tmp_path / "mounted" + copied_dir.mkdir() + mounted_dir.mkdir() - entries, bind_mounts, _staged = build_session_entries( + copied, bind_mounts = split_local_sources( [ - _source("copied", str(copied)), - _source("mounted", str(mounted), mount=True), + _source("copied", str(copied_dir)), + _source("mounted", str(mounted_dir), mount=True), ] ) - assert list(entries) == ["copied"] - assert isinstance(entries["copied"], LocalDir) + assert [c["workspace_subdir"] for c in copied] == ["copied"] assert [m["target"] for m in bind_mounts] == ["/workspace/mounted"] def test_incomplete_sources_are_skipped() -> None: - entries, bind_mounts, staged_dirs = build_session_entries( + copied, bind_mounts = split_local_sources( [ {"source_path": "", "workspace_subdir": "x"}, {"source_path": "/p", "workspace_subdir": ""}, ] ) - assert entries == {} + assert copied == [] assert bind_mounts == [] - assert staged_dirs == [] -def test_symlink_tree_is_staged(tmp_path: Path) -> None: - repo = tmp_path / "repo" - repo.mkdir() - (repo / "real.txt").write_text("content") - (repo / "link.txt").symlink_to(repo / "real.txt") +def test_workspace_subdir_slashes_are_stripped(tmp_path: Path) -> None: + copied, _ = split_local_sources([_source("/repo/", str(tmp_path))]) + assert copied[0]["workspace_subdir"] == "repo" + + +def test_tar_packs_files_under_arc_prefix(tmp_path: Path) -> None: + (tmp_path / "a.txt").write_text("a", encoding="utf-8") + nested = tmp_path / "sub" + nested.mkdir() + (nested / "b.txt").write_text("b", encoding="utf-8") + + with _built_tar(tmp_path, "repo") as (tar_bytes, added, skipped): + assert added == 2 + assert skipped == 0 + assert _tar_file_names(tar_bytes) == {"repo/a.txt", "repo/sub/b.txt"} + assert {"repo", "repo/sub"} <= _tar_dir_names(tar_bytes) + + +def test_tar_preserves_dotfiles_and_git(tmp_path: Path) -> None: + (tmp_path / ".env").write_text("SECRET=1", encoding="utf-8") + git_dir = tmp_path / ".git" + git_dir.mkdir() + (git_dir / "HEAD").write_text("ref: refs/heads/main", encoding="utf-8") + + with _built_tar(tmp_path, "repo") as (tar_bytes, added, _): + assert added == 2 + assert _tar_file_names(tar_bytes) == {"repo/.env", "repo/.git/HEAD"} + + +@pytest.mark.skipif(not _HAS_SYMLINK, reason="requires symlink support") +def test_tar_skips_file_symlinks(tmp_path: Path) -> None: + (tmp_path / "real.txt").write_text("real", encoding="utf-8") + try: + (tmp_path / "link.txt").symlink_to(tmp_path / "real.txt") + except OSError: + pytest.skip("symlink creation requires elevated privileges on this platform") + + with _built_tar(tmp_path, "repo") as (tar_bytes, added, skipped): + assert added == 1 + assert skipped == 1 + assert _tar_file_names(tar_bytes) == {"repo/real.txt"} + + +@pytest.mark.skipif(not _HAS_SYMLINK, reason="requires symlink support") +def test_tar_skips_dir_symlinks_without_descending(tmp_path: Path) -> None: + outside = tmp_path / "outside" + outside.mkdir() + (outside / "secret.txt").write_text("secret", encoding="utf-8") + + src = tmp_path / "src" + src.mkdir() + (src / "keep.txt").write_text("keep", encoding="utf-8") + try: + (src / "escape").symlink_to(outside) + except OSError: + pytest.skip("symlink creation requires elevated privileges on this platform") + + with _built_tar(src, "repo") as (tar_bytes, added, skipped): + assert added == 1 + assert skipped == 1 + assert _tar_file_names(tar_bytes) == {"repo/keep.txt"} + assert "repo/escape" not in _tar_names(tar_bytes) + + +def test_tar_preserves_empty_dirs(tmp_path: Path) -> None: + (tmp_path / "empty").mkdir() + (tmp_path / "keep.txt").write_text("k", encoding="utf-8") + + with _built_tar(tmp_path, "repo") as (tar_bytes, added, skipped): + assert added == 1 + assert skipped == 0 + assert _tar_file_names(tar_bytes) == {"repo/keep.txt"} + assert {"repo", "repo/empty"} <= _tar_dir_names(tar_bytes) + + +def test_tar_empty_root_still_packs_prefix(tmp_path: Path) -> None: + with _built_tar(tmp_path, "repo") as (tar_bytes, added, skipped): + assert added == 0 + assert skipped == 0 + assert _tar_file_names(tar_bytes) == set() + assert _tar_dir_names(tar_bytes) == {"repo"} + + +def test_unsafe_workspace_subdir_is_skipped(tmp_path: Path) -> None: + copied, bind_mounts = split_local_sources( + [ + _source("../escape", str(tmp_path)), + _source("ok/../../escape", str(tmp_path)), + ] + ) + assert copied == [] + assert bind_mounts == [] + + +def test_unsafe_workspace_subdir_skipped_for_mount(tmp_path: Path) -> None: + copied, bind_mounts = split_local_sources([_source("../escape", str(tmp_path), mount=True)]) + assert copied == [] + assert bind_mounts == [] - entries, _mounts, staged_dirs = build_session_entries([_source("repo", str(repo))]) - assert len(staged_dirs) == 1 - entry = entries["repo"] - assert isinstance(entry, LocalDir) - assert entry.src == staged_dirs[0] - assert not (staged_dirs[0] / "link.txt").is_symlink() - assert (staged_dirs[0] / "link.txt").read_text() == "content" +def test_tar_raises_on_unreadable_directory(tmp_path: Path) -> None: + """Unreadable subtrees must fail import, not silently disappear.""" + nested = tmp_path / "secret" + nested.mkdir() + (nested / "x.txt").write_text("hidden", encoding="utf-8") + + real_walk = os.walk + + def _walk_boom( + top: str | os.PathLike[str], + *args: Any, + **kwargs: Any, + ) -> Any: + onerror = kwargs.get("onerror") + # First yield root so walk has started, then simulate listdir failure + # on the nested dir by invoking onerror like os.walk would. + yield from real_walk(top, *args, **kwargs) + if onerror is not None: + onerror(PermissionError(13, "Permission denied", str(nested))) + + with ( + patch("strix.runtime.session_manager.os.walk", side_effect=_walk_boom), + pytest.raises(PermissionError, match="Permission denied"), + ): + _build_source_tar(tmp_path, "repo")