snowflake-sandbox-python 0.2.1a1__py3-none-any.whl

snowflake/sandbox/shell.py

@@ -0,0 +1,556 @@
+"""Interactive shell sessions on a sandbox.
+
+A terminal is duplex and the REST path in front of the sandbox is not: the
+gateway proxies with a synchronous send, so it never upgrades to a WebSocket.
+The sandbox therefore exposes one pty across two ordinary requests -- an SSE
+response carrying output, POSTs carrying input -- and this module drives that
+pair.
+
+**This transport is a deliberate stopgap and is meant to be deleted.** It exists
+only because the hop in front of the sandbox is half-duplex -- not because
+splitting a terminal across two requests is the right shape for one. When a
+genuine end-to-end duplex stream reaches the sandbox from an external client (a
+bidirectional session on the data plane rather than a request/response gateway),
+this module and the two server-side halves behind it should be removed rather than
+maintained alongside the real thing.
+
+The *public surface* is the part worth preserving across that change: `shell()`,
+`send`, `output`, `resize` and `attach` describe a terminal, not a transport, so a
+duplex replacement can keep them and drop everything below -- the cursor
+bookkeeping, the reconnect loop, and the base64 framing all exist to work around
+the half-duplex hop and have no reason to outlive it.
+
+`Shell` is the programmatic surface (send bytes, read output, resize). `attach()`
+is the thing a CLI wants: it puts the local terminal in raw mode, forwards
+keystrokes, and forwards window-size changes, so `python -m ...` becomes an
+actual terminal into the sandbox.
+
+**Resumption is not optional here.** The gateway bounds a streaming response at
+its exec-stream timeout (900s), so any shell used for real work will have its
+output stream cut while the shell itself is still alive and healthy. Every frame
+carries an id; on a cut we re-open from the last id seen, which the sandbox
+answers exactly (no replay, no loss) from a bounded backlog. When the backlog has
+already evicted our cursor the sandbox says so with a `gap` frame rather than
+resuming mid-escape-sequence, which would render as a corrupt screen.
+
+**Sync and async share one protocol.** Everything that is *not* I/O -- id
+validation, the request path, the resume cursor, base64 framing in both
+directions, and exit/gap frame interpretation -- lives on `_ShellCore`. The async
+`Shell` here and the blocking `SyncShell` in `sync_shell` each add only the
+transport calls (awaited on one side, blocking on the other), so the wire
+contract cannot drift between them.
+"""
+
+from __future__ import annotations
+
+import asyncio
+import base64
+import binascii
+import contextlib
+import os
+import re
+import signal
+import sys
+from collections.abc import AsyncIterator, Callable
+from typing import TYPE_CHECKING
+
+from snowflake.sandbox._ansi import _ClipboardStripper
+from snowflake.sandbox.exceptions import SandboxError, SandboxTransportError
+
+if TYPE_CHECKING:
+    from snowflake.sandbox._transport import SSEEvent, Transport
+
+__all__ = ["Shell", "ShellClosed", "open_shell"]
+
+# Frames the sandbox sends. `out` is base64 pty bytes, `control` is JSON in the
+# same shape the WebSocket transport uses, `gap` reports evicted output.
+_EVENT_OUT = "out"
+_EVENT_CONTROL = "control"
+_EVENT_GAP = "gap"
+
+# Reconnect pacing after the stream is cut. Deliberately short: the common cause
+# is the gateway's timeout on a *healthy* session, so the user is sitting at a
+# live prompt waiting for output to resume.
+_RECONNECT_DELAYS = (0.1, 0.25, 0.5, 1.0, 2.0)
+
+# Shape an id must have before it is interpolated into a request path. Applied to
+# BOTH the session id and the sandbox id: they come from the same place (a server
+# response, in the sandbox id's case often the picker's `list_sandboxes()`) and go
+# to the same place (the path of a token-bearing request), so validating one and
+# not the other was inconsistent rather than considered.
+_ID_RE = re.compile(r"[A-Za-z0-9_-]{1,64}")
+
+
+class ShellClosed(SandboxError):
+    """Raised when input is sent to a shell whose process has exited."""
+
+
+class _ShellCore:
+    """State and pure protocol logic shared by `Shell` (async) and `SyncShell`.
+
+    Everything here is transport-agnostic: id validation, the request path, the
+    resume cursor, base64 framing in both directions, and exit/gap frame
+    interpretation. Neither shell may re-implement any of it -- the two halves
+    speak the same wire protocol, and copying a piece back into one of them is the
+    drift `SyncShell` exists to avoid. The subclasses add only the actual requests
+    (awaited vs blocking) on top of these helpers.
+
+    Not constructed directly.
+    """
+
+    def __init__(
+        self,
+        sandbox_id: str,
+        session_id: str,
+        *,
+        rows: int,
+        cols: int,
+    ) -> None:
+        # Both ids are validated here, in one place, because both are
+        # server-supplied and both are interpolated into the path of a request
+        # carrying the caller's token.
+        self._sandbox_id = self._checked_id(sandbox_id, "sandbox_id")
+        self._session_id = self._checked_id(session_id, "session_id")
+        self._rows = rows
+        self._cols = cols
+        self._cursor = 0
+        self._exit_code: int | None = None
+        # Separate from `_exit_code`: the shell has ended is a different fact
+        # from we know its exit code. Deriving one from the other meant a shell
+        # that exited with a code we could not read looked like a live session,
+        # and the output loop reconnected to it forever.
+        self._ended = False
+        self._closed = False
+        self.gaps = 0
+        """How many times output was lost to backlog eviction (see module doc)."""
+
+    @property
+    def session_id(self) -> str:
+        """The sandbox-assigned id of this shell session."""
+        return self._session_id
+
+    @property
+    def exit_code(self) -> int | None:
+        """The shell's exit code once it has exited, else None."""
+        return self._exit_code
+
+    @property
+    def closed(self) -> bool:
+        """True once the shell has ended or the session has been closed."""
+        return self._closed or self._ended
+
+    def _path(self, tail: str = "") -> str:
+        return f"containers/{self._sandbox_id}/shell/sessions/{self._session_id}{tail}"
+
+    @staticmethod
+    def _checked_id(value: str, kind: str) -> str:
+        """Reject an id that could steer a token-bearing request off its path."""
+        if not _ID_RE.fullmatch(str(value)):
+            raise SandboxError(
+                f"shell: unusable {kind} {value!r} (expected an opaque token of "
+                "letters, digits, '-' or '_')"
+            )
+        return str(value)
+
+    # --- framing: shared so the two shells cannot encode the wire differently ---
+
+    @staticmethod
+    def _input_body(data: bytes | str) -> dict[str, str] | None:
+        """Base64-frame a stdin write for a POST `/input`, or None for an empty one.
+
+        Bytes, not text, because the interesting input is control characters:
+        Ctrl-C is 0x03, and a transport that treated keystrokes as text would
+        make the terminal unable to interrupt anything. A `str` is encoded UTF-8
+        for convenience. An empty write has no body -- the caller skips the POST.
+        """
+        if isinstance(data, str):
+            data = data.encode("utf-8")
+        if not data:
+            return None
+        return {"data": base64.b64encode(data).decode("ascii")}
+
+    @staticmethod
+    def _resize_body(rows: int, cols: int) -> dict[str, dict[str, int]]:
+        """The POST `/input` body that tells the pty its new window size."""
+        return {"resize": {"rows": rows, "cols": cols}}
+
+    def _stream_path(self) -> str:
+        """The output-stream path, carrying the resume cursor as a query param.
+
+        The cursor travels as a query parameter, not `Last-Event-ID`: the Snowflake
+        passthrough forwards only an allowlist of request headers (X-Request-Id
+        and Accept) and would drop the header, while query strings are forwarded
+        verbatim. The sandbox accepts either.
+        """
+        path = self._path("/stream")
+        if self._cursor:
+            path += f"?last_event_id={self._cursor}"
+        return path
+
+    def _decode_out_frame(self, evt: SSEEvent) -> bytes:
+        """Base64-decode one `out` frame, mapping a bad frame to a transport error.
+
+        Decode is the caller's first step, BEFORE advancing the cursor: a decode
+        failure raised straight out of the read loop (as a bare `binascii.Error`)
+        killed an otherwise-healthy session -- the exact failure resumption exists
+        to prevent. Converting it to the error type the reconnect loop handles, and
+        leaving the cursor put, makes the retry re-fetch the frame rather than skip
+        past it; if it is permanently corrupt the retry budget ends the session with
+        a clear error instead of losing output silently.
+        """
+        try:
+            return base64.b64decode(evt.data, validate=True)
+        except (ValueError, binascii.Error) as exc:
+            raise SandboxTransportError(
+                f"shell session {self._session_id}: undecodable output frame"
+                f"{f' (id {evt.id})' if evt.id else ''}: {exc}"
+            ) from exc
+
+    def _advance(self, evt: SSEEvent) -> None:
+        """Move the resume cursor forward. Never backwards -- a frame with no id
+        (or an out-of-order one) must not cause a replay of output already seen."""
+        if not evt.id:
+            return
+        try:
+            self._cursor = max(self._cursor, int(evt.id))
+        except ValueError:  # pragma: no cover - defensive
+            pass
+
+    def _apply_control(self, evt: SSEEvent) -> bool:
+        """Handle one control frame. Returns True when the session is over."""
+        try:
+            msg = evt.json()
+        except (ValueError, TypeError):  # pragma: no cover - defensive
+            return False
+        if isinstance(msg, dict) and msg.get("type") == "exit":
+            raw = msg.get("code")
+            # A code we cannot read (null, or a type we did not expect) still
+            # means the shell is gone -- record the end, leave the code unknown.
+            self._exit_code = raw if isinstance(raw, int) else None
+            self._ended = True
+            return True
+        return False
+
+    @staticmethod
+    def _lost_at(evt: SSEEvent) -> int:
+        try:
+            payload = evt.json()
+            return int(payload.get("resumed_at", 0)) if isinstance(payload, dict) else 0
+        except (ValueError, TypeError):  # pragma: no cover - defensive
+            return 0
+
+
+class Shell(_ShellCore):
+    """One interactive shell session on a sandbox.
+
+    Obtained from `AsyncSandbox.shell()`. Not constructed directly.
+    """
+
+    def __init__(
+        self,
+        transport: Transport,
+        sandbox_id: str,
+        session_id: str,
+        *,
+        rows: int,
+        cols: int,
+    ) -> None:
+        super().__init__(sandbox_id, session_id, rows=rows, cols=cols)
+        self._t = transport
+
+    # --- input ----------------------------------------------------------
+
+    async def send(self, data: bytes | str) -> None:
+        """Write to the shell's stdin.
+
+        Bytes, not text, because the interesting input is control characters:
+        Ctrl-C is 0x03, and a transport that treated keystrokes as text would
+        make the terminal unable to interrupt anything. A `str` is encoded UTF-8
+        for convenience.
+        """
+        if self.closed:
+            raise ShellClosed(f"shell session {self._session_id} has ended")
+        body = self._input_body(data)
+        if body is None:
+            return
+        await self._t.request("POST", self._path("/input"), json_body=body)
+
+    async def resize(self, rows: int, cols: int) -> None:
+        """Tell the pty its new window size, so curses apps re-lay-out."""
+        if self.closed:
+            return
+        self._rows, self._cols = rows, cols
+        await self._t.request("POST", self._path("/input"), json_body=self._resize_body(rows, cols))
+
+    # --- output ---------------------------------------------------------
+
+    async def output(self, *, on_gap: Callable[[int], None] | None = None) -> AsyncIterator[bytes]:
+        """Yield pty output as it arrives, across gateway stream cuts.
+
+        Ends when the shell process exits (`exit_code` is then set) or `close()`
+        is called. A cut stream is re-opened from the last frame id seen; only a
+        cut we cannot recover from raises `SandboxTransportError`.
+        """
+        attempt = 0
+        while not self.closed:
+            delivered = False
+            try:
+                async for chunk in self._read_once(on_gap):
+                    delivered = True
+                    attempt = 0  # progress resets the budget
+                    yield chunk
+                if self.closed:
+                    return
+                # A clean end with no exit frame is the ordinary 900s case: the
+                # gateway closed a response whose shell is still alive. Re-open
+                # from the cursor.
+                if not delivered and attempt >= len(_RECONNECT_DELAYS):
+                    # ...but a stream that keeps closing having sent nothing is
+                    # not that. Without this the loop reconnects forever and the
+                    # caller sees a terminal that is simply silent.
+                    raise SandboxTransportError(
+                        f"shell session {self._session_id}: stream closed without output "
+                        f"after {attempt} attempts"
+                    )
+            except SandboxTransportError:
+                if self.closed:
+                    return
+                if attempt >= len(_RECONNECT_DELAYS):
+                    raise
+            if self.closed:
+                return
+            delay = _RECONNECT_DELAYS[min(attempt, len(_RECONNECT_DELAYS) - 1)]
+            attempt += 1
+            await asyncio.sleep(delay)
+
+    async def _read_once(self, on_gap: Callable[[int], None] | None) -> AsyncIterator[bytes]:
+        """One SSE connection's worth of frames."""
+        async for evt in self._t.stream_sse("GET", self._stream_path()):
+            if evt.event == _EVENT_OUT:
+                # Decode BEFORE advancing the cursor (see `_decode_out_frame`): a
+                # malformed frame becomes a retryable transport error with the
+                # cursor left put, so the retry re-fetches it rather than skipping
+                # past lost output.
+                chunk = self._decode_out_frame(evt)
+                self._advance(evt)
+                yield chunk
+                continue
+            self._advance(evt)
+            if evt.event == _EVENT_GAP:
+                self.gaps += 1
+                if on_gap is not None:
+                    on_gap(self._lost_at(evt))
+            elif evt.event == _EVENT_CONTROL:
+                if self._apply_control(evt):
+                    return
+
+    # --- lifecycle ------------------------------------------------------
+
+    async def close(self) -> None:
+        """End the session and reap the shell. Idempotent."""
+        if self._closed:
+            return
+        self._closed = True
+        with contextlib.suppress(Exception):
+            # A close that fails is not worth raising over: the sandbox reaps
+            # idle sessions anyway, and the caller is on its way out.
+            await self._t.request("DELETE", self._path())
+
+    async def __aenter__(self) -> Shell:
+        return self
+
+    async def __aexit__(self, *exc: object) -> None:
+        await self.close()
+
+    # --- interactive attach ---------------------------------------------
+
+    async def attach(
+        self,
+        *,
+        stdin: int | None = None,
+        stdout: int | None = None,
+        sanitize: bool = True,
+    ) -> int | None:
+        """Wire the local terminal to this shell until it exits. Returns exit code.
+
+        **Trust boundary, worth stating plainly:** this hands the sandbox's raw
+        output to your *local* terminal emulator, which interprets it. A sandbox
+        runs untrusted and often agent-generated code, so treat attaching the way
+        you would treat `ssh`-ing into a machine you do not control -- the remote
+        end can emit control sequences, not just text.
+
+        `sanitize` (default on) removes OSC 52, the sequence that writes the local
+        clipboard: no display purpose, and a sandbox quietly replacing what you
+        are about to paste is a real consequence. Everything else passes through,
+        deliberately -- see `_ansi._ClipboardStripper` for what is not filtered
+        and why, including one residual risk that cannot be filtered without
+        breaking legitimate terminal programs (a sandbox can emit a capability
+        query, and the local terminal's *reply* lands as typed input at the next
+        prompt). Pass `sanitize=False` for byte-exact passthrough.
+
+        `output()` is always raw: it returns bytes to a program, and silently
+        altering them there would be the surprising choice.
+
+        Three things have to happen for this to feel like a terminal, and each
+        one is a thing users notice by its absence:
+
+        * **Raw mode** on the local tty, so keystrokes go to the *remote* shell
+          instead of being line-buffered and interpreted locally -- without it
+          Ctrl-C kills the client and nothing reaches the sandbox.
+        * **SIGWINCH** forwarded as a resize, so a window change re-lays-out the
+          remote curses app rather than leaving it drawing at the old size.
+        * The initial size sent up front, since the session was created with a
+          default that is probably not this terminal's.
+
+        Requires a real tty on both ends; raises otherwise rather than silently
+        degrading to something that looks broken.
+        """
+        import termios
+        import tty
+
+        in_fd = sys.stdin.fileno() if stdin is None else stdin
+        out_fd = sys.stdout.fileno() if stdout is None else stdout
+        if not os.isatty(in_fd):
+            raise SandboxError("attach() needs a tty on stdin; use send()/output() instead")
+
+        loop = asyncio.get_running_loop()
+        saved = termios.tcgetattr(in_fd)
+        keys: asyncio.Queue[bytes | None] = asyncio.Queue()
+        # Strong refs to in-flight resize tasks: asyncio holds only a weak ref to a
+        # bare create_task(), so a fire-and-forget resize can be garbage-collected
+        # mid-run. Hold it until it finishes, then drop it.
+        winch_tasks: set[asyncio.Task[None]] = set()
+
+        def _readable() -> None:
+            try:
+                data = os.read(in_fd, 4096)
+            except OSError:
+                data = b""
+            keys.put_nowait(data or None)
+
+        def _on_winch() -> None:
+            size = os.get_terminal_size(out_fd)
+            # Fire-and-forget: a resize is advisory, and awaiting it here would
+            # block the signal handler.
+            task = loop.create_task(self._resize_quietly(size.lines, size.columns))
+            winch_tasks.add(task)
+            task.add_done_callback(winch_tasks.discard)
+
+        def _note(resumed_at: int) -> None:
+            del resumed_at
+            os.write(out_fd, b"\r\n[output was dropped while reconnecting]\r\n")
+
+        # Acquire *inside* the try, so a failure part-way through setup still runs
+        # the matched release below. `add_signal_handler` is the realistic tripwire
+        # (RuntimeError off the main thread), but leaving a terminal in raw mode is
+        # bad enough that nothing here should sit outside the guard. Each release
+        # is a no-op when its acquire never happened -- `remove_reader` returns
+        # False for an unregistered fd, `remove_signal_handler` for an uninstalled
+        # signal, and `tcsetattr` re-applies the attributes we just read -- so one
+        # try/finally covering both setup and use is enough, and keeps acquire and
+        # release paired in a single readable order.
+        try:
+            tty.setraw(in_fd)
+            loop.add_reader(in_fd, _readable)
+            with contextlib.suppress(NotImplementedError, ValueError):
+                loop.add_signal_handler(signal.SIGWINCH, _on_winch)
+
+            if os.isatty(out_fd):
+                size = os.get_terminal_size(out_fd)
+                await self.resize(size.lines, size.columns)
+
+            pump = loop.create_task(self._pump_keys(keys))
+            stripper = _ClipboardStripper() if sanitize else None
+            try:
+                async for chunk in self.output(on_gap=_note):
+                    os.write(out_fd, stripper.feed(chunk) if stripper else chunk)
+            finally:
+                if stripper is not None:
+                    # Anything held back mid-sequence still belongs on screen.
+                    tail = stripper.flush()
+                    if tail:
+                        os.write(out_fd, tail)
+                pump.cancel()
+                with contextlib.suppress(asyncio.CancelledError):
+                    await pump
+        finally:
+            with contextlib.suppress(NotImplementedError, ValueError):
+                loop.remove_signal_handler(signal.SIGWINCH)
+            loop.remove_reader(in_fd)
+            termios.tcsetattr(in_fd, termios.TCSADRAIN, saved)
+        return self._exit_code
+
+    async def _pump_keys(self, keys: asyncio.Queue[bytes | None]) -> None:
+        while True:
+            data = await keys.get()
+            if data is None:
+                return
+            # Coalesce every keystroke already queued behind this one into a single
+            # POST. Each POST /input is a full round-trip (~150ms+ against a remote
+            # deployment) and `_readable` keeps enqueuing while one is in flight, so
+            # sending them one-at-a-time makes input latency grow with typing speed --
+            # a burst of N bytes costs N sequential round-trips. Draining the
+            # already-ready items adds no latency to a lone keystroke (nothing else is
+            # queued) and collapses a burst to one request. `None` is the stdin-EOF
+            # sentinel: send what we have, then stop.
+            chunks = [data]
+            eof = False
+            while True:
+                try:
+                    nxt = keys.get_nowait()
+                except asyncio.QueueEmpty:
+                    break
+                if nxt is None:
+                    eof = True
+                    break
+                chunks.append(nxt)
+            try:
+                await self.send(b"".join(chunks))
+            except ShellClosed:
+                return
+            except SandboxTransportError:
+                # A dropped keystroke is better than tearing down a live
+                # terminal; the output stream reports the real end of session.
+                pass
+            if eof:
+                return
+
+    async def _resize_quietly(self, rows: int, cols: int) -> None:
+        with contextlib.suppress(Exception):
+            await self.resize(rows, cols)
+
+
+async def open_shell(
+    transport: Transport,
+    sandbox_id: str,
+    *,
+    rows: int = 24,
+    cols: int = 80,
+    cwd: str | None = None,
+) -> Shell:
+    """Start a shell session on a sandbox and return it.
+
+    The pty is running by the time this returns, so output produced before the
+    caller starts reading -- the shell's greeting and first prompt -- is already
+    in the sandbox's backlog and arrives on the first read rather than being
+    lost to the gap between the two requests.
+    """
+    body: dict[str, object] = {"rows": rows, "cols": cols}
+    if cwd is not None:
+        body["cwd"] = cwd
+    resp = await transport.request(
+        "POST", f"containers/{sandbox_id}/shell/sessions", json_body=body
+    )
+    payload = resp.json() if resp.content else {}
+    session_id = payload.get("session_id") if isinstance(payload, dict) else None
+    if not session_id:
+        raise SandboxError(f"shell: sandbox returned no session_id (body: {payload!r})")
+    # Defence in depth. The id goes straight into the path of every follow-up
+    # request, all of which carry the caller's bearer token, so a value with path
+    # characters in it could point those requests somewhere else. The endpoint is
+    # authenticated and trusted, so this is not attacker-reachable today -- it
+    # bounds the blast radius of a buggy or compromised gateway for the price of
+    # one check. The charset is wider than the ids the sandbox actually mints
+    # (hex) so a future id format does not need a client release.
+    # Shell.__init__ validates both ids (see _ShellCore._checked_id).
+    return Shell(transport, sandbox_id, str(session_id), rows=rows, cols=cols)