feat(sandbox): detached PTY exec — run long commands without holding a connection

nemo_gym/sandbox/api.py

@@ -195,6 +195,8 @@ async def exec(
         rows: int = 24,
         cols: int = 80,
         pty: bool = True,
+        detach: bool = False,
+        poll_interval_s: float = 15.0,
     ) -> SandboxExecResult:
         """Run one command in a terminal session and collect its output.
 
@@ -214,13 +216,42 @@ async def exec(
         only), and a command that ends the shell (``exit``) raises
         ``SandboxPtyError``.
 
+        With ``detach=True`` the command runs without holding a connection
+        while it works: it starts in a session, the socket is dropped, and the
+        session is briefly re-attached every ``poll_interval_s`` to drain
+        output and check for completion, so a long command occupies a
+        connection for milliseconds per poll instead of its whole runtime
+        (completion latency is bounded by ``poll_interval_s``). Nothing is
+        written to the sandbox filesystem; output rides the server's retained
+        window (~1 MiB) between polls, comes back as one merged stream
+        (``stderr`` is ``None``), and exceeding the window raises rather than
+        returning truncated output — run bulk-output commands attached or via
+        the exec API instead. A detached exec never reuses the default-shell
+        session implicitly: without ``session`` it opens a private one. With
+        ``session``, the session is detached while the command works, must
+        not be used concurrently, and is attached and reusable again when
+        this returns.
+
         PTY mode returns all output on ``stdout`` and ``None`` stderr; pipe mode
         splits the two. A command that outlives ``timeout_s`` returns
         ``error_type="timeout"`` like ``sandbox.exec()`` rather than raising;
         in an explicitly passed session that command keeps running and leaves
         unread output behind, so discard the session rather than reusing it (an
         implicitly reused session is retired automatically).
         """
+        if detach:
+            return await self._exec_detached(
+                command,
+                session=session,
+                cwd=cwd,
+                env=env,
+                user=user,
+                rows=rows,
+                cols=cols,
+                pty=pty,
+                timeout_s=timeout_s,
+                poll_interval_s=poll_interval_s,
+            )
         implicit = False
         if session is None and cwd is None and env is None and user is None and pty and (rows, cols) == (24, 80):
             if self._default_session is not None and self._default_session.closed:
@@ -274,6 +305,59 @@ async def drain(read: Callable[[], Awaitable[bytes]]) -> bytes:
             return_code=return_code,
         )
 
+    async def _exec_detached(
+        self,
+        command: str,
+        *,
+        session: SandboxPtySession | None,
+        cwd: str | None,
+        env: dict[str, str] | None,
+        user: str | int | None,
+        rows: int,
+        cols: int,
+        pty: bool,
+        timeout_s: int | float | None,
+        poll_interval_s: float,
+    ) -> SandboxExecResult:
+        """``exec(detach=True)``: hand the command to the session's detached
+        runner, which holds the socket only for brief completion polls."""
+        private = session is None
+        if private:
+            session = await self.create(cwd=cwd, env=env, user=user, rows=rows, cols=cols, pty=pty)
+            if self._default_session is session:
+                # Private to this call: an implicit exec() grabbing it would
+                # collide with the detach cycle.
+                self._default_session = None
+        elif cwd is not None or env is not None or user is not None:
+            raise ValueError(
+                "cwd/env/user apply only when a detached exec opens its own session; "
+                "for an existing session they are fixed at pty.create() time"
+            )
+        if not hasattr(session, "run_detached"):
+            raise NotImplementedError(f"{type(session).__name__} does not support detached execution")
+        try:
+            async with asyncio.timeout(timeout_s):
+                # Same serialization as attached session execs: one command per
+                # sandbox at a time, for the command's whole duration.
+                async with self._session_exec_lock:
+                    output, exit_code = await session.run_detached(command, poll_interval_s=poll_interval_s)
+        except (TimeoutError, asyncio.TimeoutError):
+            return _pty_timeout_result(command, timeout_s, reusable=False)
+        finally:
+            if private:
+                await session.close()
+            else:
+                try:
+                    await session.reattach()  # no-op unless a timeout left it detached
+                except Exception:
+                    pass
+        return SandboxExecResult(
+            stdout=output.decode(errors="replace"),
+            stderr=None,
+            return_code=exit_code if exit_code is not None else SANDBOX_PTY_RUNTIME_RETURN_CODE,
+            error_type=None if exit_code is not None else "pty",
+        )
+
 
 class AsyncSandbox:
     """Async sandbox object backed by a runtime provider."""

nemo_gym/sandbox/providers/base.py

@@ -311,6 +311,13 @@ async def wait_exit(self, *, timeout_s: float | None = None) -> int:
         """Block until the process exits and return its exit code."""
         ...
 
+    async def run_detached(self, command: str, *, poll_interval_s: float = 15.0) -> tuple[bytes, int | None]:
+        """Run one command holding the transport only for brief completion
+        polls; returns ``(merged output, exit code or None)``. The server
+        retains a bounded window of output between polls, and exceeding it
+        raises rather than returning truncated output."""
+        ...
+
     async def close(self) -> None:
         """Idempotent: release local resources; a session this client created
         is also ended, while an attached one is merely detached and lives on

nemo_gym/sandbox/providers/opensandbox/pty.py

@@ -30,6 +30,7 @@
 import json
 import logging
 import shlex
+import uuid
 from collections.abc import AsyncIterator
 from typing import Any
 from urllib.parse import urlencode
@@ -110,14 +111,20 @@ def __init__(
         self._connected: asyncio.Future[None] = asyncio.get_running_loop().create_future()
         self._error: SandboxPtyError | None = None
         self._closed = False
+        self._detached = False
         self._received = 0  # bytes of the session's retained stream seen so far
+        self._replay_gap = 0  # bytes the server evicted before we could replay them
         self._pump_task = asyncio.create_task(self._pump())
 
     @property
     def closed(self) -> bool:
         """True once the session can no longer run commands: after ``close()``,
         or once the connection pump has ended (process exit, takeover eviction,
-        or connection loss). Resources are released by ``close()``."""
+        or connection loss). A ``detach()``-ed session is not closed: the
+        server side keeps running and ``reattach()`` restores it. Resources
+        are released by ``close()``."""
+        if self._detached:
+            return self._closed
         return self._closed or self._pump_task.done()
 
     async def _pump_socket(self) -> None:
@@ -137,9 +144,11 @@ async def _pump_socket(self) -> None:
                     await self._stderr.put(data[1:])
                 elif channel == CHAN_REPLAY and len(data) > REPLAY_HEADER_BYTES:
                     # Replay is one merged stream regardless of mode. The
-                    # server clamps a `since` older than its 1 MiB buffer,
-                    # so a higher offset here means output was evicted.
+                    # server clamps a `since` older than its retained window,
+                    # so a frame starting past what we saw means eviction.
                     self.replay_offset = int.from_bytes(data[1:REPLAY_HEADER_BYTES], "big")
+                    if self.replay_offset > self._received:
+                        self._replay_gap += self.replay_offset - self._received
                     self._received = self.replay_offset + len(data) - REPLAY_HEADER_BYTES
                     await self._output.put(data[REPLAY_HEADER_BYTES:])
             elif message.type == aiohttp.WSMsgType.TEXT:
@@ -203,16 +212,19 @@ async def _pump(self) -> None:
                 if barren >= 3 or not await self._reattach_socket():
                     break
         finally:
-            if not self._exit.done():
-                self._exit.set_exception(self._close_error())
-                self._exit.exception()  # retrieved; silences never-retrieved warnings
-            # A pump that ends before `connected` arrived means the session
-            # never became usable; fail the waiter.
-            if not self._connected.done():
-                self._connected.set_exception(self._close_error())
-                self._connected.exception()  # retrieved; silences never-retrieved warnings
-            await self._output.put(None)
-            await self._stderr.put(None)
+            # A detach ends the pump without ending the session: skip the
+            # finalization so reads and the exit future survive reattach().
+            if not self._detached:
+                if not self._exit.done():
+                    self._exit.set_exception(self._close_error())
+                    self._exit.exception()  # retrieved; silences never-retrieved warnings
+                # A pump that ends before `connected` arrived means the session
+                # never became usable; fail the waiter.
+                if not self._connected.done():
+                    self._connected.set_exception(self._close_error())
+                    self._connected.exception()  # retrieved; silences never-retrieved warnings
+                await self._output.put(None)
+                await self._stderr.put(None)
 
     def _close_error(self) -> SandboxPtyError:
         if self._error is not None:
@@ -229,6 +241,8 @@ async def _wait_connected(self, timeout_s: float | None) -> None:
         await asyncio.wait_for(asyncio.shield(self._connected), timeout=timeout_s)
 
     async def _read_stream(self, queue: asyncio.Queue[bytes | None], timeout_s: float | None) -> bytes:
+        if self._detached and not self._closed:
+            raise SandboxPtyError("PTY session is detached; reattach() first")
         chunk = await asyncio.wait_for(queue.get(), timeout=timeout_s)
         if chunk is None:
             # Keep the EOF observable by subsequent reads and iterators.
@@ -252,6 +266,8 @@ async def _iterate() -> AsyncIterator[bytes]:
         return _iterate()
 
     async def _send(self, frame: bytes | str) -> None:
+        if self._detached:
+            raise SandboxPtyError("PTY session is detached; reattach() first")
         if self._closed or self._ws.closed:
             raise SandboxPtyError("PTY session is closed")
         try:
@@ -277,6 +293,108 @@ async def wait_exit(self, *, timeout_s: float | None = None) -> int:
         # shield: the future is shared; a timed-out waiter must not cancel it.
         return await asyncio.wait_for(asyncio.shield(self._exit), timeout=timeout_s)
 
+    async def detach(self) -> None:
+        """Drop the WebSocket while the server-side session keeps running.
+
+        Output produced while detached lands in execd's replay buffer (a 1 MiB
+        ring; older bytes are evicted), and ``reattach()`` resumes from the
+        last byte this object saw. A detached session refuses reads and
+        writes; ``close()`` still releases it (and ends it when owned).
+        """
+        if self._detached:
+            return
+        if self.closed:
+            # Covers close() and a pump that already ended (process exit,
+            # takeover, connection loss): a dead session must not be
+            # resurrected into a not-closed, prune-evading detached state.
+            raise SandboxPtyError("PTY session is closed")
+        self._detached = True
+        self._pump_task.cancel()
+        # Let the pump observe the detach before the socket goes away.
+        await asyncio.gather(self._pump_task, return_exceptions=True)
+        await self._ws.close()
+
+    async def reattach(self) -> None:
+        """Re-dial a ``detach()``-ed session, replaying output produced since."""
+        if self._closed:
+            raise SandboxPtyError("PTY session is closed")
+        if not self._detached:
+            return
+        self._ws = await _connect_ws(
+            client=self._client,
+            base_url=self._session_url.rsplit("/pty/", 1)[0],
+            headers=self._headers,
+            session_id=self.session_id,
+            query={"takeover": "1", "since": str(self._received)},
+            request_timeout_s=self._request_timeout_s,
+        )
+        self._detached = False
+        self._pump_task = asyncio.create_task(self._pump())
+
+    async def run_detached(self, command: str, *, poll_interval_s: float = 15.0) -> tuple[bytes, int | None]:
+        """Run one command holding the socket only for brief polls.
+
+        The command is written with the same marker discipline as session
+        exec, the socket is dropped, and every ``poll_interval_s`` the session
+        re-attaches and drains output produced in the meantime from the
+        server's retained window. Nothing is written to the sandbox
+        filesystem; if the command produced more output between polls than
+        the server retains, the loss is detected and raised rather than
+        returned truncated. Returns ``(output, exit_code)`` — output is one
+        merged stream (replayed bytes carry no stdout/stderr split) and
+        ``exit_code`` is ``None`` when the marker line came back mangled.
+        The session ends attached. Callers serialize: one command per
+        session at a time, as with session exec.
+        """
+        token = f"NGPTY{uuid.uuid4().hex[:12]}"
+        needle = f"{token}:".encode()
+        # Marker from two literals so the echo cannot match it; brace group
+        # keeps shell state while putting stdin at EOF (see _run_in_pty_session
+        # in the api module for the same discipline).
+        await self.write(
+            f"{{ {command}\n}} </dev/null\nprintf '%s%s:%s\\n' '{token[:5]}' '{token[5:]}' \"$?\"\n".encode()
+        )
+        buffer = bytearray()
+        while True:
+            while needle not in buffer:
+                try:
+                    chunk = await self.read(timeout_s=1.0)
+                except (TimeoutError, asyncio.TimeoutError):
+                    break  # stream is quiet; wait detached
+                if not chunk:
+                    raise SandboxPtyError("PTY session ended before the command finished")
+                buffer.extend(chunk)
+            # Checked after draining (replay frames land asynchronously), and
+            # before accepting the marker: a mid-stream hole must not come
+            # back as silently truncated output.
+            if self._replay_gap:
+                raise SandboxPtyError(
+                    "PTY output exceeded the server's retained window while detached; "
+                    "run bulk-output commands attached or through the exec API instead"
+                )
+            if needle in buffer:
+                break
+            await self.detach()
+            await asyncio.sleep(poll_interval_s)
+            await self.reattach()
+        output, _, trailing = bytes(buffer).partition(needle)
+        while b"\n" not in trailing:
+            # The status digits can straddle the chunk that carried the marker.
+            chunk = await self.read(timeout_s=5.0)
+            if not chunk:
+                break
+            trailing += chunk
+        exit_text = trailing.split(b"\n", 1)[0].strip()
+        # Pipe mode splits live (attached) stderr onto its own queue; fold any
+        # of it into the merged result, ordering best-effort.
+        stderr = bytearray()
+        try:
+            while chunk := await self.read_stderr(timeout_s=0.05):
+                stderr.extend(chunk)
+        except (TimeoutError, asyncio.TimeoutError):
+            pass
+        return bytes(output + stderr), int(exit_text) if exit_text.isdigit() else None
+
     async def close(self) -> None:
         if self._closed:
             return
@@ -288,6 +406,10 @@ async def close(self) -> None:
         self._pump_task.cancel()
         # Let the pump's finally run before tearing the socket down.
         await asyncio.gather(self._pump_task, return_exceptions=True)
+        if self._detached:
+            # The detach-time pump skipped finalization; readers still need EOF.
+            await self._output.put(None)
+            await self._stderr.put(None)
         try:
             await self._ws.close()
             if self._owned: