arcbox 0.1.1__py3-none-any.whl

arcbox/__init__.py

@@ -0,0 +1,56 @@
+"""arcbox — run isolated microVM sandboxes on the local ArcBox daemon.
+
+Two mirrored surfaces: the sync classes (`Sandbox`, `ArcBox`, ...) and
+their `Async*` counterparts. Public shapes are hand-written and mapped
+at the transport boundary; everything under `arcbox._gen` is generated
+wire code and is deliberately NOT exported. The error hierarchy lives
+in :mod:`arcbox.errors`.
+"""
+
+from arcbox._async._client import AsyncConnectClient
+from arcbox._async.commands import AsyncCommandHandle, AsyncCommands, AsyncOutputStream
+from arcbox._async.files import AsyncFiles
+from arcbox._async.sandbox import AsyncArcBox, AsyncSandbox
+from arcbox._connection import Connection
+from arcbox._sync._client import ConnectClient
+from arcbox._sync.commands import CommandHandle, Commands, OutputStream
+from arcbox._sync.files import Files
+from arcbox._sync.sandbox import ArcBox, Sandbox
+from arcbox._types import (
+    MAX_FILE_BYTES,
+    CommandResult,
+    IdlePolicy,
+    OutputChannel,
+    OutputChunk,
+    SandboxInfo,
+    SandboxState,
+    SandboxSummary,
+    SignalName,
+)
+
+__all__ = [
+    "MAX_FILE_BYTES",
+    "ArcBox",
+    "AsyncArcBox",
+    "AsyncCommandHandle",
+    "AsyncCommands",
+    "AsyncConnectClient",
+    "AsyncFiles",
+    "AsyncOutputStream",
+    "AsyncSandbox",
+    "CommandHandle",
+    "CommandResult",
+    "Commands",
+    "ConnectClient",
+    "Connection",
+    "Files",
+    "IdlePolicy",
+    "OutputChannel",
+    "OutputChunk",
+    "OutputStream",
+    "Sandbox",
+    "SandboxInfo",
+    "SandboxState",
+    "SandboxSummary",
+    "SignalName",
+]

arcbox/_async/__init__.py

@@ -0,0 +1,5 @@
+"""One flavor of the SDK surface.
+
+The async core lives in `arcbox/_async`; `arcbox/_sync` is generated
+from it by scripts/gen_sync.py — edit only the async tree.
+"""

arcbox/_async/_client.py

@@ -0,0 +1,235 @@
+"""Connect-over-httpx transport.
+
+Unary RPCs are plain POSTs with binary-protobuf bodies
+(``application/proto``); streaming RPCs speak the Connect streaming
+envelope (``application/connect+proto``, framing in ``_envelope``).
+The local tier dials the daemon's Unix socket through httpx's ``uds``
+transport; the same code path serves a future remote tier over HTTPS.
+"""
+
+from __future__ import annotations
+
+import os
+from typing import TYPE_CHECKING, Generic, TypeVar
+
+import httpx
+
+from arcbox._connection import Connection, resolve_connection
+from arcbox._envelope import (
+    FLAG_COMPRESSED,
+    FLAG_END_STREAM,
+    EnvelopeDecoder,
+    encode_envelope,
+    end_stream_error,
+    unary_error,
+)
+from arcbox.errors import ArcBoxError, InvalidArgumentError
+
+if TYPE_CHECKING:
+    from collections.abc import AsyncIterator, Iterable, Mapping
+    from contextlib import AbstractAsyncContextManager
+    from types import TracebackType
+
+    from google.protobuf.message import Message
+
+M = TypeVar("M", bound="Message")
+
+_UNARY_HEADERS = {
+    "content-type": "application/proto",
+    "connect-protocol-version": "1",
+}
+_STREAM_HEADERS = {
+    "content-type": "application/connect+proto",
+    "connect-protocol-version": "1",
+}
+
+
+class _UseDefault:
+    """Sentinel: apply the connection's ``request_timeout``."""
+
+    def __repr__(self) -> str:
+        # Stable across the async and sync trees so signatures compare
+        # equal in the parity test (and read cleanly in help()).
+        return "USE_DEFAULT"
+
+
+_USE_DEFAULT = _UseDefault()
+
+
+class AsyncConnectClient:
+    """One resolved connection to a daemon; every handle shares it."""
+
+    def __init__(
+        self,
+        connection: Connection | None = None,
+        env: Mapping[str, str] | None = None,
+    ) -> None:
+        resolved = resolve_connection(connection, os.environ if env is None else env)
+        #: Per-unary-RPC deadline in seconds, when configured. Streams
+        #: and long-polls never receive it.
+        self.request_timeout = resolved.request_timeout
+        injected = connection.http_client if connection is not None else None
+        #: Whether this client constructed (and must eventually close)
+        #: its HTTP client; an injected one belongs to the caller.
+        self._owns_http = injected is None
+        if injected is not None:
+            required = httpx.AsyncClient
+            if not isinstance(injected, required):
+                raise InvalidArgumentError(
+                    f"Connection.http_client must be an httpx.{required.__name__} for this surface"
+                )
+            self._http = injected
+            return
+        headers = {"authorization": f"Bearer {resolved.api_key}"} if resolved.api_key else None
+        transport = (
+            httpx.AsyncHTTPTransport(uds=resolved.socket_path)
+            if resolved.socket_path is not None
+            else None
+        )
+        self._http = httpx.AsyncClient(
+            base_url=resolved.base_url,
+            transport=transport,
+            headers=headers,
+            timeout=None,
+        )
+
+    async def aclose(self) -> None:
+        """Release the HTTP client's pooled connections. A no-op when the
+        client was injected — it belongs to the caller."""
+        if self._owns_http:
+            await self._http.aclose()
+
+    async def unary(
+        self,
+        path: str,
+        request: Message,
+        response_type: type[M],
+        timeout: float | _UseDefault | None = _USE_DEFAULT,
+    ) -> M:
+        """One unary RPC. ``timeout`` overrides the configured deadline;
+        an explicit ``None`` means no deadline (checkpoint restores take
+        as long as they take)."""
+        effective = self.request_timeout if isinstance(timeout, _UseDefault) else timeout
+        headers = dict(_UNARY_HEADERS)
+        if effective is not None:
+            headers["connect-timeout-ms"] = str(int(effective * 1000))
+        response = await self._http.post(
+            path,
+            content=request.SerializeToString(),
+            headers=headers,
+            timeout=None if effective is None else httpx.Timeout(effective),
+        )
+        if response.status_code != 200:
+            raise unary_error(response.status_code, response.content)
+        content_type = response.headers.get("content-type", "")
+        if not content_type.startswith("application/proto"):
+            raise ArcBoxError(f"unexpected unary response content-type {content_type!r}")
+        message = response_type()
+        message.ParseFromString(response.content)
+        return message
+
+    def stream(self, path: str, request: Message, response_type: type[M]) -> AsyncServerStream[M]:
+        """One server-streaming RPC. Entering the returned context sends
+        the request (that is what registers a subscription server-side);
+        iterate it for messages."""
+        return AsyncServerStream(self._http, path, request, response_type)
+
+    async def client_stream(
+        self, path: str, requests: Iterable[Message], response_type: type[M]
+    ) -> M:
+        """One client-streaming RPC with the full request sequence known
+        up front: the envelopes are sent as one body and the single
+        response message is returned."""
+        body = b"".join(encode_envelope(0, m.SerializeToString()) for m in requests)
+        response = await self._http.post(path, content=body, headers=_STREAM_HEADERS, timeout=None)
+        if response.status_code != 200:
+            raise unary_error(response.status_code, response.content)
+        decoder = EnvelopeDecoder()
+        message: M | None = None
+        ended = False
+        for flags, payload in decoder.feed(response.content):
+            if flags & FLAG_COMPRESSED:
+                raise ArcBoxError("received a compressed frame without negotiating compression")
+            if flags & FLAG_END_STREAM:
+                error = end_stream_error(payload)
+                if error is not None:
+                    raise error
+                ended = True
+            else:
+                decoded = response_type()
+                decoded.ParseFromString(payload)
+                message = decoded
+        if not ended:
+            # Without the terminal frame a truncated body is
+            # indistinguishable from success — never report one as
+            # completed (WriteFile rides this path).
+            raise ArcBoxError("the stream ended without an EndStreamResponse")
+        if message is None:
+            raise ArcBoxError("the stream ended without a response message")
+        return message
+
+
+class AsyncServerStream(Generic[M]):
+    """One server-streaming call: a context manager yielding an iterator.
+
+    The stream ends deterministically — with the terminal
+    EndStreamResponse frame or a typed error, never silence.
+    """
+
+    def __init__(
+        self,
+        http: httpx.AsyncClient,
+        path: str,
+        request: Message,
+        response_type: type[M],
+    ) -> None:
+        self._http = http
+        self._path = path
+        self._body = encode_envelope(0, request.SerializeToString())
+        self._response_type = response_type
+        self._cm: AbstractAsyncContextManager[httpx.Response] | None = None
+        self._response: httpx.Response | None = None
+
+    async def __aenter__(self) -> AsyncServerStream[M]:
+        cm = self._http.stream(
+            "POST", self._path, content=self._body, headers=_STREAM_HEADERS, timeout=None
+        )
+        response = await cm.__aenter__()
+        if response.status_code != 200:
+            body = await response.aread()
+            await cm.__aexit__(None, None, None)
+            raise unary_error(response.status_code, body)
+        self._cm = cm
+        self._response = response
+        return self
+
+    async def __aexit__(
+        self,
+        exc_type: type[BaseException] | None,
+        exc: BaseException | None,
+        tb: TracebackType | None,
+    ) -> None:
+        if self._cm is not None:
+            await self._cm.__aexit__(exc_type, exc, tb)
+
+    def __aiter__(self) -> AsyncIterator[M]:
+        return self._messages()
+
+    async def _messages(self) -> AsyncIterator[M]:
+        response = self._response
+        if response is None:
+            raise ArcBoxError("stream iterated before it was entered")
+        decoder = EnvelopeDecoder()
+        async for data in response.aiter_bytes():
+            for flags, payload in decoder.feed(data):
+                if flags & FLAG_COMPRESSED:
+                    raise ArcBoxError("received a compressed frame without negotiating compression")
+                if flags & FLAG_END_STREAM:
+                    error = end_stream_error(payload)
+                    if error is not None:
+                        raise error
+                    return
+                message = self._response_type()
+                message.ParseFromString(payload)
+                yield message
+        raise ArcBoxError("the stream ended without an EndStreamResponse")

arcbox/_async/commands.py

@@ -0,0 +1,322 @@
+"""The ``sandbox.commands`` namespace: run processes inside one sandbox."""
+
+from __future__ import annotations
+
+import asyncio
+import math
+import time
+import uuid
+from typing import TYPE_CHECKING, Literal, overload
+
+from google.protobuf import empty_pb2
+
+from arcbox._boundary import wrap_errors
+from arcbox._gen import process_pb2
+from arcbox._types import (
+    SIGNAL_VALUES,
+    CommandResult,
+    OutputChunk,
+    command_result_from_execution,
+)
+from arcbox.errors import InvalidArgumentError, TimeoutError
+
+if TYPE_CHECKING:
+    from collections.abc import AsyncGenerator, AsyncIterator, Mapping, Sequence
+    from types import TracebackType
+
+    from arcbox._types import OutputChannel, SignalName
+
+    from ._client import AsyncConnectClient, AsyncServerStream
+
+_PROCESS = "/arcbox.sandbox.v1.SandboxProcessService/"
+
+#: Longest single WaitExecution long-poll slice, so a dropped daemon
+#: surfaces as an error instead of an infinite silent wait.
+_WAIT_SLICE_SECONDS = 30
+
+
+def _normalize_cmd(cmd: str | Sequence[str]) -> list[str]:
+    """``str`` is sugar for a shell command; a sequence is argv, executed
+    directly with no shell involved."""
+    return ["/bin/sh", "-lc", cmd] if isinstance(cmd, str) else list(cmd)
+
+
+def _channel_name(channel: int) -> OutputChannel:
+    if channel == process_pb2.STDIO_CHANNEL_STDERR:
+        return "stderr"
+    if channel == process_pb2.STDIO_CHANNEL_PTY:
+        return "pty"
+    return "stdout"
+
+
+def _decode_output(chunks: list[bytes]) -> str:
+    return b"".join(chunks).decode("utf-8", errors="replace")
+
+
+class AsyncOutputStream:
+    """A command's output, iterable chunk by chunk.
+
+    Consuming it to the end releases the underlying HTTP stream
+    deterministically. When exiting early (``break``), iterate inside
+    the context-manager form — ``async with handle.output as stream`` in
+    the async flavor, ``with`` in the sync one — so the stream closes at
+    the break instead of whenever the generator finalizer runs."""
+
+    def __init__(self, chunks: AsyncGenerator[OutputChunk]) -> None:
+        self._chunks = chunks
+
+    def __aiter__(self) -> AsyncIterator[OutputChunk]:
+        return self._chunks
+
+    async def aclose(self) -> None:
+        """Release the underlying HTTP stream without consuming the rest."""
+        await self._chunks.aclose()
+
+    async def __aenter__(self) -> AsyncOutputStream:
+        return self
+
+    async def __aexit__(
+        self,
+        exc_type: type[BaseException] | None,
+        exc: BaseException | None,
+        tb: TracebackType | None,
+    ) -> None:
+        await self.aclose()
+
+
+class AsyncCommandHandle:
+    """A handle to a running (or finished) command. The process is
+    decoupled from this object: dropping the handle never kills the
+    process."""
+
+    def __init__(self, client: AsyncConnectClient, sandbox_id: str, command_id: str) -> None:
+        self._client = client
+        self._sandbox_id = sandbox_id
+        #: Execution id, unique within the sandbox; addressable across clients.
+        self.command_id = command_id
+
+    @property
+    def output(self) -> AsyncOutputStream:
+        """Stream the command's output from the beginning — or from the
+        earliest byte the daemon still retains (8 MiB per channel);
+        replayed buffered output comes first, then live output follows;
+        the stream ends when the process exits (deterministic
+        termination — never silence). When breaking out early, iterate
+        via the context-manager form (see :class:`AsyncOutputStream`)."""
+        return AsyncOutputStream(self._stream_output())
+
+    async def _stream_output(self) -> AsyncGenerator[OutputChunk]:
+        with wrap_errors("commands.output"):
+            async with self._attach() as stream:
+                async for event in stream:
+                    kind = event.WhichOneof("event")
+                    if kind == "output":
+                        chunk = event.output
+                        yield OutputChunk(_channel_name(chunk.channel), chunk.data)
+                    elif kind == "exited":
+                        return
+
+    async def wait_for_exit(self, timeout: float | None = None) -> CommandResult:
+        """Wait until the command exits and return its result (server-side
+        long-poll; no client-side spinning). ``timeout`` bounds the WAIT,
+        not the process: on expiry a :class:`TimeoutError` is raised and
+        the process keeps running."""
+        deadline = None if timeout is None else time.monotonic() + timeout
+        with wrap_errors("commands.wait_for_exit"):
+            while True:
+                remaining = None if deadline is None else deadline - time.monotonic()
+                if remaining is not None and remaining <= 0:
+                    raise TimeoutError(
+                        "wait_for_exit(timeout) elapsed before the command exited",
+                        suggestion=(
+                            "increase the wait_for_exit timeout argument, or kill() the command"
+                        ),
+                        context={"command_id": self.command_id},
+                    )
+                if remaining is None:
+                    execution = await self._wait_slice(_WAIT_SLICE_SECONDS)
+                elif remaining >= 1:
+                    execution = await self._wait_slice(min(_WAIT_SLICE_SECONDS, int(remaining)))
+                else:
+                    # The wire's wait granularity is whole seconds. Poll
+                    # immediately (an already-exited command returns at
+                    # once); only if it is still running, sleep out the
+                    # remainder and take one final poll so an exit inside
+                    # the budget is still caught — past it, the deadline
+                    # check above raises.
+                    execution = await self._wait_slice(0)
+                    if execution.state != process_pb2.EXECUTION_STATE_EXITED:
+                        await asyncio.sleep(remaining)
+                        execution = await self._wait_slice(0)
+                if execution.state == process_pb2.EXECUTION_STATE_EXITED:
+                    break
+            return await self._collect_result(execution)
+
+    async def _wait_slice(self, slice_seconds: int) -> process_pb2.Execution:
+        """One WaitExecution long-poll (0 = an immediate state poll)."""
+        return await self._client.unary(
+            _PROCESS + "WaitExecution",
+            process_pb2.WaitExecutionRequest(
+                sandbox_id=self._sandbox_id,
+                execution_id=self.command_id,
+                timeout_seconds=slice_seconds,
+            ),
+            process_pb2.Execution,
+            # Exempt from request_timeout: this unary deliberately parks
+            # server-side for the slice; grant it that long plus grace.
+            timeout=float(slice_seconds + 5),
+        )
+
+    async def kill(self, signal: SignalName = "SIGTERM") -> None:
+        """Deliver a signal to the whole process group (default SIGTERM)."""
+        with wrap_errors("commands.kill"):
+            await self._client.unary(
+                _PROCESS + "SignalExecution",
+                process_pb2.SignalExecutionRequest(
+                    sandbox_id=self._sandbox_id,
+                    execution_id=self.command_id,
+                    signal=SIGNAL_VALUES[signal],
+                ),
+                empty_pb2.Empty,
+            )
+
+    def _attach(self) -> AsyncServerStream[process_pb2.ExecutionEvent]:
+        return self._client.stream(
+            _PROCESS + "AttachExecution",
+            process_pb2.AttachExecutionRequest(
+                sandbox_id=self._sandbox_id,
+                execution_id=self.command_id,
+                stdout_offset=0,
+                stderr_offset=0,
+            ),
+            process_pb2.ExecutionEvent,
+        )
+
+    async def _collect_result(self, execution: process_pb2.Execution) -> CommandResult:
+        """Assemble the result of an exited execution. Output is re-read
+        from offset 0 — the daemon retains and replays it, so the result
+        is complete even when nobody consumed the live stream, UNLESS the
+        command outgrew the daemon's per-channel retention (8 MiB): chunk
+        offsets expose the dropped head, reported as ``truncated``."""
+        stdout: list[bytes] = []
+        stderr: list[bytes] = []
+        next_stdout = 0
+        next_stderr = 0
+        truncated = False
+        async with self._attach() as stream:
+            async for event in stream:
+                kind = event.WhichOneof("event")
+                if kind == "output":
+                    chunk = event.output
+                    is_stderr = chunk.channel == process_pb2.STDIO_CHANNEL_STDERR
+                    # A chunk landing past the expected offset means
+                    # retention already dropped bytes we asked for.
+                    if chunk.offset > (next_stderr if is_stderr else next_stdout):
+                        truncated = True
+                    after = chunk.offset + len(chunk.data)
+                    if is_stderr:
+                        next_stderr = after
+                        stderr.append(chunk.data)
+                    else:
+                        next_stdout = after
+                        stdout.append(chunk.data)
+                elif kind == "exited":
+                    break
+        return command_result_from_execution(
+            execution, _decode_output(stdout), _decode_output(stderr), truncated
+        )
+
+
+class AsyncCommands:
+    """Run processes inside one sandbox."""
+
+    def __init__(self, client: AsyncConnectClient, sandbox_id: str) -> None:
+        self._client = client
+        self._sandbox_id = sandbox_id
+
+    @overload
+    async def run(
+        self,
+        cmd: str | Sequence[str],
+        *,
+        cwd: str | None = None,
+        env: Mapping[str, str] | None = None,
+        user: str | None = None,
+        timeout: float | None = None,
+        check: bool = False,
+        background: Literal[False] = False,
+    ) -> CommandResult: ...
+
+    @overload
+    async def run(
+        self,
+        cmd: str | Sequence[str],
+        *,
+        cwd: str | None = None,
+        env: Mapping[str, str] | None = None,
+        user: str | None = None,
+        timeout: float | None = None,
+        background: Literal[True],
+    ) -> AsyncCommandHandle: ...
+
+    async def run(
+        self,
+        cmd: str | Sequence[str],
+        *,
+        cwd: str | None = None,
+        env: Mapping[str, str] | None = None,
+        user: str | None = None,
+        timeout: float | None = None,
+        check: bool = False,
+        background: bool = False,
+    ) -> CommandResult | AsyncCommandHandle:
+        """Run a command. Foreground (default): returns the complete
+        :class:`CommandResult` once the process exits — non-zero exit is
+        data unless ``check=True`` (`subprocess.run` semantics).
+        Background (``background=True``): returns as soon as the process
+        is started, with a handle for streaming and waiting.
+
+        ``timeout`` kills the whole process group after that many
+        seconds; expiry surfaces as signal death in the result
+        (exit-as-data), not as a raised error."""
+        if background and check:
+            raise InvalidArgumentError(
+                "check=True applies to foreground runs; call "
+                ".wait_for_exit() and .expect() on the handle instead",
+                operation="commands.run",
+            )
+        handle = await self._start(cmd, cwd=cwd, env=env, user=user, timeout=timeout)
+        if background:
+            return handle
+        result = await handle.wait_for_exit()
+        return result.expect() if check else result
+
+    async def _start(
+        self,
+        cmd: str | Sequence[str],
+        *,
+        cwd: str | None,
+        env: Mapping[str, str] | None,
+        user: str | None,
+        timeout: float | None,
+    ) -> AsyncCommandHandle:
+        with wrap_errors("commands.run"):
+            # The execution id is minted client-side: a lost response
+            # leaves an addressable execution, and retries are idempotent
+            # by contract.
+            execution_id = str(uuid.uuid4())
+            execution = await self._client.unary(
+                _PROCESS + "StartExecution",
+                process_pb2.StartExecutionRequest(
+                    sandbox_id=self._sandbox_id,
+                    execution_id=execution_id,
+                    cmd=_normalize_cmd(cmd),
+                    env=dict(env) if env else {},
+                    working_dir=cwd or "",
+                    user=user or "",
+                    timeout_seconds=0 if timeout is None else math.ceil(timeout),
+                    stdin=False,
+                ),
+                process_pb2.Execution,
+            )
+            return AsyncCommandHandle(self._client, self._sandbox_id, execution.id)