fix(sandboxes): make live-process RPCs retry-safe via sandboxd idempotency keys

packages/prime-sandboxes/src/prime_sandboxes/_connectrpc.py

@@ -25,3 +25,20 @@ def _reject_legacy_connect_python(legacy_version: str | None) -> None:
 from connectrpc.compat import google_protobuf_binary_codec  # noqa: E402
 
 GOOGLE_PROTOBUF_BINARY_CODEC = google_protobuf_binary_codec()
+
+# Response-body read faults raised inside pyqwest's Rust core (string set
+# verified against the installed pyqwest 0.8.0 extension binary). They reach
+# fault classifiers inside ConnectError INTERNAL messages minted upstream
+# (e.g. the gateway); a client-local pyqwest fault on a unary RPC surfaces as
+# UNAVAILABLE, already retried by code. Classifiers match the markers
+# case-insensitively as substrings to tell this transport-level class apart
+# from a real server INTERNAL. "read cancelled" subsumes the longer sibling
+# under substring matching; both stay listed as the verified source strings.
+PYQWEST_BODY_READ_ERROR_MARKERS = frozenset(
+    {
+        "error reading content",
+        "error reading full content",
+        "response body read cancelled",
+        "read cancelled",
+    }
+)

packages/prime-sandboxes/src/prime_sandboxes/process.py

@@ -7,13 +7,12 @@
 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 .core import APIError
-from .rpc_command_session import parse_command_session_start_event
+from .rpc_command_session import is_recoverable_stream_fault, parse_command_session_start_event
 
 logger = logging.getLogger(__name__)
 
@@ -22,9 +21,11 @@
 _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]]
+_WriteStdin = Callable[[bytes], Awaitable[None]]
+_SendSignal = Callable[[Literal["terminate", "kill"]], Awaitable[None]]
+# The argument is whether a StartEvent has been observed yet; the callee picks
+# between retrying Start (create-or-attach) and Connect-ing to the session.
+_Reconnect = Callable[[bool], AsyncIterator[Message]]
 
 
 class _AsyncProcessStream(AsyncIterator[bytes]):
@@ -148,7 +149,7 @@ async def write_stdin(self, data: bytes) -> None:
             return
         if self._closed or self._remote_exited:
             raise BrokenPipeError("process has exited")
-        await self._write_stdin(self.pid, data)
+        await self._write_stdin(data)
 
     async def wait(self) -> int:
         """Wait for the process to exit and return its exit code."""
@@ -165,7 +166,7 @@ async def kill(self) -> None:
     async def _send_signal(self, signal: Literal["terminate", "kill"]) -> None:
         if self._closed or self._remote_exited:
             return
-        await self._send_process_signal(self.pid, signal)
+        await self._send_process_signal(signal)
         self._signals_sent.add(signal)
 
     async def aclose(self) -> None:
@@ -234,12 +235,10 @@ async def _wait_for_exit_event(self) -> bool:
         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."""
+        """Whether the stream fault is recoverable and reconnect budget remains."""
         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
+        return is_recoverable_stream_fault(error)
 
     async def _aclose_stream(self) -> None:
         close = getattr(self._stream, "aclose", None)
@@ -261,10 +260,11 @@ async def _reconnect_stream(self, reconnects: int, error: BaseException | None)
         )
         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)
+        # Re-attachment — a retried Start before the pid was seen, Connect after —
+        # never replays output emitted while detached. Both arms re-announce the
+        # StartEvent and replay the retained EndEvent of a session that exited
+        # within sandboxd's retention window, so a missed exit is still observed.
+        self._stream = reconnect(self._started.done())
 
     async def _pump(self) -> None:
         ended = False
@@ -279,7 +279,7 @@ async def _pump(self) -> None:
                             continue
                         kind, value = event
                         if kind == "start":
-                            # A reconnected (Connect) stream re-announces the pid; keep the first.
+                            # A re-attached stream re-announces the pid; keep the first.
                             if not self._started.done():
                                 self._started.set_result(value)
                         elif kind == "stdout":

packages/prime-sandboxes/src/prime_sandboxes/rpc_command_session.py

@@ -2,9 +2,12 @@
 
 from typing import Dict, List, Literal, Optional, Protocol, Sequence, cast
 
+from connectrpc.code import Code
+from connectrpc.errors import ConnectError
 from connectrpc.method import IdempotencyLevel, MethodInfo
 from google.protobuf.message import Message
 
+from ._connectrpc import PYQWEST_BODY_READ_ERROR_MARKERS
 from ._proto.command_session import command_session_pb2
 
 
@@ -18,24 +21,24 @@ def __call__(self, *, cmd: str, args: List[str], envs: Dict[str, str]) -> _Comma
 
 class _CommandSessionStartRequestFactory(Protocol):
     def __call__(
-        self, *, command: _CommandSpecLike, stdin: bool, tag: str | None = None
+        self, *, command: _CommandSpecLike, stdin: bool, session_uuid: str | None = None
     ) -> Message: ...
 
 
 class _CommandSessionSelectorFactory(Protocol):
-    def __call__(self, *, pid: int) -> Message: ...
+    def __call__(self, *, session_uuid: str) -> Message: ...
 
 
 class _CommandInputFactory(Protocol):
     def __call__(self, *, stdin: bytes) -> Message: ...
 
 
 class _CommandSessionSendInputRequestFactory(Protocol):
-    def __call__(self, *, session: Message, input: Message) -> Message: ...
+    def __call__(self, *, session: Message, input: Message, input_uuid: str) -> Message: ...
 
 
 class _CommandSessionSendSignalRequestFactory(Protocol):
-    def __call__(self, *, session: Message, signal: int) -> Message: ...
+    def __call__(self, *, session: Message, signal: int, signal_uuid: str) -> Message: ...
 
 
 class _CommandSessionConnectRequestFactory(Protocol):
@@ -44,7 +47,8 @@ def __call__(self, *, session: Message) -> Message: ...
 
 class _CommandSessionInfoLike(Protocol):
     pid: int
-    tag: str
+    session_uuid: str
+    command: _CommandSpecLike
 
 
 class _CommandSessionListResponseLike(Protocol):
@@ -164,6 +168,8 @@ def HasField(self, field_name: str) -> bool: ...
     idempotency_level=IdempotencyLevel.NO_SIDE_EFFECTS,
 )
 
+# Live-process introspection: pid, session_uuid, and command for each running
+# process. Permanent public API; exited sessions are not listed.
 COMMAND_SESSION_LIST_RPC_METHOD = MethodInfo(
     name="List",
     service_name="command_session.CommandSession",
@@ -173,13 +179,48 @@ def HasField(self, field_name: str) -> bool: ...
 )
 
 
+# The two fault predicates below classify command-session RPC failures for
+# retry, with deliberately opposite polarity. Stream re-attach (Connect, or a
+# create-or-attach Start resending the identical request) is idempotent, so
+# is_recoverable_stream_fault is a deny-list: retry everything except the codes
+# command_session.proto promises as definitive answers. A unary control RPC's
+# unknown fault may itself be a definitive answer, so is_transient_control_fault
+# is an allow-list: fail fast on everything except known link faults.
+
+# Stream faults recovery cannot fix, per command_session.proto's code promises:
+# NOT_FOUND (the session is gone or its retention expired) and
+# FAILED_PRECONDITION (a Start reusing the session_uuid with a different spec —
+# a guard for a future non-identical retry; today's reconnect resends the
+# identical request, so the server cannot answer it with a spec conflict).
+_STREAM_FATAL_CODES = frozenset({Code.NOT_FOUND, Code.FAILED_PRECONDITION})
+
+# Link faults a unary control RPC may retry; pyqwest body-read faults surface
+# as ConnectError INTERNAL and are matched by message marker instead.
+_TRANSIENT_CONTROL_CODES = frozenset({Code.DEADLINE_EXCEEDED, Code.UNAVAILABLE})
+
+
+def is_recoverable_stream_fault(error: BaseException | None) -> bool:
+    """Whether a dropped command-session stream may be re-attached (None: clean EOF)."""
+    return not (isinstance(error, ConnectError) and error.code in _STREAM_FATAL_CODES)
+
+
+def is_transient_control_fault(error: ConnectError) -> bool:
+    """Whether a unary control-RPC fault is a link hiccup rather than a definitive answer."""
+    if error.code in _TRANSIENT_CONTROL_CODES:
+        return True
+    message = (error.message or "").lower()
+    return error.code == Code.INTERNAL and any(
+        marker in message for marker in PYQWEST_BODY_READ_ERROR_MARKERS
+    )
+
+
 def build_command_session_start_request(
+    *,
     command: str,
     working_dir: Optional[str],
     env: Optional[Dict[str, str]],
-    *,
     stdin: bool = False,
-    tag: str | None = None,
+    session_uuid: str | None = None,
 ) -> Message:
     command_spec = _COMMAND_SPEC_FACTORY(
         cmd="/bin/bash",
@@ -192,42 +233,41 @@ def build_command_session_start_request(
     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)
+        session_uuid=session_uuid,
     )
 
 
 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_connect_request(*, session_uuid: str) -> Message:
+    return _COMMAND_SESSION_CONNECT_REQUEST_FACTORY(
+        session=_COMMAND_SESSION_SELECTOR_FACTORY(session_uuid=session_uuid)
+    )
 
 
-def build_command_session_send_input_request(pid: int, data: bytes) -> Message:
+def build_command_session_send_input_request(
+    *, session_uuid: str, data: bytes, input_uuid: str
+) -> Message:
     return _COMMAND_SESSION_SEND_INPUT_REQUEST_FACTORY(
-        session=_COMMAND_SESSION_SELECTOR_FACTORY(pid=pid),
+        session=_COMMAND_SESSION_SELECTOR_FACTORY(session_uuid=session_uuid),
         input=_COMMAND_INPUT_FACTORY(stdin=data),
+        input_uuid=input_uuid,
     )
 
 
 def build_command_session_send_signal_request(
-    pid: int, signal: Literal["terminate", "kill"]
+    *, session_uuid: str, signal: Literal["terminate", "kill"], signal_uuid: str
 ) -> Message:
     signal_value = getattr(
         command_session_pb2,
         "SIGNAL_SIGTERM" if signal == "terminate" else "SIGNAL_SIGKILL",
     )
     return _COMMAND_SESSION_SEND_SIGNAL_REQUEST_FACTORY(
-        session=_COMMAND_SESSION_SELECTOR_FACTORY(pid=pid),
+        session=_COMMAND_SESSION_SELECTOR_FACTORY(session_uuid=session_uuid),
         signal=signal_value,
+        signal_uuid=signal_uuid,
     )