snowflake-sandbox-python 0.2.2a3__py3-none-any.whl → 0.2.2a5__py3-none-any.whl

snowflake/cli_sandbox/_egress_flags.py

@@ -106,7 +106,12 @@ def build_egress(
 
     egress = Egress(
         external_access_integrations=integrations,
-        **({"allow_default_egress": False} if no_default_egress else {}),
+        # Uses the current parameter name. The CLI FLAG stays `--no-default-egress`: it is
+        # released, renaming it would break users' scripts, and that is a separate decision
+        # from the library parameter's name. So the flag spelling and the parameter spelling
+        # differ here deliberately — the alternative is the CLI depending on a deprecated
+        # parameter, which is worse.
+        **({"allow_public_package_registries": False} if no_default_egress else {}),
     )
     # Validate now, not at create time. `compile_egress` is pure and is what the client
     # runs anyway -- but it runs *after* connecting, so a mistyped host pair otherwise

snowflake/cli_sandbox/_lifecycle_command.py

@@ -0,0 +1,108 @@
+"""``snow sandbox suspend`` and ``snow sandbox resume``.
+
+The CLI counterparts to ``Sandbox.suspend()`` / ``Sandbox.resume()``. Snowflake
+suspends an idle sandbox on its own timer; these let an operator bank a sandbox's
+state on demand, and wake one without running anything in it.
+"""
+
+from __future__ import annotations
+
+from typing import Any
+
+import typer
+
+# From _common, where these are DEFINED. Importing them via _get_command (which
+# re-exports them) is an implicit re-export and mypy rejects it.
+from snowflake.cli_sandbox._common import (
+    _CONN_HELP,
+    _USE_SNOW_CONN_HELP,
+    _apply_connection,
+    _resolve_sandbox_id,
+)
+
+
+def _suspend_command(
+    sandbox: str = typer.Argument(..., help="Sandbox name (or id) to suspend."),
+    use_snow_connection: bool = typer.Option(
+        True,
+        "--use-snow-connection/--no-snow-connection",
+        help=_USE_SNOW_CONN_HELP,
+    ),
+    connection: str | None = typer.Option(
+        None,
+        "--connection",
+        "-c",
+        help=_CONN_HELP,
+    ),
+    **options: Any,
+) -> None:
+    """Suspend a sandbox: snapshot its filesystem and tear the container down.
+
+    The sandbox keeps its files; the next thing that touches it brings it back on
+    that snapshot with a fresh process tree. Already-suspended is success.
+
+    Example:
+        snow sandbox suspend SANDBOX
+    """
+    from snowflake.cli_sandbox._adapter import sdk_connection
+    from snowflake.sandbox.exceptions import SandboxError
+    from snowflake.sandbox.sync_client import get_sandbox
+
+    _apply_connection(use_snow_connection, connection)
+    try:
+        # The SYNC client: AsyncSandbox.connect is itself a coroutine, so a bare
+        # asyncio.run(sb.suspend()) on its result would fail on a coroutine object.
+        # get_sandbox(), not Sandbox.connect(): _resolve_sandbox_id already turned a
+        # name into a cntr_ id, and connect() resolves its argument as a NAME -- so
+        # feeding it an id fails with "no sandbox named cntr_...".
+        #
+        # The SYNC client throughout: AsyncSandbox.connect is itself a coroutine, so a
+        # bare asyncio.run(sb.suspend()) on its result fails on a coroutine object.
+        get_sandbox(_resolve_sandbox_id(sandbox), connection=sdk_connection()).suspend()
+    except SandboxError as exc:
+        typer.echo(f"error: {exc}", err=True)
+        raise typer.Exit(1) from exc
+    typer.echo(f"suspended {sandbox}")
+
+
+def _resume_command(
+    sandbox: str = typer.Argument(..., help="Sandbox name (or id) to resume."),
+    use_snow_connection: bool = typer.Option(
+        True,
+        "--use-snow-connection/--no-snow-connection",
+        help=_USE_SNOW_CONN_HELP,
+    ),
+    connection: str | None = typer.Option(
+        None,
+        "--connection",
+        "-c",
+        help=_CONN_HELP,
+    ),
+    **options: Any,
+) -> None:
+    """Wake a suspended sandbox without running anything in it.
+
+    Rarely needed -- anything that touches a sandbox resumes it anyway -- so this
+    is for paying the cold start up front. Returns once the resume is under way,
+    not once the sandbox is serving.
+
+    Example:
+        snow sandbox resume SANDBOX
+    """
+    from snowflake.cli_sandbox._adapter import sdk_connection
+    from snowflake.sandbox.exceptions import SandboxError
+    from snowflake.sandbox.sync_client import get_sandbox
+
+    _apply_connection(use_snow_connection, connection)
+    try:
+        # get_sandbox(), not Sandbox.connect(): _resolve_sandbox_id already turned a
+        # name into a cntr_ id, and connect() resolves its argument as a NAME.
+        get_sandbox(_resolve_sandbox_id(sandbox), connection=sdk_connection()).resume()
+    except SandboxError as exc:
+        typer.echo(f"error: {exc}", err=True)
+        raise typer.Exit(1) from exc
+    typer.echo(f"resuming {sandbox}")
+
+
+_SUSPEND_HELP = "Suspend a sandbox, snapshotting its filesystem.\n\nExample:\n    snow sandbox suspend SANDBOX\n"
+_RESUME_HELP = "Resume a suspended sandbox.\n\nExample:\n    snow sandbox resume SANDBOX\n"

snowflake/cli_sandbox/_run_command.py

@@ -312,7 +312,13 @@ def _print_inferred_plan(plan: _RunPlan) -> None:
     typer.echo(f"command:    {plan.command or '(image default)'}")
     if plan.egress is not None:
         eais = getattr(plan.egress, "external_access_integrations", ())
-        closed = getattr(plan.egress, "allow_default_egress", None) is False
+        # Read the RESOLVER, not any one spelling. There are three accepted names for this
+        # control and reading a specific field goes silently blind when a caller (or this
+        # CLI) uses a different one — which is exactly what happened: switching the flag's
+        # construction to allow_public_package_registries made this summary stop reporting
+        # "default-egress=closed" while the sandbox was still correctly closed. A dry run
+        # that under-reports the egress posture is worse than no dry run.
+        closed = getattr(plan.egress, "default_egress_allowed", None) is False
         bits = [f"eai={','.join(eais)}" if eais else "", "default-egress=closed" if closed else ""]
         if plan.secrets:
             bits.append(f"secrets={len(plan.secrets)}")

snowflake/cli_sandbox/_upload_plan.py

@@ -2,7 +2,7 @@
 
 The *selection* engine (the walk, the skip reasons, `UploadPlan`) lives in
 `snowflake.sandbox._upload_plan`, because the library needs it too: `Sandbox.upload_dir`
-and the `from_local` bundle are the same decision about the same trees. What stays here
+and create-time bundle assembly are the same decision about the same trees. What stays here
 is the part that is genuinely about a command line — turning argv tokens into targets,
 and the errors that name CLI flags in their remediation.
 

snowflake/cli_sandbox/commands.py

@@ -26,6 +26,12 @@ from snowflake.cli_sandbox._common import (
     _resolve_sandbox_id,
 )
 from snowflake.cli_sandbox._get_command import _CANCEL_HELP, _cancel_command, get
+from snowflake.cli_sandbox._lifecycle_command import (
+    _RESUME_HELP,
+    _SUSPEND_HELP,
+    _resume_command,
+    _suspend_command,
+)
 from snowflake.cli_sandbox._run_command import _RUN_HELP, _run_command
 from snowflake.cli_sandbox._shell_command import (  # noqa: F401  (re-export: import surface)
     _CREATE_RETRY_ATTEMPTS,
@@ -279,10 +285,15 @@ async def _stream_process_logs(sb: Any, tail: int) -> None:
 
 _LOGS_HELP = (
     "Print or follow a sandbox's output.\n\n"
-    "By default this is the captured stdout/stderr of the sandbox's managed process, "
-    "which only exists for command containers (created with `command=`). Snapshot by "
-    "default; --follow tails the SAME source live until the process finishes, printing "
-    "stdout to stdout and stderr to stderr so `2>` separates them.\n\n"
+    "By default this is the server's `all`: the container's own output plus the managed "
+    "process's, where a command container has one. Snapshot by "
+    "default; --follow instead tails the MANAGED process live until it finishes, printing "
+    "stdout to stdout and stderr to stderr so `2>` separates them. That is the only source "
+    "it can follow, which is why --source is snapshot-only.\n\n"
+    "--source reads a different captured stream: `exec` is what `exec()` calls printed, "
+    "kept per producer by guest-init; `base` is the container's own output; `all` is base "
+    "+ managed, and deliberately NOT exec, since an exec's output is already returned by "
+    "the call that produced it. --pid narrows `exec` to one producer.\n\n"
     "--file reads a file inside the sandbox instead. That is how you read a "
     "`run --detach` local-code run: its command is backgrounded rather than being the "
     "managed process, so its output is tee'd to a file and the managed-process log is "
@@ -291,6 +302,8 @@ _LOGS_HELP = (
     "    snow sandbox logs SANDBOX_ID\n"
     "    snow sandbox logs SANDBOX_ID --tail 50\n"
     "    snow sandbox logs SANDBOX_ID --follow\n"
+    "    snow sandbox logs SANDBOX_ID --source exec\n"
+    "    snow sandbox logs SANDBOX_ID --source exec --pid 295\n"
     "    snow sandbox logs SANDBOX_ID --file /tmp/sandbox-run.out --follow"
 )
 
@@ -318,6 +331,22 @@ def _logs_command(
         metavar="MS",
         help="Only lines at or after this server timestamp (Unix epoch ms). Ignored with --follow.",
     ),
+    source: str | None = typer.Option(
+        None,
+        "--source",
+        metavar="WHICH",
+        help="Which output to read: all (the default when omitted: base + managed, and it does "
+        "NOT include exec), managed (just the command= process), exec (what `exec()` calls "
+        "printed), or base (the container's own stdout/stderr). Snapshot only — --follow tails "
+        "the managed process and cannot follow another source.",
+    ),
+    pid: int = typer.Option(
+        0,
+        "--pid",
+        metavar="N",
+        help="With --source exec, read only this producer instead of every exec merged "
+        "by capture time. The pid is shown per line in the merged output.",
+    ),
     file: str | None = typer.Option(
         None,
         "--file",
@@ -347,6 +376,8 @@ def _logs_command(
     Example:
         snow sandbox logs SANDBOX_ID --tail 50
         snow sandbox logs SANDBOX_ID --follow
+        snow sandbox logs SANDBOX_ID --source exec
+        snow sandbox logs SANDBOX_ID --source exec --pid 295
         snow sandbox logs SANDBOX_ID --file /tmp/sandbox-run.out
     """
     import shlex
@@ -360,6 +391,19 @@ def _logs_command(
             "--since applies to the managed-process log, not to --file (a file has no "
             "server timestamps). Use --tail to limit the output instead."
         )
+    # Each of these would otherwise be accepted and then silently ignored, which is worse
+    # than a rejection: the caller reads output and believes it is what they asked for.
+    if source is not None and file is not None:
+        raise typer.BadParameter(
+            "--source selects a captured stream; --file reads a path. Pass one."
+        )
+    if source is not None and follow and source not in ("managed", "process"):
+        raise typer.BadParameter(
+            f"--follow tails the managed process and cannot follow --source {source}. "
+            "Drop --follow for a snapshot, or drop --source to follow the managed process."
+        )
+    if pid and source != "exec":
+        raise typer.BadParameter("--pid selects one exec producer, so it needs --source exec.")
 
     _apply_connection(use_snow_connection, connection)
 
@@ -373,7 +417,7 @@ def _logs_command(
         elif follow:
             await _stream_process_logs(sb, tail)
         else:
-            typer.echo(await sb.logs(tail=tail, since_ts_ms=since))
+            typer.echo(await sb.logs(tail=tail, since_ts_ms=since, source=source, pid=pid))
 
     async def _read_sandbox_file(sb: Any, path: str, *, tail: int, follow: bool) -> None:
         """Print, or tail, a file inside the sandbox.
@@ -554,3 +598,7 @@ app.command(name="delete", requires_connection=False, help=_DELETE_HELP)(_delete
 # so these are plain top-level verbs alongside `list`/`logs`/`delete`.
 app.command(name="get", requires_connection=False)(get)
 app.command(name="cancel", requires_connection=False, help=_CANCEL_HELP)(_cancel_command)
+# The two halves of the suspend/resume pair. Snowflake suspends an idle sandbox on its
+# own timer; these make both directions available on demand.
+app.command(name="suspend", requires_connection=False, help=_SUSPEND_HELP)(_suspend_command)
+app.command(name="resume", requires_connection=False, help=_RESUME_HELP)(_resume_command)

snowflake/sandbox/__init__.py

@@ -28,6 +28,7 @@ from typing import TYPE_CHECKING, Any
 from snowflake.sandbox._version import __version__
 from snowflake.sandbox.exceptions import (
     SandboxAuthError,
+    SandboxBuildError,
     SandboxConflictError,
     SandboxContractWarning,
     SandboxCreateError,
@@ -124,6 +125,7 @@ __all__ = [
     "SandboxCreateError",
     "SandboxFailure",
     "SandboxValidationError",
+    "SandboxBuildError",
     "SandboxConflictError",
     "SandboxNotFoundError",
     "SandboxNotReadyError",

snowflake/sandbox/_bundle.py

@@ -75,7 +75,7 @@ def _collect_tree(
     include: list[str] | None = None,
 ) -> tuple[list[tuple[str, bytes]], UploadPlan]:
     """Sorted ``(relpath, content)`` for every file under *src*, plus the plan that
-    says what was left out and why. Used by ``from_local``.
+    says what was left out and why. Used by create-time bundle assembly.
 
     Selection is delegated to `_upload_plan`, which is the one walk the CLI and
     `Sandbox.upload_dir` also use. This function used to own a second, weaker walk:
@@ -123,7 +123,7 @@ def _collect_tree(
         raise SandboxValidationError(
             f"no files to bundle from {src}: all {len(plan.skipped)} path(s) were "
             f"filtered out ({_reason_summary(plan)}). Pass include= to re-admit a "
-            f"filtered path, or point from_local at the directory holding your code."
+            f"filtered path, or point create(code=...) at the directory holding your code."
         )
     return files, plan
 

snowflake/sandbox/_code_delivery.py

@@ -0,0 +1,88 @@
+"""Create-time code delivery via ``upload_dir`` (no stage / no bootstrap).
+
+Deploy and ``Sandbox.create(code=)`` create an empty sandbox, upload the local
+tree into ``CODE_DIR``, then signal a ready marker so a wait-then-exec main
+command (detach / jobs / ``create(code=, command=…)``) can start the entry.
+"""
+
+from __future__ import annotations
+
+import shlex
+from collections.abc import Sequence
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from snowflake.sandbox.client import AsyncSandbox
+    from snowflake.sandbox.sync_client import Sandbox
+
+__all__ = [
+    "CODE_DIR",
+    "READY_MARKER",
+    "cd_into_code_dir",
+    "entry_in_code_dir",
+    "signal_bundle_ready",
+    "signal_bundle_ready_sync",
+    "upload_bundle",
+    "upload_bundle_sync",
+    "wait_then_exec",
+]
+
+# Where ``upload_dir`` lands the project. Relative entries ``cd`` here before exec.
+CODE_DIR = "/tmp/_sandbox_code"
+READY_MARKER = f"{CODE_DIR}/.sandbox_bundle_ready"
+
+_CD_INTO_CODE_DIR = f"cd {shlex.quote(CODE_DIR)} && "
+# A copy of this name inside the uploaded tree must not count as "ready" —
+# the marker is written only after upload_dir returns.
+_MARKER_EXCLUDE = [".sandbox_bundle_ready", "**/.sandbox_bundle_ready"]
+
+
+def cd_into_code_dir() -> str:
+    """Shell fragment that cds into ``CODE_DIR`` and stops the command if that fails."""
+    return _CD_INTO_CODE_DIR
+
+
+def entry_in_code_dir(entry: Sequence[str]) -> list[str]:
+    """Wrap *entry* so it runs with cwd=``CODE_DIR`` after ``upload_dir``."""
+    inner = " ".join(shlex.quote(str(a)) for a in entry)
+    return ["/bin/bash", "-c", _CD_INTO_CODE_DIR + "exec " + inner]
+
+
+def wait_then_exec(entry: Sequence[str]) -> list[str]:
+    """Container main command: block until the ready marker exists, then exec *entry*.
+
+    Used when the entry must be the managed process (detach / jobs / ``create(code=)``)
+    but the bundle can only be uploaded after the sandbox is up.
+    """
+    inner = " ".join(shlex.quote(str(a)) for a in entry)
+    ready = shlex.quote(READY_MARKER)
+    body = f"while [ ! -f {ready} ]; do sleep 0.1; done; " + _CD_INTO_CODE_DIR + "exec " + inner
+    return ["/bin/bash", "-c", body]
+
+
+def _signal_script() -> str:
+    return f"mkdir -p {shlex.quote(CODE_DIR)} && touch {shlex.quote(READY_MARKER)}"
+
+
+async def upload_bundle(sb: AsyncSandbox, local_dir: str) -> None:
+    """Upload *local_dir* into ``CODE_DIR`` on a live sandbox.
+
+    The ready-marker name is excluded so a file in the tree cannot release a
+    ``wait_then_exec`` main process before the upload finishes.
+    """
+    await sb.upload_dir(local_dir, CODE_DIR, exclude=_MARKER_EXCLUDE)
+
+
+def upload_bundle_sync(sb: Sandbox, local_dir: str) -> None:
+    """Sync counterpart of `upload_bundle`."""
+    sb.upload_dir(local_dir, CODE_DIR, exclude=_MARKER_EXCLUDE)
+
+
+async def signal_bundle_ready(sb: AsyncSandbox) -> None:
+    """Write the ready marker so a `wait_then_exec` main process can proceed."""
+    await sb.exec(["/bin/bash", "-c", _signal_script()], timeout=60.0)
+
+
+def signal_bundle_ready_sync(sb: Sandbox) -> None:
+    """Sync counterpart of `signal_bundle_ready`."""
+    sb.exec(["/bin/bash", "-c", _signal_script()], timeout=60.0)

snowflake/sandbox/_diagnostics.py

@@ -127,19 +127,19 @@ def preflight_checks(spec: DeploySpec, bundle_root: Path | None = None) -> list[
         problems.append(Problem("error", str(exc)))
 
     # A secret's `host` and Snowflake reachability are different controls: the first is
-    # enforced at the credential swap, the second by allow_internet / the group
-    # booleans. A secret scoped to a host the sandbox cannot reach on Snowflake is
+    # enforced at the credential swap, the second by the egress level
+    # (allow_public_package_registries) / the group booleans. A secret scoped to a host the sandbox cannot reach on Snowflake is
     # silently useless, so say so. (eai_hosts is the SPCS perimeter and does not
     # affect Snowflake reachability, so it is deliberately not consulted here.)
     # Only warn when egress was *explicitly* closed AND no EAI was granted:
-    #   - allow_internet is tri-state: left unset the platform applies its permissive
+    #   - the egress level is tri-state: left unset the platform applies its permissive
     #     default, and actual reachability is then an Istio question we cannot answer
     #     from here — warning on it would be noise on the common path.
     #   - allow_github / allow_pypi are NOT credited here, and now never could be: both
     #     are inert, so neither makes a host reachable in any mode. The retired
     #     allowed_egress_hosts is likewise not consulted — hosts named there are ignored
     #     platform-side.
-    #   - Under allow_internet=False the only grant left is an EAI, whose network rules
+    #   - Under allow_public_package_registries=False the only grant left is an EAI, whose rules
     #     live in Snowflake and cannot be resolved from here — so an EAI being present is
     #     taken as "the caller has said how" (loop-invariant, hence hoisted out below).
     eg = spec.egress
@@ -153,13 +153,15 @@ def preflight_checks(spec: DeploySpec, bundle_root: Path | None = None) -> list[
                 hint = (
                     f"grant {h!r} with an External Access Integration and pass "
                     f'external_access_integrations=("MY_EAI",) — under '
-                    f"allow_internet=False that is the only grant the platform applies"
+                    f"allow_public_package_registries=False that is the only grant the "
+                    f"platform applies"
                 )
                 problems.append(
                     Problem(
                         "warn",
                         f"secret {s.fqn!r} is scoped to host {h!r}, which the sandbox cannot "
-                        f"reach with allow_internet=False — the credential could never be used",
+                        f"reach with allow_public_package_registries=False — the credential "
+                        f"could never be used",
                         hint,
                     )
                 )

snowflake/sandbox/_env.py

@@ -117,7 +117,7 @@ def validate_env_mapping(
     ``str(False)`` coercion is a footgun — ``env(DEBUG=False)`` would bake the
     string ``"False"``, which is truthy). *allow* names keys that are exempt
     from the reserved check because the SDK sets them itself (e.g.
-    ``SNOWFLAKE_CODE_URL``, the presigned-download URL ``from_local`` injects).
+    legacy ``SNOWFLAKE_CODE_URL`` stage bootstrap, now unused).
     Returns a plain ``dict[str, str]``. Raises `SandboxError`.
     """
     if not env:

snowflake/sandbox/_preview.py

@@ -0,0 +1,164 @@
+"""Resolve the stable browser URL for a port inside a sandbox.
+
+The URL is ``<handle>-<port>-sandbox.<suffix>``, where ``handle`` is the sandbox id
+without its ``cntr_`` prefix and ``suffix`` is the account's ``*.snowflake.app``
+host suffix. Both halves of the leftmost label are opaque to the edge: the CNG
+ingress proxy parses them out of the hostname, so any port works with no
+registration step.
+
+The suffix is NOT assembled here. It comes from
+``SYSTEM$GET_SNOWFLAKE_APP_HOST_SUFFIX()``, which resolves the region id and the
+TLD the way CNG app hosts do -- including ``.snowflakeapp.cn`` on a China
+deployment and a deployment-level form for accounts without per-account app URLs.
+A client-side guess does not fail loudly on those; it produces a confident wrong
+hostname.
+
+SQL, not REST: the sandbox REST surface accepts the connector's session token, but
+``/api/v2/statements`` requires ``Bearer`` and rejects it (390146). So this takes
+the same route ``list_mcp_servers`` does -- a connector connection resolved from
+the transport -- and is a plain sync function for the same reason: both the async
+and sync clients call it directly.
+"""
+
+from __future__ import annotations
+
+from typing import Any
+
+from snowflake.sandbox.exceptions import SandboxError, SandboxValidationError
+
+__all__ = ["app_host_suffix", "preview_url", "validate_preview_port"]
+
+# A DNS label caps at 63 octets and the first label is "<handle>-<port>-sandbox",
+# so the proxy bounds app_key at 63 - len("-sandbox"). Mirrors sandbox-api's
+# ingress/resolver.go maxAppKey.
+_MAX_APP_KEY = 63 - len("-sandbox")
+
+_CONTAINER_ID_PREFIX = "cntr_"
+
+# One statement, both values: the suffix alone cannot be checked for the
+# per-account form without knowing what that form would be for this account.
+_SUFFIX_SQL = (
+    "SELECT SYSTEM$GET_SNOWFLAKE_APP_HOST_SUFFIX(), "
+    "LOWER(CURRENT_ORGANIZATION_NAME() || '-' || CURRENT_ACCOUNT_NAME())"
+)
+
+# Keyed by base_url, which identifies the account. The suffix is an account
+# property that does not change for the life of a process.
+_SUFFIX_CACHE: dict[str, str] = {}
+
+
+def validate_preview_port(port: int) -> int:
+    """Return ``port`` if the proxy could parse it out of a hostname.
+
+    The bounds are sandbox-api's, not TCP's in general: ``parseAppKey`` reads the
+    port with ParseUint and rejects 0 and a leading zero, so a port this refuses
+    would come back as an opaque 400 from the edge rather than anything the caller
+    could act on.
+    """
+    if isinstance(port, bool) or not isinstance(port, int):
+        raise SandboxValidationError(f"port must be an int, got {type(port).__name__}")
+    if port < 1 or port > 65535:
+        raise SandboxValidationError(f"port must be in 1..65535, got {port}")
+    return port
+
+
+def preview_url(sandbox_id: str, port: int, suffix: str) -> str:
+    """Build the preview URL. Pure; the suffix is already resolved."""
+    handle = sandbox_id.removeprefix(_CONTAINER_ID_PREFIX)
+    app_key = f"{handle}-{port}"
+    if len(app_key) > _MAX_APP_KEY:
+        raise SandboxValidationError(
+            f"sandbox id and port do not fit in a DNS label: {app_key!r} "
+            f"is {len(app_key)} chars, limit {_MAX_APP_KEY}"
+        )
+    return f"https://{app_key}-sandbox.{suffix}"
+
+
+def app_host_suffix(transport: Any = None, connection: Any = None) -> str:
+    """Return the account's ``*.snowflake.app`` host suffix, cached per account.
+
+    Blocking, and deliberately so: the connector cursor is blocking, so this is a
+    sync function that each client adapts -- the async client hands it to
+    ``asyncio.to_thread`` so it does not stall the loop, the sync client calls it
+    directly. Same split ``list_mcp_servers`` uses for its SQL.
+    """
+    cache_key = _cache_key(transport, connection)
+    cached = _SUFFIX_CACHE.get(cache_key) if cache_key else None
+    if cached:
+        return cached
+
+    own_connection = connection is None
+    if connection is None:
+        from snowflake.sandbox._stage import _connect_for_stage
+
+        try:
+            connection, own_connection = _connect_for_stage(transport)
+        except Exception as exc:
+            raise SandboxError(
+                "preview needs a Snowflake connection to resolve this account's "
+                f"app host suffix, and none could be opened: {exc}"
+            ) from exc
+
+    try:
+        cur = connection.cursor()
+        try:
+            cur.execute(_SUFFIX_SQL)
+            row = cur.fetchone()
+        finally:
+            cur.close()
+    except SandboxError:
+        raise
+    except Exception as exc:
+        # The sysfunc raises a SQL-compilation error naming itself when the
+        # account does not have it; that message is the useful one, so pass it on.
+        raise SandboxError(f"could not resolve this account's app host suffix: {exc}") from exc
+    finally:
+        if own_connection:
+            try:
+                connection.close()
+            except Exception:  # noqa: BLE001 - closing our own connection is best-effort
+                pass
+
+    suffix = _require_per_account_suffix(row)
+    if cache_key:
+        _SUFFIX_CACHE[cache_key] = suffix
+    return suffix
+
+
+def _require_per_account_suffix(row: Any) -> str:
+    """Return the suffix, refusing the deployment-level form.
+
+    An account without per-account app URLs (ENABLE_PER_ACCOUNT_APP_SERVICE_URL,
+    itself gated on a per-account certificate) gets ``<deployment>.snowflake.app``
+    with no ``<org>-<account>`` label. The sandbox edge route only matches the
+    per-account form, and no certificate covers the other one -- so a URL built
+    from it fails in the TLS handshake, with nothing naming the cause. Refuse here
+    instead.
+    """
+    if not row or not row[0]:
+        raise SandboxError("this account reported no app host suffix")
+    suffix = str(row[0]).strip().lower()
+    # `_` -> `-`: an account name may legitimately contain an underscore, which is
+    # not a legal DNS label character, so the host form carries a hyphen instead
+    # (the connector does the same mapping for the URL it dials -- see the
+    # ACCOUNT_REGION note in config.py). Comparing the raw name would make this
+    # guard reject every such account.
+    org_acct = str(row[1]).strip().lower().replace("_", "-") if len(row) > 1 and row[1] else ""
+    if org_acct and not suffix.startswith(f"{org_acct}."):
+        raise SandboxError(
+            "preview needs per-account app URLs, which this account does not have: "
+            f"its app host suffix is {suffix!r}, not {org_acct!r} followed by the "
+            "region. Enable ENABLE_PER_ACCOUNT_APP_SERVICE_URL for the account."
+        )
+    return suffix
+
+
+def _cache_key(transport: Any, connection: Any) -> str:
+    """Identify the account. Empty means do not cache."""
+    if connection is not None:
+        return ""
+    config = getattr(transport, "config", None)
+    try:
+        return str(config.base_url) if config is not None else ""
+    except Exception:  # noqa: BLE001 - an unresolvable config just means no caching
+        return ""

snowflake/sandbox/_runtime/_shims.py

@@ -34,7 +34,7 @@ __all__ = [
 ]
 
 
-def _session_runner_shim(module: str, name: str) -> str:
+def _session_runner_shim(module: str, name: str, mod_basename: str = "") -> str:
     """Generate the daemon shim for a ``@app.session`` class.
 
     The shim runs as the container's main process (the manifest entry is
@@ -47,7 +47,9 @@ def _session_runner_shim(module: str, name: str) -> str:
     Behaviour (mirrors ``session_loop`` but with the user's class as warm state):
 
     1. Insert the runner's own dir on ``sys.path`` and import the user's class
-       by ``module:name`` qualname.
+       by ``module:name`` qualname — falling back to loading *mod_basename*
+       from the bundle root when the package path is not preserved (same
+       contract as `_function_runner_source`).
     2. Instantiate it once — **the instance IS the warm state across turns**.
     3. Run its cold-start hook once: the method flagged by ``@enter``
        (``_ENTER_FLAG``), else a method named ``enter`` or ``setup`` if present.
@@ -67,6 +69,7 @@ def _session_runner_shim(module: str, name: str) -> str:
     return textwrap.dedent(f"""\
         import base64
         import importlib
+        import importlib.util
         import json
         import os
         import signal
@@ -84,9 +87,34 @@ def _session_runner_shim(module: str, name: str) -> str:
         # reply. Empty when the deploying SDK predates the nonce.
         _REPLY_NONCE = os.environ.pop("SANDBOX_SESSION_NONCE", "")
 
-        sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
+        _here = os.path.dirname(os.path.abspath(__file__))
+        sys.path.insert(0, _here)
+
+        _MODULE = {module!r}
+        _MODFILE = {mod_basename!r}
 
-        _mod = importlib.import_module({module!r})
+        def _load_user_module():
+            # Prefer import-by-name (packages / sibling imports). Fall back to
+            # loading the bundled source file — the session bundle root is the
+            # class's source directory, so a dotted package name like
+            # ``tests.sandbox.fixtures.live_workload`` is NOT on sys.path even
+            # though ``live_workload.py`` is next to this shim.
+            if _MODULE and _MODULE != "__main__":
+                try:
+                    return importlib.import_module(_MODULE)
+                except Exception:
+                    pass
+            _path = os.path.join(_here, _MODFILE)
+            if _MODFILE and os.path.isfile(_path):
+                _spec = importlib.util.spec_from_file_location(
+                    "__sandbox_user_module__", _path
+                )
+                _m = importlib.util.module_from_spec(_spec)
+                _spec.loader.exec_module(_m)
+                return _m
+            return importlib.import_module(_MODULE)
+
+        _mod = _load_user_module()
         _cls = getattr(_mod, {name!r})
         _instance = _cls()