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

snowflake/sandbox/_wire.py

@@ -0,0 +1,348 @@
+"""The server's side of the contract: its vocabulary and its budgets.
+
+What the platform says (container statuses, memory tiers, pagination markers,
+streaming-exec frames) and how long the client waits for it. Kept in one module
+because these values are only correct in relation to the server, so they must be
+identical for the async and sync clients — a status added on one side only made a
+resumed sandbox read as ``"unknown"`` on the other, with a warning blaming the
+caller's SDK version.
+"""
+
+from __future__ import annotations
+
+import re
+from collections.abc import Mapping, Sequence
+from datetime import timedelta
+from typing import TYPE_CHECKING, Any, get_args
+
+from snowflake.sandbox.exceptions import SandboxContractWarning, SandboxError
+from snowflake.sandbox.types import MemoryTier, SandboxStatus
+
+if TYPE_CHECKING:
+    from snowflake.sandbox._transport import SSEEvent
+
+__all__ = [
+    "_parse_status",
+    "_warn_if_truncated",
+    "_exec_session_id",
+    "_resolve_exec_budget",
+    "_idle_suspend_minutes",
+]
+
+
+# Server status vocabulary -> SDK `SandboxStatus`.
+#
+# `Container.status` (sandbox-api containers/models.go ContainerFromInfo) carries
+# the managed process's state for a command container -- "starting" | "running" |
+# "crashed" -- and is normalized to "running" when the registry tracks no
+# liveness. The Snowflake app statuses ("suspended" | "stopped") can also surface. A
+# suspended container is alive from the caller's side: Snowflake resumes it
+# transparently on the next exec.
+#
+# Anything NOT in this table maps to "unknown", never to "ready". Coercing an
+# unrecognized status to the most optimistic value is what made a field rename
+# invisible: `LogStream` waited for a "dead" the server never says, and a reaper
+# filtering on "ready" silently reaped nothing.
+_SERVER_STATUS: dict[str, SandboxStatus] = {
+    "starting": "pending",
+    "pending": "pending",
+    # Mid-transition, so NOT "ready": the exec would work (every exec re-drives StartApp,
+    # which is what completes a resume) but it can block for seconds, and a status a caller
+    # checks before acting must not promise a sandbox is servable when it is not yet.
+    # CNG's ActiveStatuses lists both alongside `starting` -- that list means "consuming host
+    # resources", not "usable", which is why it cannot be the readiness rule.
+    "resuming": "pending",
+    "suspending": "pending",
+    # Teardown has started and does not come back, so terminal rather than pending.
+    "stopping": "dead",
+    # A pre-run state the process-status probe can report (e.g. right after a
+    # ?refresh=true resume, before the managed process relaunches). Same bucket
+    # as "starting": the process is not up yet, but this is not a failure.
+    "not_started": "pending",
+    "running": "ready",
+    "ready": "ready",
+    "suspended": "ready",
+    "crashed": "dead",
+    "stopped": "dead",
+    "dead": "dead",
+    # async create (202 path): CNG StartApp failed server-side
+    "failed": "failed",
+    # The server's own "I do not know". Needs a row, or it takes the unrecognized path and
+    # advises an SDK upgrade for a word the server chose deliberately.
+    "unknown": "unknown",
+}
+
+
+# The closed set of memory tiers, mirroring `types.MemoryTier` at runtime and
+# sandbox-api's `tierMap` on the wire. Derived from the Literal with ``get_args``
+# rather than written out, so the two cannot drift: when they did, the SDK warned
+# "unrecognized memory tier" on every 8g/32g sandbox the server had served fine.
+_MEMORY_TIERS: frozenset[str] = frozenset(get_args(MemoryTier))
+
+
+# Keys a paginated response could use to say "there is more". None of these is in
+# today's contract; they are checked so that if one appears, the SDK says so
+# instead of quietly returning a first page as the whole answer.
+_TRUNCATION_KEYS: tuple[str, ...] = (
+    "next_page_token",
+    "nextPageToken",
+    "next_cursor",
+    "nextCursor",
+    "next_page",
+    "next",
+    "has_more",
+    "hasMore",
+    "truncated",
+)
+
+
+# Default exec budget the SDK applies when the caller names no timeout. Sent as
+# the exec's timeout_s / timeout_seconds; the server enforces it and reports its
+# own expiry (408), the only report that reflects what actually happened in the
+# sandbox. 1 hour — the server's exec ceiling — so a long-running exec started
+# without an explicit timeout is not cut short at the old 10-minute default.
+_DEFAULT_EXEC_BUDGET_S = 3600.0
+
+
+# How much longer than the server's exec budget the client waits for the
+# response. The HTTP read timeout must be the LONGER of the two: when it was the
+# shorter one (an unset `timeout` left it at config.timeout_s = 120s against a
+# 600s server budget) every exec between 2 and 10 minutes looked like a
+# transport failure to the caller while still running in the sandbox — and, with
+# a method-agnostic retry loop, was then re-POSTed on top of the live one.
+_EXEC_TIMEOUT_GRACE_S = 30.0
+
+
+# Reconnect pacing after a streaming exec's transport is cut. The common cause is
+# the customer-facing ingress GOAWAY at ~180s (and the 900s Istio cap) on a
+# healthy exec whose command is still running, so the first retries are fast. The
+# length of this tuple is also the budget for a stream that keeps closing having
+# produced NOTHING: a genuinely silent command survives across cuts because the
+# container emits its terminal exit within timeout_s, but a stream that reopens
+# repeatedly with zero output is treated as dead rather than reconnected forever.
+_EXEC_RECONNECT_DELAYS = (0.1, 0.25, 0.5, 1.0, 2.0)
+
+
+def _resolve_exec_budget(timeout: float | None, *, caller: str) -> float:
+    """Validate a caller's ``timeout`` and resolve the exec budget in seconds.
+
+    Both halves belong together because the old inline form
+    (``float(timeout) if timeout else _DEFAULT_EXEC_BUDGET_S``) got the resolution
+    wrong in a way only a guard can fix: ``0`` is FALSY, so ``timeout=0`` did not
+    mean "no time", it silently selected the 600s default. A caller asking for the
+    shortest possible budget got the longest one.
+
+    The three rejected shapes, and what each used to do instead:
+
+    * ``inf`` / ``-inf`` -- ``math.ceil`` raised a bare ``OverflowError``, which is
+      not a `SandboxError`, so ``except SandboxError`` around the exec missed it.
+    * ``nan`` -- ``math.ceil`` raised a bare ``ValueError``, same problem.
+    * ``<= 0`` -- ``max(1, ceil(-9999))`` silently clamped the wire value to 1s
+      while the client's read timeout became ``-9999 + 30`` -- a negative httpx
+      timeout. The caller asked for one thing and got two different wrong ones.
+
+    ``0`` is rejected rather than honored because it is not representable: the wire
+    field is an int and the server reads ``0`` as "unset" and substitutes its own
+    default (which is why the sub-second case rounds UP to 1). There is no way to
+    express "expire immediately", so accepting ``0`` could only lie.
+    """
+    if timeout is None:
+        return _DEFAULT_EXEC_BUDGET_S
+    import math
+
+    try:
+        value = float(timeout)
+    except (TypeError, ValueError) as exc:
+        raise SandboxError(
+            f"{caller}() timeout must be a number of seconds, got "
+            f"{type(timeout).__name__} ({timeout!r})."
+        ) from exc
+    if not math.isfinite(value):
+        raise SandboxError(
+            f"{caller}() timeout must be a finite number of seconds, got {value!r}. "
+            f"There is no 'wait forever' budget: the server enforces its own "
+            f"deadline and reports the expiry. Omit timeout= to take the default "
+            f"({_DEFAULT_EXEC_BUDGET_S:.0f}s), or pass the seconds you want."
+        )
+    if value <= 0:
+        raise SandboxError(
+            f"{caller}() timeout must be greater than 0 seconds, got {value!r}. "
+            f"0 is not 'expire immediately' -- the wire field is an integer and the "
+            f"server reads 0 as unset, replacing it with its own default, so this "
+            f"used to silently become {_DEFAULT_EXEC_BUDGET_S:.0f}s. Pass a positive "
+            f"number of seconds, or omit timeout= to take the default."
+        )
+    return value
+
+
+# What `poll()` returns for a terminal container whose exit code the server never
+# reported (an older server without the cheap `/status` route, or a process lost
+# before it could report one). Nonzero on purpose: the container HAS finished, so
+# `None` would misread as "still running", and it did NOT exit cleanly, so `0`
+# would misread as "success". `1` is the generic-failure convention.
+_POLL_UNKNOWN_EXIT_CODE = 1
+
+
+def _parse_status(raw: object) -> SandboxStatus | None:
+    """Map a server ``status`` string onto `types.SandboxStatus`.
+
+    Returns ``None`` when the field is absent (nothing to learn), and
+    ``"unknown"`` -- with a `exceptions.SandboxContractWarning` -- when the
+    server sent a state this SDK does not know.
+    """
+    if not isinstance(raw, str) or not raw:
+        return None
+    mapped = _SERVER_STATUS.get(raw.strip().lower())
+    if mapped is not None:
+        return mapped
+    import warnings
+
+    warnings.warn(
+        f"server reported an unrecognized container status {raw!r}; treating it as "
+        "'unknown'. Upgrade snowflake-sandbox-python if the server is newer than "
+        "this SDK.",
+        SandboxContractWarning,
+        stacklevel=3,
+    )
+    return "unknown"
+
+
+def _warn_if_truncated(payload: Mapping[str, Any], returned: int) -> None:
+    """Warn when a list response carries any marker of a further page."""
+    import warnings
+
+    for key in _TRUNCATION_KEYS:
+        marker = payload.get(key)
+        if marker:
+            warnings.warn(
+                f"the list response carries {key}={marker!r}: only the first "
+                f"{returned} sandboxes are returned. This SDK does not follow the "
+                "cursor, so do not treat this as the complete set.",
+                SandboxContractWarning,
+                stacklevel=3,
+            )
+            return
+
+
+def _exec_session_id(evt: SSEEvent) -> str | None:
+    """Read the ``exec_session`` id from a streaming exec ``session`` frame.
+
+    Returns None for a malformed frame rather than raising: a session id we
+    cannot read just means this stream is not resumable, which degrades to the
+    old behaviour (a cut ends the exec) instead of crashing a live run.
+    """
+    try:
+        payload = evt.json()
+    except (ValueError, TypeError):
+        return None
+    if isinstance(payload, dict):
+        sid = payload.get("exec_session")
+        if isinstance(sid, str) and sid:
+            return sid
+    return None
+
+
+_DURATION_TOKEN = re.compile(r"(\d+)\s*([smhd])")
+_DURATION_UNIT_SECONDS = {"s": 1, "m": 60, "h": 3600, "d": 86400}
+
+
+def _parse_duration_secs(text: str) -> int:
+    """Parse a duration string (``"30m"``, ``"2h"``, ``"1h30m"``, ``"1d"``) to seconds.
+
+    One or more ``<int><unit>`` tokens with units ``s``/``m``/``h``/``d`` and nothing
+    else; anything else is a ``ValueError``.
+    """
+    s = text.strip().lower()
+    total = 0
+    pos = 0
+    for m in _DURATION_TOKEN.finditer(s):
+        if m.start() != pos:
+            break
+        total += int(m.group(1)) * _DURATION_UNIT_SECONDS[m.group(2)]
+        pos = m.end()
+    if not s or pos != len(s):
+        raise ValueError(f"invalid duration {text!r}; use forms like '30m', '2h', '1h30m'")
+    return total
+
+
+def _idle_suspend_minutes(value: timedelta | str | None) -> int | None:
+    """Normalize a ``create(idle_suspend=)`` argument to whole minutes for the
+    ``lifecycle.idle_suspend_minutes`` create field.
+
+    Accepts a ``datetime.timedelta`` or a duration string (``"30m"``, ``"2h"``);
+    ``None`` leaves the field unset so CNG applies its global idle-suspend default.
+    CNG's resolution is minutes, so a value that is not a whole number of minutes
+    (e.g. ``"90s"`` / ``timedelta(seconds=90)``) raises ``ValueError`` rather than
+    being silently rounded — the API does not accept a precision it cannot honor. A
+    bare ``int`` is rejected with ``TypeError``: an unqualified number reads as
+    seconds in Python, so callers must be explicit via a timedelta or unit-suffixed
+    string.
+    """
+    if value is None:
+        return None
+    if isinstance(value, timedelta):
+        secs = value.total_seconds()
+    elif isinstance(value, str):
+        secs = _parse_duration_secs(value)
+    else:
+        raise TypeError(
+            "idle_suspend must be a datetime.timedelta or a duration string like "
+            f"'30m'/'2h' (not a bare number); got {type(value).__name__}"
+        )
+    if secs < 0:
+        raise ValueError("idle_suspend must be non-negative")
+    if secs % 60:
+        raise ValueError(
+            f"idle_suspend must be a whole number of minutes (CNG's resolution); got {value!r}"
+        )
+    return int(secs // 60)
+
+
+def _fresh_exec_body(
+    cmd: Sequence[str],
+    env: Mapping[str, str] | None,
+    working_dir: str | None,
+    exec_budget_s: float,
+) -> dict[str, Any]:
+    """Build the body for the FIRST (fresh-run) request of a streaming exec.
+
+    This is the request that actually runs the command; every later reopen
+    resumes from the cursor with a different body and must NOT carry ``cmd``
+    (which would re-run it). ``env`` / ``working_dir`` are included only when set
+    so an unset field is absent from the wire rather than sent as ``null``.
+    """
+    body: dict[str, Any] = {"cmd": list(cmd), "timeout_s": exec_budget_s}
+    if env:
+        body["env"] = dict(env)
+    if working_dir:
+        body["working_dir"] = working_dir
+    return body
+
+
+def _advance_exec_cursor(cursor: int, evt: SSEEvent) -> int:
+    """Advance the resume cursor past a frame's ``id``.
+
+    The cursor is the ``last_event_id`` a reconnect resends so the server
+    replays from just after it. A frame with no readable integer ``id`` (none
+    present, or one that will not parse) leaves the cursor where it was --
+    resuming from the last id we could trust is safe, resuming from a bad one is
+    not.
+    """
+    if not evt.id:
+        return cursor
+    try:
+        return max(cursor, int(evt.id))
+    except ValueError:  # pragma: no cover - defensive
+        return cursor
+
+
+def _exec_reconnect_exhausted(attempt: int, *, delivered: bool) -> bool:
+    """Whether a stream that keeps reopening should be given up as dead.
+
+    A reopen that has produced no output (``delivered`` is False) and has
+    already used every entry in ``_EXEC_RECONNECT_DELAYS`` is treated as dead
+    rather than retried forever -- the caller reports the truncation. Any
+    delivered output resets the attempt count, so a long but productive stream
+    never trips this.
+    """
+    return not delivered and attempt >= len(_EXEC_RECONNECT_DELAYS)

snowflake/sandbox/app.py

@@ -0,0 +1,256 @@
+"""``App`` — the registry that ``@app.function`` / ``@app.session`` attach to.
+
+Modal-style authoring entry point: the module is the contract, as in Modal;
+there is no manifest file.
+
+    from snowflake.sandbox import App, Bundle, Image, Secret
+
+    app = App("jira-triage")
+
+    @app.function(
+        image=Image.from_catalog("sandbox-base"),
+        bundle=Bundle.from_dir("..", include=["tools/jira-triage/**"]),
+        secrets=[Secret.from_name("MYDB.SECRETS.JIRA_CREDS",
+                                  env_var="JIRA_CREDS", host="jira.example.com")],
+        entry=["/bin/bash", "setup.sh"],
+    )
+    def triage(team: str, ticket: str) -> None:
+        ...  # local impl; body is the Python entry for .local() / direct call
+
+    import asyncio
+    result = asyncio.run(triage.remote(TRIAGE_TEAM="platform"))
+    assert result.exit_code == 0
+
+`Function` lives in ``snowflake.sandbox.function``; the ``@app.session`` surface
+(`SessionApp`, ``enter``) in ``snowflake.sandbox.session_app``.
+"""
+
+from __future__ import annotations
+
+from collections.abc import Callable, Mapping, Sequence
+from typing import Any
+
+from snowflake.sandbox._assemble import Bundle
+from snowflake.sandbox._deploy_spec import DeploySpec
+from snowflake.sandbox.egress import Egress
+from snowflake.sandbox.function import Function, FunctionSpec
+from snowflake.sandbox.image import Image
+from snowflake.sandbox.mount import StageMount
+from snowflake.sandbox.secret import Secret
+from snowflake.sandbox.session_app import SessionApp, SessionSpec
+from snowflake.sandbox.types import MemoryTier
+
+__all__ = ["App"]
+
+
+class App:
+    """A named collection of deployable functions.
+
+    Decorate functions with `@app.function(...)` to register them, then call
+    `.remote()` to run them in a sandbox.
+
+    Example:
+        app = App("my-agent")
+
+        @app.function(
+            image=Image.from_catalog("sandbox-base"),
+            egress=Egress(),
+        )
+        def run() -> None:
+            ...
+
+        result = asyncio.run(run.remote())
+    """
+
+    def __init__(self, name: str) -> None:
+        if not name:
+            raise ValueError("App name must be non-empty")
+        self._name = name
+        self._functions: dict[str, Function] = {}
+        self._sessions: dict[str, SessionApp] = {}
+
+    @property
+    def name(self) -> str:
+        return self._name
+
+    def function(
+        self,
+        *,
+        image: Image | str,
+        memory: MemoryTier = "4g",
+        cpu: float | None = None,
+        bundle: Bundle | None = None,
+        secrets: Sequence[Secret] = (),
+        egress: Egress | None = None,
+        env: Mapping[str, str] | None = None,
+        timeout: float = 3600.0,
+        entry: Sequence[str] | None = None,
+        code_stage: str | None = None,
+        stage_mounts: Sequence[StageMount] = (),
+    ) -> Callable[[Callable[..., Any]], Function]:
+        """Decorator that registers ``fn`` as a deployable `Function`.
+
+        Parameters
+        ----------
+        image:
+            Base image name or `Image` object.  Bare strings are
+            coerced to ``Image(name=str)``.
+        memory:
+            Memory tier: ``"1g"``, ``"4g"`` (default), ``"8g"``, ``"16g"``,
+            ``"32g"``, ``"64g"``.
+        cpu:
+            CPU in cores, overriding what the memory tier would imply
+            (``1g``:1, ``4g``:2, ``8g``:3, ``16g``:4, ``32g``:6, ``64g``:8).  Fractional values are
+            allowed — ``0.5`` is half a core.  Omit to inherit the tier default.
+            Sets only CPU; memory stays tier-selected.
+        bundle:
+            Source-tree slice to upload.  ``None`` uses the temp directory
+            itself as the bundle root (function-mode shim only).
+        secrets:
+            Snowflake SECRETs to broker in (``Secret.from_name(...)``). The
+            sandbox receives a dummy; the egress proxy swaps in the real value.
+        egress:
+            Outbound network scope (``Egress()``,
+            ``Egress(allow_internet=True)``,
+            ``Egress(external_access_integrations=("MY_EAI",))``).
+            Defaults to no egress.
+        env:
+            Static env vars set in the container.
+        timeout:
+            Wall-clock timeout in seconds for the remote exec (default 3600).
+        entry:
+            Shell entry command run in the container (process-mode).  ``None``
+            (default) selects function-mode — the function body runs via the
+            auto-generated ``__app_runner__.py`` shim.
+        code_stage:
+            Snowflake stage path for presigned code delivery, e.g.
+            ``"TEMP.DB.SANDBOX_CODE_SSE"``.  Optional.
+        stage_mounts:
+            Snowflake stages/workspaces to mount into the container's
+            filesystem (``StageMount(...)``).  Defaults to no mounts.
+        """
+        resolved_image = image if isinstance(image, Image) else Image(name=str(image))
+        resolved_entry = tuple(entry) if entry is not None else None
+
+        spec = FunctionSpec(
+            image=resolved_image,
+            memory=memory,
+            cpu=cpu,
+            bundle=bundle,
+            secrets=tuple(secrets),
+            egress=egress,
+            env=dict(env) if env is not None else {},
+            timeout_s=float(timeout),
+            entry=resolved_entry,
+            code_stage=code_stage,
+            stage_mounts=tuple(stage_mounts),
+        )
+
+        def decorator(fn: Callable[..., Any]) -> Function:
+            if hasattr(fn, "__app_function_spec__"):
+                raise ValueError(f"function {fn!r} is already registered with an App")
+            func = Function(fn, spec, self)
+            try:
+                fn.__app_function_spec__ = spec  # type: ignore[attr-defined]
+            except AttributeError:
+                pass
+            self._functions[fn.__name__] = func
+            return func
+
+        return decorator
+
+    def entrypoint(
+        self,
+        **kwargs: Any,
+    ) -> Callable[[Callable[..., Any]], Function]:
+        """Alias for `function()` — marks the primary entry function.
+
+        The first ``@app.entrypoint`` (or ``@app.function``) registered is
+        the default target of `resolve()` / `deploy_spec()` called with no
+        arguments.
+        """
+        return self.function(**kwargs)
+
+    def session(
+        self,
+        *,
+        image: Image | str,
+        memory: MemoryTier = "4g",
+        cpu: float | None = None,
+        bundle: Bundle | None = None,
+        secrets: Sequence[Secret] = (),
+        egress: Egress | None = None,
+        env: Mapping[str, str] | None = None,
+        timeout: float = 3600.0,
+        code_stage: str | None = None,
+        stage_mounts: Sequence[StageMount] = (),
+    ) -> Callable[[type], SessionApp]:
+        """Decorator that registers a class as a stateful keep-alive Session.
+
+        The ``@app.cls``-style sibling of `function()`. The decorated class
+        becomes a `SessionApp` whose ``.session(key=…)`` deploys a
+        long-running daemon that holds an instance of the class as warm state
+        across turns (a Slack-bot-shaped agent).
+
+        The class should define ``on_message(self, msg: str) -> str`` (called
+        per turn) and may mark a cold-start hook with ``@enter`` (or name it
+        ``enter`` / ``setup``) to build warm state once on cold start.
+
+        Parameters mirror `function()`, minus ``entry`` (a Session's entry
+        is always the generated daemon shim).
+        """
+        resolved_image = image if isinstance(image, Image) else Image(name=str(image))
+        spec = SessionSpec(
+            image=resolved_image,
+            memory=memory,
+            cpu=cpu,
+            bundle=bundle,
+            secrets=tuple(secrets),
+            egress=egress,
+            env=dict(env) if env is not None else {},
+            timeout_s=float(timeout),
+            code_stage=code_stage,
+            stage_mounts=tuple(stage_mounts),
+        )
+
+        def decorator(cls: type) -> SessionApp:
+            sess = SessionApp(cls, spec, self)
+            self._sessions[getattr(cls, "__name__", "Session")] = sess
+            return sess
+
+        return decorator
+
+    def resolve(self, fn_name: str | None = None) -> Function:
+        """Look up a registered function — the CLI's entry point into an App.
+
+        Parameters
+        ----------
+        fn_name:
+            Name of the registered function.  ``None`` (default) uses the
+            first registered function.
+
+        Raises
+        ------
+        ValueError:
+            If no functions are registered, or ``fn_name`` is unknown.
+        """
+        if not self._functions:
+            raise ValueError(f"App {self._name!r} has no registered functions")
+        if fn_name is None:
+            return next(iter(self._functions.values()))
+        func = self._functions.get(fn_name)
+        if func is None:
+            known = list(self._functions)
+            raise ValueError(
+                f"unknown function {fn_name!r} in App {self._name!r}; registered: {known}"
+            )
+        return func
+
+    def deploy_spec(self, fn_name: str | None = None, **env_overrides: str) -> DeploySpec:
+        """The `DeploySpec` for a registered function."""
+        return self.resolve(fn_name).deploy_spec(**env_overrides)
+
+    @property
+    def functions(self) -> Mapping[str, Function]:
+        """The registered functions, by name."""
+        return dict(self._functions)