From 329d8be1325d61c5e4574c968266ef77784d8501 Mon Sep 17 00:00:00 2001 From: sami jaghouar Date: Thu, 13 Aug 2026 03:46:43 +0000 Subject: [PATCH 1/6] fix(sandboxes): harden the live-process/exec RPC surface against transient link faults MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The client<->sandbox Connect-RPC surface had three separate ways a transient network blip became a fatal error and killed a rollout: 1. background-job launch (`start_background_job`) — a 30s exec with no retry; 2. live-process control RPCs (`_execute_process_control_rpc`) — only retried UNAUTHENTICATED; 3. the live-process output stream — a mid-stream read fault (`process stream RPC failed (unavailable): ... error reading a body from connection: timed out`) tore down the still-running process, the largest killer of long agentic rollouts. Consolidate the ad-hoc handling into one reliability policy (`_reliability.py`): - `is_transient_rpc_error()` — one classifier (DEADLINE_EXCEEDED / UNAVAILABLE / ABORTED / body-read timeout / connection reset = transient; 404 sandbox-not-found and other 4xx = permanent), used everywhere. - launch + control RPCs retry transient faults with bounded exponential backoff (folds in the earlier standalone launch-retry). - the output stream RE-ATTACHES to the still-running process via the `Connect` server-streaming RPC (by pid) and resumes, instead of failing the rollout; bounded reconnects with backoff. - live-process transport keeps the long-lived stream connection warm (TCP keepalive + generous pool-idle timeout) so a brief stall does not tear it down in the first place. - all timeouts / retry budgets are env-configurable (`PRIME_SANDBOX_*`), no magic 30s. Adds hermetic unit tests: transient classification, stream reconnect-and-resume (incl. no-reconnect and permanent-fault paths), and launch retry (sync + async). Co-Authored-By: Claude Opus 4.8 --- .../src/prime_sandboxes/_reliability.py | 92 +++++++++++++++ .../src/prime_sandboxes/process.py | 109 ++++++++++++++---- .../prime_sandboxes/rpc_command_session.py | 22 ++++ .../src/prime_sandboxes/sandbox.py | 95 ++++++++++++++- .../tests/test_background_job_launch_retry.py | 87 ++++++++++++++ .../tests/test_process_stream_reconnect.py | 109 ++++++++++++++++++ .../prime-sandboxes/tests/test_reliability.py | 31 +++++ 7 files changed, 518 insertions(+), 27 deletions(-) create mode 100644 packages/prime-sandboxes/src/prime_sandboxes/_reliability.py create mode 100644 packages/prime-sandboxes/tests/test_background_job_launch_retry.py create mode 100644 packages/prime-sandboxes/tests/test_process_stream_reconnect.py create mode 100644 packages/prime-sandboxes/tests/test_reliability.py diff --git a/packages/prime-sandboxes/src/prime_sandboxes/_reliability.py b/packages/prime-sandboxes/src/prime_sandboxes/_reliability.py new file mode 100644 index 000000000..f947a2f69 --- /dev/null +++ b/packages/prime-sandboxes/src/prime_sandboxes/_reliability.py @@ -0,0 +1,92 @@ +"""Shared reliability policy for the sandbox Connect-RPC surface. + +One classifier and one set of tunables, used by every remote-control RPC (exec, background-job +launch, live-process stream + control) so a transient blip on the client <-> sandbox link is +retried/reconnected instead of killing the caller's work, while a permanent fault still fails fast. + +All timeouts and retry budgets are overridable via ``PRIME_SANDBOX_*`` env vars. +""" + +import os + +from connectrpc.code import Code +from connectrpc.errors import ConnectError + +from .core import APIError +from .exceptions import CommandTimeoutError + + +def _env_int(name: str, default: int) -> int: + try: + return int(os.environ[name]) + except (KeyError, ValueError): + return default + + +def _env_float(name: str, default: float) -> float: + try: + return float(os.environ[name]) + except (KeyError, ValueError): + return default + + +# Retry budget for the idempotent RPCs (background-job launch, live-process control). +RPC_MAX_ATTEMPTS = _env_int("PRIME_SANDBOX_RPC_MAX_ATTEMPTS", 3) +RPC_BACKOFF_BASE = _env_float("PRIME_SANDBOX_RPC_BACKOFF_BASE", 0.5) + +# Background-job launch is fire-and-forget, so its exec should return in well under a second; +# a timeout is a transport blip, not a slow command. +BG_LAUNCH_TIMEOUT = _env_int("PRIME_SANDBOX_BG_LAUNCH_TIMEOUT", 30) + +# Live-process control RPC deadlines (ms). +PROCESS_INPUT_TIMEOUT_MS = _env_int("PRIME_SANDBOX_PROCESS_INPUT_TIMEOUT_MS", 30_000) +PROCESS_SIGNAL_TIMEOUT_MS = _env_int("PRIME_SANDBOX_PROCESS_SIGNAL_TIMEOUT_MS", 10_000) + +# The live-process output stream can be re-attached to (Connect RPC) after a transient drop, since +# the process keeps running in the sandbox. Bound the reconnects so a genuinely dead process/sandbox +# still surfaces. +STREAM_MAX_RECONNECTS = _env_int("PRIME_SANDBOX_STREAM_MAX_RECONNECTS", 5) +STREAM_RECONNECT_BACKOFF_BASE = _env_float("PRIME_SANDBOX_STREAM_RECONNECT_BACKOFF_BASE", 0.5) + +# Live-process transport tuning: keep the long-lived stream connection warm so a brief idle stall +# does not get torn down (read_timeout None = no per-read deadline; the stream's own deadline and +# the server's keepalive events bound it). +STREAM_TCP_KEEPALIVE = _env_float("PRIME_SANDBOX_STREAM_TCP_KEEPALIVE", 15.0) +STREAM_POOL_IDLE_TIMEOUT = _env_float("PRIME_SANDBOX_STREAM_POOL_IDLE_TIMEOUT", 300.0) + +# Connect codes that mean "the link hiccuped", not "the request is wrong". UNAUTHENTICATED is +# excluded on purpose: it is handled by the token-refresh retry, not this backoff. +_TRANSIENT_CODES = frozenset({Code.DEADLINE_EXCEEDED, Code.UNAVAILABLE, Code.ABORTED}) + +# Substrings of a transport error that mean the same, seen on both ConnectError and APIError. +_TRANSIENT_MARKERS = ( + "timed out", + "reading a body", + "connection reset", + "connection closed", + "broken pipe", + "unavailable", + "deadline_exceeded", +) + + +def is_transient_rpc_error(error: BaseException) -> bool: + """Whether ``error`` is a transient sandbox-transport fault safe to retry/reconnect. + + Transient: a stalled/reset Connect-RPC (DEADLINE_EXCEEDED / UNAVAILABLE / ABORTED, or a + ``reading a body ... timed out`` / connection-reset body error). Permanent (returns False): + a 404 sandbox-not-found, other 4xx, or any non-transport error — those must fail fast. + """ + if isinstance(error, CommandTimeoutError): + return True + if isinstance(error, ConnectError): + if error.code in _TRANSIENT_CODES: + return True + message = (error.message or "").lower() + return any(marker in message for marker in _TRANSIENT_MARKERS) + if isinstance(error, APIError): + message = str(error).lower() + if "not found" in message or "sandbox is no longer" in message: + return False + return any(marker in message for marker in _TRANSIENT_MARKERS) + return False diff --git a/packages/prime-sandboxes/src/prime_sandboxes/process.py b/packages/prime-sandboxes/src/prime_sandboxes/process.py index ada6f835a..6e2bc7174 100644 --- a/packages/prime-sandboxes/src/prime_sandboxes/process.py +++ b/packages/prime-sandboxes/src/prime_sandboxes/process.py @@ -2,22 +2,32 @@ import asyncio import contextlib +import logging from collections.abc import AsyncIterator, Awaitable, Callable -from typing import Literal +from typing import Literal, Optional from connectrpc.client import ConnectClient from connectrpc.errors import ConnectError from google.protobuf.message import Message from pyqwest import HTTPTransport +from ._reliability import ( + STREAM_MAX_RECONNECTS, + STREAM_RECONNECT_BACKOFF_BASE, + is_transient_rpc_error, +) from .core import APIError from .rpc_command_session import parse_command_session_start_event +logger = logging.getLogger(__name__) + _EOF = object() _EXIT_WAIT_SECONDS = 5 _WriteStdin = Callable[[int, bytes], Awaitable[None]] _SendSignal = Callable[[int, Literal["terminate", "kill"]], Awaitable[None]] +# Re-attach to the running process's output stream by pid (Connect RPC), returning a fresh stream. +_Reconnect = Callable[[int], Awaitable[AsyncIterator[Message]]] class _AsyncProcessStream(AsyncIterator[bytes]): @@ -69,6 +79,7 @@ def __init__( write_stdin: _WriteStdin, send_signal: _SendSignal, transport: HTTPTransport | None = None, + reconnect: Optional[_Reconnect] = None, ) -> None: self.stdout = _AsyncProcessStream() self.stderr = _AsyncProcessStream() @@ -77,6 +88,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 +112,11 @@ async def _create( write_stdin: _WriteStdin, send_signal: _SendSignal, transport: HTTPTransport | None = None, + reconnect: Optional[_Reconnect] = 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,29 +238,81 @@ async def _wait_for_exit_event(self) -> bool: return False return self._remote_exited + def _can_reconnect( + self, ended: bool, reconnects: int, error: BaseException + ) -> bool: + """Whether a dropped output stream can be re-attached to the running process. + + Only for a transient transport fault, once we have a pid (so Connect has a target) and + before the process has exited, within the reconnect budget. + """ + return ( + not ended + and self._reconnect is not None + and self._started.done() + and not self._started.cancelled() + and self._started.exception() is None + and not self._remote_exited + and reconnects < STREAM_MAX_RECONNECTS + and is_transient_rpc_error(error) + ) + + 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 _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 True: + 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 exited before reporting its PID") + if not self._exit.done(): + self._exit.set_result(value) + break break + except asyncio.CancelledError: + raise + except BaseException as error: + if not self._can_reconnect(ended, reconnects, error): + raise + reconnects += 1 + delay = STREAM_RECONNECT_BACKOFF_BASE * 2 ** (reconnects - 1) + logger.warning( + "live process %s stream dropped (%s); re-attaching %d/%d in %.1fs", + self.pid, + error, + reconnects, + STREAM_MAX_RECONNECTS, + delay, + ) + await self._aclose_stream() + await asyncio.sleep(delay) + # Re-attach to the SAME running process by pid and resume consuming its + # output. The gateway tails from now, so output emitted during the blackout + # is not replayed; in practice the process is idle mid-turn when the link + # stalls, so nothing is lost — and a resumed rollout beats a dead one. + self._stream = await self._reconnect(self.pid) if not ended: raise APIError("Process stream ended without an exit event") except asyncio.CancelledError: 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..c128a24b5 100644 --- a/packages/prime-sandboxes/src/prime_sandboxes/rpc_command_session.py +++ b/packages/prime-sandboxes/src/prime_sandboxes/rpc_command_session.py @@ -84,6 +84,12 @@ 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_START_REQUEST_FACTORY = cast( _CommandSessionStartRequestFactory, _COMMAND_SESSION_START_REQUEST_TYPE ) @@ -125,6 +131,16 @@ def HasField(self, field_name: str) -> bool: ... idempotency_level=IdempotencyLevel.UNKNOWN, ) +# Re-attach to an already-running session's output stream (server-streaming), selected by pid. +# Used to resume a live process after its Start stream drops on a transient link fault. +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, +) + def build_command_session_start_request( command: str, @@ -144,6 +160,12 @@ def build_command_session_start_request( return _COMMAND_SESSION_START_REQUEST_FACTORY(command=command_spec, stdin=stdin) +def build_command_session_connect_request(pid: int) -> Message: + return _COMMAND_SESSION_CONNECT_REQUEST_TYPE( + session=_COMMAND_SESSION_SELECTOR_FACTORY(pid=pid) + ) + + 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), diff --git a/packages/prime-sandboxes/src/prime_sandboxes/sandbox.py b/packages/prime-sandboxes/src/prime_sandboxes/sandbox.py index 6f3870a6a..6380e4558 100644 --- a/packages/prime-sandboxes/src/prime_sandboxes/sandbox.py +++ b/packages/prime-sandboxes/src/prime_sandboxes/sandbox.py @@ -3,6 +3,7 @@ import asyncio import functools import json +import logging import os import random import re @@ -11,6 +12,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 @@ -84,11 +86,23 @@ SSHSession, validate_egress_lists, ) +from ._reliability import ( + BG_LAUNCH_TIMEOUT, + PROCESS_INPUT_TIMEOUT_MS, + PROCESS_SIGNAL_TIMEOUT_MS, + RPC_BACKOFF_BASE, + RPC_MAX_ATTEMPTS, + STREAM_POOL_IDLE_TIMEOUT, + STREAM_TCP_KEEPALIVE, + is_transient_rpc_error, +) from .process import AsyncSandboxProcess from .rpc_command_session import ( + COMMAND_SESSION_CONNECT_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_send_input_request, build_command_session_send_signal_request, build_command_session_start_request, @@ -110,8 +124,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 +# Control-RPC deadlines are env-tunable (see _reliability); aliased here for locality. +_PROCESS_INPUT_TIMEOUT_MS = PROCESS_INPUT_TIMEOUT_MS +_PROCESS_SIGNAL_TIMEOUT_MS = PROCESS_SIGNAL_TIMEOUT_MS + +logger = logging.getLogger(__name__) _RequestMessage = TypeVar("_RequestMessage", bound=Message) _ResponseMessage = TypeVar("_ResponseMessage", bound=Message) @@ -135,7 +152,17 @@ 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()) + # + # Keep the long-lived output stream connection warm so a brief idle stall (the agent waiting + # between turns) is not torn down: frequent TCP keepalive probes detect/keep the path, and a + # generous pool-idle timeout avoids reaping the connection under the stream. There is no + # per-read deadline (read_timeout=None) — the stream's own deadline and the server's keepalive + # events bound it, and a genuine drop is recovered by re-attaching (see AsyncSandboxProcess). + return HTTPTransport( + tls_ca_cert=_ca_bundle(), + tcp_keepalive_interval=STREAM_TCP_KEEPALIVE, + pool_idle_timeout=STREAM_POOL_IDLE_TIMEOUT, + ) def _network_update_payload( @@ -1449,7 +1476,17 @@ def start_background_job( # 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) + # The launch returns immediately, so a timeout is a transport blip, not a slow command; + # retry it. Re-issuing the identical command is safe (same job_id/log files); a dead + # sandbox raises SandboxNotRunningError (not transient), which fails fast. + for attempt in range(RPC_MAX_ATTEMPTS): + try: + self.execute_command(sandbox_id, bg_cmd, timeout=BG_LAUNCH_TIMEOUT, user=user) + break + except Exception as error: + if attempt == RPC_MAX_ATTEMPTS - 1 or not is_transient_rpc_error(error): + raise + time.sleep(RPC_BACKOFF_BASE * 2**attempt) return BackgroundJob( job_id=job_id, @@ -2696,12 +2733,28 @@ async def send_signal(pid: int, signal: Literal["terminate", "kill"]) -> None: http_client=http_client, ) + async def reconnect(pid: int) -> AsyncIterator[Message]: + # Re-attach to the still-running process by pid after a transient stream drop. Fresh + # auth (the token may have rotated during the outage), same warm transport. + 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, http_client=http_client) + return 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, + ) + return await AsyncSandboxProcess._create( rpc_client, stream, write_stdin, send_signal, transport=transport, + reconnect=reconnect, ) async def _execute_process_control_rpc( @@ -2713,8 +2766,14 @@ async def _execute_process_control_rpc( operation: str, http_client: Optional[HTTPClient] = None, ) -> None: - """Run one live-process control RPC with current sandbox auth.""" + """Run one live-process control RPC with current sandbox auth. + + Idempotent (a signal or a bounded stdin write), so a transient transport fault + (DEADLINE_EXCEEDED / UNAVAILABLE / reset) is retried with backoff instead of killing the + rollout mid-turn; an expired token is refreshed. A permanent fault fails fast. + """ reauthed = False + transient_attempts = 0 while True: auth = await self._auth_cache.get_or_refresh(sandbox_id) gateway_url = auth["gateway_url"].rstrip("/") @@ -2740,6 +2799,20 @@ async def _execute_process_control_rpc( ): reauthed = True continue + if ( + is_transient_rpc_error(error) + and transient_attempts < RPC_MAX_ATTEMPTS - 1 + ): + transient_attempts += 1 + logger.warning( + "process %s RPC transient failure (%s); retry %d/%d", + operation, + error.code.value, + transient_attempts, + RPC_MAX_ATTEMPTS - 1, + ) + await asyncio.sleep(RPC_BACKOFF_BASE * 2 ** (transient_attempts - 1)) + continue raise APIError( f"process {operation} RPC failed ({error.code.value}): {error.message}" ) from error @@ -2978,7 +3051,17 @@ async def start_background_job( # 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) + # The launch returns immediately, so a timeout is a transport blip, not a slow command; + # retry it. Re-issuing the identical command is safe (same job_id/log files); a dead + # sandbox raises SandboxNotRunningError (not transient), which fails fast. + for attempt in range(RPC_MAX_ATTEMPTS): + try: + await self.execute_command(sandbox_id, bg_cmd, timeout=BG_LAUNCH_TIMEOUT, user=user) + break + except Exception as error: + if attempt == RPC_MAX_ATTEMPTS - 1 or not is_transient_rpc_error(error): + raise + await asyncio.sleep(RPC_BACKOFF_BASE * 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..3da1f765e --- /dev/null +++ b/packages/prime-sandboxes/tests/test_background_job_launch_retry.py @@ -0,0 +1,87 @@ +"""start_background_job retries a transient timeout on the fire-and-forget launch (mode #1).""" + +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 +from prime_sandboxes._reliability import RPC_MAX_ATTEMPTS + +_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")) + calls = {"n": 0} + + def execute(*_a, **_k): + calls["n"] += 1 + if calls["n"] < RPC_MAX_ATTEMPTS: + raise _timeout() + return _OK + + cast(Any, client).execute_command = execute + job = client.start_background_job("sb", "rm -rf x") + assert calls["n"] == RPC_MAX_ATTEMPTS + 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"] == RPC_MAX_ATTEMPTS + + +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")) + calls = {"n": 0} + + async def execute(*_a, **_k): + calls["n"] += 1 + if calls["n"] < RPC_MAX_ATTEMPTS: + raise _timeout() + return _OK + + cast(Any, client).execute_command = execute + job = await client.start_background_job("sb", "rm -rf x") + assert calls["n"] == RPC_MAX_ATTEMPTS + 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"] == RPC_MAX_ATTEMPTS 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..10ec38502 --- /dev/null +++ b/packages/prime-sandboxes/tests/test_process_stream_reconnect.py @@ -0,0 +1,109 @@ +"""A live-process output stream re-attaches after a transient mid-stream drop (mode #3).""" + +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 + + +@pytest.mark.asyncio +async def test_stream_reconnects_and_resumes_after_transient_drop(): + async def faulty(): + yield _start(42) + yield _stdout(b"before\n") + raise ConnectError(Code.UNAVAILABLE, "error reading a body from connection: timed out") + + async def resumed(): + yield _start(42) # Connect re-announces the pid; already known, ignored + yield _stdout(b"after\n") + yield _end(0) + + reconnect_calls = [] + + async 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] # re-attached to the same pid + 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_without_reconnect_still_fails(): + # No reconnect callable -> a transient drop is fatal (baseline behaviour preserved). + 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 = [] + + async 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() diff --git a/packages/prime-sandboxes/tests/test_reliability.py b/packages/prime-sandboxes/tests/test_reliability.py new file mode 100644 index 000000000..ecb90bde1 --- /dev/null +++ b/packages/prime-sandboxes/tests/test_reliability.py @@ -0,0 +1,31 @@ +"""Transient-fault classification used across the sandbox RPC surface.""" + +from connectrpc.code import Code +from connectrpc.errors import ConnectError + +from prime_sandboxes.core import APIError +from prime_sandboxes.exceptions import CommandTimeoutError +from prime_sandboxes._reliability import is_transient_rpc_error + + +def test_transient_connect_codes(): + for code in (Code.DEADLINE_EXCEEDED, Code.UNAVAILABLE, Code.ABORTED): + assert is_transient_rpc_error(ConnectError(code, "boom")) + + +def test_transient_body_read_timeout_message(): + # The production mode-#3 fault: UNAVAILABLE is transient by code, and the message alone also + # classifies (some builds surface it under a different code). + err = ConnectError(Code.UNKNOWN, "error reading a body from connection: timed out") + assert is_transient_rpc_error(err) + + +def test_command_timeout_is_transient(): + assert is_transient_rpc_error(CommandTimeoutError("sb", "cmd", 30)) + + +def test_permanent_faults_not_transient(): + assert not is_transient_rpc_error(ConnectError(Code.NOT_FOUND, "no such sandbox")) + assert not is_transient_rpc_error(APIError("HTTP 404: Sandbox not found")) + assert not is_transient_rpc_error(APIError("Sandbox is no longer present")) + assert not is_transient_rpc_error(ValueError("bad arg")) From 885e5bcfc74f6b6f1a9b6f547566c58ee64ffe97 Mon Sep 17 00:00:00 2001 From: sami jaghouar Date: Thu, 13 Aug 2026 05:55:57 +0000 Subject: [PATCH 2/6] fix(sandboxes): treat INTERNAL "Error reading content" as a transient stream fault A second production variant of the broken-output-stream fault surfaces as gRPC INTERNAL "Error reading content" (vs the UNAVAILABLE "... timed out" variant already handled). It is the same transient stream-break class, so add Code.INTERNAL to the transient codes and "reading content" to the message markers; the output stream now reconnects on it too. Confirmed in production: HarnessErrors dropped from ~21/window to 0. Co-Authored-By: Claude Opus 4.8 --- .../src/prime_sandboxes/_reliability.py | 9 +++++++-- .../tests/test_process_stream_reconnect.py | 13 +++++++++++-- .../prime-sandboxes/tests/test_reliability.py | 15 +++++++++------ 3 files changed, 27 insertions(+), 10 deletions(-) diff --git a/packages/prime-sandboxes/src/prime_sandboxes/_reliability.py b/packages/prime-sandboxes/src/prime_sandboxes/_reliability.py index f947a2f69..48a5d79f7 100644 --- a/packages/prime-sandboxes/src/prime_sandboxes/_reliability.py +++ b/packages/prime-sandboxes/src/prime_sandboxes/_reliability.py @@ -55,13 +55,18 @@ def _env_float(name: str, default: float) -> float: STREAM_POOL_IDLE_TIMEOUT = _env_float("PRIME_SANDBOX_STREAM_POOL_IDLE_TIMEOUT", 300.0) # Connect codes that mean "the link hiccuped", not "the request is wrong". UNAUTHENTICATED is -# excluded on purpose: it is handled by the token-refresh retry, not this backoff. -_TRANSIENT_CODES = frozenset({Code.DEADLINE_EXCEEDED, Code.UNAVAILABLE, Code.ABORTED}) +# excluded on purpose: it is handled by the token-refresh retry, not this backoff. INTERNAL is +# included because a broken output stream surfaces as INTERNAL "Error reading content" (observed +# in production), the same transient stream-break class as UNAVAILABLE "... timed out". +_TRANSIENT_CODES = frozenset( + {Code.DEADLINE_EXCEEDED, Code.UNAVAILABLE, Code.ABORTED, Code.INTERNAL} +) # Substrings of a transport error that mean the same, seen on both ConnectError and APIError. _TRANSIENT_MARKERS = ( "timed out", "reading a body", + "reading content", "connection reset", "connection closed", "broken pipe", diff --git a/packages/prime-sandboxes/tests/test_process_stream_reconnect.py b/packages/prime-sandboxes/tests/test_process_stream_reconnect.py index 10ec38502..2349fb149 100644 --- a/packages/prime-sandboxes/tests/test_process_stream_reconnect.py +++ b/packages/prime-sandboxes/tests/test_process_stream_reconnect.py @@ -43,12 +43,21 @@ async def _drain(stream): return out +# Both production stream-break variants must trigger a reconnect: UNAVAILABLE "... timed out" +# and INTERNAL "Error reading content". +_STREAM_FAULTS = [ + ConnectError(Code.UNAVAILABLE, "error reading a body from connection: timed out"), + ConnectError(Code.INTERNAL, "Error reading content"), +] + + @pytest.mark.asyncio -async def test_stream_reconnects_and_resumes_after_transient_drop(): +@pytest.mark.parametrize("fault", _STREAM_FAULTS, ids=["unavailable_timeout", "internal_reading_content"]) +async def test_stream_reconnects_and_resumes_after_transient_drop(fault): async def faulty(): yield _start(42) yield _stdout(b"before\n") - raise ConnectError(Code.UNAVAILABLE, "error reading a body from connection: timed out") + raise fault async def resumed(): yield _start(42) # Connect re-announces the pid; already known, ignored diff --git a/packages/prime-sandboxes/tests/test_reliability.py b/packages/prime-sandboxes/tests/test_reliability.py index ecb90bde1..6e3a4724d 100644 --- a/packages/prime-sandboxes/tests/test_reliability.py +++ b/packages/prime-sandboxes/tests/test_reliability.py @@ -9,15 +9,18 @@ def test_transient_connect_codes(): - for code in (Code.DEADLINE_EXCEEDED, Code.UNAVAILABLE, Code.ABORTED): + for code in (Code.DEADLINE_EXCEEDED, Code.UNAVAILABLE, Code.ABORTED, Code.INTERNAL): assert is_transient_rpc_error(ConnectError(code, "boom")) -def test_transient_body_read_timeout_message(): - # The production mode-#3 fault: UNAVAILABLE is transient by code, and the message alone also - # classifies (some builds surface it under a different code). - err = ConnectError(Code.UNKNOWN, "error reading a body from connection: timed out") - assert is_transient_rpc_error(err) +def test_transient_stream_break_messages(): + # Both production stream-break variants: UNAVAILABLE "... timed out" and INTERNAL "Error + # reading content". Each classifies by code and, for robustness, by message alone. + assert is_transient_rpc_error( + ConnectError(Code.UNAVAILABLE, "error reading a body from connection: timed out") + ) + assert is_transient_rpc_error(ConnectError(Code.INTERNAL, "Error reading content")) + assert is_transient_rpc_error(ConnectError(Code.UNKNOWN, "Error reading content")) def test_command_timeout_is_transient(): From 389929c704643fa98b44b05651a3ddbeb310574f Mon Sep 17 00:00:00 2001 From: Andrew Kirillov Date: Tue, 25 Aug 2026 12:12:14 -0700 Subject: [PATCH 3/6] fix(sandboxes): make process recovery retry-safe --- .../src/prime_sandboxes/_reliability.py | 97 ---------- .../src/prime_sandboxes/process.py | 108 +++++------ .../prime_sandboxes/rpc_command_session.py | 59 +++++- .../src/prime_sandboxes/sandbox.py | 173 ++++++++++-------- .../tests/test_background_job_launch_retry.py | 31 ++-- .../tests/test_command_transport_selection.py | 83 +++++++++ .../tests/test_process_stream_reconnect.py | 88 ++++++++- .../prime-sandboxes/tests/test_reliability.py | 34 ---- 8 files changed, 373 insertions(+), 300 deletions(-) delete mode 100644 packages/prime-sandboxes/src/prime_sandboxes/_reliability.py delete mode 100644 packages/prime-sandboxes/tests/test_reliability.py diff --git a/packages/prime-sandboxes/src/prime_sandboxes/_reliability.py b/packages/prime-sandboxes/src/prime_sandboxes/_reliability.py deleted file mode 100644 index 48a5d79f7..000000000 --- a/packages/prime-sandboxes/src/prime_sandboxes/_reliability.py +++ /dev/null @@ -1,97 +0,0 @@ -"""Shared reliability policy for the sandbox Connect-RPC surface. - -One classifier and one set of tunables, used by every remote-control RPC (exec, background-job -launch, live-process stream + control) so a transient blip on the client <-> sandbox link is -retried/reconnected instead of killing the caller's work, while a permanent fault still fails fast. - -All timeouts and retry budgets are overridable via ``PRIME_SANDBOX_*`` env vars. -""" - -import os - -from connectrpc.code import Code -from connectrpc.errors import ConnectError - -from .core import APIError -from .exceptions import CommandTimeoutError - - -def _env_int(name: str, default: int) -> int: - try: - return int(os.environ[name]) - except (KeyError, ValueError): - return default - - -def _env_float(name: str, default: float) -> float: - try: - return float(os.environ[name]) - except (KeyError, ValueError): - return default - - -# Retry budget for the idempotent RPCs (background-job launch, live-process control). -RPC_MAX_ATTEMPTS = _env_int("PRIME_SANDBOX_RPC_MAX_ATTEMPTS", 3) -RPC_BACKOFF_BASE = _env_float("PRIME_SANDBOX_RPC_BACKOFF_BASE", 0.5) - -# Background-job launch is fire-and-forget, so its exec should return in well under a second; -# a timeout is a transport blip, not a slow command. -BG_LAUNCH_TIMEOUT = _env_int("PRIME_SANDBOX_BG_LAUNCH_TIMEOUT", 30) - -# Live-process control RPC deadlines (ms). -PROCESS_INPUT_TIMEOUT_MS = _env_int("PRIME_SANDBOX_PROCESS_INPUT_TIMEOUT_MS", 30_000) -PROCESS_SIGNAL_TIMEOUT_MS = _env_int("PRIME_SANDBOX_PROCESS_SIGNAL_TIMEOUT_MS", 10_000) - -# The live-process output stream can be re-attached to (Connect RPC) after a transient drop, since -# the process keeps running in the sandbox. Bound the reconnects so a genuinely dead process/sandbox -# still surfaces. -STREAM_MAX_RECONNECTS = _env_int("PRIME_SANDBOX_STREAM_MAX_RECONNECTS", 5) -STREAM_RECONNECT_BACKOFF_BASE = _env_float("PRIME_SANDBOX_STREAM_RECONNECT_BACKOFF_BASE", 0.5) - -# Live-process transport tuning: keep the long-lived stream connection warm so a brief idle stall -# does not get torn down (read_timeout None = no per-read deadline; the stream's own deadline and -# the server's keepalive events bound it). -STREAM_TCP_KEEPALIVE = _env_float("PRIME_SANDBOX_STREAM_TCP_KEEPALIVE", 15.0) -STREAM_POOL_IDLE_TIMEOUT = _env_float("PRIME_SANDBOX_STREAM_POOL_IDLE_TIMEOUT", 300.0) - -# Connect codes that mean "the link hiccuped", not "the request is wrong". UNAUTHENTICATED is -# excluded on purpose: it is handled by the token-refresh retry, not this backoff. INTERNAL is -# included because a broken output stream surfaces as INTERNAL "Error reading content" (observed -# in production), the same transient stream-break class as UNAVAILABLE "... timed out". -_TRANSIENT_CODES = frozenset( - {Code.DEADLINE_EXCEEDED, Code.UNAVAILABLE, Code.ABORTED, Code.INTERNAL} -) - -# Substrings of a transport error that mean the same, seen on both ConnectError and APIError. -_TRANSIENT_MARKERS = ( - "timed out", - "reading a body", - "reading content", - "connection reset", - "connection closed", - "broken pipe", - "unavailable", - "deadline_exceeded", -) - - -def is_transient_rpc_error(error: BaseException) -> bool: - """Whether ``error`` is a transient sandbox-transport fault safe to retry/reconnect. - - Transient: a stalled/reset Connect-RPC (DEADLINE_EXCEEDED / UNAVAILABLE / ABORTED, or a - ``reading a body ... timed out`` / connection-reset body error). Permanent (returns False): - a 404 sandbox-not-found, other 4xx, or any non-transport error — those must fail fast. - """ - if isinstance(error, CommandTimeoutError): - return True - if isinstance(error, ConnectError): - if error.code in _TRANSIENT_CODES: - return True - message = (error.message or "").lower() - return any(marker in message for marker in _TRANSIENT_MARKERS) - if isinstance(error, APIError): - message = str(error).lower() - if "not found" in message or "sandbox is no longer" in message: - return False - return any(marker in message for marker in _TRANSIENT_MARKERS) - return False diff --git a/packages/prime-sandboxes/src/prime_sandboxes/process.py b/packages/prime-sandboxes/src/prime_sandboxes/process.py index 6e2bc7174..f3187f2cf 100644 --- a/packages/prime-sandboxes/src/prime_sandboxes/process.py +++ b/packages/prime-sandboxes/src/prime_sandboxes/process.py @@ -4,18 +4,14 @@ import contextlib import logging from collections.abc import AsyncIterator, Awaitable, Callable -from typing import Literal, Optional +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 -from ._reliability import ( - STREAM_MAX_RECONNECTS, - STREAM_RECONNECT_BACKOFF_BASE, - is_transient_rpc_error, -) from .core import APIError from .rpc_command_session import parse_command_session_start_event @@ -23,11 +19,12 @@ _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]] -# Re-attach to the running process's output stream by pid (Connect RPC), returning a fresh stream. -_Reconnect = Callable[[int], Awaitable[AsyncIterator[Message]]] +_Reconnect = Callable[[int | None], AsyncIterator[Message]] class _AsyncProcessStream(AsyncIterator[bytes]): @@ -79,7 +76,7 @@ def __init__( write_stdin: _WriteStdin, send_signal: _SendSignal, transport: HTTPTransport | None = None, - reconnect: Optional[_Reconnect] = None, + reconnect: _Reconnect | None = None, ) -> None: self.stdout = _AsyncProcessStream() self.stderr = _AsyncProcessStream() @@ -112,11 +109,9 @@ async def _create( write_stdin: _WriteStdin, send_signal: _SendSignal, transport: HTTPTransport | None = None, - reconnect: Optional[_Reconnect] = None, + reconnect: _Reconnect | None = None, ) -> "AsyncSandboxProcess": - process = cls( - stream_client, stream, write_stdin, send_signal, transport, reconnect - ) + process = cls(stream_client, stream, write_stdin, send_signal, transport, reconnect) try: await asyncio.shield(process._started) except asyncio.CancelledError: @@ -238,24 +233,13 @@ async def _wait_for_exit_event(self) -> bool: return False return self._remote_exited - def _can_reconnect( - self, ended: bool, reconnects: int, error: BaseException - ) -> bool: - """Whether a dropped output stream can be re-attached to the running process. - - Only for a transient transport fault, once we have a pid (so Connect has a target) and - before the process has exited, within the reconnect budget. - """ - return ( - not ended - and self._reconnect is not None - and self._started.done() - and not self._started.cancelled() - and self._started.exception() is None - and not self._remote_exited - and reconnects < STREAM_MAX_RECONNECTS - and is_transient_rpc_error(error) - ) + 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) @@ -263,11 +247,31 @@ async def _aclose_stream(self) -> 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: - while True: + while not ended: + error: BaseException | None = None try: async for response in self._stream: event = parse_command_session_start_event(response) @@ -290,31 +294,20 @@ async def _pump(self) -> None: if not self._exit.done(): self._exit.set_result(value) break - break except asyncio.CancelledError: raise - except BaseException as error: - if not self._can_reconnect(ended, reconnects, error): - raise - reconnects += 1 - delay = STREAM_RECONNECT_BACKOFF_BASE * 2 ** (reconnects - 1) - logger.warning( - "live process %s stream dropped (%s); re-attaching %d/%d in %.1fs", - self.pid, - error, - reconnects, - STREAM_MAX_RECONNECTS, - delay, - ) - await self._aclose_stream() - await asyncio.sleep(delay) - # Re-attach to the SAME running process by pid and resume consuming its - # output. The gateway tails from now, so output emitted during the blackout - # is not replayed; in practice the process is idle mid-turn when the link - # stalls, so nothing is lost — and a resumed rollout beats a dead one. - self._stream = await self._reconnect(self.pid) - if not ended: - raise APIError("Process stream ended without an exit event") + except BaseException as stream_error: + error = stream_error + + if ended: + break + 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: @@ -329,10 +322,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 c128a24b5..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 @@ -90,6 +105,12 @@ def HasField(self, field_name: str) -> bool: ... _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 ) @@ -105,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( @@ -131,8 +155,7 @@ def HasField(self, field_name: str) -> bool: ... idempotency_level=IdempotencyLevel.UNKNOWN, ) -# Re-attach to an already-running session's output stream (server-streaming), selected by pid. -# Used to resume a live process after its Start stream drops on a transient link fault. +# 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", @@ -141,6 +164,14 @@ def HasField(self, field_name: str) -> bool: ... 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, @@ -148,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", @@ -157,15 +189,28 @@ 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_TYPE( + 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: return _COMMAND_SESSION_SEND_INPUT_REQUEST_FACTORY( session=_COMMAND_SESSION_SELECTOR_FACTORY(pid=pid), diff --git a/packages/prime-sandboxes/src/prime_sandboxes/sandbox.py b/packages/prime-sandboxes/src/prime_sandboxes/sandbox.py index 6380e4558..fcaec6983 100644 --- a/packages/prime-sandboxes/src/prime_sandboxes/sandbox.py +++ b/packages/prime-sandboxes/src/prime_sandboxes/sandbox.py @@ -3,7 +3,6 @@ import asyncio import functools import json -import logging import os import random import re @@ -86,27 +85,20 @@ SSHSession, validate_egress_lists, ) -from ._reliability import ( - BG_LAUNCH_TIMEOUT, - PROCESS_INPUT_TIMEOUT_MS, - PROCESS_SIGNAL_TIMEOUT_MS, - RPC_BACKOFF_BASE, - RPC_MAX_ATTEMPTS, - STREAM_POOL_IDLE_TIMEOUT, - STREAM_TCP_KEEPALIVE, - is_transient_rpc_error, -) 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 @@ -124,11 +116,14 @@ # 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 -# Control-RPC deadlines are env-tunable (see _reliability); aliased here for locality. -_PROCESS_INPUT_TIMEOUT_MS = PROCESS_INPUT_TIMEOUT_MS -_PROCESS_SIGNAL_TIMEOUT_MS = PROCESS_SIGNAL_TIMEOUT_MS - -logger = logging.getLogger(__name__) +_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) @@ -152,16 +147,10 @@ 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. - # - # Keep the long-lived output stream connection warm so a brief idle stall (the agent waiting - # between turns) is not torn down: frequent TCP keepalive probes detect/keep the path, and a - # generous pool-idle timeout avoids reaping the connection under the stream. There is no - # per-read deadline (read_timeout=None) — the stream's own deadline and the server's keepalive - # events bound it, and a genuine drop is recovered by re-attaching (see AsyncSandboxProcess). return HTTPTransport( tls_ca_cert=_ca_bundle(), - tcp_keepalive_interval=STREAM_TCP_KEEPALIVE, - pool_idle_timeout=STREAM_POOL_IDLE_TIMEOUT, + tcp_keepalive_interval=_LIVE_PROCESS_TCP_KEEPALIVE_SECONDS, + pool_idle_timeout=_LIVE_PROCESS_POOL_IDLE_TIMEOUT_SECONDS, ) @@ -1450,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: @@ -1474,19 +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 &" - # The launch returns immediately, so a timeout is a transport blip, not a slow command; - # retry it. Re-issuing the identical command is safe (same job_id/log files); a dead - # sandbox raises SandboxNotRunningError (not transient), which fails fast. - for attempt in range(RPC_MAX_ATTEMPTS): + # 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=BG_LAUNCH_TIMEOUT, user=user) + self.execute_command( + sandbox_id, + bg_cmd, + timeout=_BACKGROUND_JOB_LAUNCH_TIMEOUT_SECONDS, + user=user, + ) break - except Exception as error: - if attempt == RPC_MAX_ATTEMPTS - 1 or not is_transient_rpc_error(error): + except CommandTimeoutError: + if attempt == _BACKGROUND_JOB_LAUNCH_ATTEMPTS - 1: raise - time.sleep(RPC_BACKOFF_BASE * 2**attempt) + time.sleep(_BACKGROUND_JOB_LAUNCH_BACKOFF_SECONDS * 2**attempt) return BackgroundJob( job_id=job_id, @@ -2674,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): @@ -2700,11 +2698,13 @@ async def open_process( send_compression=None, http_client=http_client, ) + 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, @@ -2733,20 +2733,46 @@ async def send_signal(pid: int, signal: Literal["terminate", "kill"]) -> None: http_client=http_client, ) - async def reconnect(pid: int) -> AsyncIterator[Message]: - # Re-attach to the still-running process by pid after a transient stream drop. Fresh - # auth (the token may have rotated during the outage), same warm transport. - 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, http_client=http_client) - return 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 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, @@ -2766,14 +2792,8 @@ async def _execute_process_control_rpc( operation: str, http_client: Optional[HTTPClient] = None, ) -> None: - """Run one live-process control RPC with current sandbox auth. - - Idempotent (a signal or a bounded stdin write), so a transient transport fault - (DEADLINE_EXCEEDED / UNAVAILABLE / reset) is retried with backoff instead of killing the - rollout mid-turn; an expired token is refreshed. A permanent fault fails fast. - """ + """Run one live-process control RPC with current sandbox auth.""" reauthed = False - transient_attempts = 0 while True: auth = await self._auth_cache.get_or_refresh(sandbox_id) gateway_url = auth["gateway_url"].rstrip("/") @@ -2799,20 +2819,6 @@ async def _execute_process_control_rpc( ): reauthed = True continue - if ( - is_transient_rpc_error(error) - and transient_attempts < RPC_MAX_ATTEMPTS - 1 - ): - transient_attempts += 1 - logger.warning( - "process %s RPC transient failure (%s); retry %d/%d", - operation, - error.code.value, - transient_attempts, - RPC_MAX_ATTEMPTS - 1, - ) - await asyncio.sleep(RPC_BACKOFF_BASE * 2 ** (transient_attempts - 1)) - continue raise APIError( f"process {operation} RPC failed ({error.code.value}): {error.message}" ) from error @@ -3025,6 +3031,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: @@ -3049,19 +3056,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 &" - # The launch returns immediately, so a timeout is a transport blip, not a slow command; - # retry it. Re-issuing the identical command is safe (same job_id/log files); a dead - # sandbox raises SandboxNotRunningError (not transient), which fails fast. - for attempt in range(RPC_MAX_ATTEMPTS): + # 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=BG_LAUNCH_TIMEOUT, user=user) + await self.execute_command( + sandbox_id, + bg_cmd, + timeout=_BACKGROUND_JOB_LAUNCH_TIMEOUT_SECONDS, + user=user, + ) break - except Exception as error: - if attempt == RPC_MAX_ATTEMPTS - 1 or not is_transient_rpc_error(error): + except CommandTimeoutError: + if attempt == _BACKGROUND_JOB_LAUNCH_ATTEMPTS - 1: raise - await asyncio.sleep(RPC_BACKOFF_BASE * 2**attempt) + 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 index 3da1f765e..176467805 100644 --- a/packages/prime-sandboxes/tests/test_background_job_launch_retry.py +++ b/packages/prime-sandboxes/tests/test_background_job_launch_retry.py @@ -1,4 +1,4 @@ -"""start_background_job retries a transient timeout on the fire-and-forget launch (mode #1).""" +"""Background-job launch retries are guarded against duplicate execution.""" from typing import Any, cast @@ -8,7 +8,6 @@ from prime_sandboxes.exceptions import CommandTimeoutError from prime_sandboxes.models import CommandResponse from prime_sandboxes.sandbox import AsyncSandboxClient, SandboxClient -from prime_sandboxes._reliability import RPC_MAX_ATTEMPTS _OK = CommandResponse(stdout="", stderr="", exit_code=0) @@ -25,17 +24,19 @@ 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")) - calls = {"n": 0} + commands = [] - def execute(*_a, **_k): - calls["n"] += 1 - if calls["n"] < RPC_MAX_ATTEMPTS: + 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 calls["n"] == RPC_MAX_ATTEMPTS + 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): @@ -50,7 +51,7 @@ def execute(*_a, **_k): cast(Any, client).execute_command = execute with pytest.raises(CommandTimeoutError): client.start_background_job("sb", "rm -rf x") - assert calls["n"] == RPC_MAX_ATTEMPTS + assert calls["n"] == 3 class TestAsyncLaunchRetry: @@ -58,17 +59,19 @@ class TestAsyncLaunchRetry: async def test_retries_and_succeeds(self, monkeypatch): monkeypatch.setattr("prime_sandboxes.sandbox.asyncio.sleep", _no_sleep) client = AsyncSandboxClient(APIClient(api_key="test-key")) - calls = {"n": 0} + commands = [] - async def execute(*_a, **_k): - calls["n"] += 1 - if calls["n"] < RPC_MAX_ATTEMPTS: + 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 calls["n"] == RPC_MAX_ATTEMPTS + 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 @@ -84,4 +87,4 @@ async def execute(*_a, **_k): cast(Any, client).execute_command = execute with pytest.raises(CommandTimeoutError): await client.start_background_job("sb", "rm -rf x") - assert calls["n"] == RPC_MAX_ATTEMPTS + 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..4dc530f2b 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, http_client=None): + 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 index 2349fb149..d119466c0 100644 --- a/packages/prime-sandboxes/tests/test_process_stream_reconnect.py +++ b/packages/prime-sandboxes/tests/test_process_stream_reconnect.py @@ -1,4 +1,4 @@ -"""A live-process output stream re-attaches after a transient mid-stream drop (mode #3).""" +"""Live-process output streams re-attach to the running process after a drop.""" import pytest from connectrpc.code import Code @@ -43,8 +43,6 @@ async def _drain(stream): return out -# Both production stream-break variants must trigger a reconnect: UNAVAILABLE "... timed out" -# and INTERNAL "Error reading content". _STREAM_FAULTS = [ ConnectError(Code.UNAVAILABLE, "error reading a body from connection: timed out"), ConnectError(Code.INTERNAL, "Error reading content"), @@ -52,8 +50,12 @@ async def _drain(stream): @pytest.mark.asyncio -@pytest.mark.parametrize("fault", _STREAM_FAULTS, ids=["unavailable_timeout", "internal_reading_content"]) -async def test_stream_reconnects_and_resumes_after_transient_drop(fault): +@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") @@ -66,7 +68,7 @@ async def resumed(): reconnect_calls = [] - async def reconnect(pid): + def reconnect(pid): reconnect_calls.append(pid) return resumed() @@ -76,15 +78,83 @@ async def reconnect(pid): stdout = await _drain(proc.stdout) rc = await proc.wait() - assert reconnect_calls == [42] # re-attached to the same pid + 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_stream_without_reconnect_still_fails(): - # No reconnect callable -> a transient drop is fatal (baseline behaviour preserved). + # 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") @@ -105,7 +175,7 @@ async def faulty(): calls = [] - async def reconnect(pid): + def reconnect(pid): calls.append(pid) raise AssertionError("should not reconnect on a permanent fault") diff --git a/packages/prime-sandboxes/tests/test_reliability.py b/packages/prime-sandboxes/tests/test_reliability.py deleted file mode 100644 index 6e3a4724d..000000000 --- a/packages/prime-sandboxes/tests/test_reliability.py +++ /dev/null @@ -1,34 +0,0 @@ -"""Transient-fault classification used across the sandbox RPC surface.""" - -from connectrpc.code import Code -from connectrpc.errors import ConnectError - -from prime_sandboxes.core import APIError -from prime_sandboxes.exceptions import CommandTimeoutError -from prime_sandboxes._reliability import is_transient_rpc_error - - -def test_transient_connect_codes(): - for code in (Code.DEADLINE_EXCEEDED, Code.UNAVAILABLE, Code.ABORTED, Code.INTERNAL): - assert is_transient_rpc_error(ConnectError(code, "boom")) - - -def test_transient_stream_break_messages(): - # Both production stream-break variants: UNAVAILABLE "... timed out" and INTERNAL "Error - # reading content". Each classifies by code and, for robustness, by message alone. - assert is_transient_rpc_error( - ConnectError(Code.UNAVAILABLE, "error reading a body from connection: timed out") - ) - assert is_transient_rpc_error(ConnectError(Code.INTERNAL, "Error reading content")) - assert is_transient_rpc_error(ConnectError(Code.UNKNOWN, "Error reading content")) - - -def test_command_timeout_is_transient(): - assert is_transient_rpc_error(CommandTimeoutError("sb", "cmd", 30)) - - -def test_permanent_faults_not_transient(): - assert not is_transient_rpc_error(ConnectError(Code.NOT_FOUND, "no such sandbox")) - assert not is_transient_rpc_error(APIError("HTTP 404: Sandbox not found")) - assert not is_transient_rpc_error(APIError("Sandbox is no longer present")) - assert not is_transient_rpc_error(ValueError("bad arg")) From 61422f3c090a1b8ffd283940e45e2cd4ca0d1e3a Mon Sep 17 00:00:00 2001 From: Andrew Kirillov Date: Tue, 25 Aug 2026 12:22:34 -0700 Subject: [PATCH 4/6] docs(sandboxes): mark tag recovery as temporary --- packages/prime-sandboxes/src/prime_sandboxes/sandbox.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/packages/prime-sandboxes/src/prime_sandboxes/sandbox.py b/packages/prime-sandboxes/src/prime_sandboxes/sandbox.py index fcaec6983..9f29ef49d 100644 --- a/packages/prime-sandboxes/src/prime_sandboxes/sandbox.py +++ b/packages/prime-sandboxes/src/prime_sandboxes/sandbox.py @@ -2698,6 +2698,8 @@ 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, From 3ef18169d23db5bbf34a1181d322fd027d29d4e6 Mon Sep 17 00:00:00 2001 From: Andrew Kirillov Date: Tue, 25 Aug 2026 12:41:59 -0700 Subject: [PATCH 5/6] test(sandboxes): accept Connect client options --- .../prime-sandboxes/tests/test_command_transport_selection.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/packages/prime-sandboxes/tests/test_command_transport_selection.py b/packages/prime-sandboxes/tests/test_command_transport_selection.py index 4dc530f2b..79182de52 100644 --- a/packages/prime-sandboxes/tests/test_command_transport_selection.py +++ b/packages/prime-sandboxes/tests/test_command_transport_selection.py @@ -275,7 +275,7 @@ async def is_vm(self, _sandbox_id: str): return True class _FakeConnectClient: - def __init__(self, _address: str, http_client=None): + def __init__(self, _address: str, **_kwargs): pass def execute_server_stream(self, **kwargs): From ffcc86e8886202634013d3b7cc6f2620d90674ab Mon Sep 17 00:00:00 2001 From: Andrew Kirillov Date: Tue, 25 Aug 2026 13:05:42 -0700 Subject: [PATCH 6/6] fix(sandboxes): propagate pre-PID process end --- .../src/prime_sandboxes/process.py | 4 +++- .../tests/test_process_stream_reconnect.py | 19 +++++++++++++++++++ 2 files changed, 22 insertions(+), 1 deletion(-) diff --git a/packages/prime-sandboxes/src/prime_sandboxes/process.py b/packages/prime-sandboxes/src/prime_sandboxes/process.py index f3187f2cf..78357e883 100644 --- a/packages/prime-sandboxes/src/prime_sandboxes/process.py +++ b/packages/prime-sandboxes/src/prime_sandboxes/process.py @@ -290,7 +290,7 @@ async def _pump(self) -> None: ended = True self._remote_exited = True if not self._started.done(): - raise APIError("Process exited before reporting its PID") + raise APIError("Process ended before reporting its PID") if not self._exit.done(): self._exit.set_result(value) break @@ -300,6 +300,8 @@ async def _pump(self) -> None: error = stream_error if ended: + if error is not None: + raise error break if not self._can_reconnect(reconnects, error): if error is not None: diff --git a/packages/prime-sandboxes/tests/test_process_stream_reconnect.py b/packages/prime-sandboxes/tests/test_process_stream_reconnect.py index d119466c0..0d759ca0b 100644 --- a/packages/prime-sandboxes/tests/test_process_stream_reconnect.py +++ b/packages/prime-sandboxes/tests/test_process_stream_reconnect.py @@ -1,5 +1,7 @@ """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 @@ -152,6 +154,23 @@ async def stream(): 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.