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

snowflake/sandbox/warm_session.py

@@ -0,0 +1,700 @@
+"""Warm session engine — the internal machinery behind ``Sandbox.send``.
+
+Deploy a project's daemon **detached** (a long-running loop, never
+run-to-completion), get a `Session` handle back, and drive it with
+``send()`` — a round-trip that injects a message into the live container and
+reads the daemon's reply. It is the stateful sibling of the detached run
+(``jobs.py``): a detached run executes to completion and exits (stateless); a
+session stays alive across many turns with context warm in process memory.
+
+The daemon stays warm and pays for idle. ``session.id == container_id``;
+``key`` → container is tracked in a module-level dict in-process.
+
+**The mailbox protocol:**
+
+* The daemon (inside the container) runs `session_loop()`, which tails
+  ``/tmp/sandbox_mailbox.jsonl``. Each new line is
+  ``{"id": "<msg_id>", "msg": "<text>"}``. For each, the loop calls the user's
+  ``handle(msg, state)`` and prints to **stdout**:
+
+      __SANDBOX_REPLY__<msg_id>__<base64(json(reply))>
+
+* `Session.send()` mints a ``msg_id`` (uuid hex), execs a shell-quoted
+  ``printf … >> /tmp/sandbox_mailbox.jsonl`` to append the message, then polls
+  ``logs()`` until the matching ``__SANDBOX_REPLY__<msg_id>__<b64>`` line
+  appears, base64-decodes it, and returns the reply string. A timeout raises
+  `SandboxExecTimeoutError`.
+
+``state`` is a plain dict that the daemon keeps in process memory across calls —
+this is THE warm context mechanism: it stays warm across turns.
+"""
+
+from __future__ import annotations
+
+import asyncio
+import base64
+import binascii
+import json
+import shlex
+import time
+import uuid
+from collections.abc import Callable
+from contextlib import AbstractContextManager, nullcontext
+from pathlib import Path
+from typing import TYPE_CHECKING, Any
+
+from snowflake.sandbox._deploy_spec import DeploySpec
+from snowflake.sandbox._runtime._protocol import _MAILBOX_PATH, _REPLY_SENTINEL
+from snowflake.sandbox.client import get_sandbox
+from snowflake.sandbox.exceptions import SandboxError, SandboxExecTimeoutError
+from snowflake.sandbox.types import TERMINAL_STATUSES
+
+if TYPE_CHECKING:
+    from snowflake.sandbox.config import ConnectionLike
+
+__all__ = ["Session", "agent_session", "agent_session_sync", "session_loop"]
+
+# Env var carrying the per-session reply nonce into the daemon. The daemon loop
+# reads it once at startup and unsets it before running any user handler, so the
+# untrusted ``on_message`` code cannot read it and forge an authenticated reply.
+#
+# Defined here rather than in ``_runtime._protocol`` with its siblings: no
+# generated in-container program reads this name from Python (the session shim
+# bakes the literal into its source), while ``client``/``sync`` import it from
+# this module to set the env var on create — and under mypy's
+# no-implicit-reexport a pass-through alias would not satisfy them.
+_SESSION_NONCE_ENV = "SANDBOX_SESSION_NONCE"
+
+# In-process key -> container_id tracking. A backend-side registry would let
+# create-or-resume survive process restarts; today this is per-process.
+_SESSIONS: dict[str, str] = {}
+# Parallel key -> reply-nonce map, so a same-process resume re-attaches with
+# the nonce that authenticates the daemon's replies. Lost on process restart, like
+# _SESSIONS itself; a rehydrated handle then falls back to unauthenticated matching.
+_SESSION_NONCES: dict[str, str] = {}
+
+
+def _parse_reply_sentinel(log_text: str, msg_id: str, *, nonce: str | None = None) -> str | None:
+    """Scan *log_text* for the LAST reply sentinel matching *msg_id* and return the
+    decoded reply string. Pure + unit-testable.
+
+    The reply line is ``__SANDBOX_REPLY__<msg_id>__<nonce>__<b64>`` — the daemon
+    reads the per-session ``nonce`` from its environment once at startup and
+    unsets it before running any handler, so the (untrusted) ``on_message`` code
+    cannot learn it and forge a reply. When *nonce* is given, only a sentinel
+    carrying that exact nonce is accepted. The older
+    ``__SANDBOX_REPLY__<msg_id>__<b64>`` form (a base image shipping a pre-nonce
+    SDK) is still decoded when the caller holds no nonce, but is rejected when a
+    nonce is required — its reply cannot be authenticated. ``msg_id`` and the
+    nonce are uuid hex and standard base64 has no ``_``, so ``__`` splits the
+    fields unambiguously.
+
+    Returns ``None`` if no matching well-formed sentinel is present (a malformed
+    base64/json payload is skipped, not raised, so a partial log line never masks
+    an earlier good one)."""
+    prefix = _REPLY_SENTINEL
+    found: str | None = None
+    for line in log_text.splitlines():
+        idx = line.find(prefix)
+        if idx == -1:
+            continue
+        rest = line[idx + len(prefix) :].strip()
+        parts = rest.split("__")
+        if len(parts) == 3:
+            r_msg_id, r_nonce, b64 = parts
+        elif len(parts) == 2:
+            r_msg_id, r_nonce, b64 = parts[0], "", parts[1]
+        else:
+            continue
+        if r_msg_id != msg_id or not b64:
+            continue
+        if nonce is not None and r_nonce != nonce:
+            # A reply that cannot prove it came from the trusted daemon loop.
+            continue
+        try:
+            decoded = base64.b64decode(b64, validate=True)
+            reply = json.loads(decoded)
+        except (binascii.Error, ValueError, UnicodeDecodeError):
+            continue
+        if not isinstance(reply, str):
+            continue
+        found = reply
+    return found
+
+
+def _mailbox_append_cmd(msg_id: str, message: str) -> list[str]:
+    """Build the shell argv that appends one mailbox line for *msg_id*/*message*.
+
+    The line is ``{"id": "<msg_id>", "msg": "<text>"}`` (compact JSON). The JSON
+    is shell-quoted with `shlex.quote()` and emitted with ``printf '%s\\n'``
+    so embedded quotes/newlines/backslashes survive the round-trip into the
+    container's mailbox file."""
+    payload = json.dumps({"id": msg_id, "msg": message}, separators=(",", ":"))
+    quoted = shlex.quote(payload)
+    return [
+        "/bin/bash",
+        "-c",
+        f"printf '%s\\n' {quoted} >> {_MAILBOX_PATH}",
+    ]
+
+
+def _ended_before_reply(status: str, msg_id: str) -> str:
+    """The message for a container that went terminal mid-round-trip."""
+    return (
+        f"the sandbox ended (status {status!r}) before replying to message {msg_id}. "
+        f"The daemon is gone, so this is not a timeout and retrying send() on this "
+        f"session will not help -- start a new one with force_fresh=True, or check "
+        f"the container's logs for why it exited."
+    )
+
+
+class Session:
+    """A handle to a warm session.
+
+    ``id == container_id``. Construct via `agent_session()` (which
+    deploys the daemon detached) or `from_id()` to rehydrate later.
+    """
+
+    def __init__(
+        self,
+        container_id: str,
+        key: str = "",
+        *,
+        nonce: str | None = None,
+        connection: ConnectionLike | None = None,
+    ) -> None:
+        self.id = container_id
+        self.key = key
+        # The Snowflake connection this handle reads its container over. Every method
+        # re-fetches by id (``get_sandbox``), and a handle rehydrated in a new process
+        # has no ambient connection to inherit.
+        self._connection = connection
+        # Per-session reply nonce: set by `agent_session` at create time and
+        # held only in this process, so `send()` can authenticate that a reply came
+        # from the trusted daemon loop rather than a forged stdout line. A handle
+        # rehydrated by `from_id()` has none — its replies cannot be authenticated.
+        self.nonce = nonce
+
+    @classmethod
+    def from_id(
+        cls, container_id: str, key: str = "", *, connection: ConnectionLike | None = None
+    ) -> Session:
+        """Rehydrate a `Session` handle from a container id.
+
+        The per-session reply nonce lives only in the process that created the
+        session, so a rehydrated handle has none: `send()` matches replies by
+        message id but cannot authenticate them.
+
+        ``connection`` names the Snowflake connection the session lives in. Omit it to
+        use the enclosing ``using()`` block, else the default connection."""
+        return cls(container_id, key=key, connection=connection)
+
+    async def send(self, message: str, timeout: float = 120.0, poll_s: float = 1.0) -> str:
+        """Send *message* to the daemon and return its reply (the mailbox
+        round-trip).
+
+        Mints a ``msg_id``, execs a shell-quoted append into
+        ``/tmp/sandbox_mailbox.jsonl`` inside the live container, then polls
+        ``logs()`` until the daemon prints ``__SANDBOX_REPLY__<msg_id>__<b64>``.
+        Base64-decodes the payload and returns the reply string.
+
+        Raises `SandboxError` if the container reaches a terminal state without
+        replying, and `SandboxExecTimeoutError` if it is still alive but silent for
+        *timeout* seconds."""
+        msg_id = uuid.uuid4().hex
+        sb = await get_sandbox(self.id, connection=self._connection)
+        await sb.exec(_mailbox_append_cmd(msg_id, message))
+
+        deadline = time.monotonic() + timeout
+        while True:
+            log_text = await sb.logs()
+            reply = _parse_reply_sentinel(log_text, msg_id, nonce=self.nonce)
+            if reply is not None:
+                return reply
+            # A dead daemon otherwise costs the whole timeout and is then reported as
+            # one, which points the caller at their message rather than at the gone
+            # container. refresh(), not `sb.status`: that is cached from create time
+            # and nothing on the container's lifecycle mutates it (see LogStream).
+            if await sb.refresh() in TERMINAL_STATUSES:
+                # One last read first -- the daemon may have replied just before
+                # exiting, and that reply is the answer, not an error.
+                reply = _parse_reply_sentinel(await sb.logs(), msg_id, nonce=self.nonce)
+                if reply is not None:
+                    return reply
+                raise SandboxError(_ended_before_reply(sb.status, msg_id))
+            if time.monotonic() + poll_s >= deadline:
+                raise SandboxExecTimeoutError(f"no reply for message {msg_id} within {timeout}s")
+            await asyncio.sleep(poll_s)
+
+    def send_sync(self, message: str, timeout: float = 120.0, poll_s: float = 1.0) -> str:
+        """Send *message* to the daemon and return its reply — synchronous
+        counterpart of `send`.
+
+        Same mailbox round-trip as `send`, blocking instead of awaiting: mints a
+        ``msg_id``, execs the shell-quoted append into
+        ``/tmp/sandbox_mailbox.jsonl`` inside the live container (via the
+        synchronous `Sandbox`), then polls ``logs()`` until the daemon prints
+        ``__SANDBOX_REPLY__<msg_id>__<b64>``, base64-decodes it, and returns the
+        reply string. The request built and the reply parsed are byte-identical
+        to `send`'s — both go through the shared ``_mailbox_append_cmd`` /
+        ``_parse_reply_sentinel``. Safe to call with no running event loop and
+        never starts one, so it does not deadlock inside a notebook.
+
+        Raises `SandboxError` if the container reaches a terminal state without
+        replying, and `SandboxExecTimeoutError` if it is still alive but silent for
+        *timeout* seconds.
+
+        Example:
+            session = TriageSession.session_sync(key="triage")
+            reply = session.send_sync("what changed in prod yesterday?")
+        """
+        from snowflake.sandbox.sync_client import get_sandbox as _get_sandbox_sync
+
+        msg_id = uuid.uuid4().hex
+        sb = _get_sandbox_sync(self.id, connection=self._connection)
+        sb.exec(_mailbox_append_cmd(msg_id, message))
+
+        deadline = time.monotonic() + timeout
+        while True:
+            log_text = sb.logs()
+            reply = _parse_reply_sentinel(log_text, msg_id, nonce=self.nonce)
+            if reply is not None:
+                return reply
+            # As in `send`: a terminal container is not a timeout, and refresh() is
+            # the only live read of its status.
+            if sb.refresh() in TERMINAL_STATUSES:
+                reply = _parse_reply_sentinel(sb.logs(), msg_id, nonce=self.nonce)
+                if reply is not None:
+                    return reply
+                raise SandboxError(_ended_before_reply(sb.status, msg_id))
+            if time.monotonic() + poll_s >= deadline:
+                raise SandboxExecTimeoutError(f"no reply for message {msg_id} within {timeout}s")
+            time.sleep(poll_s)
+
+    async def status_async(self) -> str:
+        """The lifecycle status from a real server read: ``cold`` | ``running`` |
+        ``ended``.
+
+        ``ended`` when the container is terminal (``dead`` or ``failed``); otherwise ``running``
+        (a daemon the platform has checkpointed still reports ``running`` — it is
+        restored warm on the next ``send()``), or ``cold`` before an id exists.
+        Safe from async code,
+        including from inside a running event loop; prefer it over the sync
+        `status` in any ``async def``."""
+        sb = await get_sandbox(self.id, connection=self._connection)
+        if sb.status in TERMINAL_STATUSES:
+            return "ended"
+        return "running" if sb._id else "cold"
+
+    def status(self) -> str:
+        """The lifecycle status from a real server read: ``cold`` | ``running`` |
+        ``ended`` (``ended`` = the container is terminal: ``dead`` or ``failed``).
+
+        A synchronous convenience wrapper around `status_async`: it drives the
+        read on a fresh event loop, so it may be called **only when no event loop
+        is running**. Called from *inside* a running loop it cannot block on the
+        network read — and fabricating a status from the id alone can never observe
+        ``ended`` (it would report a dead session as ``running``, the silent-wrong
+        result this method must not produce) — so it raises `SandboxError` instead.
+        From async code, ``await session.status_async()``."""
+        try:
+            asyncio.get_running_loop()
+        except RuntimeError:
+            return asyncio.run(self.status_async())
+        raise SandboxError(
+            "Session.status() cannot read the session status from inside a running "
+            "event loop: it would have to block on a network read, and reporting a "
+            "status from the id alone can never return 'ended' for a dead session. "
+            "Use `await session.status_async()` from async code."
+        )
+
+    async def end(self) -> None:
+        """End the Session by destroying its container.
+
+        Also drops the key→container mapping so a later ``agent_session(key=…)``
+        cold-starts a fresh daemon rather than resuming a dead one."""
+        sb = await get_sandbox(self.id, connection=self._connection)
+        await sb.destroy()
+        if self.key and _SESSIONS.get(self.key) == self.id:
+            _SESSIONS.pop(self.key, None)
+            _SESSION_NONCES.pop(self.key, None)
+
+    def end_sync(self) -> None:
+        """End the Session by destroying its container — synchronous counterpart
+        of `end`.
+
+        Blocks instead of awaiting; the same teardown as `end`, and it likewise
+        drops the key→container mapping so a later ``agent_session(key=…)`` /
+        `agent_session_sync()` cold-starts a fresh daemon rather than resuming a
+        dead one. Safe to call with no running event loop and never starts one.
+
+        Example:
+            session = TriageSession.session_sync(key="triage")
+            session.end_sync()
+        """
+        from snowflake.sandbox.sync_client import get_sandbox as _get_sandbox_sync
+
+        sb = _get_sandbox_sync(self.id, connection=self._connection)
+        sb.destroy()
+        if self.key and _SESSIONS.get(self.key) == self.id:
+            _SESSIONS.pop(self.key, None)
+            _SESSION_NONCES.pop(self.key, None)
+
+
+def _layer_shim(spec: DeploySpec, shim_dir: Path) -> DeploySpec:
+    """Return *spec* with the generated shim files copied into its bundle root.
+
+    ``@app.session`` writes its mailbox-loop shim to a temp dir; the daemon also
+    needs the user's module. Copying the shim into a throwaway copy of the bundle
+    root keeps both without mutating the user's tree.
+    """
+    import shutil
+    import tempfile
+    from dataclasses import replace
+
+    from snowflake.sandbox._assemble import _DEFAULT_EXCLUDES, Bundle, _copy_tree_into
+
+    root = Path(spec.bundle.root) if spec.bundle is not None else shim_dir
+    exclude = list(spec.bundle.exclude) if spec.bundle is not None else []
+    staged = Path(tempfile.mkdtemp(prefix="sbx-session-"))
+    if root.is_dir() and root != shim_dir:
+        # Route the user's tree through the credential + symlink filters: a bare
+        # copytree would ship `.env` / `id_rsa` / `~/.aws` / `*.pem` and ignore the
+        # bundle's own excludes. _copy_tree_into applies both.
+        _copy_tree_into(root, staged, list(_DEFAULT_EXCLUDES) + exclude)
+    for f in shim_dir.iterdir():
+        if f.is_file():
+            shutil.copy2(f, staged / f.name)
+    include = list(spec.bundle.include) if spec.bundle is not None else []
+    return replace(
+        spec, bundle=Bundle(root=str(staged), include=tuple(include), exclude=tuple(exclude))
+    )
+
+
+def _connection_scope(connection: ConnectionLike | None) -> AbstractContextManager[object]:
+    """``using(connection)`` when one was named, else a no-op.
+
+    Lets an entry point accept ``connection=`` and have it apply to a whole pipeline of
+    calls that take no connection argument themselves.
+    ``connection`` picks the Snowflake connection for this call; the handle carries it
+    onward. See `config.using` for the resolution rules (names, live connections,
+    memoisation, and how an enclosing ``using()`` block interacts).
+
+    """
+    if connection is None:
+        return nullcontext()
+    from snowflake.sandbox.config import using
+
+    return using(connection)
+
+
+async def agent_session(
+    spec: DeploySpec,
+    *,
+    key: str,
+    shim_dir: str | Path | None = None,
+    force_fresh: bool = False,
+    timeout: float = 3600.0,
+    connection: ConnectionLike | None = None,
+    git_repo: str | None = None,
+    git_ref: str | None = None,
+    git_subdir: str | None = None,
+) -> Session:
+    """Create-or-resume a keep-alive Session keyed by *key*.
+
+    Deploys the project's daemon **detached** — reusing ``deploy_spec(...,
+    detach=True)`` internals exactly like ``deploy_async`` reuses the bundle
+    path: the spec's ``entry`` is the user's long-running loop (it should call
+    `session_loop()`), and the container's main process is that loop. Returns a
+    `Session` whose ``id`` is the daemon's container id.
+
+    **Create-or-resume by ``key``:** an in-process ``dict[str, str]`` tracks
+    ``key → container_id``. If *key* is already
+    tracked and that container is alive, the existing Session is returned
+    (resume); otherwise a fresh daemon is deployed and ``key → id`` recorded.
+    Set ``force_fresh=True`` to bypass resume and always cold-start.
+    ``connection`` picks the Snowflake connection for this call; the handle carries it
+    onward. See `config.using` for the resolution rules (names, live connections,
+    memoisation, and how an enclosing ``using()`` block interacts).
+
+    """
+    # --- resume path: reuse a live daemon already tracked for this key ---
+    if not force_fresh:
+        existing_id = _SESSIONS.get(key)
+        if existing_id is not None:
+            try:
+                # Same connection the caller named: checking liveness on the DEFAULT
+                # connection instead would raise for a container that is alive on
+                # theirs, cold-starting a second daemon and leaking the first.
+                sb = await get_sandbox(existing_id, connection=connection)
+            except SandboxError:
+                sb = None
+            if sb is not None and sb.status not in TERMINAL_STATUSES:
+                return Session(
+                    existing_id, key=key, nonce=_SESSION_NONCES.get(key), connection=connection
+                )
+            # Stale/dead -> drop it and cold-start below.
+            _SESSIONS.pop(key, None)
+            _SESSION_NONCES.pop(key, None)
+
+    # --- create path: deploy the daemon detached ---
+    # Imported here (not at module import) to keep the lazy __init__ contract:
+    # bare ``import snowflake.sandbox`` must not pull deploy.py / httpx.
+    from snowflake.sandbox.deploy import deploy_spec
+
+    # The generated daemon shim lives in shim_dir; layer it over the declared
+    # bundle so the user's module stays importable alongside it.
+    if shim_dir is not None:
+        spec = _layer_shim(spec, Path(shim_dir))
+
+    # Per-session reply nonce: delivered as the SANDBOX_SESSION_NONCE env var
+    # the daemon loop reads once at startup and unsets before running any handler,
+    # so an untrusted ``on_message`` cannot forge an authenticated reply. Held
+    # in-process so ``send()`` can verify it.
+    #
+    # It rides the `sandbox_env` PLATFORM channel (deploy_spec -> from_local ->
+    # Sandbox._add_platform_env), NOT the user `env` map: SANDBOX_ is a
+    # server-reserved prefix, and the create path rejects (400) a reserved-prefix
+    # key on the `env` map. Enforcement is unchanged: the container still receives
+    # SANDBOX_SESSION_NONCE in its environment (the server merges sandbox_env into
+    # the app env), and session_loop still pops it before any handler runs.
+    nonce = uuid.uuid4().hex
+
+    # ``connection=`` has to reach the whole deploy pipeline (deploy_spec -> from_local
+    # -> create), not just this function, and threading a kwarg through every layer of
+    # it would mean every future layer remembering to forward it. A ``using()`` block
+    # carries it through code that takes no connection argument at all -- the case the
+    # scope exists for. ``contextlib.nullcontext`` when nothing was named, so the
+    # ambient/enclosing binding is left exactly as it was.
+    with _connection_scope(connection):
+        result = await deploy_spec(
+            spec,
+            git_repo=git_repo,
+            git_ref=git_ref,
+            git_subdir=git_subdir,
+            detach=True,
+            force_fresh=force_fresh,
+            timeout=timeout,
+            platform_env={_SESSION_NONCE_ENV: nonce},
+        )
+    if not result.deployed or not result.sandbox_id:
+        problems = "; ".join(p.message for p in result.problems if p.severity == "error")
+        raise SandboxError(
+            "agent_session failed to deploy the daemon" + (f": {problems}" if problems else "")
+        )
+
+    _SESSIONS[key] = result.sandbox_id
+    _SESSION_NONCES[key] = nonce
+    return Session(result.sandbox_id, key=key, nonce=nonce, connection=connection)
+
+
+def agent_session_sync(
+    spec: DeploySpec,
+    *,
+    key: str,
+    shim_dir: str | Path | None = None,
+    force_fresh: bool = False,
+    timeout: float = 3600.0,
+    connection: ConnectionLike | None = None,
+    git_repo: str | None = None,
+    git_ref: str | None = None,
+    git_subdir: str | None = None,
+) -> Session:
+    """Create-or-resume a keep-alive Session keyed by *key* — synchronous
+    counterpart of `agent_session`.
+
+    Same create-or-resume contract, blocking instead of awaiting: an in-process
+    ``key -> container_id`` map is consulted first (resume a live daemon), and
+    otherwise the daemon is deployed detached via the synchronous deploy core.
+    It shares the same `_SESSIONS` registry as `agent_session`, so a session
+    opened on either side resumes from the other. Returns a `Session` whose
+    ``id`` is the daemon's container id; drive it with ``s.send_sync(msg)``. Safe
+    with no running event loop and never starts one.
+
+    ``connection`` picks the Snowflake connection for this call; the handle carries it
+    onward. See `config.using` for the resolution rules (names, live connections,
+    memoisation, and how an enclosing ``using()`` block interacts).
+
+    Example:
+        session = agent_session_sync(app_session.deploy_spec(), key="triage")
+        reply = session.send_sync("hello")
+    """
+    from snowflake.sandbox.sync_client import get_sandbox as _get_sandbox_sync
+
+    # --- resume path: reuse a live daemon already tracked for this key ---
+    if not force_fresh:
+        existing_id = _SESSIONS.get(key)
+        if existing_id is not None:
+            try:
+                # Same connection the caller named — see the async twin.
+                sb = _get_sandbox_sync(existing_id, connection=connection)
+            except SandboxError:
+                sb = None
+            if sb is not None and sb.status not in TERMINAL_STATUSES:
+                return Session(
+                    existing_id, key=key, nonce=_SESSION_NONCES.get(key), connection=connection
+                )
+            # Stale/dead -> drop it and cold-start below.
+            _SESSIONS.pop(key, None)
+            _SESSION_NONCES.pop(key, None)
+
+    # --- create path: deploy the daemon detached (sync deploy core) ---
+    # Imported lazily (not at module scope) to keep the lazy __init__ contract:
+    # a bare ``import snowflake.sandbox`` must not pull deploy.py / httpx.
+    from snowflake.sandbox.deploy import deploy_spec_sync
+
+    if shim_dir is not None:
+        spec = _layer_shim(spec, Path(shim_dir))
+
+    # Per-session reply nonce: rides the `sandbox_env` PLATFORM channel, not the user
+    # `env` map (SANDBOX_ is a server-reserved prefix the create path rejects on `env`).
+    # See `agent_session` for the full rationale.
+    nonce = uuid.uuid4().hex
+
+    # ``connection=`` has to reach the whole deploy pipeline (deploy_spec -> from_local
+    # -> create), not just this function, and threading a kwarg through every layer of
+    # it would mean every future layer remembering to forward it. A ``using()`` block
+    # carries it through code that takes no connection argument at all -- the case the
+    # scope exists for. ``contextlib.nullcontext`` when nothing was named, so the
+    # ambient/enclosing binding is left exactly as it was.
+    with _connection_scope(connection):
+        result = deploy_spec_sync(
+            spec,
+            git_repo=git_repo,
+            git_ref=git_ref,
+            git_subdir=git_subdir,
+            detach=True,
+            force_fresh=force_fresh,
+            timeout=timeout,
+            platform_env={_SESSION_NONCE_ENV: nonce},
+        )
+    if not result.deployed or not result.sandbox_id:
+        problems = "; ".join(p.message for p in result.problems if p.severity == "error")
+        raise SandboxError(
+            "agent_session failed to deploy the daemon" + (f": {problems}" if problems else "")
+        )
+
+    _SESSIONS[key] = result.sandbox_id
+    _SESSION_NONCES[key] = nonce
+    return Session(result.sandbox_id, key=key, nonce=nonce, connection=connection)
+
+
+def session_loop(
+    handle: Callable[[str, dict[str, Any]], str],
+    *,
+    mailbox_path: str = _MAILBOX_PATH,
+    poll_s: float = 0.5,
+) -> None:
+    """Run the daemon-side message loop (the in-container half of a warm agent).
+
+    The daemon's entry program imports this so you don't hand-write the
+    mailbox loop:
+
+        from snowflake.sandbox import session_loop
+        session_loop(lambda msg, state: ...)
+
+    *handle* is ``handle(msg, state) -> reply``: *msg* is the incoming message
+    text, *state* is a plain dict that **persists across calls in process
+    memory** (THE warm context). The reply string is returned to the controller.
+
+    The loop tails *mailbox_path* for new ``{"id", "msg"}`` lines; for each it
+    calls ``handle(msg, state)`` and prints
+    ``__SANDBOX_REPLY__<id>__<base64(json(reply))>`` to **stdout** (flushed). It
+    runs forever and is robust to malformed lines (a bad line is skipped, not
+    fatal) and to handler exceptions (the error string is returned as the reply
+    so ``send()`` doesn't hang).
+
+    **Surviving the platform's idle checkpoint.** Snowflake's idle auto-suspend
+    (``runsc checkpoint``/``restore``) sends **SIGUSR1** to the app immediately
+    *before* the memory checkpoint and **SIGUSR2** immediately *after* restore.
+    SIGUSR1's default action is *terminate*, so this loop sets both to
+    ``SIG_IGN`` — without that a checkpoint kills the daemon. In-process *memory*
+    (``state``, counters, loaded objects) survives the checkpoint automatically;
+    open network sockets (Cortex/Snowflake/TCP) go stale and must be reopened
+    lazily by the handler. Ignoring the signals is harmless when the platform
+    never checkpoints — they simply never arrive.
+    """
+    import os
+    import signal
+
+    # Per-session reply nonce: read once and remove it from the environment
+    # before any handler runs, so the (untrusted) handler cannot read it and forge
+    # an authenticated reply. Empty when the caller's SDK predates the nonce. Uses
+    # the module constant (not a literal) so it stays in sync with the producer.
+    _reply_nonce = os.environ.pop(_SESSION_NONCE_ENV, "")
+
+    state: dict[str, Any] = {}
+
+    # Ignore the SIGUSR1 (pre-checkpoint) / SIGUSR2 (post-restore) signals the
+    # platform's idle auto-suspend sends: SIGUSR1's default action is terminate, so
+    # a daemon that installs nothing is killed the first time it is checkpointed.
+    # Nothing is wired to them beyond survival — every reply is already printed
+    # with flush=True, so no output is buffered across the checkpoint. Must run on
+    # the main thread (Python requirement); if the loop is driven off-thread we
+    # degrade to the default disposition rather than crash.
+    try:
+        signal.signal(signal.SIGUSR1, signal.SIG_IGN)
+        signal.signal(signal.SIGUSR2, signal.SIG_IGN)
+    except (ValueError, OSError, AttributeError):
+        # Off the main thread, or a platform without SIGUSR1/2 — checkpoint
+        # survival is unavailable, but the keep-alive loop still works.
+        pass
+
+    # Track how many bytes we've consumed so we only process new lines, even
+    # across file truncation/rotation.
+    offset = 0
+    # Ensure the mailbox exists so the first send's append has a target and our
+    # read doesn't race a missing file.
+    try:
+        open(mailbox_path, "a").close()
+    except OSError:
+        pass
+
+    while True:
+        try:
+            size = os.path.getsize(mailbox_path)
+        except OSError:
+            time.sleep(poll_s)
+            continue
+        if size < offset:
+            # File was truncated/rotated — restart from the top.
+            offset = 0
+        if size == offset:
+            time.sleep(poll_s)
+            continue
+        # Read in BINARY so `offset` is a true byte count: it is compared against
+        # os.path.getsize() (bytes) and seeked with, and a text-mode tell() is an
+        # opaque cookie that must not be offset-arithmetic'd (doing so seeks to an
+        # invalid position on some CPython builds).
+        with open(mailbox_path, "rb") as fh:
+            fh.seek(offset)
+            chunk = fh.read()
+        offset += len(chunk)
+        # A trailing line with no newline is a partial append still in flight (a
+        # large `printf ... >> mailbox` is not always one atomic write). Rewind the
+        # byte offset past it and drop it from this batch, so it is re-read whole
+        # next poll rather than split into two fragments that each fail json.loads
+        # and silently lose the message.
+        if chunk and not chunk.endswith(b"\n"):
+            partial = chunk.rpartition(b"\n")[2]
+            offset -= len(partial)
+            chunk = chunk[: len(chunk) - len(partial)]
+        new_lines = chunk.decode("utf-8", "replace").splitlines()
+        for line in new_lines:
+            line = line.strip()
+            if not line:
+                continue
+            try:
+                rec = json.loads(line)
+                msg_id = rec["id"]
+                msg = rec["msg"]
+            except (ValueError, KeyError, TypeError):
+                continue  # malformed line — skip, never fatal
+            try:
+                reply = handle(msg, state)
+            except Exception as exc:  # noqa: BLE001 — never let one turn kill the daemon
+                reply = f"handler error: {type(exc).__name__}: {exc}"
+            if not isinstance(reply, str):
+                reply = str(reply)
+            b64 = base64.b64encode(json.dumps(reply).encode()).decode()
+            print(f"{_REPLY_SENTINEL}{msg_id}__{_reply_nonce}__{b64}", flush=True)