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

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

@@ -22,9 +22,15 @@
 _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]]
+
+# Stream faults that recovery cannot fix: the session is gone (NOT_FOUND) or the
+# retried Start conflicts with the session's original spec (FAILED_PRECONDITION).
+_STREAM_FATAL_CODES = frozenset({Code.NOT_FOUND, Code.FAILED_PRECONDITION})
 
 
 class _AsyncProcessStream(AsyncIterator[bytes]):
@@ -148,7 +154,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 +171,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 +240,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 not (isinstance(error, ConnectError) and error.code in _STREAM_FATAL_CODES)
 
     async def _aclose_stream(self) -> None:
         close = getattr(self._stream, "aclose", None)
@@ -261,10 +265,9 @@ 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 never replays output emitted while detached; only the
+        # StartEvent and a retained EndEvent are, so a missed exit is still observed.
+        self._stream = reconnect(self._started.done())
 
     async def _pump(self) -> None:
         ended = False
@@ -279,7 +282,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":