diff --git a/packages/prime-sandboxes/src/prime_sandboxes/process.py b/packages/prime-sandboxes/src/prime_sandboxes/process.py index ada6f835a..78357e883 100644 --- a/packages/prime-sandboxes/src/prime_sandboxes/process.py +++ b/packages/prime-sandboxes/src/prime_sandboxes/process.py @@ -2,10 +2,12 @@ import asyncio import contextlib +import logging from collections.abc import AsyncIterator, Awaitable, Callable from typing import Literal from connectrpc.client import ConnectClient +from connectrpc.code import Code from connectrpc.errors import ConnectError from google.protobuf.message import Message from pyqwest import HTTPTransport @@ -13,11 +15,16 @@ from .core import APIError from .rpc_command_session import parse_command_session_start_event +logger = logging.getLogger(__name__) + _EOF = object() _EXIT_WAIT_SECONDS = 5 +_STREAM_MAX_RECONNECTS = 5 +_STREAM_RECONNECT_BACKOFF_SECONDS = 0.5 _WriteStdin = Callable[[int, bytes], Awaitable[None]] _SendSignal = Callable[[int, Literal["terminate", "kill"]], Awaitable[None]] +_Reconnect = Callable[[int | None], AsyncIterator[Message]] class _AsyncProcessStream(AsyncIterator[bytes]): @@ -69,6 +76,7 @@ def __init__( write_stdin: _WriteStdin, send_signal: _SendSignal, transport: HTTPTransport | None = None, + reconnect: _Reconnect | None = None, ) -> None: self.stdout = _AsyncProcessStream() self.stderr = _AsyncProcessStream() @@ -77,6 +85,7 @@ def __init__( self._transport = transport self._write_stdin = write_stdin self._send_process_signal = send_signal + self._reconnect = reconnect self._remote_exited = False self._signals_sent: set[Literal["terminate", "kill"]] = set() self._closed = False @@ -100,8 +109,9 @@ async def _create( write_stdin: _WriteStdin, send_signal: _SendSignal, transport: HTTPTransport | None = None, + reconnect: _Reconnect | None = None, ) -> "AsyncSandboxProcess": - process = cls(stream_client, stream, write_stdin, send_signal, transport) + process = cls(stream_client, stream, write_stdin, send_signal, transport, reconnect) try: await asyncio.shield(process._started) except asyncio.CancelledError: @@ -223,31 +233,83 @@ async def _wait_for_exit_event(self) -> bool: return False return self._remote_exited + def _can_reconnect(self, reconnects: int, error: BaseException | None) -> bool: + """Whether the process has enough identity and budget for another attach.""" + if self._reconnect is None or self._remote_exited or reconnects >= _STREAM_MAX_RECONNECTS: + return False + if not isinstance(error, ConnectError) or error.code != Code.NOT_FOUND: + return True + return not self._started.done() and reconnects > 0 + + async def _aclose_stream(self) -> None: + close = getattr(self._stream, "aclose", None) + if close is not None: + with contextlib.suppress(BaseException): + await close() + + async def _reconnect_stream(self, reconnects: int, error: BaseException | None) -> None: + reconnect = self._reconnect + assert reconnect is not None + + delay = _STREAM_RECONNECT_BACKOFF_SECONDS * 2 ** (reconnects - 1) + logger.warning( + "live process stream dropped (%s); re-attaching %d/%d in %.1fs", + error or "ended without an exit event", + reconnects, + _STREAM_MAX_RECONNECTS, + delay, + ) + await self._aclose_stream() + await asyncio.sleep(delay) + # Connect tails the same process from re-attachment time; output emitted while detached + # is not replayed. + pid = self.pid if self._started.done() else None + self._stream = reconnect(pid) + async def _pump(self) -> None: ended = False + reconnects = 0 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 - self._remote_exited = True - if not self._started.done(): - raise APIError("Process exited before reporting its PID") - if not self._exit.done(): - self._exit.set_result(value) + while not ended: + error: BaseException | None = None + 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": + # A reconnected (Connect) stream re-announces the pid; keep the first. + 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 + self._remote_exited = True + if not self._started.done(): + raise APIError("Process ended before reporting its PID") + if not self._exit.done(): + self._exit.set_result(value) + break + except asyncio.CancelledError: + raise + except BaseException as stream_error: + error = stream_error + + if ended: + if error is not None: + raise error break - if not ended: - raise APIError("Process stream ended without an exit event") + if not self._can_reconnect(reconnects, error): + if error is not None: + raise error + raise APIError("Process stream ended without an exit event") + + reconnects += 1 + await self._reconnect_stream(reconnects, error) except asyncio.CancelledError: raise except BaseException as error: @@ -262,10 +324,7 @@ async def _pump(self) -> None: 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._aclose_stream() await self._stream_client.close() # Callers that only consume the streams never reach aclose(), so a # process-owned transport is released here too once the stream ends. 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 e4ef3bb83..1d23b412c 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.""" +"""Command-session RPC helpers.""" -from typing import Dict, List, Literal, Optional, Protocol, cast +from typing import Dict, List, Literal, Optional, Protocol, Sequence, cast from connectrpc.method import IdempotencyLevel, MethodInfo from google.protobuf.message import Message @@ -17,7 +17,9 @@ def __call__(self, *, cmd: str, args: List[str], envs: Dict[str, str]) -> _Comma class _CommandSessionStartRequestFactory(Protocol): - def __call__(self, *, command: _CommandSpecLike, stdin: bool) -> Message: ... + def __call__( + self, *, command: _CommandSpecLike, stdin: bool, tag: str | None = None + ) -> Message: ... class _CommandSessionSelectorFactory(Protocol): @@ -36,6 +38,19 @@ class _CommandSessionSendSignalRequestFactory(Protocol): def __call__(self, *, session: Message, signal: int) -> Message: ... +class _CommandSessionConnectRequestFactory(Protocol): + def __call__(self, *, session: Message) -> Message: ... + + +class _CommandSessionInfoLike(Protocol): + pid: int + tag: str + + +class _CommandSessionListResponseLike(Protocol): + sessions: Sequence[_CommandSessionInfoLike] + + class _CommandSessionDataEventLike(Protocol): stdout: bytes stderr: bytes @@ -84,6 +99,18 @@ def HasField(self, field_name: str) -> bool: ... _COMMAND_SESSION_SEND_SIGNAL_RESPONSE_TYPE = cast( type[Message], getattr(command_session_pb2, "SendSignalResponse") ) +_COMMAND_SESSION_CONNECT_REQUEST_TYPE = cast( + type[Message], getattr(command_session_pb2, "ConnectRequest") +) +_COMMAND_SESSION_CONNECT_RESPONSE_TYPE = cast( + type[Message], getattr(command_session_pb2, "ConnectResponse") +) +_COMMAND_SESSION_LIST_REQUEST_TYPE = cast( + type[Message], getattr(command_session_pb2, "ListRequest") +) +_COMMAND_SESSION_LIST_RESPONSE_TYPE = cast( + type[Message], getattr(command_session_pb2, "ListResponse") +) _COMMAND_SESSION_START_REQUEST_FACTORY = cast( _CommandSessionStartRequestFactory, _COMMAND_SESSION_START_REQUEST_TYPE ) @@ -99,6 +126,9 @@ def HasField(self, field_name: str) -> bool: ... _COMMAND_SESSION_SEND_SIGNAL_REQUEST_FACTORY = cast( _CommandSessionSendSignalRequestFactory, _COMMAND_SESSION_SEND_SIGNAL_REQUEST_TYPE ) +_COMMAND_SESSION_CONNECT_REQUEST_FACTORY = cast( + _CommandSessionConnectRequestFactory, _COMMAND_SESSION_CONNECT_REQUEST_TYPE +) COMMAND_SESSION_START_RPC_METHOD = MethodInfo( @@ -125,6 +155,23 @@ def HasField(self, field_name: str) -> bool: ... idempotency_level=IdempotencyLevel.UNKNOWN, ) +# Re-attach to an already-running session's output stream by its session selector. +COMMAND_SESSION_CONNECT_RPC_METHOD = MethodInfo( + name="Connect", + service_name="command_session.CommandSession", + input=_COMMAND_SESSION_CONNECT_REQUEST_TYPE, + output=_COMMAND_SESSION_CONNECT_RESPONSE_TYPE, + idempotency_level=IdempotencyLevel.NO_SIDE_EFFECTS, +) + +COMMAND_SESSION_LIST_RPC_METHOD = MethodInfo( + name="List", + service_name="command_session.CommandSession", + input=_COMMAND_SESSION_LIST_REQUEST_TYPE, + output=_COMMAND_SESSION_LIST_RESPONSE_TYPE, + idempotency_level=IdempotencyLevel.NO_SIDE_EFFECTS, +) + def build_command_session_start_request( command: str, @@ -132,6 +179,7 @@ def build_command_session_start_request( env: Optional[Dict[str, str]], *, stdin: bool = False, + tag: str | None = None, ) -> Message: command_spec = _COMMAND_SPEC_FACTORY( cmd="/bin/bash", @@ -141,7 +189,26 @@ 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=stdin) + return _COMMAND_SESSION_START_REQUEST_FACTORY( + command=command_spec, + stdin=stdin, + tag=tag, + ) + + +def build_command_session_connect_request(pid: int) -> Message: + return _COMMAND_SESSION_CONNECT_REQUEST_FACTORY( + session=_COMMAND_SESSION_SELECTOR_FACTORY(pid=pid) + ) + + +def build_command_session_list_request() -> Message: + return _COMMAND_SESSION_LIST_REQUEST_TYPE() + + +def find_command_session_pid(response: Message, tag: str) -> int | None: + sessions = cast(_CommandSessionListResponseLike, response).sessions + return next((int(session.pid) for session in sessions if session.tag == tag), None) def build_command_session_send_input_request(pid: int, data: bytes) -> Message: diff --git a/packages/prime-sandboxes/src/prime_sandboxes/sandbox.py b/packages/prime-sandboxes/src/prime_sandboxes/sandbox.py index 6f3870a6a..9f29ef49d 100644 --- a/packages/prime-sandboxes/src/prime_sandboxes/sandbox.py +++ b/packages/prime-sandboxes/src/prime_sandboxes/sandbox.py @@ -11,6 +11,7 @@ import threading import time import uuid +from collections.abc import AsyncIterator from concurrent.futures import Future from dataclasses import dataclass from datetime import datetime, timedelta, timezone @@ -86,13 +87,18 @@ ) from .process import AsyncSandboxProcess from .rpc_command_session import ( + COMMAND_SESSION_CONNECT_RPC_METHOD, + COMMAND_SESSION_LIST_RPC_METHOD, COMMAND_SESSION_SEND_INPUT_RPC_METHOD, COMMAND_SESSION_SEND_SIGNAL_RPC_METHOD, COMMAND_SESSION_START_RPC_METHOD, + build_command_session_connect_request, + build_command_session_list_request, build_command_session_send_input_request, build_command_session_send_signal_request, build_command_session_start_request, collect_command_session_start_event, + find_command_session_pid, ) # Connection-level errors: request never reached the server, so retry is safe @@ -112,6 +118,12 @@ _LIVE_PROCESS_TIMEOUT_MS = 24 * 60 * 60 * 1000 _PROCESS_INPUT_TIMEOUT_MS = 30_000 _PROCESS_SIGNAL_TIMEOUT_MS = 10_000 +_PROCESS_DISCOVERY_TIMEOUT_MS = 10_000 +_LIVE_PROCESS_TCP_KEEPALIVE_SECONDS = 15.0 +_LIVE_PROCESS_POOL_IDLE_TIMEOUT_SECONDS = 300.0 +_BACKGROUND_JOB_LAUNCH_ATTEMPTS = 3 +_BACKGROUND_JOB_LAUNCH_BACKOFF_SECONDS = 0.5 +_BACKGROUND_JOB_LAUNCH_TIMEOUT_SECONDS = 30 _RequestMessage = TypeVar("_RequestMessage", bound=Message) _ResponseMessage = TypeVar("_ResponseMessage", bound=Message) @@ -135,7 +147,11 @@ def _ca_bundle() -> bytes: def _live_process_transport() -> HTTPTransport: # A bare HTTPTransport carries no trust roots on some pyqwest versions # (only the default singleton does), so pass certifi's bundle explicitly. - return HTTPTransport(tls_ca_cert=_ca_bundle()) + return HTTPTransport( + tls_ca_cert=_ca_bundle(), + tcp_keepalive_interval=_LIVE_PROCESS_TCP_KEEPALIVE_SECONDS, + pool_idle_timeout=_LIVE_PROCESS_POOL_IDLE_TIMEOUT_SECONDS, + ) def _network_update_payload( @@ -1423,6 +1439,7 @@ def start_background_job( stdout_log_file = f"/tmp/job_{job_id}.stdout.log" stderr_log_file = f"/tmp/job_{job_id}.stderr.log" exit_file = f"/tmp/job_{job_id}.exit" + launch_dir = f"/tmp/job_{job_id}.launch" env_prefix = "" if env: @@ -1447,9 +1464,25 @@ def start_background_job( ) quoted_sh_command = shlex.quote(sh_command) - # Outer nohup redirects to /dev/null since output goes to log files inside sh -c - bg_cmd = f"nohup sh -c {quoted_sh_command} < /dev/null > /dev/null 2>&1 &" - self.execute_command(sandbox_id, bg_cmd, timeout=30, user=user) + # mkdir is the launch's idempotency guard: after an ambiguous timeout, only one attempt + # can create it and run the user command. + bg_cmd = ( + f"mkdir {shlex.quote(launch_dir)} && " + f"nohup sh -c {quoted_sh_command} < /dev/null > /dev/null 2>&1 &" + ) + for attempt in range(_BACKGROUND_JOB_LAUNCH_ATTEMPTS): + try: + self.execute_command( + sandbox_id, + bg_cmd, + timeout=_BACKGROUND_JOB_LAUNCH_TIMEOUT_SECONDS, + user=user, + ) + break + except CommandTimeoutError: + if attempt == _BACKGROUND_JOB_LAUNCH_ATTEMPTS - 1: + raise + time.sleep(_BACKGROUND_JOB_LAUNCH_BACKOFF_SECONDS * 2**attempt) return BackgroundJob( job_id=job_id, @@ -2637,7 +2670,9 @@ async def open_process( 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. + do not expose this transport and fail fast. If the initial stream drops + before reporting a PID, recovery can find the session only while it is + still running because sandboxd does not retain completed sessions. """ await self._auth_cache.get_or_refresh(sandbox_id) if not await self._auth_cache.is_vm(sandbox_id): @@ -2663,11 +2698,15 @@ async def open_process( send_compression=None, http_client=http_client, ) + # The SDK tag lets reconnect find a Start whose PID response was lost. Remove tag/List + # recovery once sandboxd provides idempotent Start/create-or-attach semantics. + process_tag = f"prime-sdk-{uuid.uuid4().hex}" request = build_command_session_start_request( command, working_dir, env, stdin=True, + tag=process_tag, ) stream = rpc_client.execute_server_stream( request=request, @@ -2696,12 +2735,54 @@ async def send_signal(pid: int, signal: Literal["terminate", "kill"]) -> None: http_client=http_client, ) + async def reconnect(pid: int | None) -> AsyncIterator[Message]: + reauthed = False + while True: + auth = await self._auth_cache.get_or_refresh(sandbox_id) + base_url = f"{auth['gateway_url'].rstrip('/')}/{auth['user_ns']}/{auth['job_id']}" + client = ConnectClient( + base_url, + codec=GOOGLE_PROTOBUF_BINARY_CODEC, + send_compression=None, + http_client=http_client, + ) + try: + if pid is None: + sessions = await client.execute_unary( + request=build_command_session_list_request(), + method=COMMAND_SESSION_LIST_RPC_METHOD, + headers={"Authorization": f"Bearer {auth['token']}"}, + timeout_ms=_PROCESS_DISCOVERY_TIMEOUT_MS, + ) + pid = find_command_session_pid(sessions, process_tag) + if pid is None: + raise ConnectError(Code.NOT_FOUND, "live process not found") + stream = client.execute_server_stream( + request=build_command_session_connect_request(pid), + method=COMMAND_SESSION_CONNECT_RPC_METHOD, + headers={"Authorization": f"Bearer {auth['token']}"}, + timeout_ms=_LIVE_PROCESS_TIMEOUT_MS, + ) + async for response in stream: + yield response + return + except ConnectError as error: + if error.code == Code.UNAUTHENTICATED and await self._should_retry_401( + sandbox_id, reauthed + ): + reauthed = True + continue + raise + finally: + await client.close() + return await AsyncSandboxProcess._create( rpc_client, stream, write_stdin, send_signal, transport=transport, + reconnect=reconnect, ) async def _execute_process_control_rpc( @@ -2952,6 +3033,7 @@ async def start_background_job( stdout_log_file = f"/tmp/job_{job_id}.stdout.log" stderr_log_file = f"/tmp/job_{job_id}.stderr.log" exit_file = f"/tmp/job_{job_id}.exit" + launch_dir = f"/tmp/job_{job_id}.launch" env_prefix = "" if env: @@ -2976,9 +3058,25 @@ async def start_background_job( ) quoted_sh_command = shlex.quote(sh_command) - # Outer nohup redirects to /dev/null since output goes to log files inside sh -c - bg_cmd = f"nohup sh -c {quoted_sh_command} < /dev/null > /dev/null 2>&1 &" - await self.execute_command(sandbox_id, bg_cmd, timeout=30, user=user) + # mkdir is the launch's idempotency guard: after an ambiguous timeout, only one attempt + # can create it and run the user command. + bg_cmd = ( + f"mkdir {shlex.quote(launch_dir)} && " + f"nohup sh -c {quoted_sh_command} < /dev/null > /dev/null 2>&1 &" + ) + for attempt in range(_BACKGROUND_JOB_LAUNCH_ATTEMPTS): + try: + await self.execute_command( + sandbox_id, + bg_cmd, + timeout=_BACKGROUND_JOB_LAUNCH_TIMEOUT_SECONDS, + user=user, + ) + break + except CommandTimeoutError: + if attempt == _BACKGROUND_JOB_LAUNCH_ATTEMPTS - 1: + raise + await asyncio.sleep(_BACKGROUND_JOB_LAUNCH_BACKOFF_SECONDS * 2**attempt) return BackgroundJob( job_id=job_id, diff --git a/packages/prime-sandboxes/tests/test_background_job_launch_retry.py b/packages/prime-sandboxes/tests/test_background_job_launch_retry.py new file mode 100644 index 000000000..176467805 --- /dev/null +++ b/packages/prime-sandboxes/tests/test_background_job_launch_retry.py @@ -0,0 +1,90 @@ +"""Background-job launch retries are guarded against duplicate execution.""" + +from typing import Any, cast + +import pytest + +from prime_sandboxes.core.client import APIClient +from prime_sandboxes.exceptions import CommandTimeoutError +from prime_sandboxes.models import CommandResponse +from prime_sandboxes.sandbox import AsyncSandboxClient, SandboxClient + +_OK = CommandResponse(stdout="", stderr="", exit_code=0) + + +def _timeout(): + return CommandTimeoutError("sb", "nohup ...", 30) + + +async def _no_sleep(_): + return None + + +class TestSyncLaunchRetry: + def test_retries_and_succeeds(self, monkeypatch): + monkeypatch.setattr("prime_sandboxes.sandbox.time.sleep", lambda _: None) + client = SandboxClient(APIClient(api_key="test-key")) + commands = [] + + def execute(_sandbox_id, command, **_kwargs): + commands.append(command) + if len(commands) == 1: + raise _timeout() + return _OK + + cast(Any, client).execute_command = execute + job = client.start_background_job("sb", "rm -rf x") + assert len(commands) == 2 + assert commands[0] == commands[1] + assert commands[0].startswith(f"mkdir /tmp/job_{job.job_id}.launch && nohup") + assert job.job_id + + def test_gives_up_after_max_attempts(self, monkeypatch): + monkeypatch.setattr("prime_sandboxes.sandbox.time.sleep", lambda _: None) + client = SandboxClient(APIClient(api_key="test-key")) + calls = {"n": 0} + + def execute(*_a, **_k): + calls["n"] += 1 + raise _timeout() + + cast(Any, client).execute_command = execute + with pytest.raises(CommandTimeoutError): + client.start_background_job("sb", "rm -rf x") + assert calls["n"] == 3 + + +class TestAsyncLaunchRetry: + @pytest.mark.asyncio + async def test_retries_and_succeeds(self, monkeypatch): + monkeypatch.setattr("prime_sandboxes.sandbox.asyncio.sleep", _no_sleep) + client = AsyncSandboxClient(APIClient(api_key="test-key")) + commands = [] + + async def execute(_sandbox_id, command, **_kwargs): + commands.append(command) + if len(commands) == 1: + raise _timeout() + return _OK + + cast(Any, client).execute_command = execute + job = await client.start_background_job("sb", "rm -rf x") + assert len(commands) == 2 + assert commands[0] == commands[1] + assert commands[0].startswith(f"mkdir /tmp/job_{job.job_id}.launch && nohup") + assert job.job_id + + @pytest.mark.asyncio + async def test_gives_up_after_max_attempts(self, monkeypatch): + monkeypatch.setattr("prime_sandboxes.sandbox.asyncio.sleep", _no_sleep) + client = AsyncSandboxClient(APIClient(api_key="test-key")) + calls = {"n": 0} + + async def execute(*_a, **_k): + calls["n"] += 1 + raise _timeout() + + cast(Any, client).execute_command = execute + with pytest.raises(CommandTimeoutError): + await client.start_background_job("sb", "rm -rf x") + assert calls["n"] == 3 diff --git a/packages/prime-sandboxes/tests/test_command_transport_selection.py b/packages/prime-sandboxes/tests/test_command_transport_selection.py index de3610a5f..79182de52 100644 --- a/packages/prime-sandboxes/tests/test_command_transport_selection.py +++ b/packages/prime-sandboxes/tests/test_command_transport_selection.py @@ -253,6 +253,89 @@ async def test_async_open_process_rejects_container_sandbox(): await client.aclose() +@pytest.mark.asyncio +async def test_process_reconnect_discovers_pid_by_tag_and_refreshes_rejected_auth(monkeypatch): + discovery_tokens = [] + connected_pids = [] + process_tags = [] + + class _RejectedTokenCache: + def __init__(self): + self.invalidations = 0 + + async def get_or_refresh(self, _sandbox_id: str): + auth = _auth_payload() + auth["token"] = "stale" if self.invalidations == 0 else "fresh" + return auth + + async def invalidate(self, _sandbox_id: str): + self.invalidations += 1 + + async def is_vm(self, _sandbox_id: str): + return True + + class _FakeConnectClient: + def __init__(self, _address: str, **_kwargs): + pass + + def execute_server_stream(self, **kwargs): + method = kwargs["method"].name + if method == "Start": + process_tags.append(kwargs["request"].tag) + else: + connected_pids.append(kwargs["request"].session.pid) + + async def events(): + if method == "Start": + raise ConnectError(Code.UNAVAILABLE, "stream dropped") + yield command_session_pb2.ConnectResponse( + event=command_session_pb2.CommandSessionEvent( + start=command_session_pb2.CommandSessionEvent.StartEvent(pid=42) + ) + ) + yield command_session_pb2.ConnectResponse( + event=command_session_pb2.CommandSessionEvent( + end=command_session_pb2.CommandSessionEvent.EndEvent(exit_code=0) + ) + ) + + return events() + + async def execute_unary(self, **kwargs): + token = kwargs["headers"]["Authorization"] + discovery_tokens.append(token) + if token == "Bearer stale": + raise ConnectError(Code.UNAUTHENTICATED, "expired token") + return command_session_pb2.ListResponse( + sessions=[ + command_session_pb2.CommandSessionInfo(pid=7, tag="other-process"), + command_session_pb2.CommandSessionInfo(pid=42, tag=process_tags[0]), + ] + ) + + async def close(self): + pass + + monkeypatch.setattr("prime_sandboxes.sandbox.ConnectClient", _FakeConnectClient) + monkeypatch.setattr("prime_sandboxes.process._STREAM_RECONNECT_BACKOFF_SECONDS", 0) + + client = AsyncSandboxClient(api_key="test-key") + cache = _RejectedTokenCache() + cast(Any, client)._auth_cache = cache + try: + process = await client.open_process("sbx-vm", "sleep 1") + + assert await process.wait() == 0 + assert len(process_tags) == 1 + assert process_tags[0].startswith("prime-sdk-") + assert discovery_tokens == ["Bearer stale", "Bearer fresh"] + assert connected_pids == [42] + assert cache.invalidations == 1 + await process.aclose() + finally: + await client.aclose() + + def test_auth_cache_stores_vm_flag_for_reuse(tmp_path): class _FakeAPIClient: def __init__(self): diff --git a/packages/prime-sandboxes/tests/test_process_stream_reconnect.py b/packages/prime-sandboxes/tests/test_process_stream_reconnect.py new file mode 100644 index 000000000..0d759ca0b --- /dev/null +++ b/packages/prime-sandboxes/tests/test_process_stream_reconnect.py @@ -0,0 +1,207 @@ +"""Live-process output streams re-attach to the running process after a drop.""" + +import asyncio + +import pytest +from connectrpc.code import Code +from connectrpc.errors import ConnectError + +from prime_sandboxes._proto.command_session import command_session_pb2 as pb +from prime_sandboxes.core import APIError +from prime_sandboxes.process import AsyncSandboxProcess + +_EV = pb.CommandSessionEvent + + +def _start(pid): + return pb.StartResponse(event=_EV(start=_EV.StartEvent(pid=pid))) + + +def _stdout(data): + return pb.StartResponse(event=_EV(data=_EV.DataEvent(stdout=data))) + + +def _end(code): + return pb.StartResponse(event=_EV(end=_EV.EndEvent(exit_code=code))) + + +class _FakeStreamClient: + async def close(self): + pass + + +async def _noop_stdin(pid, data): + pass + + +async def _noop_signal(pid, sig): + pass + + +async def _drain(stream): + out = b"" + async for chunk in stream: + out += chunk + return out + + +_STREAM_FAULTS = [ + ConnectError(Code.UNAVAILABLE, "error reading a body from connection: timed out"), + ConnectError(Code.INTERNAL, "Error reading content"), +] + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "fault", _STREAM_FAULTS, ids=["unavailable_timeout", "internal_reading_content"] +) +async def test_stream_reconnects_and_resumes_after_drop(fault, monkeypatch): + monkeypatch.setattr("prime_sandboxes.process._STREAM_RECONNECT_BACKOFF_SECONDS", 0) + + async def faulty(): + yield _start(42) + yield _stdout(b"before\n") + raise fault + + async def resumed(): + yield _start(42) # Connect re-announces the pid; already known, ignored + yield _stdout(b"after\n") + yield _end(0) + + reconnect_calls = [] + + def reconnect(pid): + reconnect_calls.append(pid) + return resumed() + + proc = await AsyncSandboxProcess._create( + _FakeStreamClient(), faulty(), _noop_stdin, _noop_signal, reconnect=reconnect + ) + stdout = await _drain(proc.stdout) + rc = await proc.wait() + + assert reconnect_calls == [42] + assert rc == 0 # exit observed on the resumed stream + assert stdout == b"before\nafter\n" # output from both segments + await proc.aclose() + + +@pytest.mark.asyncio +async def test_stream_reconnects_after_clean_eof(monkeypatch): + monkeypatch.setattr("prime_sandboxes.process._STREAM_RECONNECT_BACKOFF_SECONDS", 0) + + async def ended_without_exit(): + yield _start(42) + yield _stdout(b"before\n") + + async def resumed(): + yield _start(42) + yield _stdout(b"after\n") + yield _end(0) + + reconnect_calls = [] + + def reconnect(pid): + reconnect_calls.append(pid) + return resumed() + + proc = await AsyncSandboxProcess._create( + _FakeStreamClient(), + ended_without_exit(), + _noop_stdin, + _noop_signal, + reconnect=reconnect, + ) + + assert await _drain(proc.stdout) == b"before\nafter\n" + assert await proc.wait() == 0 + assert reconnect_calls == [42] + await proc.aclose() + + +@pytest.mark.asyncio +async def test_stream_reconnects_before_pid_is_received(monkeypatch): + monkeypatch.setattr("prime_sandboxes.process._STREAM_RECONNECT_BACKOFF_SECONDS", 0) + + async def dropped_before_start(): + raise ConnectError(Code.UNAVAILABLE, "stream dropped") + yield _start(42) + + reconnect_calls = [] + + def reconnect(pid): + reconnect_calls.append(pid) + + async def stream(): + if len(reconnect_calls) == 1: + raise ConnectError(Code.NOT_FOUND, "process not registered yet") + yield _start(42) + yield _end(0) + + return stream() + + proc = await AsyncSandboxProcess._create( + _FakeStreamClient(), + dropped_before_start(), + _noop_stdin, + _noop_signal, + reconnect=reconnect, + ) + + assert proc.pid == 42 + assert await proc.wait() == 0 + assert reconnect_calls == [None, None] + await proc.aclose() + + +@pytest.mark.asyncio +async def test_end_before_pid_fails_instead_of_hanging(): + async def ended_before_start(): + yield _end(0) + + with pytest.raises(APIError, match="ended before reporting its PID"): + await asyncio.wait_for( + AsyncSandboxProcess._create( + _FakeStreamClient(), + ended_before_start(), + _noop_stdin, + _noop_signal, + ), + timeout=1, + ) + + +@pytest.mark.asyncio +async def test_stream_without_reconnect_still_fails(): + # No reconnect callable preserves the previous fatal behavior. + async def faulty(): + yield _start(7) + raise ConnectError(Code.UNAVAILABLE, "error reading a body from connection: timed out") + + proc = await AsyncSandboxProcess._create( + _FakeStreamClient(), faulty(), _noop_stdin, _noop_signal, reconnect=None + ) + with pytest.raises(APIError, match="process stream RPC failed"): + await proc.wait() + await proc.aclose() + + +@pytest.mark.asyncio +async def test_permanent_fault_is_not_reconnected(): + async def faulty(): + yield _start(9) + raise ConnectError(Code.NOT_FOUND, "session gone") + + calls = [] + + def reconnect(pid): + calls.append(pid) + raise AssertionError("should not reconnect on a permanent fault") + + proc = await AsyncSandboxProcess._create( + _FakeStreamClient(), faulty(), _noop_stdin, _noop_signal, reconnect=reconnect + ) + with pytest.raises(APIError, match="process stream RPC failed"): + await proc.wait() + assert calls == [] + await proc.aclose()