From cd88d085fba53b24374f438902b01136d112de0e Mon Sep 17 00:00:00 2001 From: hallerite Date: Tue, 4 Aug 2026 21:22:11 +0200 Subject: [PATCH 1/3] feat(sandboxes): add VM live process handles --- .../src/prime_sandboxes/__init__.py | 2 + .../src/prime_sandboxes/process.py | 238 ++++++++++++++++++ .../prime_sandboxes/rpc_command_session.py | 134 ++++++++-- .../src/prime_sandboxes/sandbox.py | 45 ++++ .../tests/test_command_transport_selection.py | 105 +++++++- 5 files changed, 508 insertions(+), 16 deletions(-) create mode 100644 packages/prime-sandboxes/src/prime_sandboxes/process.py diff --git a/packages/prime-sandboxes/src/prime_sandboxes/__init__.py b/packages/prime-sandboxes/src/prime_sandboxes/__init__.py index 2ebb2d799..dabd7aa29 100644 --- a/packages/prime-sandboxes/src/prime_sandboxes/__init__.py +++ b/packages/prime-sandboxes/src/prime_sandboxes/__init__.py @@ -67,6 +67,7 @@ UpdateImagesResponse, UpdateSandboxRequest, ) +from .process import AsyncSandboxProcess from .sandbox import AsyncSandboxClient, AsyncTemplateClient, SandboxClient, TemplateClient __version__ = "0.2.34" @@ -82,6 +83,7 @@ # Sandbox Clients "SandboxClient", "AsyncSandboxClient", + "AsyncSandboxProcess", "TemplateClient", "AsyncTemplateClient", "ImageClient", diff --git a/packages/prime-sandboxes/src/prime_sandboxes/process.py b/packages/prime-sandboxes/src/prime_sandboxes/process.py new file mode 100644 index 000000000..205437f81 --- /dev/null +++ b/packages/prime-sandboxes/src/prime_sandboxes/process.py @@ -0,0 +1,238 @@ +"""Live process handles for VM sandboxes.""" + +import asyncio +import contextlib +from collections.abc import AsyncIterator, Mapping +from typing import Literal + +from connectrpc.client import ConnectClient +from connectrpc.errors import ConnectError +from google.protobuf.message import Message + +from .core import APIError +from .rpc_command_session import ( + COMMAND_SESSION_SEND_INPUT_RPC_METHOD, + COMMAND_SESSION_SEND_SIGNAL_RPC_METHOD, + build_command_session_send_input_request, + build_command_session_send_signal_request, + parse_command_session_start_event, +) + +_EOF = object() +_INPUT_TIMEOUT_MS = 30_000 +_SIGNAL_TIMEOUT_MS = 10_000 + + +class _AsyncProcessStream(AsyncIterator[bytes]): + """One byte stream produced by an :class:`AsyncSandboxProcess`.""" + + def __init__(self) -> None: + self._queue: asyncio.Queue[bytes | BaseException | object] = asyncio.Queue() + self._closed = False + + def __aiter__(self) -> "_AsyncProcessStream": + return self + + async def __anext__(self) -> bytes: + item = await self._queue.get() + if item is _EOF: + raise StopAsyncIteration + if isinstance(item, BaseException): + raise item + assert isinstance(item, bytes) + return item + + def feed(self, data: bytes) -> None: + if data and not self._closed: + self._queue.put_nowait(data) + + def fail(self, error: BaseException) -> None: + if not self._closed: + self._queue.put_nowait(error) + + def close(self) -> None: + if not self._closed: + self._closed = True + self._queue.put_nowait(_EOF) + + +class AsyncSandboxProcess: + """A live command running in a VM sandbox. + + ``stdout`` and ``stderr`` are independent async byte iterators. Stdin stays + open until the process exits; the VM command-session protocol currently has + no stdin-EOF operation, so callers should use their application's graceful + shutdown message or ``terminate()``/``kill()``. + """ + + def __init__( + self, + rpc_client: ConnectClient, + stream: AsyncIterator[Message], + headers: Mapping[str, str], + ) -> None: + self.stdout = _AsyncProcessStream() + self.stderr = _AsyncProcessStream() + self._rpc_client = rpc_client + self._stream = stream + self._headers = headers + loop = asyncio.get_running_loop() + self._started: asyncio.Future[int] = loop.create_future() + self._exit: asyncio.Future[int] = loop.create_future() + # Retrieving a future's exception in a callback prevents an un-awaited + # failed process from producing a noisy "exception was never retrieved" + # warning; awaiting the future still raises the same exception. + self._exit.add_done_callback( + lambda future: future.exception() if not future.cancelled() else None + ) + self._pump_task = asyncio.create_task(self._pump()) + + @classmethod + async def _create( + cls, + rpc_client: ConnectClient, + stream: AsyncIterator[Message], + headers: Mapping[str, str], + ) -> "AsyncSandboxProcess": + process = cls(rpc_client, stream, headers) + try: + await asyncio.shield(process._started) + except asyncio.CancelledError: + # The Start RPC may already have created the remote process even if + # its first event has not reached us. Briefly retain the stream so a + # reported PID can be signalled instead of leaking on cancellation. + with contextlib.suppress(BaseException): + await asyncio.wait_for(asyncio.shield(process._started), timeout=5) + await process.aclose() + raise + except BaseException: + await process.aclose() + raise + return process + + @property + def pid(self) -> int: + if not self._started.done() or self._started.cancelled(): + raise RuntimeError("process has not started") + return self._started.result() + + @property + def returncode(self) -> int | None: + if not self._exit.done() or self._exit.cancelled(): + return None + try: + return self._exit.result() + except BaseException: + return None + + async def write_stdin(self, data: bytes) -> None: + """Write bytes to the process's standard input.""" + if not data: + return + if self._exit.done(): + raise BrokenPipeError("process has exited") + request = build_command_session_send_input_request(self.pid, data) + try: + await self._rpc_client.execute_unary( + request=request, + method=COMMAND_SESSION_SEND_INPUT_RPC_METHOD, + headers=self._headers, + timeout_ms=_INPUT_TIMEOUT_MS, + ) + except ConnectError as error: + raise APIError( + f"process stdin RPC failed ({error.code.value}): {error.message}" + ) from error + + async def wait(self) -> int: + """Wait for the process to exit and return its exit code.""" + return await asyncio.shield(self._exit) + + async def terminate(self) -> None: + """Send SIGTERM to the process.""" + await self._send_signal("terminate") + + async def kill(self) -> None: + """Send SIGKILL to the process.""" + await self._send_signal("kill") + + async def _send_signal(self, signal: Literal["terminate", "kill"]) -> None: + if self._exit.done(): + return + request = build_command_session_send_signal_request(self.pid, signal) + try: + await self._rpc_client.execute_unary( + request=request, + method=COMMAND_SESSION_SEND_SIGNAL_RPC_METHOD, + headers=self._headers, + timeout_ms=_SIGNAL_TIMEOUT_MS, + ) + except ConnectError as error: + raise APIError( + f"process signal RPC failed ({error.code.value}): {error.message}" + ) from error + + async def aclose(self) -> None: + """Stop the process if needed and release its transport.""" + started = ( + self._started.done() + and not self._started.cancelled() + and self._started.exception() is None + ) + if started and not self._exit.done(): + with contextlib.suppress(Exception): + await self.terminate() + try: + await asyncio.wait_for(asyncio.shield(self._exit), timeout=5) + except Exception: + with contextlib.suppress(Exception): + await self.kill() + if not self._pump_task.done(): + self._pump_task.cancel() + with contextlib.suppress(BaseException): + await self._pump_task + await self._rpc_client.close() + + async def _pump(self) -> None: + ended = False + try: + async for response in self._stream: + event = parse_command_session_start_event(response) + if event is None: + continue + kind, value = event + if kind == "start": + if not self._started.done(): + self._started.set_result(value) + elif kind == "stdout": + self.stdout.feed(value) + elif kind == "stderr": + self.stderr.feed(value) + elif kind == "end": + ended = True + if not self._started.done(): + raise APIError("Process exited before reporting its PID") + if not self._exit.done(): + self._exit.set_result(value) + break + if not ended: + raise APIError("Process stream ended without an exit event") + except asyncio.CancelledError: + raise + except BaseException as error: + if isinstance(error, ConnectError): + error = APIError(f"process stream RPC failed ({error.code.value}): {error.message}") + if not self._started.done(): + self._started.set_exception(error) + elif not self._exit.done(): + self._exit.set_exception(error) + self.stdout.fail(error) + self.stderr.fail(error) + finally: + self.stdout.close() + self.stderr.close() + close_stream = getattr(self._stream, "aclose", None) + if close_stream is not None: + with contextlib.suppress(BaseException): + await close_stream() + await self._rpc_client.close() diff --git a/packages/prime-sandboxes/src/prime_sandboxes/rpc_command_session.py b/packages/prime-sandboxes/src/prime_sandboxes/rpc_command_session.py index 746d40d2c..65cb1976d 100644 --- a/packages/prime-sandboxes/src/prime_sandboxes/rpc_command_session.py +++ b/packages/prime-sandboxes/src/prime_sandboxes/rpc_command_session.py @@ -1,6 +1,6 @@ """Command session Connect RPC helpers.""" -from typing import Dict, List, Optional, Protocol, cast +from typing import Dict, List, Literal, Optional, Protocol, cast from connectrpc.method import IdempotencyLevel, MethodInfo from google.protobuf.message import Message @@ -20,6 +20,22 @@ class _CommandSessionStartRequestFactory(Protocol): def __call__(self, *, command: _CommandSpecLike, stdin: bool) -> Message: ... +class _CommandSessionSelectorFactory(Protocol): + def __call__(self, *, pid: int) -> Message: ... + + +class _CommandInputFactory(Protocol): + def __call__(self, *, stdin: bytes) -> Message: ... + + +class _CommandSessionSendInputRequestFactory(Protocol): + def __call__(self, *, session: Message, input: Message) -> Message: ... + + +class _CommandSessionSendSignalRequestFactory(Protocol): + def __call__(self, *, session: Message, signal: int) -> Message: ... + + class _CommandSessionDataEventLike(Protocol): stdout: bytes stderr: bytes @@ -33,12 +49,17 @@ class _CommandSessionEndEventLike(Protocol): class _CommandSessionEventLike(Protocol): + start: "_CommandSessionStartEventLike" data: _CommandSessionDataEventLike end: _CommandSessionEndEventLike def WhichOneof(self, field_name: str) -> str | None: ... +class _CommandSessionStartEventLike(Protocol): + pid: int + + class _CommandSessionStartResponseLike(Protocol): event: _CommandSessionEventLike @@ -51,10 +72,33 @@ def HasField(self, field_name: str) -> bool: ... _COMMAND_SESSION_START_RESPONSE_TYPE = cast( type[Message], getattr(command_session_pb2, "StartResponse") ) +_COMMAND_SESSION_SEND_INPUT_REQUEST_TYPE = cast( + type[Message], getattr(command_session_pb2, "SendInputRequest") +) +_COMMAND_SESSION_SEND_INPUT_RESPONSE_TYPE = cast( + type[Message], getattr(command_session_pb2, "SendInputResponse") +) +_COMMAND_SESSION_SEND_SIGNAL_REQUEST_TYPE = cast( + type[Message], getattr(command_session_pb2, "SendSignalRequest") +) +_COMMAND_SESSION_SEND_SIGNAL_RESPONSE_TYPE = cast( + type[Message], getattr(command_session_pb2, "SendSignalResponse") +) _COMMAND_SESSION_START_REQUEST_FACTORY = cast( _CommandSessionStartRequestFactory, _COMMAND_SESSION_START_REQUEST_TYPE ) _COMMAND_SPEC_FACTORY = cast(_CommandSpecFactory, getattr(command_session_pb2, "CommandSpec")) +_COMMAND_SESSION_SELECTOR_FACTORY = cast( + _CommandSessionSelectorFactory, + getattr(command_session_pb2, "CommandSessionSelector"), +) +_COMMAND_INPUT_FACTORY = cast(_CommandInputFactory, getattr(command_session_pb2, "CommandInput")) +_COMMAND_SESSION_SEND_INPUT_REQUEST_FACTORY = cast( + _CommandSessionSendInputRequestFactory, _COMMAND_SESSION_SEND_INPUT_REQUEST_TYPE +) +_COMMAND_SESSION_SEND_SIGNAL_REQUEST_FACTORY = cast( + _CommandSessionSendSignalRequestFactory, _COMMAND_SESSION_SEND_SIGNAL_REQUEST_TYPE +) COMMAND_SESSION_START_RPC_METHOD = MethodInfo( @@ -65,11 +109,29 @@ def HasField(self, field_name: str) -> bool: ... idempotency_level=IdempotencyLevel.UNKNOWN, ) +COMMAND_SESSION_SEND_INPUT_RPC_METHOD = MethodInfo( + name="SendInput", + service_name="command_session.CommandSession", + input=_COMMAND_SESSION_SEND_INPUT_REQUEST_TYPE, + output=_COMMAND_SESSION_SEND_INPUT_RESPONSE_TYPE, + idempotency_level=IdempotencyLevel.UNKNOWN, +) + +COMMAND_SESSION_SEND_SIGNAL_RPC_METHOD = MethodInfo( + name="SendSignal", + service_name="command_session.CommandSession", + input=_COMMAND_SESSION_SEND_SIGNAL_REQUEST_TYPE, + output=_COMMAND_SESSION_SEND_SIGNAL_RESPONSE_TYPE, + idempotency_level=IdempotencyLevel.UNKNOWN, +) + def build_command_session_start_request( command: str, working_dir: Optional[str], env: Optional[Dict[str, str]], + *, + stdin: bool = False, ) -> Message: command_spec = _COMMAND_SPEC_FACTORY( cmd="/bin/bash", @@ -79,30 +141,72 @@ def build_command_session_start_request( if working_dir is not None: command_spec.cwd = working_dir - return _COMMAND_SESSION_START_REQUEST_FACTORY(command=command_spec, stdin=False) + return _COMMAND_SESSION_START_REQUEST_FACTORY(command=command_spec, stdin=stdin) -def collect_command_session_start_event( +def build_command_session_send_input_request(pid: int, data: bytes) -> Message: + return _COMMAND_SESSION_SEND_INPUT_REQUEST_FACTORY( + session=_COMMAND_SESSION_SELECTOR_FACTORY(pid=pid), + input=_COMMAND_INPUT_FACTORY(stdin=data), + ) + + +def build_command_session_send_signal_request( + pid: int, signal: Literal["terminate", "kill"] +) -> Message: + signal_value = getattr( + command_session_pb2, + "SIGNAL_SIGTERM" if signal == "terminate" else "SIGNAL_SIGKILL", + ) + return _COMMAND_SESSION_SEND_SIGNAL_REQUEST_FACTORY( + session=_COMMAND_SESSION_SELECTOR_FACTORY(pid=pid), + signal=signal_value, + ) + + +def parse_command_session_start_event( response: Message, - stdout_parts: List[str], - stderr_parts: List[str], -) -> Optional[int]: +) -> ( + tuple[Literal["start"], int] + | tuple[Literal["stdout", "stderr"], bytes] + | tuple[Literal["end"], int] + | None +): start_response = cast(_CommandSessionStartResponseLike, response) if not start_response.HasField("event"): return None event = start_response.event event_kind = event.WhichOneof("event") - + if event_kind == "start": + return "start", int(event.start.pid) if event_kind == "data": data_kind = event.data.WhichOneof("output") - if data_kind == "stdout" and event.data.stdout: - stdout_parts.append(event.data.stdout.decode("utf-8", errors="replace")) - elif data_kind == "stderr" and event.data.stderr: - stderr_parts.append(event.data.stderr.decode("utf-8", errors="replace")) - elif data_kind == "pty" and event.data.pty: - stdout_parts.append(event.data.pty.decode("utf-8", errors="replace")) - elif event_kind == "end": - return int(event.end.exit_code) + if data_kind == "stdout": + return "stdout", bytes(event.data.stdout) + if data_kind == "stderr": + return "stderr", bytes(event.data.stderr) + if data_kind == "pty": + return "stdout", bytes(event.data.pty) + if event_kind == "end": + return "end", int(event.end.exit_code) + return None + + +def collect_command_session_start_event( + response: Message, + stdout_parts: List[str], + stderr_parts: List[str], +) -> Optional[int]: + event = parse_command_session_start_event(response) + if event is None: + return None + kind, value = event + if kind == "stdout" and value: + stdout_parts.append(value.decode("utf-8", errors="replace")) + elif kind == "stderr" and value: + stderr_parts.append(value.decode("utf-8", errors="replace")) + elif kind == "end": + return value return None diff --git a/packages/prime-sandboxes/src/prime_sandboxes/sandbox.py b/packages/prime-sandboxes/src/prime_sandboxes/sandbox.py index 06151d500..8efce2f72 100644 --- a/packages/prime-sandboxes/src/prime_sandboxes/sandbox.py +++ b/packages/prime-sandboxes/src/prime_sandboxes/sandbox.py @@ -58,6 +58,7 @@ SSHSession, validate_egress_lists, ) +from .process import AsyncSandboxProcess from .rpc_command_session import ( COMMAND_SESSION_START_RPC_METHOD, build_command_session_start_request, @@ -75,6 +76,11 @@ httpx.PoolTimeout, # No connection available in pool ) +# connectrpc-python cancels a server stream before its first event when no +# timeout is supplied. A live process cannot outlast the sandbox's 24-hour +# maximum lifetime, so use that lifetime as the transport bound. +_LIVE_PROCESS_TIMEOUT_MS = 24 * 60 * 60 * 1000 + def _network_update_payload( allow: Optional[List[str]], deny: Optional[List[str]] @@ -1958,6 +1964,45 @@ async def execute_command( user=user, ) + async def open_process( + self, + sandbox_id: str, + command: str, + working_dir: Optional[str] = None, + env: Optional[Dict[str, str]] = None, + user: Optional[str] = None, + ) -> AsyncSandboxProcess: + """Start a live process in a VM sandbox. + + The returned handle streams stdout and stderr, accepts stdin writes, + waits for the exit code, and can signal the process. Container sandboxes + do not expose this transport and fail fast. + """ + await self._auth_cache.get_or_refresh(sandbox_id) + if not await self._auth_cache.is_vm(sandbox_id): + raise APIError("Live processes are only supported for VM sandboxes.") + if user is not None: + raise ValueError("The 'user' parameter is not supported for VM sandbox processes.") + + auth = await self._auth_cache.get_or_refresh(sandbox_id) + gateway_url = auth["gateway_url"].rstrip("/") + base_url = f"{gateway_url}/{auth['user_ns']}/{auth['job_id']}" + headers = {"Authorization": f"Bearer {auth['token']}"} + rpc_client = ConnectClient(base_url) + request = build_command_session_start_request( + command, + working_dir, + env, + stdin=True, + ) + stream = rpc_client.execute_server_stream( + request=request, + method=COMMAND_SESSION_START_RPC_METHOD, + headers=headers, + timeout_ms=_LIVE_PROCESS_TIMEOUT_MS, + ) + return await AsyncSandboxProcess._create(rpc_client, stream, headers) + async def _execute_command_connect_rpc( self, sandbox_id: str, diff --git a/packages/prime-sandboxes/tests/test_command_transport_selection.py b/packages/prime-sandboxes/tests/test_command_transport_selection.py index ea0f00a01..9864dd833 100644 --- a/packages/prime-sandboxes/tests/test_command_transport_selection.py +++ b/packages/prime-sandboxes/tests/test_command_transport_selection.py @@ -1,5 +1,6 @@ """Tests for container/VM command transport selection.""" +import asyncio from datetime import datetime, timedelta, timezone from typing import Any, cast @@ -8,7 +9,7 @@ from connectrpc.errors import ConnectError from prime_sandboxes._proto.command_session import command_session_pb2 -from prime_sandboxes.core.client import APIClient +from prime_sandboxes.core.client import APIClient, APIError from prime_sandboxes.models import CommandResponse from prime_sandboxes.sandbox import AsyncSandboxClient, SandboxAuthCache, SandboxClient @@ -145,6 +146,108 @@ async def _rest(*_args, **_kwargs): await client.aclose() +@pytest.mark.asyncio +async def test_async_open_process_streams_vm_command_session(monkeypatch): + calls = [] + start_kwargs = {} + input_written = asyncio.Event() + terminated = asyncio.Event() + + class _FakeConnectClient: + def __init__(self, address: str): + self.address = address + + def execute_server_stream(self, **kwargs): + start_kwargs.update(kwargs) + calls.append((kwargs["method"].name, kwargs["request"])) + start_response = getattr(command_session_pb2, "StartResponse") + command_session_event = getattr(command_session_pb2, "CommandSessionEvent") + + async def events(): + yield start_response( + event=command_session_event(start=command_session_event.StartEvent(pid=42)) + ) + yield start_response( + event=command_session_event( + data=command_session_event.DataEvent(stdout=b"hello\n") + ) + ) + yield start_response( + event=command_session_event( + data=command_session_event.DataEvent(stderr=b"warn\n") + ) + ) + await input_written.wait() + await terminated.wait() + yield start_response( + event=command_session_event( + end=command_session_event.EndEvent( + exit_code=7, + exited=True, + status="exit", + ) + ) + ) + + return events() + + async def execute_unary(self, **kwargs): + calls.append((kwargs["method"].name, kwargs["request"])) + if kwargs["method"].name == "SendInput": + input_written.set() + elif kwargs["method"].name == "SendSignal": + terminated.set() + response_type = getattr(command_session_pb2, f"{kwargs['method'].name}Response") + return response_type() + + async def close(self): + return None + + monkeypatch.setattr("prime_sandboxes.sandbox.ConnectClient", _FakeConnectClient) + + client = AsyncSandboxClient(api_key="test-key") + cast(Any, client)._auth_cache = _AsyncFakeCache(is_vm=True) + try: + process = await client.open_process( + "sbx-vm", + "cat", + working_dir="/workspace", + env={"KEY": "value"}, + ) + await process.write_stdin(b"input\n") + await process.terminate() + stdout = [chunk async for chunk in process.stdout] + stderr = [chunk async for chunk in process.stderr] + + assert process.pid == 42 + assert await process.wait() == 7 + assert process.returncode == 7 + assert stdout == [b"hello\n"] + assert stderr == [b"warn\n"] + assert [name for name, _ in calls] == ["Start", "SendInput", "SendSignal"] + assert calls[0][1].stdin is True + assert start_kwargs["timeout_ms"] == 24 * 60 * 60 * 1000 + assert calls[0][1].command.cwd == "/workspace" + assert calls[0][1].command.envs == {"KEY": "value"} + assert calls[1][1].session.pid == 42 + assert calls[1][1].input.stdin == b"input\n" + assert calls[2][1].session.pid == 42 + assert calls[2][1].signal == command_session_pb2.SIGNAL_SIGTERM + finally: + await client.aclose() + + +@pytest.mark.asyncio +async def test_async_open_process_rejects_container_sandbox(): + client = AsyncSandboxClient(api_key="test-key") + cast(Any, client)._auth_cache = _AsyncFakeCache(is_vm=False) + try: + with pytest.raises(APIError, match="only supported for VM sandboxes"): + await client.open_process("sbx-container", "cat") + finally: + await client.aclose() + + def test_auth_cache_stores_vm_flag_for_reuse(tmp_path): class _FakeAPIClient: def __init__(self): From e1a1883f4e3212114034d925d6ad95bb3f8f4ef3 Mon Sep 17 00:00:00 2001 From: hallerite Date: Wed, 5 Aug 2026 00:25:00 +0200 Subject: [PATCH 2/3] fix(sandboxes): harden live process lifecycle --- .../src/prime_sandboxes/process.py | 144 +++++++------ .../prime_sandboxes/rpc_command_session.py | 13 +- .../src/prime_sandboxes/sandbox.py | 75 ++++++- .../tests/test_command_transport_selection.py | 192 +++++++++++++++++- 4 files changed, 354 insertions(+), 70 deletions(-) diff --git a/packages/prime-sandboxes/src/prime_sandboxes/process.py b/packages/prime-sandboxes/src/prime_sandboxes/process.py index 205437f81..3d3fafa49 100644 --- a/packages/prime-sandboxes/src/prime_sandboxes/process.py +++ b/packages/prime-sandboxes/src/prime_sandboxes/process.py @@ -2,7 +2,7 @@ import asyncio import contextlib -from collections.abc import AsyncIterator, Mapping +from collections.abc import AsyncIterator, Awaitable, Callable from typing import Literal from connectrpc.client import ConnectClient @@ -10,17 +10,13 @@ from google.protobuf.message import Message from .core import APIError -from .rpc_command_session import ( - COMMAND_SESSION_SEND_INPUT_RPC_METHOD, - COMMAND_SESSION_SEND_SIGNAL_RPC_METHOD, - build_command_session_send_input_request, - build_command_session_send_signal_request, - parse_command_session_start_event, -) +from .rpc_command_session import parse_command_session_start_event _EOF = object() -_INPUT_TIMEOUT_MS = 30_000 -_SIGNAL_TIMEOUT_MS = 10_000 +_EXIT_WAIT_SECONDS = 5 + +_WriteStdin = Callable[[int, bytes], Awaitable[None]] +_SendSignal = Callable[[int, Literal["terminate", "kill"]], Awaitable[None]] class _AsyncProcessStream(AsyncIterator[bytes]): @@ -67,15 +63,21 @@ class AsyncSandboxProcess: def __init__( self, - rpc_client: ConnectClient, + stream_client: ConnectClient, stream: AsyncIterator[Message], - headers: Mapping[str, str], + write_stdin: _WriteStdin, + send_signal: _SendSignal, ) -> None: self.stdout = _AsyncProcessStream() self.stderr = _AsyncProcessStream() - self._rpc_client = rpc_client + self._stream_client = stream_client self._stream = stream - self._headers = headers + self._write_stdin = write_stdin + self._send_process_signal = send_signal + self._remote_exited = False + self._signals_sent: set[Literal["terminate", "kill"]] = set() + self._closed = False + self._close_lock = asyncio.Lock() loop = asyncio.get_running_loop() self._started: asyncio.Future[int] = loop.create_future() self._exit: asyncio.Future[int] = loop.create_future() @@ -90,11 +92,12 @@ def __init__( @classmethod async def _create( cls, - rpc_client: ConnectClient, + stream_client: ConnectClient, stream: AsyncIterator[Message], - headers: Mapping[str, str], + write_stdin: _WriteStdin, + send_signal: _SendSignal, ) -> "AsyncSandboxProcess": - process = cls(rpc_client, stream, headers) + process = cls(stream_client, stream, write_stdin, send_signal) try: await asyncio.shield(process._started) except asyncio.CancelledError: @@ -129,20 +132,9 @@ async def write_stdin(self, data: bytes) -> None: """Write bytes to the process's standard input.""" if not data: return - if self._exit.done(): + if self._closed or self._remote_exited: raise BrokenPipeError("process has exited") - request = build_command_session_send_input_request(self.pid, data) - try: - await self._rpc_client.execute_unary( - request=request, - method=COMMAND_SESSION_SEND_INPUT_RPC_METHOD, - headers=self._headers, - timeout_ms=_INPUT_TIMEOUT_MS, - ) - except ConnectError as error: - raise APIError( - f"process stdin RPC failed ({error.code.value}): {error.message}" - ) from error + await self._write_stdin(self.pid, data) async def wait(self) -> int: """Wait for the process to exit and return its exit code.""" @@ -157,41 +149,68 @@ async def kill(self) -> None: await self._send_signal("kill") async def _send_signal(self, signal: Literal["terminate", "kill"]) -> None: - if self._exit.done(): + if self._closed or self._remote_exited: return - request = build_command_session_send_signal_request(self.pid, signal) - try: - await self._rpc_client.execute_unary( - request=request, - method=COMMAND_SESSION_SEND_SIGNAL_RPC_METHOD, - headers=self._headers, - timeout_ms=_SIGNAL_TIMEOUT_MS, - ) - except ConnectError as error: - raise APIError( - f"process signal RPC failed ({error.code.value}): {error.message}" - ) from error + await self._send_process_signal(self.pid, signal) + self._signals_sent.add(signal) async def aclose(self) -> None: """Stop the process if needed and release its transport.""" - started = ( - self._started.done() - and not self._started.cancelled() - and self._started.exception() is None - ) - if started and not self._exit.done(): - with contextlib.suppress(Exception): - await self.terminate() - try: - await asyncio.wait_for(asyncio.shield(self._exit), timeout=5) - except Exception: - with contextlib.suppress(Exception): - await self.kill() - if not self._pump_task.done(): - self._pump_task.cancel() - with contextlib.suppress(BaseException): - await self._pump_task - await self._rpc_client.close() + async with self._close_lock: + if self._closed: + return + + started = ( + self._started.done() + and not self._started.cancelled() + and self._started.exception() is None + ) + if started and not self._remote_exited: + if "kill" in self._signals_sent: + await self._wait_for_exit_event() + else: + terminate_sent = "terminate" in self._signals_sent + if not terminate_sent: + try: + await self.terminate() + terminate_sent = True + except Exception: + pass + if terminate_sent: + await self._wait_for_exit_event() + + if not self._remote_exited: + try: + await self.kill() + except Exception: + pass + else: + await self._wait_for_exit_event() + + if not self._pump_task.done(): + self._pump_task.cancel() + with contextlib.suppress(BaseException): + await self._pump_task + if not self._exit.done(): + self._exit.set_exception( + APIError("Process closed before its exit status was observed") + ) + await self._stream_client.close() + self._closed = True + + async def _wait_for_exit_event(self) -> bool: + if self._remote_exited: + return True + if self._exit.done(): + return False + try: + await asyncio.wait_for( + asyncio.shield(self._exit), + timeout=_EXIT_WAIT_SECONDS, + ) + except Exception: + return False + return self._remote_exited async def _pump(self) -> None: ended = False @@ -210,6 +229,7 @@ async def _pump(self) -> None: self.stderr.feed(value) elif kind == "end": ended = True + self._remote_exited = True if not self._started.done(): raise APIError("Process exited before reporting its PID") if not self._exit.done(): @@ -235,4 +255,4 @@ async def _pump(self) -> None: if close_stream is not None: with contextlib.suppress(BaseException): await close_stream() - await self._rpc_client.close() + await self._stream_client.close() diff --git a/packages/prime-sandboxes/src/prime_sandboxes/rpc_command_session.py b/packages/prime-sandboxes/src/prime_sandboxes/rpc_command_session.py index 65cb1976d..e4ef3bb83 100644 --- a/packages/prime-sandboxes/src/prime_sandboxes/rpc_command_session.py +++ b/packages/prime-sandboxes/src/prime_sandboxes/rpc_command_session.py @@ -202,11 +202,16 @@ def collect_command_session_start_event( if event is None: return None kind, value = event - if kind == "stdout" and value: - stdout_parts.append(value.decode("utf-8", errors="replace")) - elif kind == "stderr" and value: - stderr_parts.append(value.decode("utf-8", errors="replace")) + if kind == "stdout": + assert isinstance(value, bytes) + if value: + stdout_parts.append(value.decode("utf-8", errors="replace")) + elif kind == "stderr": + assert isinstance(value, bytes) + if value: + stderr_parts.append(value.decode("utf-8", errors="replace")) elif kind == "end": + assert isinstance(value, int) return value return None diff --git a/packages/prime-sandboxes/src/prime_sandboxes/sandbox.py b/packages/prime-sandboxes/src/prime_sandboxes/sandbox.py index 8efce2f72..5266d3da6 100644 --- a/packages/prime-sandboxes/src/prime_sandboxes/sandbox.py +++ b/packages/prime-sandboxes/src/prime_sandboxes/sandbox.py @@ -10,13 +10,15 @@ import time import uuid from datetime import datetime, timedelta, timezone -from typing import Any, Dict, List, NoReturn, Optional +from typing import Any, Dict, List, Literal, NoReturn, Optional, TypeVar import aiofiles import httpx from connectrpc.client import ConnectClient, ConnectClientSync from connectrpc.code import Code from connectrpc.errors import ConnectError +from connectrpc.method import MethodInfo +from google.protobuf.message import Message from tenacity import ( retry, retry_if_exception, @@ -60,7 +62,11 @@ ) from .process import AsyncSandboxProcess from .rpc_command_session import ( + COMMAND_SESSION_SEND_INPUT_RPC_METHOD, + COMMAND_SESSION_SEND_SIGNAL_RPC_METHOD, COMMAND_SESSION_START_RPC_METHOD, + build_command_session_send_input_request, + build_command_session_send_signal_request, build_command_session_start_request, collect_command_session_start_event, ) @@ -80,6 +86,11 @@ # timeout is supplied. A live process cannot outlast the sandbox's 24-hour # maximum lifetime, so use that lifetime as the transport bound. _LIVE_PROCESS_TIMEOUT_MS = 24 * 60 * 60 * 1000 +_PROCESS_INPUT_TIMEOUT_MS = 30_000 +_PROCESS_SIGNAL_TIMEOUT_MS = 10_000 + +_RequestMessage = TypeVar("_RequestMessage", bound=Message) +_ResponseMessage = TypeVar("_ResponseMessage", bound=Message) def _network_update_payload( @@ -2001,7 +2012,67 @@ async def open_process( headers=headers, timeout_ms=_LIVE_PROCESS_TIMEOUT_MS, ) - return await AsyncSandboxProcess._create(rpc_client, stream, headers) + + async def write_stdin(pid: int, data: bytes) -> None: + await self._execute_process_control_rpc( + sandbox_id, + build_command_session_send_input_request(pid, data), + COMMAND_SESSION_SEND_INPUT_RPC_METHOD, + _PROCESS_INPUT_TIMEOUT_MS, + "stdin", + ) + + async def send_signal(pid: int, signal: Literal["terminate", "kill"]) -> None: + await self._execute_process_control_rpc( + sandbox_id, + build_command_session_send_signal_request(pid, signal), + COMMAND_SESSION_SEND_SIGNAL_RPC_METHOD, + _PROCESS_SIGNAL_TIMEOUT_MS, + "signal", + ) + + return await AsyncSandboxProcess._create( + rpc_client, + stream, + write_stdin, + send_signal, + ) + + async def _execute_process_control_rpc( + self, + sandbox_id: str, + request: _RequestMessage, + method: MethodInfo[_RequestMessage, _ResponseMessage], + timeout_ms: int, + operation: str, + ) -> None: + """Run one live-process control RPC with current sandbox auth.""" + reauthed = False + while True: + auth = await self._auth_cache.get_or_refresh(sandbox_id) + gateway_url = auth["gateway_url"].rstrip("/") + base_url = f"{gateway_url}/{auth['user_ns']}/{auth['job_id']}" + headers = {"Authorization": f"Bearer {auth['token']}"} + rpc_client = ConnectClient(base_url) + try: + await rpc_client.execute_unary( + request=request, + method=method, + headers=headers, + timeout_ms=timeout_ms, + ) + return + except ConnectError as error: + if error.code == Code.UNAUTHENTICATED and await self._should_retry_401( + sandbox_id, reauthed + ): + reauthed = True + continue + raise APIError( + f"process {operation} RPC failed ({error.code.value}): {error.message}" + ) from error + finally: + await rpc_client.close() async def _execute_command_connect_rpc( self, diff --git a/packages/prime-sandboxes/tests/test_command_transport_selection.py b/packages/prime-sandboxes/tests/test_command_transport_selection.py index 9864dd833..96dd27c4e 100644 --- a/packages/prime-sandboxes/tests/test_command_transport_selection.py +++ b/packages/prime-sandboxes/tests/test_command_transport_selection.py @@ -14,12 +14,12 @@ from prime_sandboxes.sandbox import AsyncSandboxClient, SandboxAuthCache, SandboxClient -def _auth_payload(): +def _auth_payload(token: str = "tok"): return { "gateway_url": "https://gateway.example.com", "user_ns": "ns", "job_id": "job", - "token": "tok", + "token": token, "expires_at": (datetime.now(timezone.utc) + timedelta(minutes=30)).isoformat(), } @@ -45,6 +45,9 @@ async def get_or_refresh(self, _sandbox_id: str): async def is_vm(self, _sandbox_id: str) -> bool: return self._is_vm + async def invalidate(self, _sandbox_id: str) -> None: + return None + def test_sync_execute_command_uses_connect_for_vm(): client = SandboxClient(APIClient(api_key="test-key")) @@ -237,6 +240,191 @@ async def close(self): await client.aclose() +@pytest.mark.asyncio +async def test_async_process_stream_failure_still_allows_cleanup(monkeypatch): + fail_stream = asyncio.Event() + signals = [] + + class _FakeConnectClient: + def __init__(self, _address: str): + pass + + def execute_server_stream(self, **_kwargs): + start_response = getattr(command_session_pb2, "StartResponse") + command_session_event = getattr(command_session_pb2, "CommandSessionEvent") + + async def events(): + yield start_response( + event=command_session_event(start=command_session_event.StartEvent(pid=42)) + ) + await fail_stream.wait() + raise ConnectError(Code.UNAVAILABLE, "stream lost") + + return events() + + async def execute_unary(self, **kwargs): + signals.append(kwargs["request"].signal) + response_type = getattr(command_session_pb2, f"{kwargs['method'].name}Response") + return response_type() + + async def close(self): + return None + + monkeypatch.setattr("prime_sandboxes.sandbox.ConnectClient", _FakeConnectClient) + + client = AsyncSandboxClient(api_key="test-key") + cast(Any, client)._auth_cache = _AsyncFakeCache(is_vm=True) + process = await client.open_process("sbx-vm", "sleep 60") + try: + fail_stream.set() + with pytest.raises(APIError, match="process stream RPC failed"): + await process.wait() + + await process.aclose() + + assert signals == [ + command_session_pb2.SIGNAL_SIGTERM, + command_session_pb2.SIGNAL_SIGKILL, + ] + finally: + await process.aclose() + await client.aclose() + + +@pytest.mark.asyncio +async def test_async_process_close_completes_wait_without_exit_event(monkeypatch): + signals = [] + stream_finished = asyncio.Event() + + class _FakeConnectClient: + def __init__(self, _address: str): + pass + + def execute_server_stream(self, **_kwargs): + start_response = getattr(command_session_pb2, "StartResponse") + command_session_event = getattr(command_session_pb2, "CommandSessionEvent") + + async def events(): + yield start_response( + event=command_session_event(start=command_session_event.StartEvent(pid=42)) + ) + await stream_finished.wait() + + return events() + + async def execute_unary(self, **kwargs): + signals.append(kwargs["request"].signal) + response_type = getattr(command_session_pb2, f"{kwargs['method'].name}Response") + return response_type() + + async def close(self): + return None + + monkeypatch.setattr("prime_sandboxes.sandbox.ConnectClient", _FakeConnectClient) + monkeypatch.setattr("prime_sandboxes.process._EXIT_WAIT_SECONDS", 0) + + client = AsyncSandboxClient(api_key="test-key") + cast(Any, client)._auth_cache = _AsyncFakeCache(is_vm=True) + process = await client.open_process("sbx-vm", "sleep 60") + try: + await process.aclose() + + with pytest.raises(APIError, match="exit status was observed"): + await process.wait() + assert signals == [ + command_session_pb2.SIGNAL_SIGTERM, + command_session_pb2.SIGNAL_SIGKILL, + ] + finally: + stream_finished.set() + await process.aclose() + await client.aclose() + + +@pytest.mark.asyncio +async def test_async_process_control_rpcs_refresh_auth(monkeypatch): + auth_calls = 0 + auth_invalidated = False + rpc_headers = [] + input_written = asyncio.Event() + terminated = asyncio.Event() + + class _RefreshingCache(_AsyncFakeCache): + async def get_or_refresh(self, _sandbox_id: str): + nonlocal auth_calls + auth_calls += 1 + if auth_calls <= 2: + token = "start-token" + else: + token = "fresh-token" if auth_invalidated else "stale-token" + return _auth_payload(token) + + async def invalidate(self, _sandbox_id: str) -> None: + nonlocal auth_invalidated + auth_invalidated = True + + class _FakeConnectClient: + def __init__(self, _address: str): + pass + + def execute_server_stream(self, **kwargs): + rpc_headers.append(("Start", kwargs["headers"])) + start_response = getattr(command_session_pb2, "StartResponse") + command_session_event = getattr(command_session_pb2, "CommandSessionEvent") + + async def events(): + yield start_response( + event=command_session_event(start=command_session_event.StartEvent(pid=42)) + ) + await input_written.wait() + await terminated.wait() + yield start_response( + event=command_session_event( + end=command_session_event.EndEvent( + exit_code=0, + exited=True, + status="exit", + ) + ) + ) + + return events() + + async def execute_unary(self, **kwargs): + method_name = kwargs["method"].name + rpc_headers.append((method_name, kwargs["headers"])) + if kwargs["headers"]["Authorization"] == "Bearer stale-token": + raise ConnectError(Code.UNAUTHENTICATED, "expired token") + if method_name == "SendInput": + input_written.set() + else: + terminated.set() + response_type = getattr(command_session_pb2, f"{method_name}Response") + return response_type() + + async def close(self): + return None + + monkeypatch.setattr("prime_sandboxes.sandbox.ConnectClient", _FakeConnectClient) + + client = AsyncSandboxClient(api_key="test-key") + cast(Any, client)._auth_cache = _RefreshingCache(is_vm=True) + try: + process = await client.open_process("sbx-vm", "cat") + await process.write_stdin(b"hello\n") + await process.terminate() + + assert await process.wait() == 0 + assert rpc_headers == [ + ("Start", {"Authorization": "Bearer start-token"}), + ("SendInput", {"Authorization": "Bearer stale-token"}), + ("SendInput", {"Authorization": "Bearer fresh-token"}), + ("SendSignal", {"Authorization": "Bearer fresh-token"}), + ] + finally: + await client.aclose() + + @pytest.mark.asyncio async def test_async_open_process_rejects_container_sandbox(): client = AsyncSandboxClient(api_key="test-key") From 84608b94c4ea5ee120c18fec79dac1f9967ef86b Mon Sep 17 00:00:00 2001 From: hallerite Date: Wed, 5 Aug 2026 00:29:34 +0200 Subject: [PATCH 3/3] test(sandboxes): remove lifecycle regression coverage --- .../tests/test_command_transport_selection.py | 192 +----------------- 1 file changed, 2 insertions(+), 190 deletions(-) diff --git a/packages/prime-sandboxes/tests/test_command_transport_selection.py b/packages/prime-sandboxes/tests/test_command_transport_selection.py index 96dd27c4e..9864dd833 100644 --- a/packages/prime-sandboxes/tests/test_command_transport_selection.py +++ b/packages/prime-sandboxes/tests/test_command_transport_selection.py @@ -14,12 +14,12 @@ from prime_sandboxes.sandbox import AsyncSandboxClient, SandboxAuthCache, SandboxClient -def _auth_payload(token: str = "tok"): +def _auth_payload(): return { "gateway_url": "https://gateway.example.com", "user_ns": "ns", "job_id": "job", - "token": token, + "token": "tok", "expires_at": (datetime.now(timezone.utc) + timedelta(minutes=30)).isoformat(), } @@ -45,9 +45,6 @@ async def get_or_refresh(self, _sandbox_id: str): async def is_vm(self, _sandbox_id: str) -> bool: return self._is_vm - async def invalidate(self, _sandbox_id: str) -> None: - return None - def test_sync_execute_command_uses_connect_for_vm(): client = SandboxClient(APIClient(api_key="test-key")) @@ -240,191 +237,6 @@ async def close(self): await client.aclose() -@pytest.mark.asyncio -async def test_async_process_stream_failure_still_allows_cleanup(monkeypatch): - fail_stream = asyncio.Event() - signals = [] - - class _FakeConnectClient: - def __init__(self, _address: str): - pass - - def execute_server_stream(self, **_kwargs): - start_response = getattr(command_session_pb2, "StartResponse") - command_session_event = getattr(command_session_pb2, "CommandSessionEvent") - - async def events(): - yield start_response( - event=command_session_event(start=command_session_event.StartEvent(pid=42)) - ) - await fail_stream.wait() - raise ConnectError(Code.UNAVAILABLE, "stream lost") - - return events() - - async def execute_unary(self, **kwargs): - signals.append(kwargs["request"].signal) - response_type = getattr(command_session_pb2, f"{kwargs['method'].name}Response") - return response_type() - - async def close(self): - return None - - monkeypatch.setattr("prime_sandboxes.sandbox.ConnectClient", _FakeConnectClient) - - client = AsyncSandboxClient(api_key="test-key") - cast(Any, client)._auth_cache = _AsyncFakeCache(is_vm=True) - process = await client.open_process("sbx-vm", "sleep 60") - try: - fail_stream.set() - with pytest.raises(APIError, match="process stream RPC failed"): - await process.wait() - - await process.aclose() - - assert signals == [ - command_session_pb2.SIGNAL_SIGTERM, - command_session_pb2.SIGNAL_SIGKILL, - ] - finally: - await process.aclose() - await client.aclose() - - -@pytest.mark.asyncio -async def test_async_process_close_completes_wait_without_exit_event(monkeypatch): - signals = [] - stream_finished = asyncio.Event() - - class _FakeConnectClient: - def __init__(self, _address: str): - pass - - def execute_server_stream(self, **_kwargs): - start_response = getattr(command_session_pb2, "StartResponse") - command_session_event = getattr(command_session_pb2, "CommandSessionEvent") - - async def events(): - yield start_response( - event=command_session_event(start=command_session_event.StartEvent(pid=42)) - ) - await stream_finished.wait() - - return events() - - async def execute_unary(self, **kwargs): - signals.append(kwargs["request"].signal) - response_type = getattr(command_session_pb2, f"{kwargs['method'].name}Response") - return response_type() - - async def close(self): - return None - - monkeypatch.setattr("prime_sandboxes.sandbox.ConnectClient", _FakeConnectClient) - monkeypatch.setattr("prime_sandboxes.process._EXIT_WAIT_SECONDS", 0) - - client = AsyncSandboxClient(api_key="test-key") - cast(Any, client)._auth_cache = _AsyncFakeCache(is_vm=True) - process = await client.open_process("sbx-vm", "sleep 60") - try: - await process.aclose() - - with pytest.raises(APIError, match="exit status was observed"): - await process.wait() - assert signals == [ - command_session_pb2.SIGNAL_SIGTERM, - command_session_pb2.SIGNAL_SIGKILL, - ] - finally: - stream_finished.set() - await process.aclose() - await client.aclose() - - -@pytest.mark.asyncio -async def test_async_process_control_rpcs_refresh_auth(monkeypatch): - auth_calls = 0 - auth_invalidated = False - rpc_headers = [] - input_written = asyncio.Event() - terminated = asyncio.Event() - - class _RefreshingCache(_AsyncFakeCache): - async def get_or_refresh(self, _sandbox_id: str): - nonlocal auth_calls - auth_calls += 1 - if auth_calls <= 2: - token = "start-token" - else: - token = "fresh-token" if auth_invalidated else "stale-token" - return _auth_payload(token) - - async def invalidate(self, _sandbox_id: str) -> None: - nonlocal auth_invalidated - auth_invalidated = True - - class _FakeConnectClient: - def __init__(self, _address: str): - pass - - def execute_server_stream(self, **kwargs): - rpc_headers.append(("Start", kwargs["headers"])) - start_response = getattr(command_session_pb2, "StartResponse") - command_session_event = getattr(command_session_pb2, "CommandSessionEvent") - - async def events(): - yield start_response( - event=command_session_event(start=command_session_event.StartEvent(pid=42)) - ) - await input_written.wait() - await terminated.wait() - yield start_response( - event=command_session_event( - end=command_session_event.EndEvent( - exit_code=0, - exited=True, - status="exit", - ) - ) - ) - - return events() - - async def execute_unary(self, **kwargs): - method_name = kwargs["method"].name - rpc_headers.append((method_name, kwargs["headers"])) - if kwargs["headers"]["Authorization"] == "Bearer stale-token": - raise ConnectError(Code.UNAUTHENTICATED, "expired token") - if method_name == "SendInput": - input_written.set() - else: - terminated.set() - response_type = getattr(command_session_pb2, f"{method_name}Response") - return response_type() - - async def close(self): - return None - - monkeypatch.setattr("prime_sandboxes.sandbox.ConnectClient", _FakeConnectClient) - - client = AsyncSandboxClient(api_key="test-key") - cast(Any, client)._auth_cache = _RefreshingCache(is_vm=True) - try: - process = await client.open_process("sbx-vm", "cat") - await process.write_stdin(b"hello\n") - await process.terminate() - - assert await process.wait() == 0 - assert rpc_headers == [ - ("Start", {"Authorization": "Bearer start-token"}), - ("SendInput", {"Authorization": "Bearer stale-token"}), - ("SendInput", {"Authorization": "Bearer fresh-token"}), - ("SendSignal", {"Authorization": "Bearer fresh-token"}), - ] - finally: - await client.aclose() - - @pytest.mark.asyncio async def test_async_open_process_rejects_container_sandbox(): client = AsyncSandboxClient(api_key="test-key")