wehub-resource-sync
291 lines
10 KiB
Python
291 lines
10 KiB
Python
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import signal
|
|
from pathlib import Path
|
|
from types import SimpleNamespace
|
|
from typing import cast
|
|
|
|
import pytest
|
|
|
|
from agents.sandbox.errors import PtySessionNotFoundError
|
|
from agents.sandbox.manifest import Manifest
|
|
from agents.sandbox.sandboxes.unix_local import (
|
|
UnixLocalSandboxClient,
|
|
UnixLocalSandboxSession,
|
|
UnixLocalSandboxSessionState,
|
|
_UnixPtyProcessEntry,
|
|
)
|
|
from agents.sandbox.snapshot import NoopSnapshot
|
|
from agents.sandbox.types import ExecResult, User
|
|
|
|
|
|
class _RecordingUnixLocalSession(UnixLocalSandboxSession):
|
|
def __init__(self, root: Path) -> None:
|
|
super().__init__(
|
|
state=UnixLocalSandboxSessionState(
|
|
manifest=Manifest(root=str(root)),
|
|
snapshot=NoopSnapshot(id="noop"),
|
|
)
|
|
)
|
|
self.exec_commands: list[tuple[str, ...]] = []
|
|
|
|
async def _exec_internal(
|
|
self,
|
|
*command: str | Path,
|
|
timeout: float | None = None,
|
|
) -> ExecResult:
|
|
_ = timeout
|
|
self.exec_commands.append(tuple(str(part) for part in command))
|
|
return ExecResult(stdout=b"", stderr=b"", exit_code=0)
|
|
|
|
|
|
class TestUnixLocalPty:
|
|
@pytest.mark.asyncio
|
|
async def test_tty_fd_close_is_owned_without_blocking_termination(
|
|
self,
|
|
tmp_path: Path,
|
|
monkeypatch: pytest.MonkeyPatch,
|
|
) -> None:
|
|
session = _RecordingUnixLocalSession(tmp_path)
|
|
close_started = asyncio.Event()
|
|
release_close = asyncio.Event()
|
|
|
|
async def blocked_to_thread(*args: object, **kwargs: object) -> None:
|
|
_ = (args, kwargs)
|
|
close_started.set()
|
|
await release_close.wait()
|
|
|
|
monkeypatch.setattr(asyncio, "to_thread", blocked_to_thread)
|
|
process = cast(
|
|
asyncio.subprocess.Process,
|
|
SimpleNamespace(returncode=0, pid=None),
|
|
)
|
|
entry = _UnixPtyProcessEntry(process=process, tty=True, primary_fd=123)
|
|
|
|
await asyncio.wait_for(session._terminate_pty_entry(entry), timeout=0.5)
|
|
await close_started.wait()
|
|
|
|
assert len(session._fd_close_tasks) == 1
|
|
await asyncio.wait_for(session._after_stop(), timeout=0.5)
|
|
assert len(session._fd_close_tasks) == 1
|
|
|
|
release_close.set()
|
|
await asyncio.gather(*session._fd_close_tasks)
|
|
await asyncio.sleep(0)
|
|
|
|
assert session._fd_close_tasks == set()
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_pty_exec_write_poll_and_unknown_session_errors(self, tmp_path: Path) -> None:
|
|
client = UnixLocalSandboxClient()
|
|
manifest = Manifest(root=str(tmp_path / "workspace"))
|
|
|
|
async with await client.create(manifest=manifest, snapshot=None, options=None) as session:
|
|
started = await session.pty_exec_start(
|
|
"sh",
|
|
"-c",
|
|
"IFS= read -r line; printf '%s\\n' \"$line\"",
|
|
shell=False,
|
|
tty=True,
|
|
yield_time_s=0.05,
|
|
)
|
|
|
|
assert started.process_id is not None
|
|
assert started.exit_code is None
|
|
|
|
written = await session.pty_write_stdin(
|
|
session_id=started.process_id,
|
|
chars="hello from pty\n",
|
|
yield_time_s=0.25,
|
|
)
|
|
assert written.process_id is None
|
|
assert written.exit_code == 0
|
|
assert "hello from pty" in written.output.decode("utf-8", errors="replace")
|
|
|
|
with pytest.raises(PtySessionNotFoundError):
|
|
await session.pty_write_stdin(session_id=started.process_id, chars="")
|
|
|
|
with pytest.raises(PtySessionNotFoundError):
|
|
await session.pty_write_stdin(session_id=999_999, chars="")
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_pty_ctrl_c_interrupts_long_running_process(self, tmp_path: Path) -> None:
|
|
client = UnixLocalSandboxClient()
|
|
manifest = Manifest(root=str(tmp_path / "workspace"))
|
|
|
|
async with await client.create(manifest=manifest, snapshot=None, options=None) as session:
|
|
started = await session.pty_exec_start(
|
|
"sleep",
|
|
"30",
|
|
shell=False,
|
|
tty=True,
|
|
yield_time_s=0.05,
|
|
)
|
|
|
|
assert started.process_id is not None
|
|
assert started.exit_code is None
|
|
|
|
first_interrupt = await session.pty_write_stdin(
|
|
session_id=started.process_id,
|
|
chars="\x03",
|
|
yield_time_s=0.25,
|
|
)
|
|
if first_interrupt.process_id is None:
|
|
interrupted = first_interrupt
|
|
else:
|
|
interrupted = await session.pty_write_stdin(
|
|
session_id=started.process_id,
|
|
chars="",
|
|
yield_time_s=5.5,
|
|
)
|
|
|
|
assert interrupted.process_id is None
|
|
assert interrupted.exit_code is not None
|
|
|
|
with pytest.raises(PtySessionNotFoundError):
|
|
await session.pty_write_stdin(session_id=started.process_id, chars="")
|
|
|
|
@pytest.mark.parametrize(
|
|
("signum", "chars"),
|
|
[
|
|
pytest.param(signal.SIGINT, "\x03", id="sigint"),
|
|
pytest.param(signal.SIGQUIT, "\x1c", id="sigquit"),
|
|
],
|
|
)
|
|
@pytest.mark.asyncio
|
|
async def test_pty_terminal_signals_interrupt_even_if_parent_ignores_signal(
|
|
self, tmp_path: Path, signum: signal.Signals, chars: str
|
|
) -> None:
|
|
client = UnixLocalSandboxClient()
|
|
manifest = Manifest(root=str(tmp_path / "workspace"))
|
|
previous_handler = signal.getsignal(signum)
|
|
|
|
signal.signal(signum, signal.SIG_IGN)
|
|
try:
|
|
async with await client.create(
|
|
manifest=manifest, snapshot=None, options=None
|
|
) as session:
|
|
started = await session.pty_exec_start(
|
|
"sleep",
|
|
"30",
|
|
shell=False,
|
|
tty=True,
|
|
yield_time_s=0.05,
|
|
)
|
|
assert started.process_id is not None
|
|
|
|
interrupted = await session.pty_write_stdin(
|
|
session_id=started.process_id,
|
|
chars=chars,
|
|
yield_time_s=5.5,
|
|
)
|
|
|
|
assert interrupted.process_id is None
|
|
assert interrupted.exit_code == -signum
|
|
finally:
|
|
signal.signal(signum, previous_handler)
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_non_tty_pty_session_rejects_stdin_and_can_still_be_polled(
|
|
self, tmp_path: Path
|
|
) -> None:
|
|
client = UnixLocalSandboxClient()
|
|
manifest = Manifest(root=str(tmp_path / "workspace"))
|
|
|
|
async with await client.create(manifest=manifest, snapshot=None, options=None) as session:
|
|
started = await session.pty_exec_start(
|
|
"sh",
|
|
"-c",
|
|
"printf 'stdout\\n'; printf 'stderr\\n' >&2; sleep 1",
|
|
shell=False,
|
|
tty=False,
|
|
yield_time_s=0.05,
|
|
)
|
|
|
|
assert started.process_id is not None
|
|
assert started.exit_code is None
|
|
started_text = started.output.decode("utf-8", errors="replace")
|
|
assert "stdout" in started_text
|
|
assert "stderr" in started_text
|
|
|
|
with pytest.raises(RuntimeError, match="stdin is not available for this process"):
|
|
await session.pty_write_stdin(session_id=started.process_id, chars="hello")
|
|
|
|
finished = await session.pty_write_stdin(
|
|
session_id=started.process_id,
|
|
chars="",
|
|
yield_time_s=5.5,
|
|
)
|
|
text = finished.output.decode("utf-8", errors="replace")
|
|
assert finished.process_id is None
|
|
assert finished.exit_code == 0
|
|
assert text == ""
|
|
|
|
with pytest.raises(PtySessionNotFoundError):
|
|
await session.pty_write_stdin(session_id=started.process_id, chars="")
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_stop_terminates_active_pty_sessions(self, tmp_path: Path) -> None:
|
|
client = UnixLocalSandboxClient()
|
|
manifest = Manifest(root=str(tmp_path / "workspace"))
|
|
|
|
session = await client.create(manifest=manifest, snapshot=None, options=None)
|
|
await session.start()
|
|
started = await session.pty_exec_start(
|
|
"sh",
|
|
"-c",
|
|
"printf 'ready\\n'; sleep 30",
|
|
shell=False,
|
|
tty=True,
|
|
yield_time_s=0.25,
|
|
)
|
|
|
|
assert started.process_id is not None
|
|
assert "ready" in started.output.decode("utf-8", errors="replace")
|
|
|
|
await session.stop()
|
|
|
|
with pytest.raises(PtySessionNotFoundError):
|
|
await session.pty_write_stdin(session_id=started.process_id, chars="")
|
|
|
|
|
|
class TestUnixLocalUserScopedFilesystem:
|
|
@pytest.mark.asyncio
|
|
async def test_mkdir_as_user_checks_permissions_then_uses_local_fs(
|
|
self,
|
|
tmp_path: Path,
|
|
) -> None:
|
|
workspace = tmp_path / "workspace"
|
|
workspace.mkdir()
|
|
session = _RecordingUnixLocalSession(workspace)
|
|
|
|
await session.mkdir("nested", user=User(name="sandbox-user"))
|
|
|
|
assert (workspace / "nested").is_dir()
|
|
assert len(session.exec_commands) == 1
|
|
assert session.exec_commands[0][:4] == ("sudo", "-u", "sandbox-user", "--")
|
|
assert session.exec_commands[0][4:6] == ("sh", "-lc")
|
|
assert session.exec_commands[0][-2:] == (str(workspace / "nested"), "0")
|
|
assert not any(part.startswith("mkdir ") for part in session.exec_commands[0])
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_rm_as_user_checks_permissions_then_uses_local_fs(
|
|
self,
|
|
tmp_path: Path,
|
|
) -> None:
|
|
workspace = tmp_path / "workspace"
|
|
workspace.mkdir()
|
|
target = workspace / "stale.txt"
|
|
target.write_text("stale", encoding="utf-8")
|
|
session = _RecordingUnixLocalSession(workspace)
|
|
|
|
await session.rm("stale.txt", user=User(name="sandbox-user"))
|
|
|
|
assert not target.exists()
|
|
assert len(session.exec_commands) == 1
|
|
assert session.exec_commands[0][:4] == ("sudo", "-u", "sandbox-user", "--")
|
|
assert session.exec_commands[0][4:6] == ("sh", "-lc")
|
|
assert session.exec_commands[0][-2:] == (str(target), "0")
|
|
assert not any(part.startswith("rm ") for part in session.exec_commands[0])
|