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

snowflake/sandbox/sync_log_stream.py

@@ -0,0 +1,142 @@
+"""Iterable reader over a command container's captured output — sync counterpart.
+
+`SyncLogStream` is what `Sandbox.stdout` and `Sandbox.stderr` return, giving a
+command container the same ``.stdout`` / ``.stderr`` shape a streamed exec already
+has via `sync_exec_stream.SyncExecStream`.
+
+It is a line-for-line port of `log_stream.LogStream`: same polling of
+``GET containers/{id}/logs?since=<ts>``, same watermark handling, same
+positional dedupe of untimestamped lines, same drain-after-terminal rule. The
+async version's comments explain WHY each of those exists and are not repeated
+here — read `log_stream` for the reasoning, and change the two together.
+
+The dedupe in particular is load-bearing rather than tidy: the server filters on
+``since``, so a line it never timestamped cannot be filtered out and returns on
+every poll. Counting them positionally is what stops `read()` growing without
+bound, and unlike matching on text it keeps a genuinely repeated log line.
+"""
+
+from __future__ import annotations
+
+import time
+from collections.abc import Iterator
+from typing import TYPE_CHECKING, Self
+
+from snowflake.sandbox._runtime._protocol import _RESULT_SENTINEL
+from snowflake.sandbox.log_stream import (
+    _DRAIN_AFTER_TERMINAL_S,
+    _POLL_INTERVAL_S,
+)
+from snowflake.sandbox.types import TERMINAL_STATUSES, StreamName
+
+if TYPE_CHECKING:  # pragma: no cover - typing only
+    from snowflake.sandbox.sync_client import Sandbox
+
+__all__ = ["SyncLogStream"]
+
+
+class SyncLogStream(Iterator[str]):
+    """Iterator over one channel of a command container's output.
+
+    Yields line text with the trailing newline stripped, oldest first. Iteration
+    ends once the container is in a terminal state and no further lines arrive, so
+    ``for line in sb.stdout`` over a finished container terminates rather than
+    hanging.
+
+    Only meaningful for command containers (created with ``command=``); an
+    ``exec``-only sandbox produces no managed-process output and the stream ends
+    immediately.
+    """
+
+    def __init__(self, sandbox: Sandbox, channel: StreamName, *, since_ts_ms: int = 0) -> None:
+        self._sandbox = sandbox
+        self._channel = channel
+        self._since_ts_ms = since_ts_ms
+        self._pending: list[str] = []
+        self._done = False
+        # Set once the managed process's exit marker is seen in the log stream.
+        self._process_done = False
+        self._terminal_since: float | None = None
+        self._untimestamped_seen = 0
+
+    def __iter__(self) -> Self:
+        return self
+
+    def __next__(self) -> str:
+        while True:
+            if self._pending:
+                return self._pending.pop(0)
+            if self._done:
+                raise StopIteration
+            self._fill()
+
+    def _fill(self) -> None:
+        """Fetch any lines newer than the last one seen, or decide we are done."""
+        lines = self._sandbox._fetch_log_lines(since_ts_ms=self._since_ts_ms)
+
+        # Snapshot the skip threshold once; advancing it mid-loop would drop later
+        # lines in the batch that share a `ts` (same millisecond). See log_stream.
+        start_ts = self._since_ts_ms
+        next_ts = self._since_ts_ms
+        got = False
+        untimestamped = 0
+        for ln in lines:
+            text = str(ln.get("text") or "")
+            # End-of-stream marker the injected runner prints on stdout when the
+            # managed process exits; checked before the watermark/dedup logic and on
+            # every channel, and suppressed from output. See log_stream._fill for the
+            # full reasoning (a command container stays `ready` after its process
+            # ends, so the refresh() fallback below would otherwise poll forever).
+            if _RESULT_SENTINEL in text:
+                self._process_done = True
+                continue
+            ts = int(ln.get("ts") or 0)
+            if ts:
+                if ts < start_ts:
+                    continue
+                next_ts = max(next_ts, ts + 1)
+            else:
+                untimestamped += 1
+                if untimestamped <= self._untimestamped_seen:
+                    continue
+                self._untimestamped_seen = untimestamped
+            if ln.get("stream") != self._channel:
+                continue
+            self._pending.append(text)
+            got = True
+
+        # Advance the watermark once, after the whole batch (see log_stream).
+        self._since_ts_ms = next_ts
+
+        if got:
+            self._terminal_since = None
+            return
+
+        # Managed process has exited (its result marker was seen) and everything
+        # newer than the watermark is drained -- stop now, rather than waiting for
+        # the container lifecycle to reach a terminal status (it never will: a
+        # command container stays `ready` after its process exits).
+        if self._process_done:
+            self._done = True
+            return
+
+        # A live read, not the cached status: `Sandbox.status` is set at create
+        # time and nothing on the container's lifecycle mutates it, so trusting it
+        # would poll a finished container forever.
+        if self._sandbox.refresh() in TERMINAL_STATUSES:
+            now = time.monotonic()
+            if self._terminal_since is None:
+                self._terminal_since = now
+            elif now - self._terminal_since >= _DRAIN_AFTER_TERMINAL_S:
+                self._done = True
+                return
+
+        time.sleep(_POLL_INTERVAL_S)
+
+    def read(self) -> str:
+        """Drain the stream and return it as one newline-joined string.
+
+        Blocks until the stream ends, so only use it on a container that will
+        finish.
+        """
+        return "\n".join(list(self))

snowflake/sandbox/sync_shell.py

@@ -0,0 +1,413 @@
+"""Synchronous interactive shell sessions on a sandbox.
+
+The blocking counterpart to `shell.Shell`. `SyncShell` is returned by
+`Sandbox.shell()`; `AsyncSandbox.shell()` returns the async `Shell`.
+
+There is no separate sync *transport* problem to solve here, contrary to how the
+half-duplex hop reads at first: a shell is not one duplex socket but an SSE
+response carrying output (`stream_sse`, a streaming read) paired with POSTs
+carrying input (`request`). `_sync_transport.SyncTransport` already does both, so
+the sync shell rides the exact same two-request protocol as the async one -- see
+the `shell` module docstring for why that protocol exists and why it is a stopgap.
+
+Everything that is not I/O -- id validation, the request path, the resume cursor,
+base64 framing, exit/gap interpretation -- is inherited from `_ShellCore`, so the
+two shells cannot encode the wire differently. What differs is only the shape of
+the I/O:
+
+* `output()` is a blocking generator over `SyncTransport.stream_sse`, with the
+  same resume-across-a-cut loop as the async side but `time.sleep` for its
+  backoff instead of `await asyncio.sleep`.
+* `attach()` cannot lean on an event loop's reader/signal registration, so it runs
+  a **worker thread** that pumps keystrokes and applies resizes. All blocking work
+  (both the input POSTs and the resize POSTs) happens on that thread; the SIGWINCH
+  handler and the output loop stay on the main thread, and the handler does no I/O
+  -- it records the new size and wakes the worker. The worker blocks in
+  `select()` on stdin plus a wake pipe rather than a bare `os.read`, which a
+  finished session could not interrupt: a leaked reader thread would sit blocked
+  on the real terminal and steal the next keystroke after `attach()` returned.
+"""
+
+from __future__ import annotations
+
+import contextlib
+import os
+import select
+import signal
+import sys
+import threading
+import time
+from collections.abc import Callable, Iterator
+from typing import TYPE_CHECKING, Any
+
+from snowflake.sandbox._ansi import _ClipboardStripper
+from snowflake.sandbox.exceptions import SandboxError, SandboxTransportError
+from snowflake.sandbox.shell import (
+    _EVENT_CONTROL,
+    _EVENT_GAP,
+    _EVENT_OUT,
+    _RECONNECT_DELAYS,
+    ShellClosed,
+    _ShellCore,
+)
+
+if TYPE_CHECKING:
+    from snowflake.sandbox._sync_transport import SyncTransport
+
+__all__ = ["SyncShell", "open_shell_sync"]
+
+# How long `attach()` waits for its worker thread to unwind on teardown. The
+# thread is blocked in `select()` and wakes the instant the stop byte is written,
+# so this is only a backstop against a wedged fd, not a normal-path delay.
+_WORKER_JOIN_TIMEOUT_S = 2.0
+
+
+class SyncShell(_ShellCore):
+    """One interactive shell session on a sandbox (blocking API).
+
+    Obtained from `Sandbox.shell()`. Not constructed directly.
+    """
+
+    def __init__(
+        self,
+        transport: SyncTransport,
+        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 ----------------------------------------------------------
+
+    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
+        self._t.request("POST", self._path("/input"), json_body=body)
+
+    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
+        self._t.request("POST", self._path("/input"), json_body=self._resize_body(rows, cols))
+
+    # --- output ---------------------------------------------------------
+
+    def output(self, *, on_gap: Callable[[int], None] | None = None) -> Iterator[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:
+                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
+            time.sleep(delay)
+
+    def _read_once(self, on_gap: Callable[[int], None] | None) -> Iterator[bytes]:
+        """One SSE connection's worth of frames."""
+        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 ------------------------------------------------------
+
+    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.
+            self._t.request("DELETE", self._path())
+
+    def __enter__(self) -> SyncShell:
+        return self
+
+    def __exit__(self, *exc: object) -> None:
+        self.close()
+
+    # --- interactive attach ---------------------------------------------
+
+    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:
+
+        * **Raw mode** on the local tty, so keystrokes go to the *remote* shell
+          instead of being line-buffered and interpreted locally.
+        * **SIGWINCH** forwarded as a resize, so a window change re-lays-out the
+          remote curses app. Installed on the main thread only; a `SyncShell`
+          driven from a worker thread simply loses live resize (the handler cannot
+          be registered off the main thread), which is advisory, not fatal -- the
+          async `attach()` treats the analogous failure the same way.
+        * The initial size sent up front, since the session was created with a
+          default that is probably not this terminal's.
+
+        A worker thread reads keystrokes and applies resizes so every blocking
+        request happens off the main thread; it is torn down cleanly on exit and
+        is not left blocked on the terminal after this returns. Requires a real tty
+        on stdin; raises otherwise rather than silently degrading.
+        """
+        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")
+
+        saved = termios.tcgetattr(in_fd)
+
+        def _note(resumed_at: int) -> None:
+            del resumed_at
+            os.write(out_fd, b"\r\n[output was dropped while reconnecting]\r\n")
+
+        # The worker thread + wake pipe + resize mailbox are one unit -- the sync
+        # stand-in for the event loop's reader/signal registration the async attach
+        # gets for free. See `_ShellInputForwarder`.
+        forwarder = _ShellInputForwarder(self, in_fd, out_fd)
+        prev_winch: Any = None
+        winch_installed = False
+        stripper = _ClipboardStripper() if sanitize else None
+        # Acquire inside the try, exactly as the async attach does: a failure
+        # part-way through setup must still restore the terminal and stop the worker.
+        # setraw lands first, so a raise on any later line is the tightest window,
+        # and it must not escape over a terminal left in raw mode or a live thread.
+        try:
+            tty.setraw(in_fd)
+            forwarder.start()
+            if hasattr(signal, "SIGWINCH"):
+                # signal.signal raises ValueError off the main thread (and the
+                # attribute is absent on Windows). Resize-on-winch is advisory, so
+                # this is survivable -- the async side suppresses the analogous
+                # error rather than failing the attach.
+                with contextlib.suppress(ValueError, OSError):
+                    prev_winch = signal.signal(signal.SIGWINCH, forwarder.on_winch)
+                    winch_installed = True
+
+            if os.isatty(out_fd):
+                size = os.get_terminal_size(out_fd)
+                self.resize(size.lines, size.columns)
+
+            try:
+                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)
+        finally:
+            if winch_installed:
+                with contextlib.suppress(ValueError, OSError, TypeError):
+                    signal.signal(signal.SIGWINCH, prev_winch)
+            forwarder.stop_and_join()
+            termios.tcsetattr(in_fd, termios.TCSADRAIN, saved)
+        return self._exit_code
+
+    def _resize_quietly(self, rows: int, cols: int) -> None:
+        with contextlib.suppress(Exception):
+            self.resize(rows, cols)
+
+
+class _ShellInputForwarder:
+    """Pumps local stdin to the remote shell, off the main thread.
+
+    The sync stand-in for what the async `attach()` gets from the event loop --
+    a reader on stdin plus a SIGWINCH handler. A worker thread blocks in
+    ``select()`` on stdin and a wake pipe rather than a bare ``os.read``: writing
+    one byte to the pipe (on resize or on teardown) breaks the block at once, so
+    the thread never sits uninterruptibly on the real terminal -- the failure
+    mode that would leak a thread eating the next keystroke after `attach()`
+    returns.
+
+    ``on_winch`` is the SIGWINCH handler the caller installs. It runs on the main
+    thread between bytecode instructions and does NO I/O -- a blocking POST from
+    inside a signal handler could reenter the httpx client -- so it records the
+    new size in a locked mailbox and wakes the worker, which sends the resize.
+    """
+
+    def __init__(self, shell: SyncShell, in_fd: int, out_fd: int) -> None:
+        self._shell = shell
+        self._in_fd = in_fd
+        self._out_fd = out_fd
+        self._stop = threading.Event()
+        self._wake_r, self._wake_w = os.pipe()
+        # Single-slot mailbox for the latest window size, written by on_winch (main
+        # thread) and drained by the worker. A lock because the handler can fire
+        # while the worker is mid-read.
+        self._pending: list[tuple[int, int]] = []
+        self._pending_lock = threading.Lock()
+        self._worker = threading.Thread(target=self._pump, name="sync-shell-input", daemon=True)
+
+    def start(self) -> None:
+        self._worker.start()
+
+    def on_winch(self, signum: int, frame: object) -> None:
+        try:
+            size = os.get_terminal_size(self._out_fd)
+        except OSError:
+            return
+        with self._pending_lock:
+            self._pending.append((size.lines, size.columns))
+        with contextlib.suppress(OSError):
+            os.write(self._wake_w, b"\0")
+
+    def _pump(self) -> None:
+        while not self._stop.is_set():
+            try:
+                readable, _, _ = select.select([self._in_fd, self._wake_r], [], [])
+            except OSError:
+                return
+            if self._stop.is_set():
+                return
+            if self._wake_r in readable:
+                with contextlib.suppress(OSError):
+                    os.read(self._wake_r, 4096)
+                with self._pending_lock:
+                    size = self._pending.pop() if self._pending else None
+                    self._pending.clear()
+                if size is not None:
+                    self._shell._resize_quietly(*size)
+            if self._in_fd in readable:
+                try:
+                    data = os.read(self._in_fd, 4096)
+                except OSError:
+                    data = b""
+                if not data:  # EOF on stdin
+                    return
+                try:
+                    self._shell.send(data)
+                except ShellClosed:
+                    return
+                except SandboxTransportError:
+                    # A dropped keystroke is better than tearing down a live
+                    # terminal; the output stream reports the real end.
+                    continue
+
+    def stop_and_join(self) -> None:
+        """Stop the worker and reclaim the wake pipe. Safe if start() never ran."""
+        self._stop.set()
+        with contextlib.suppress(OSError):
+            os.write(self._wake_w, b"\0")  # break the worker out of select()
+        if self._worker.is_alive():
+            self._worker.join(timeout=_WORKER_JOIN_TIMEOUT_S)
+        with contextlib.suppress(OSError):
+            os.close(self._wake_r)
+        with contextlib.suppress(OSError):
+            os.close(self._wake_w)
+
+
+def open_shell_sync(
+    transport: SyncTransport,
+    sandbox_id: str,
+    *,
+    rows: int = 24,
+    cols: int = 80,
+    cwd: str | None = None,
+) -> SyncShell:
+    """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 = 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. SyncShell.__init__ validates both ids (see _ShellCore._checked_id).
+    return SyncShell(transport, sandbox_id, str(session_id), rows=rows, cols=cols)