diff --git a/opencollab/adapters/_docker_snapshot.py b/opencollab/adapters/_docker_snapshot.py index e08bef6f..f233dc15 100644 --- a/opencollab/adapters/_docker_snapshot.py +++ b/opencollab/adapters/_docker_snapshot.py @@ -5,6 +5,7 @@ import posixpath import shlex import shutil +import stat import tempfile from contextlib import asynccontextmanager from pathlib import Path @@ -14,8 +15,10 @@ from opencollab.adapters._env_docker import DockerEnvironment from opencollab.adapters._env_local import LocalEnvironment from opencollab.adapters._env_vm import VMEnvironment +from opencollab.adapters._env_worktree import WorktreeEnvironment from opencollab.adapters._git_patch_bytes import _complete, _patch_temporary from opencollab.adapters._team_candidate_environment import _TeamCandidateEnvironment +from opencollab.adapters.safe_anchored_files import open_regular_read_at from opencollab.application.async_timeout import await_owned_operation @@ -40,6 +43,38 @@ async def _copy_to_host(environment: Any, source: str, target: Path, *, director raise OSError("Docker source workspace copy failed") +async def _copy_artifact_to_host(environment: Any, source: str, target: Path) -> None: + relative = ( + posixpath.relpath(source, environment.workspace) if posixpath.isabs(source) else posixpath.normpath(source) + ) + if relative == ".." or relative.startswith("../"): + raise PermissionError("artifact path escapes workspace") + source = posixpath.join(environment.workspace, relative) + delegate = environment._delegate if isinstance(environment, _TeamCandidateEnvironment) else environment + if isinstance(delegate, WorktreeEnvironment): + delegate._ensure_active() + if delegate._local_env is None: + await delegate.setup() + delegate = delegate._local_env + if isinstance(delegate, LocalEnvironment): + def copy() -> None: + with open_regular_read_at( + delegate._root_fd(), delegate.workspace, delegate._relative_path(source), + ) as handle, target.open("wb") as output: + shutil.copyfileobj(handle, output) + + await delegate._run_file_operation(copy) + return + if isinstance(delegate, DockerEnvironment): + source = await delegate._contained_file_path(source) + if isinstance(delegate, (DockerEnvironment, VMEnvironment)): + await _copy_to_host(delegate, source, target) + if not stat.S_ISREG(target.lstat().st_mode): + raise OSError("artifact transfer requires a regular file") + return + target.write_text(await environment.read_file(source), encoding="utf-8") + + def _remove_marker(path: Path) -> None: marker = path / ".git" if marker.is_dir() and not marker.is_symlink(): diff --git a/opencollab/adapters/_env_worktree.py b/opencollab/adapters/_env_worktree.py index 39161643..3c8905f5 100644 --- a/opencollab/adapters/_env_worktree.py +++ b/opencollab/adapters/_env_worktree.py @@ -679,12 +679,12 @@ async def _ensure_directory_copy_changes_exported(self) -> None: "call get_diff() and deliver its artifact first" ) - async def exec_cmd(self, cmd: str, timeout: float = 120.0) -> ExecResult: + async def exec_cmd(self, cmd: str, timeout: float = 120.0, *, input_bytes: bytes | None = None) -> ExecResult: self._ensure_active() if self._local_env is None: await self.setup() assert self._local_env is not None - return await self._local_env.exec_cmd(cmd, timeout) + return await self._local_env.exec_cmd(cmd, timeout, input_bytes=input_bytes) async def read_file(self, path: str) -> str: self._ensure_active() diff --git a/opencollab/adapters/safe_anchored_files.py b/opencollab/adapters/safe_anchored_files.py index 6625bfeb..aeff6f22 100644 --- a/opencollab/adapters/safe_anchored_files.py +++ b/opencollab/adapters/safe_anchored_files.py @@ -10,7 +10,8 @@ import os import stat import uuid -from collections.abc import Callable +from collections.abc import Callable, Iterator +from contextlib import contextmanager from typing import BinaryIO from opencollab.adapters._env_base import TextFileRange @@ -116,15 +117,13 @@ def open_parent_beneath( os.close(current_fd) raise -def read_regular_bytes_at( +@contextmanager +def open_regular_read_at( root_fd: int, root_path: str | os.PathLike[str], path: str | os.PathLike[str], - *, - max_bytes: int, -) -> bytes: - """Read a bounded regular file beneath a pinned directory descriptor.""" - limit = _require_limit(max_bytes) +) -> Iterator[BinaryIO]: + """Open a regular file beneath a pinned directory without following symlinks.""" parent_fd, name = open_parent_beneath( root_fd, root_path, @@ -141,6 +140,28 @@ def read_regular_bytes_at( opened = os.fstat(fd) if not stat.S_ISREG(opened.st_mode): raise OSError(f"read target is not a regular file: {path}") + handle = os.fdopen(fd, "rb") + fd = -1 + with handle: + yield handle + finally: + if fd >= 0: + os.close(fd) + os.close(parent_fd) + + +def read_regular_bytes_at( + root_fd: int, + root_path: str | os.PathLike[str], + path: str | os.PathLike[str], + *, + max_bytes: int, +) -> bytes: + """Read a bounded regular file beneath a pinned directory descriptor.""" + limit = _require_limit(max_bytes) + with open_regular_read_at(root_fd, root_path, path) as handle: + fd = handle.fileno() + opened = os.fstat(fd) if opened.st_size > limit: raise ValueError(f"read target exceeds {limit}-byte limit: {path}") chunks: list[bytes] = [] @@ -155,10 +176,6 @@ def read_regular_bytes_at( if len(payload) > limit: raise ValueError(f"read target exceeds {limit}-byte limit: {path}") return payload - finally: - if fd >= 0: - os.close(fd) - os.close(parent_fd) def read_regular_text_range_at( root_fd: int, diff --git a/opencollab/bootstrap/authored_workflow.py b/opencollab/bootstrap/authored_workflow.py index 725bfc45..fd67ab6e 100644 --- a/opencollab/bootstrap/authored_workflow.py +++ b/opencollab/bootstrap/authored_workflow.py @@ -16,6 +16,8 @@ from pathlib import Path from typing import TYPE_CHECKING, Any, get_args +from opencollab.adapters._docker_snapshot import _copy_artifact_to_host +from opencollab.adapters._env_process import ProcessCleanupError from opencollab.adapters.session_environments import SessionEnvironmentPool from opencollab.adapters.tools._paths import checked_path from opencollab.adapters.tools.base import Tool @@ -320,15 +322,16 @@ def _account_usage(self, child: _Child, tokens: Any) -> None: child.tokens = tokens async def _execute(self, child: _Child, command: str, request_bytes: bytes) -> None: - execution = asyncio.create_task(child.environment.exec_cmd( - command, timeout=math.inf, input_bytes=request_bytes, - )) - child.execution = execution - if child.stop_requested: - execution.cancel() + execution = None logs: dict[str, Any] = {} status = "running" try: + execution = asyncio.create_task(child.environment.exec_cmd( + command, timeout=math.inf, input_bytes=request_bytes, + )) + child.execution = execution + if child.stop_requested: + execution.cancel() while not execution.done(): await asyncio.wait({execution}, timeout=0.1) assert child.bridge is not None @@ -340,7 +343,12 @@ async def _execute(self, child: _Child, command: str, request_bytes: bytes) -> N logs = {"returncode": outcome.returncode, "stdout": outcome.stdout, "stderr": outcome.stderr} if outcome.returncode: raise RuntimeError(f"Workflow runner exited {outcome.returncode}: {outcome.stderr[-2000:]}") - child.result = json.loads(await child.environment.read_file(posixpath.join(child.artifacts, "result.json"))) + with tempfile.TemporaryDirectory(prefix="opencollab-child-result-") as temporary: + result_path = Path(temporary) / "result.json" + await _copy_artifact_to_host( + child.environment, posixpath.join(child.artifacts, "result.json"), result_path, + ) + child.result = json.loads(result_path.read_text(encoding="utf-8")) self._account_usage(child, child.result["tokens"]) status = child.result["status"] except asyncio.CancelledError: @@ -350,14 +358,16 @@ async def _execute(self, child: _Child, command: str, request_bytes: bytes) -> N status = "failed" child.result = {"status": "failed", "output": None, "reason": f"{type(exc).__name__}: {exc}"} finally: - if not execution.done(): - execution.cancel() - settled = await asyncio.gather(execution, return_exceptions=True) - if isinstance(settled[0], Exception): - status = "failed" - child.cleanup_error = settled[0] - child.result = {"status": "failed", "output": None, - "reason": f"child process cleanup failed: {settled[0]}"} + if execution is not None: + cancelling = not execution.done() + if cancelling: + execution.cancel() + settled = await asyncio.gather(execution, return_exceptions=True) + if isinstance(settled[0], ProcessCleanupError) or (cancelling and isinstance(settled[0], Exception)): + status = "failed" + child.cleanup_error = settled[0] + child.result = {"status": "failed", "output": None, + "reason": f"child process cleanup failed: {settled[0]}"} try: assert child.bridge is not None await await_owned_operation(child.bridge.aclose(), propagate_cancellation=False) @@ -388,10 +398,11 @@ async def _execute(self, child: _Child, command: str, request_bytes: bytes) -> N ): raise ValueError("invalid child artifact name") try: - content = await child.environment.read_file(posixpath.join(child.artifacts, name)) + await _copy_artifact_to_host( + child.environment, posixpath.join(child.artifacts, name), directory / name, + ) except FileNotFoundError: continue - (directory / name).write_text(content, encoding="utf-8") child.saved_artifacts = str(directory) (directory / "result.json").write_text(json.dumps(self._snapshot(child, include_result=True)), encoding="utf-8") diff --git a/tests/interfaces/test_authored_workflow_execution.py b/tests/interfaces/test_authored_workflow_execution.py new file mode 100644 index 00000000..d4b1faf9 --- /dev/null +++ b/tests/interfaces/test_authored_workflow_execution.py @@ -0,0 +1,133 @@ +from __future__ import annotations + +import json +import os +import subprocess +from pathlib import Path +from types import SimpleNamespace + +import pytest + +from opencollab.adapters._env_process import ProcessCleanupError +from opencollab.adapters._team_candidate_environment import _TeamCandidateEnvironment +from opencollab.adapters.env import LocalEnvironment, WorktreeEnvironment +from tests.interfaces.test_member_authored_workflow import _launch, _manager +from tests.support.paths import REPO_ROOT + + +@pytest.fixture(autouse=True) +def _child_import_path(monkeypatch): + monkeypatch.setenv("PYTHONPATH", str(REPO_ROOT)) + monkeypatch.setenv("OPENCOLLAB_API_USAGE_LOG", os.devnull) + + +@pytest.mark.parametrize("kind", ["local", "copy-worktree", "git-worktree", "team-candidate"]) +@pytest.mark.parametrize("size", [1024, 4 * 1024 * 1024 + 1024]) +async def test_authored_result_and_native_artifacts_cross_file_read_limit(tmp_path, kind, size): + workspace = tmp_path / "workspace" + workspace.mkdir() + (workspace / "authored.py").write_text( + "async def main(ctx, args):\n return {'payload': args['payload']}\n" + ) + (workspace / ".gitignore").write_text(".opencollab/\n") + if kind == "git-worktree": + for args in (("init", "-q"), ("add", "."), ( + "-c", "user.name=Test", "-c", "user.email=test@example.invalid", + "commit", "-qm", "test: prepare workflow fixture", + )): + subprocess.run(["git", "-C", str(workspace), *args], check=True, capture_output=True) + environment = (LocalEnvironment if kind in {"local", "team-candidate"} else WorktreeEnvironment)(str(workspace)) + if kind == "team-candidate": + environment = _TeamCandidateEnvironment( + SimpleNamespace(environment=environment, cleanup=environment.cleanup), str(workspace), + ) + await environment.setup() + parent, tool, runtime = _manager(environment) + saved = tmp_path / "saved" + saved.mkdir() + tool._save_dir = str(saved) + payload = "x" * size + try: + child = await _launch(tool, runtime, args={"payload": payload}) + result = json.loads(await tool.execute_with_runtime({"action": "wait", "id": child["id"]}, runtime)) + native_path = Path(child["artifacts"]) / "result.json" + assert json.loads(native_path.read_text())["status"] == "completed" + assert result["status"] == "completed", result + assert result["result"]["output"] == {"payload": payload} + assert tool._reserved_tokens() == 0 + persisted = Path(result["artifacts"]) + assert json.loads((persisted / "result.json").read_text())["result"]["output"] == {"payload": payload} + assert json.loads((persisted / "workflow.json").read_text())["args"] == {"payload": payload} + if size > 4 * 1024 * 1024: + assert native_path.stat().st_size > 4 * 1024 * 1024 + assert (persisted / "workflow.json").stat().st_size > 4 * 1024 * 1024 + with pytest.raises(ValueError, match="read target exceeds"): + await environment.read_file(str(native_path)) + finally: + try: + await parent.aclose() + finally: + if isinstance(environment, WorktreeEnvironment): + await environment.get_diff() + await environment.cleanup() + + +@pytest.mark.parametrize("failure", ["signature", "async-start"]) +async def test_executor_start_failure_finishes_child_and_releases_budget(tmp_path, monkeypatch, failure): + (tmp_path / "authored.py").write_text("async def main(ctx, args):\n return 1\n") + environment = LocalEnvironment(str(tmp_path)) + parent, tool, runtime = _manager(environment) + native_exec = environment.exec_cmd + + if failure == "signature": + def fail_to_start(command, timeout=120): + raise AssertionError("the input_bytes argument must fail during binding") + else: + async def fail_to_start(command, timeout=120, *, input_bytes=None): + raise OSError("executor startup failed") + + monkeypatch.setattr(environment, "exec_cmd", fail_to_start) + try: + child = await _launch(tool, runtime) + assert tool._reserved_tokens() == 500 + result = json.loads(await tool.execute_with_runtime({"action": "wait", "id": child["id"]}, runtime)) + assert result["status"] == "failed" + assert ("TypeError" if failure == "signature" else "executor startup failed") in result["result"]["reason"] + assert tool._reserved_tokens() == 0 + owned = tool._children[child["id"]] + assert owned.task.done() and owned.task.exception() is None + assert owned.bridge._pool._closed + assert owned.cleanup_error is None + assert not environment._processes._processes + finally: + monkeypatch.setattr(environment, "exec_cmd", native_exec) + try: + await parent.aclose() + finally: + await environment.cleanup() + + +async def test_executor_cleanup_failure_retains_recovery_error(tmp_path, monkeypatch): + (tmp_path / "authored.py").write_text("async def main(ctx, args):\n return 1\n") + environment = LocalEnvironment(str(tmp_path)) + parent, tool, runtime = _manager(environment) + native_exec = environment.exec_cmd + + async def cleanup_failed(command, timeout=120, *, input_bytes=None): + raise ProcessCleanupError("process group still alive") + + monkeypatch.setattr(environment, "exec_cmd", cleanup_failed) + try: + child = await _launch(tool, runtime) + result = json.loads(await tool.execute_with_runtime({"action": "wait", "id": child["id"]}, runtime)) + assert result["status"] == "failed" + assert "process group still alive" in result["result"]["reason"] + owned = tool._children[child["id"]] + assert isinstance(owned.cleanup_error, ProcessCleanupError) + assert tool._reserved_tokens() == 0 + with pytest.raises(RuntimeError, match="process group still alive"): + await parent.aclose() + assert not owned.artifacts_removed + finally: + monkeypatch.setattr(environment, "exec_cmd", native_exec) + await environment.cleanup() diff --git a/tests/interfaces/test_member_authored_workflow.py b/tests/interfaces/test_member_authored_workflow.py index 330ce25c..c74f0b57 100644 --- a/tests/interfaces/test_member_authored_workflow.py +++ b/tests/interfaces/test_member_authored_workflow.py @@ -266,12 +266,12 @@ async def test_stop_grace_leaves_finished_runner_result_persistence_owned( environment = local_environment_factory(str(tmp_path)) parent, tool, runtime = _manager(environment) entered, finish = asyncio.Event(), asyncio.Event() - native_read = environment.read_file + native_copy = authored_workflow._copy_artifact_to_host read_cancelled = False - async def delayed_result(path): + async def delayed_result(environment, path, target): nonlocal read_cancelled - value = await native_read(path) + await native_copy(environment, path, target) if path.endswith("/result.json"): entered.set() try: @@ -279,9 +279,8 @@ async def delayed_result(path): except asyncio.CancelledError: read_cancelled = True raise - return value - monkeypatch.setattr(environment, "read_file", delayed_result) + monkeypatch.setattr(authored_workflow, "_copy_artifact_to_host", delayed_result) child = await _launch(tool, runtime) stopping = None try: diff --git a/tests/runtime/test_authored_artifact_transfer.py b/tests/runtime/test_authored_artifact_transfer.py new file mode 100644 index 00000000..2fa7d6c6 --- /dev/null +++ b/tests/runtime/test_authored_artifact_transfer.py @@ -0,0 +1,171 @@ +from __future__ import annotations + +import asyncio +import os +import shutil +from pathlib import Path +from unittest.mock import AsyncMock + +import pytest + +from opencollab.adapters import _docker_snapshot as snapshot_module +from opencollab.adapters._docker_snapshot import _copy_artifact_to_host +from opencollab.adapters._env_vm import VMEnvironment +from opencollab.adapters.env import DockerEnvironment, ExecResult, LocalEnvironment, WorktreeEnvironment +from opencollab.adapters.session_environments import SessionEnvironmentPool +from opencollab.application.tool_execution import ToolRuntime +from opencollab.bootstrap.authored_workflow_bridge import AuthoredWorkflowBridge, BridgeClient, BridgeEnvironment + + +@pytest.mark.parametrize("kind", ["local", "worktree"]) +@pytest.mark.parametrize("target_kind", ["relative-escape", "absolute-escape", "file-link", "directory-link", "fifo"]) +async def test_artifact_transfer_preserves_local_path_and_regular_file_checks(tmp_path, kind, target_kind): + workspace = tmp_path / "workspace" + workspace.mkdir() + environment = (LocalEnvironment if kind == "local" else WorktreeEnvironment)(str(workspace)) + await environment.setup() + root = Path(environment.workspace) + outside = root.parent / f"{root.name}-outside" + outside.mkdir() + (outside / "result.json").write_text('{"outside": true}') + if target_kind == "relative-escape": + source = f"../{outside.name}/result.json" + elif target_kind == "absolute-escape": + source = str(outside / "result.json") + elif target_kind == "file-link": + (root / "result.json").symlink_to(outside / "result.json") + source = str(root / "result.json") + elif target_kind == "directory-link": + (root / "artifacts").symlink_to(outside, target_is_directory=True) + source = str(root / "artifacts" / "result.json") + else: + os.mkfifo(root / "result.json") + source = str(root / "result.json") + target = tmp_path / "copied.json" + try: + with pytest.raises((OSError, PermissionError)): + await _copy_artifact_to_host(environment, source, target) + assert not target.exists() + finally: + for name in ("result.json", "artifacts"): + if os.path.lexists(root / name): + (root / name).unlink() + shutil.rmtree(outside) + if isinstance(environment, WorktreeEnvironment): + await environment.get_diff() + await environment.cleanup() + + +async def test_artifact_transfer_rejects_replaced_workspace(tmp_path): + workspace = tmp_path / "workspace" + workspace.mkdir() + environment = LocalEnvironment(str(workspace)) + workspace.rename(tmp_path / "original") + workspace.mkdir() + (workspace / "result.json").write_text('{"replacement": true}') + try: + with pytest.raises(OSError, match="workspace was replaced"): + await _copy_artifact_to_host(environment, str(workspace / "result.json"), tmp_path / "copied.json") + assert not (tmp_path / "copied.json").exists() + finally: + await environment.cleanup() + + +@pytest.mark.parametrize("kind", ["docker", "vm"]) +@pytest.mark.parametrize("symlink", [False, True]) +async def test_remote_artifact_transfer_uses_existing_transport_and_requires_regular_output( + tmp_path, monkeypatch, kind, symlink, +): + environment = object.__new__(DockerEnvironment if kind == "docker" else VMEnvironment) + environment.workspace = "/workspace" + target = tmp_path / "copied.json" + native = tmp_path / "native.json" + native.write_bytes(b'{"payload":"' + b"x" * (4 * 1024 * 1024) + b'"}\n') + + async def copy(destination): + if symlink: + destination.symlink_to(native) + else: + shutil.copyfile(native, destination) + + if kind == "docker": + environment._container_id = "owned-test-container" + monkeypatch.setattr(environment, "setup", AsyncMock()) + contained = AsyncMock(return_value="/workspace/physical/result.json") + monkeypatch.setattr(environment, "_contained_file_path", contained) + + async def docker(*args): + assert args[:3] == ("cp", "--", "owned-test-container:/workspace/physical/result.json") + await copy(Path(args[3])) + return ExecResult(0, "", "") + + monkeypatch.setattr(environment, "_docker", docker) + else: + async def download(source, destination, *, directory=False): + assert source == "/workspace/result.json" and not directory + await copy(destination) + + monkeypatch.setattr(environment, "copy_to_host", download) + if symlink: + with pytest.raises(OSError, match="regular file"): + await _copy_artifact_to_host(environment, "/workspace/result.json", target) + else: + await _copy_artifact_to_host(environment, "/workspace/result.json", target) + assert target.read_bytes() == native.read_bytes() + if kind == "docker": + contained.assert_awaited_once_with("/workspace/result.json") + with pytest.raises(PermissionError, match="escapes workspace"): + await _copy_artifact_to_host(environment, "/outside/result.json", target) + + +async def test_artifact_copy_failure_closes_source_descriptor(tmp_path, monkeypatch): + (tmp_path / "result.json").write_text('{"ready": true}') + environment = LocalEnvironment(str(tmp_path)) + descriptor = None + + def fail_copy(source, target): + nonlocal descriptor + descriptor = source.fileno() + raise OSError("destination write failed") + + monkeypatch.setattr(snapshot_module.shutil, "copyfileobj", fail_copy) + try: + with pytest.raises(OSError, match="destination write failed"): + await _copy_artifact_to_host(environment, "result.json", tmp_path / "copied.json") + assert descriptor is not None + with pytest.raises(OSError): + os.fstat(descriptor) + finally: + await environment.cleanup() + + +async def test_artifact_transfer_retains_bridge_environment_fallback(tmp_path, monkeypatch): + monkeypatch.chdir(tmp_path) + artifacts = tmp_path / "bridge-artifacts" + artifacts.mkdir() + (tmp_path / "result.json").write_text('{"ready": true}') + environment = LocalEnvironment(str(tmp_path)) + bridge = AuthoredWorkflowBridge( + environment, str(artifacts), + SessionEnvironmentPool(str(tmp_path), isolation="shared", base_environment=environment), + {}, ToolRuntime(environment, None, None), allow_unisolated_shell=True, + ) + proxy = BridgeEnvironment(BridgeClient(artifacts, bridge.options()), { + "handle": 0, "workspace": str(tmp_path), "process_isolated": False, + }) + copying = asyncio.create_task(_copy_artifact_to_host(proxy, "result.json", tmp_path / "copied.json")) + + async def poll_until_copied(): + while not copying.done(): + await bridge.poll() + await asyncio.sleep(0.01) + + try: + await asyncio.wait_for(poll_until_copied(), 5) + await copying + assert (tmp_path / "copied.json").read_text() == (tmp_path / "result.json").read_text() + finally: + copying.cancel() + await asyncio.gather(copying, return_exceptions=True) + await bridge.aclose() + await environment.cleanup() diff --git a/tests/runtime/test_worktree_pool.py b/tests/runtime/test_worktree_pool.py index dac98e46..75b732d2 100644 --- a/tests/runtime/test_worktree_pool.py +++ b/tests/runtime/test_worktree_pool.py @@ -47,6 +47,20 @@ def _repo(path): def _branch_exists(repo, branch: str) -> bool: return _git(repo, "show-ref", "--verify", f"refs/heads/{branch}", check=False).returncode == 0 + +async def test_worktree_exec_forwards_binary_stdin(tmp_path): + environment = WorktreeEnvironment(str(_repo(tmp_path / "repo"))) + payload = b"\x00workflow\n\xff" + try: + result = await environment.exec_cmd(shlex.join([ + sys.executable, "-c", "import sys; print(sys.stdin.buffer.read().hex())", + ]), timeout=10, input_bytes=payload) + assert result.returncode == 0, result.stderr + assert result.stdout.strip() == payload.hex() + finally: + await environment.cleanup() + + @pytest.mark.parametrize("suffix", ["", " ", "\n"]) @pytest.mark.parametrize("subdirectory", [False, True]) async def test_git_worktree_keeps_whitespace_in_repository_root(