wehub-resource-sync
153 lines
4.6 KiB
Python
153 lines
4.6 KiB
Python
from __future__ import annotations
|
|
|
|
import io
|
|
import uuid
|
|
from pathlib import Path
|
|
|
|
from agents.sandbox import Manifest
|
|
from agents.sandbox.errors import WorkspaceReadNotFoundError
|
|
from agents.sandbox.session.base_sandbox_session import BaseSandboxSession
|
|
from agents.sandbox.snapshot import NoopSnapshot
|
|
from agents.sandbox.types import ExecResult, User
|
|
from tests.utils.factories import TestSessionState
|
|
|
|
|
|
class ApplyPatchSession(BaseSandboxSession):
|
|
def __init__(self, manifest: Manifest | None = None) -> None:
|
|
self.state = TestSessionState(
|
|
manifest=manifest or Manifest(root="/workspace"),
|
|
snapshot=NoopSnapshot(id=str(uuid.uuid4())),
|
|
)
|
|
self.files: dict[Path, bytes] = {}
|
|
self.mkdir_calls: list[tuple[Path, bool]] = []
|
|
self.rm_calls: list[tuple[Path, bool]] = []
|
|
|
|
async def start(self) -> None:
|
|
return None
|
|
|
|
async def stop(self) -> None:
|
|
return None
|
|
|
|
async def shutdown(self) -> None:
|
|
return None
|
|
|
|
async def running(self) -> bool:
|
|
return True
|
|
|
|
async def read(self, path: Path, *, user: str | User | None = None) -> io.BytesIO:
|
|
_ = user
|
|
normalized = self.normalize_path(path)
|
|
if normalized not in self.files:
|
|
raise FileNotFoundError(normalized)
|
|
return io.BytesIO(self.files[normalized])
|
|
|
|
async def write(
|
|
self,
|
|
path: Path,
|
|
data: io.IOBase,
|
|
*,
|
|
user: str | User | None = None,
|
|
) -> None:
|
|
_ = user
|
|
normalized = self.normalize_path(path)
|
|
payload = data.read()
|
|
if isinstance(payload, str):
|
|
self.files[normalized] = payload.encode("utf-8")
|
|
else:
|
|
self.files[normalized] = bytes(payload)
|
|
|
|
async def _exec_internal(
|
|
self,
|
|
*command: str | Path,
|
|
timeout: float | None = None,
|
|
) -> ExecResult:
|
|
_ = (command, timeout)
|
|
raise AssertionError("_exec_internal() should not be called")
|
|
|
|
async def persist_workspace(self) -> io.IOBase:
|
|
return io.BytesIO()
|
|
|
|
async def hydrate_workspace(self, data: io.IOBase) -> None:
|
|
_ = data
|
|
|
|
async def mkdir(
|
|
self,
|
|
path: Path | str,
|
|
*,
|
|
parents: bool = False,
|
|
user: str | User | None = None,
|
|
) -> None:
|
|
_ = user
|
|
normalized = self.normalize_path(path)
|
|
self.mkdir_calls.append((normalized, parents))
|
|
|
|
async def rm(
|
|
self,
|
|
path: Path | str,
|
|
*,
|
|
recursive: bool = False,
|
|
user: str | User | None = None,
|
|
) -> None:
|
|
_ = user
|
|
normalized = self.normalize_path(path)
|
|
self.rm_calls.append((normalized, recursive))
|
|
self.files.pop(normalized, None)
|
|
|
|
|
|
class ProviderNotFoundApplyPatchSession(ApplyPatchSession):
|
|
async def read(self, path: Path, *, user: str | User | None = None) -> io.BytesIO:
|
|
try:
|
|
return await super().read(path, user=user)
|
|
except FileNotFoundError as exc:
|
|
workspace_path = self.normalize_path(path).relative_to("/")
|
|
raise WorkspaceReadNotFoundError(
|
|
path=Path("/provider/private/root") / workspace_path
|
|
) from exc
|
|
|
|
|
|
class UserRecordingApplyPatchSession(ApplyPatchSession):
|
|
def __init__(self, manifest: Manifest | None = None) -> None:
|
|
super().__init__(manifest)
|
|
self.read_users: list[str | None] = []
|
|
self.write_users: list[str | None] = []
|
|
self.mkdir_users: list[str | None] = []
|
|
self.rm_users: list[str | None] = []
|
|
|
|
@staticmethod
|
|
def _user_name(user: str | User | None) -> str | None:
|
|
return user.name if isinstance(user, User) else user
|
|
|
|
async def read(self, path: Path, *, user: str | User | None = None) -> io.BytesIO:
|
|
self.read_users.append(self._user_name(user))
|
|
return await super().read(path)
|
|
|
|
async def write(
|
|
self,
|
|
path: Path,
|
|
data: io.IOBase,
|
|
*,
|
|
user: str | User | None = None,
|
|
) -> None:
|
|
self.write_users.append(self._user_name(user))
|
|
await super().write(path, data)
|
|
|
|
async def mkdir(
|
|
self,
|
|
path: Path | str,
|
|
*,
|
|
parents: bool = False,
|
|
user: str | User | None = None,
|
|
) -> None:
|
|
self.mkdir_users.append(self._user_name(user))
|
|
await super().mkdir(path, parents=parents)
|
|
|
|
async def rm(
|
|
self,
|
|
path: Path | str,
|
|
*,
|
|
recursive: bool = False,
|
|
user: str | User | None = None,
|
|
) -> None:
|
|
self.rm_users.append(self._user_name(user))
|
|
await super().rm(path, recursive=recursive)
|