diff --git a/sdk/src/beta9/abstractions/sandbox.py b/sdk/src/beta9/abstractions/sandbox.py index 1c759ba16..23e83aa87 100644 --- a/sdk/src/beta9/abstractions/sandbox.py +++ b/sdk/src/beta9/abstractions/sandbox.py @@ -1,6 +1,7 @@ import asyncio import atexit import io +import posixpath import shlex import time from dataclasses import dataclass, field @@ -1730,6 +1731,108 @@ class SandboxFileSystem: def __init__(self, sandbox_instance: SandboxInstance): self.sandbox_instance = sandbox_instance + @staticmethod + def _normalize_sandbox_path(sandbox_path: str) -> str: + if not sandbox_path or "\x00" in sandbox_path: + raise ValueError("sandbox_path must be a non-empty POSIX path") + + normalized = posixpath.normpath(sandbox_path) + if not normalized.startswith("/"): + normalized = f"/{normalized}" + + return normalized + + @staticmethod + def _parent_directories(sandbox_path: str) -> List[str]: + parents = [] + parent = posixpath.dirname(sandbox_path) + while parent and parent != "/": + parents.append(parent) + parent = posixpath.dirname(parent) + return list(reversed(parents)) + + def create_file( + self, + sandbox_path: str, + contents: Union[str, bytes] = "", + *, + mode: int = 0o644, + parents: bool = False, + overwrite: bool = False, + encoding: str = "utf-8", + ): + """ + Create a file in the sandbox from in-memory contents. + + Parameters: + sandbox_path (str): The destination path within the sandbox. + contents (Union[str, bytes]): File contents. Strings are encoded with + the provided encoding. + mode (int): File permissions. Default is 0o644. + parents (bool): Create missing parent directories first. + overwrite (bool): Allow replacing an existing file. + encoding (str): Encoding used when contents is a string. + + Raises: + SandboxFileSystemError: If file creation fails. + FileExistsError: If the file exists and overwrite is False. + ValueError: If the path or contents are invalid. + + Example: + ```python + fs.create_file("/workspace/app.py", "print('hello')\n", parents=True) + fs.create_file("/workspace/model.bin", b"\x00\x01", overwrite=True) + ``` + """ + sandbox_path = self._normalize_sandbox_path(sandbox_path) + + if not isinstance(contents, (str, bytes)): + raise ValueError("contents must be str or bytes") + + if not 0 <= mode <= 0o777: + raise ValueError("mode must be between 0o000 and 0o777") + + if parents: + for parent in self._parent_directories(sandbox_path): + try: + info = self.stat_file(parent) + except SandboxFileSystemError: + self.create_directory(parent) + else: + if not info.is_dir: + raise SandboxFileSystemError( + message=f"Parent path is not a directory: {parent}", + operation="create_file", + path=sandbox_path, + container_id=self.sandbox_instance.container_id, + ) + + if not overwrite: + try: + self.stat_file(sandbox_path) + except SandboxFileSystemError: + pass + else: + raise FileExistsError(f"Sandbox file already exists: {sandbox_path}") + + data = contents.encode(encoding) if isinstance(contents, str) else contents + response = self.sandbox_instance.stub.sandbox_upload_file( + PodSandboxUploadFileRequest( + container_id=self.sandbox_instance.container_id, + container_path=sandbox_path, + data=data, + mode=mode, + ) + ) + + if not response.ok: + raise SandboxFileSystemError( + message=response.error_msg, + operation="create_file", + path=sandbox_path, + container_id=self.sandbox_instance.container_id, + ) + def upload_file(self, local_path: str, sandbox_path: str): """ Upload a local file to the sandbox. @@ -3567,6 +3670,43 @@ async def upload_file(self, local_path: str, sandbox_path: str): """ return await asyncio.to_thread(self._sync.upload_file, local_path, sandbox_path) + async def create_file( + self, + sandbox_path: str, + contents: Union[str, bytes] = "", + *, + mode: int = 0o644, + parents: bool = False, + overwrite: bool = False, + encoding: str = "utf-8", + ): + """ + Create a file in the sandbox from in-memory contents asynchronously. + + Parameters: + sandbox_path (str): The destination path within the sandbox. + contents (Union[str, bytes]): File contents. Strings are encoded with + the provided encoding. + mode (int): File permissions. Default is 0o644. + parents (bool): Create missing parent directories first. + overwrite (bool): Allow replacing an existing file. + encoding (str): Encoding used when contents is a string. + + Raises: + SandboxFileSystemError: If file creation fails. + FileExistsError: If the file exists and overwrite is False. + ValueError: If the path or contents are invalid. + """ + return await asyncio.to_thread( + self._sync.create_file, + sandbox_path, + contents, + mode=mode, + parents=parents, + overwrite=overwrite, + encoding=encoding, + ) + async def write_bytes(self, sandbox_path: str, data: bytes, mode: int = 644): """ Write bytes to a sandbox file asynchronously. diff --git a/sdk/tests/test_sandbox_filesystem.py b/sdk/tests/test_sandbox_filesystem.py new file mode 100644 index 000000000..3df1237d5 --- /dev/null +++ b/sdk/tests/test_sandbox_filesystem.py @@ -0,0 +1,134 @@ +import asyncio +import unittest +from dataclasses import dataclass + +from beta9.abstractions.sandbox import SandboxFileInfo, SandboxFileSystem +from beta9.clients.pod import ( + PodSandboxCreateDirectoryResponse, + PodSandboxStatFileResponse, + PodSandboxUploadFileResponse, +) +from beta9.exceptions import SandboxFileSystemError + + +@dataclass +class FakeSandboxInstance: + container_id: str = "sandbox-123" + + def __post_init__(self): + self.stub = FakePodStub() + + +class FakePodStub: + def __init__(self): + self.files = {} + self.dirs = {"/"} + self.uploads = [] + self.created_dirs = [] + + def sandbox_stat_file(self, request): + path = request.container_path + if path in self.dirs: + return PodSandboxStatFileResponse( + ok=True, + file_info=SandboxFileInfo( + name=path, + is_dir=True, + size=0, + mode=0o755, + mod_time=0, + owner="root", + group="root", + permissions=0o755, + ), + ) + + if path in self.files: + return PodSandboxStatFileResponse( + ok=True, + file_info=SandboxFileInfo( + name=path, + is_dir=False, + size=len(self.files[path]), + mode=0o644, + mod_time=0, + owner="root", + group="root", + permissions=0o644, + ), + ) + + return PodSandboxStatFileResponse(ok=False, error_msg="not found") + + def sandbox_create_directory(self, request): + self.created_dirs.append(request.container_path) + self.dirs.add(request.container_path) + return PodSandboxCreateDirectoryResponse(ok=True) + + def sandbox_upload_file(self, request): + self.uploads.append(request) + self.files[request.container_path] = request.data + return PodSandboxUploadFileResponse(ok=True) + + +class TestSandboxFileSystemCreateFile(unittest.TestCase): + def setUp(self): + self.instance = FakeSandboxInstance() + self.fs = SandboxFileSystem(self.instance) + + def test_create_file_uploads_string_contents(self): + self.fs.create_file("/workspace/app.py", "print('hi')\n") + + upload = self.instance.stub.uploads[-1] + self.assertEqual(upload.container_id, "sandbox-123") + self.assertEqual(upload.container_path, "/workspace/app.py") + self.assertEqual(upload.data, b"print('hi')\n") + self.assertEqual(upload.mode, 0o644) + + def test_create_file_accepts_bytes_and_mode(self): + self.fs.create_file("tmp/model.bin", b"\x00\x01", mode=0o600) + + upload = self.instance.stub.uploads[-1] + self.assertEqual(upload.container_path, "/tmp/model.bin") + self.assertEqual(upload.data, b"\x00\x01") + self.assertEqual(upload.mode, 0o600) + + def test_create_file_creates_missing_parents_in_order(self): + self.fs.create_file("/workspace/src/app.py", "print(1)", parents=True) + + self.assertEqual(self.instance.stub.created_dirs, ["/workspace", "/workspace/src"]) + self.assertIn("/workspace/src/app.py", self.instance.stub.files) + + def test_create_file_refuses_to_overwrite_by_default(self): + self.instance.stub.files["/workspace/app.py"] = b"old" + + with self.assertRaises(FileExistsError): + self.fs.create_file("/workspace/app.py", "new") + + self.assertEqual(self.instance.stub.files["/workspace/app.py"], b"old") + + def test_create_file_allows_explicit_overwrite(self): + self.instance.stub.files["/workspace/app.py"] = b"old" + + self.fs.create_file("/workspace/app.py", "new", overwrite=True) + + self.assertEqual(self.instance.stub.files["/workspace/app.py"], b"new") + + def test_create_file_rejects_parent_that_is_file(self): + self.instance.stub.files["/workspace"] = b"not a dir" + + with self.assertRaises(SandboxFileSystemError): + self.fs.create_file("/workspace/app.py", "print(1)", parents=True) + + def test_create_file_rejects_invalid_contents(self): + with self.assertRaises(ValueError): + self.fs.create_file("/workspace/app.py", object()) # type: ignore[arg-type] + + def test_async_create_file_delegates_to_sync_filesystem(self): + asyncio.run(self.fs.aio.create_file("/workspace/async.py", "print('async')")) + + self.assertEqual(self.instance.stub.files["/workspace/async.py"], b"print('async')") + + +if __name__ == "__main__": + unittest.main()