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

snowflake/cli_sandbox/_shell_command.py

@@ -0,0 +1,666 @@
+"""``snow sandbox shell`` / ``ssh`` — run a command, open a terminal, or attach.
+
+Split out of ``commands.py`` unchanged: the shell command callback, its
+create-retry policy, and the streaming/PTY helpers it drives. ``commands``
+imports the callback and registers it onto the shared ``app`` (as both ``shell``
+and ``ssh``); the shared helpers this reads come from ``_common``.
+"""
+
+from __future__ import annotations
+
+import asyncio
+from collections.abc import Awaitable, Callable, Mapping, Sequence
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, Any
+
+import typer
+
+from snowflake.cli_sandbox._common import (
+    _CONN_HELP,
+    _ENV_HELP,
+    _USE_SNOW_CONN_HELP,
+    _apply_connection,
+    _parse_env,
+    _resolve_sandbox_id,
+)
+from snowflake.cli_sandbox._egress_flags import (
+    EAI_HELP,
+    NO_DEFAULT_EGRESS_HELP,
+    SECRET_HELP,
+    build_egress,
+)
+
+if TYPE_CHECKING:
+    from snowflake.sandbox.client import AsyncSandbox
+    from snowflake.sandbox.egress import Egress
+    from snowflake.sandbox.secret import Secret
+
+_SHELL_HELP = """Run a command, open a terminal, or attach to a running sandbox.
+
+What it does depends on how you invoke it:
+  * --cmd, no SANDBOX_ID  -> create a fresh ephemeral sandbox, run the command
+    (streaming stdout/stderr split, exiting with its code), then tear it down.
+    --keep leaves it running and prints its name.
+  * --cmd + SANDBOX_ID    -> run the command once in that existing sandbox, which
+    is left running (not yours to destroy).
+  * no --cmd, no SANDBOX_ID -> open an interactive terminal: attach to the one
+    running sandbox, or pick from a list when there are several. When NONE are
+    running, spin up a fresh ephemeral box interactively (rather than dead-end on
+    an empty picker). Pass --new to always create a fresh box even when others are
+    running, or a create option (e.g. --image) to size the fresh box.
+  * no --cmd + SANDBOX_ID -> open an interactive terminal in that sandbox.
+
+The create-sizing options (--new/--image/--memory/--cpu/--keep) apply
+only when creating a fresh sandbox; passing them with a SANDBOX_ID is an error
+rather than a silent no-op.
+
+Note:
+    The interactive terminal is a PTY over the ordinary Snowflake REST path, not
+    the SSH protocol -- no server, no port 22, no keys, no scp. --no-sanitize
+    passes the sandbox's output through byte-exact, including clipboard-write
+    sequences (see Shell.attach's trust note); it affects the interactive
+    terminal only.
+
+Example:
+    snow sandbox shell --cmd "python -c 'print(40 + 2)'"   # fresh box: run, tear down
+    snow sandbox shell --cmd "python train.py" --memory 16g
+    snow sandbox shell                                     # pick a running box (or make one if none)
+    snow sandbox shell --new                               # always a fresh box, interactive
+    snow sandbox shell --image sandbox-base                # fresh box, interactive
+    snow sandbox shell cntr_9f2a1b                         # attach to an existing box
+    snow sandbox shell cntr_9f2a1b --cmd "ls -la"          # run once in an existing box
+"""
+
+
+# How many times to (re)create a fresh sandbox when a create hits a transient
+# scheduling miss, and how long to pause between attempts (a fresh create reschedules,
+# so a short pause lets capacity free up).
+_CREATE_RETRY_ATTEMPTS = 3
+_CREATE_RETRY_BACKOFF_S = 2.0
+
+# Substrings that mark a create / first-use failure worth repeating with a NEW sandbox.
+# The CNG scheduler can momentarily have no placement ("no hosts match the
+# requirements"); because the server mints a brand-new sandbox per create and refuses
+# to re-drive a failed one, the only correct retry is a fresh create. These are the
+# FAST scheduling / not-yet-routable signatures (each fails in ~seconds, so a bounded
+# retry is cheap even when the cause turns out to be permanent, e.g. a cpu/memory
+# request larger than any host). Deliberately NOT included:
+#   * a create TIMEOUT ("did not become ready") — that is the 300s strict-wait giving
+#     up on a hung StartApp; retrying it would stack 5-minute waits into a many-minute
+#     hang, and a hang is not the fast scheduling-miss this targets.
+#   * the generic "container create failed" prefix — it wraps ANY reason (e.g. an
+#     image-pull/auth failure), so matching it would retry deterministic errors; the
+#     transient reasons below already match inside such a message when it IS transient.
+# Anything not listed is surfaced on the first attempt rather than hammered.
+_TRANSIENT_CREATE_MARKERS = (
+    "no hosts match the requirements",
+    "scheduling failed",
+    "failed to reach sandbox exec stream",
+    "container not found",
+)
+
+
+def _looks_transient_create_error(exc: BaseException) -> bool:
+    """Whether *exc* from a create (or the first exec on a fresh sandbox) is the kind of
+    transient scheduling failure that a repeat create is likely to clear."""
+    msg = str(exc).lower()
+    return any(marker in msg for marker in _TRANSIENT_CREATE_MARKERS)
+
+
+def _wants_fresh_sandbox(
+    *,
+    new: bool,
+    cmd: str | None,
+    image: str | None,
+    memory: str | None,
+    cpu: float | None,
+    keep: bool,
+    env_overrides: dict[str, str],
+) -> bool:
+    """True when the invocation asks to CREATE a sandbox rather than attach/pick.
+
+    An explicit ``--new``, a ``--cmd`` to run, any create-sizing option, or
+    ``--keep`` all force the create path. A bare ``shell``/``ssh`` with none of
+    these keeps the attach-or-pick behavior (and won't eagerly bill a fresh box).
+    """
+    return (
+        new
+        or cmd is not None
+        or image is not None
+        or memory is not None
+        or cpu is not None
+        or keep
+        or bool(env_overrides)
+    )
+
+
+def _reject_incompatible_attach_flags(
+    *,
+    sandbox_id: str,
+    new: bool,
+    image: str | None,
+    eai: list[str] | None = None,
+    secret: list[str] | None = None,
+    no_default_egress: bool = False,
+    memory: str | None,
+    cpu: float | None,
+    keep: bool,
+    env_overrides: dict[str, str],
+    cmd: str | None,
+) -> None:
+    """Refuse create-sizing knobs when attaching to an existing sandbox.
+
+    They describe a sandbox we would make; against one that already exists they
+    cannot take effect, so we reject loudly rather than accept-and-ignore (the
+    silent-no-op trap this SDK keeps hitting).
+    """
+    create_only = [
+        name
+        for name, given in (
+            ("--new", new),
+            ("--image", image is not None),
+            ("--memory", memory is not None),
+            ("--cpu", cpu is not None),
+            ("--keep", keep),
+            # Egress is fixed at create, so these cannot apply to a box that exists.
+            ("--eai", bool(eai)),
+            ("--secret", bool(secret)),
+            ("--no-default-egress", no_default_egress),
+        )
+        if given
+    ]
+    if create_only:
+        raise typer.BadParameter(
+            f"{', '.join(create_only)} only appl{'ies' if len(create_only) == 1 else 'y'} "
+            f"when creating a fresh sandbox; {sandbox_id} already exists."
+        )
+    if env_overrides and cmd is None:
+        raise typer.BadParameter(
+            "--env needs --cmd when attaching: env is injected per command "
+            "run, and an interactive shell in an existing sandbox cannot take it."
+        )
+
+
+def _try_interactive_picker(
+    *, sandbox_id: str | None, cwd: str | None, no_sanitize: bool, quiet: bool
+) -> bool:
+    """Attach to *sandbox_id*, or pick from the running set (delegated to ``ssh``).
+
+    Returns True if it handled the session (the caller should return); False to
+    fall through to creating a fresh ephemeral sandbox — the case where no id was
+    given and nothing is running, so there is nothing to pick. Delegates the
+    terminal wiring and exit codes to the SDK module command so they cannot drift.
+    """
+    from snowflake.sandbox.exceptions import SandboxError
+
+    use_picker = True
+    if not sandbox_id:
+        # Only the no-id case can be empty; with an id we always attach. Probe the
+        # running set cheaply -- one running box is enough. A probe failure is not a
+        # reason to skip the proven path, so fall back to opening the picker.
+        from snowflake.cli_sandbox._adapter import sdk_connection
+        from snowflake.sandbox.client import list_sandboxes
+
+        async def _any_running() -> bool:
+            async for _sb in list_sandboxes(connection=sdk_connection()):
+                return True
+            return False
+
+        try:
+            use_picker = asyncio.run(_any_running())
+        except SandboxError:
+            use_picker = True
+
+    if not use_picker:
+        _emit_status("No running sandboxes -- creating a fresh one...", quiet=quiet)
+        return False
+
+    from snowflake.sandbox.__main__ import main as sandbox_main
+
+    argv = ["ssh"]
+    if sandbox_id:
+        argv.append(sandbox_id)
+    if cwd:
+        argv += ["--cwd", cwd]
+    if no_sanitize:
+        argv.append("--no-sanitize")
+    try:
+        code = sandbox_main(argv)
+    except SandboxError as exc:
+        typer.echo(f"error: {exc}", err=True)
+        raise typer.Exit(1) from exc
+    if code:
+        raise typer.Exit(code)
+    return True
+
+
+@dataclass(frozen=True)
+class _ShellPlan:
+    """The resolved shell invocation, after flag parsing and validation.
+
+    Threads the settled request from `_shell_command` (which owns the typer
+    signature) to the async workers `_run_shell` / `_use_sandbox`, so the two
+    stay module-level and independently testable rather than closures capturing
+    the command's many flags.
+    """
+
+    attaching: bool
+    sandbox_id: str | None
+    cmd: str | None
+    cwd: str | None
+    no_sanitize: bool
+    quiet: bool
+    keep: bool
+    image: str | None
+    egress: Egress | None
+    secrets: tuple[Secret, ...]
+    memory: str | None
+    cpu: float | None
+    env_overrides: dict[str, str]
+
+
+async def _use_sandbox(sb: Any, plan: _ShellPlan) -> int:
+    """Run the command in *sb* (streamed) or open an interactive PTY."""
+    if plan.cmd is not None:
+        # A blank line sets the status block apart from the command's own output
+        # that follows, so the two never read as one stream.
+        if not plan.quiet:
+            typer.echo("", err=True)
+        # env rides the exec only when attaching; a fresh sandbox already carries
+        # it from create.
+        return await _exec_streaming(
+            sb, plan.cmd, plan.cwd, plan.env_overrides if plan.attaching else None
+        )
+    return await _attach_pty(sb, cwd=plan.cwd, sanitize=not plan.no_sanitize, quiet=plan.quiet)
+
+
+async def _run_shell(plan: _ShellPlan) -> int:
+    from snowflake.cli_sandbox._adapter import sdk_connection
+    from snowflake.sandbox.client import get_sandbox
+
+    if plan.attaching:
+        # An attached sandbox already exists -- no create, so nothing to retry, and
+        # its lifecycle is not ours to end.
+        sb = await get_sandbox(
+            plan.sandbox_id,  # type: ignore[arg-type]  # attaching => not None
+            connection=sdk_connection(),
+        )
+        if plan.cmd is not None:
+            _emit_status(f"Running command in {sb.name or sb.id}...", quiet=plan.quiet)
+        return await _use_sandbox(sb, plan)
+
+    async def _use(sb: AsyncSandbox) -> int:
+        return await _use_sandbox(sb, plan)
+
+    return await _run_on_fresh_sandbox(
+        use=_use,
+        keep=plan.keep,
+        quiet=plan.quiet,
+        # The "ready" line only makes sense before a command runs; an interactive
+        # attach prints its own "Connected to..." instead.
+        announce_ready=plan.cmd is not None,
+        keep_notice=True,
+        image=plan.image,
+        memory=plan.memory,
+        cpu=plan.cpu,
+        env=plan.env_overrides or None,
+        egress=plan.egress,
+        secrets=plan.secrets,
+    )
+
+
+async def _run_on_fresh_sandbox(
+    *,
+    use: Callable[[AsyncSandbox], Awaitable[int]],
+    keep: bool,
+    quiet: bool,
+    announce_ready: bool = False,
+    keep_notice: bool = False,
+    image: str | None = None,
+    memory: str | None = None,
+    cpu: float | None = None,
+    env: Mapping[str, str] | None = None,
+    name: str | None = None,
+    command: Sequence[str] | None = None,
+    idle_suspend: str | None = None,
+    egress: Egress | None = None,
+    secrets: tuple[Secret, ...] = (),
+) -> int:
+    """Create a fresh sandbox, run ``use(sandbox)`` on it, and return its result —
+    the shared "spin up, run, (optionally) tear down" path behind ``snow sandbox
+    shell`` and ``snow sandbox run``, so the startup status lines and the transient
+    retry live in one place.
+
+    Emits the "Spinning up..." status. A create can hit a transient CNG scheduling
+    miss ("no hosts match the requirements"), and because the server marks a sandbox
+    ready before StartApp finishes it can also surface a beat later as a 502 on the
+    first use; either way the fix is a brand-new sandbox (the server refuses to
+    re-drive a failed one), so both are retried with a fresh create. A real error is
+    surfaced on the first attempt.
+
+    Tears the sandbox down afterward unless ``keep``. With ``keep`` and
+    ``keep_notice`` it prints the surviving sandbox's id so the caller can find it
+    (``shell --keep``); a caller that prints its own handle passes ``keep_notice``
+    False (``run --detach``). ``command``/``idle_suspend``/``name`` are forwarded
+    to create for the long-running (``--detach``) callers; the ephemeral callers
+    leave them unset.
+    """
+    from snowflake.cli_sandbox._adapter import sdk_connection
+    from snowflake.sandbox.client import AsyncSandbox
+    from snowflake.sandbox.egress import Egress  # noqa: F401
+    from snowflake.sandbox.exceptions import SandboxError
+    from snowflake.sandbox.secret import Secret  # noqa: F401
+
+    last_exc: SandboxError | None = None
+    for attempt in range(1, _CREATE_RETRY_ATTEMPTS + 1):
+        # Creating a sandbox pulls the image and boots the container, which takes a
+        # few seconds with no output -- say so, so it never looks hung.
+        _emit_status(f"Spinning up sandbox ({image or 'default image'})...", quiet=quiet)
+        created: AsyncSandbox | None = None
+        retrying = False
+        handed_off = False
+        try:
+            created = await AsyncSandbox.create(
+                # Empty string -> the deployment's default runtime image, resolved
+                # server-side. Never hardcode a runtime: the fleet default moves, and
+                # a pinned name spins up a stale/unavailable image (its server 502s).
+                image=image or "",
+                memory=memory or "4g",  # type: ignore[arg-type]  # validated server-side
+                cpu=cpu,
+                env=env or None,
+                name=name,
+                command=command,
+                idle_suspend=idle_suspend,
+                # Omitted entirely when unset so the create body is byte-identical to
+                # what it was before these flags existed.
+                egress=egress,
+                secrets=list(secrets) or None,
+                connection=sdk_connection(),
+            )
+            if announce_ready:
+                _emit_status(f"Sandbox {created.name or created.id} ready", quiet=quiet)
+            result = await use(created)
+            handed_off = True  # `use` returned, so a `keep` sandbox is the caller's now
+            return result
+        except SandboxError as exc:
+            last_exc = exc
+            if not _looks_transient_create_error(exc):
+                raise  # a real error (bad image, invalid request): surface it now
+            if attempt >= _CREATE_RETRY_ATTEMPTS:
+                break  # exhausted: fall through to the capacity message below
+            retrying = True
+            _emit_status(
+                f"transient scheduling error ({exc}); retrying "
+                f"[{attempt}/{_CREATE_RETRY_ATTEMPTS - 1}]...",
+                quiet=quiet,
+            )
+        finally:
+            # A failed attempt is never "kept"; when we own the lifecycle (not keep)
+            # we always tear down. `keep` only earns the sandbox once `use` *returned*:
+            # if it raised, nobody was ever handed the id, so keeping it would leave a
+            # billed, running, unnamed-to-the-user box -- and burn a `--name` so the
+            # obvious retry 409s. Best effort: a failed teardown must not mask the use
+            # result (e.g. a command's own exit code).
+            if created is not None and (retrying or not keep or not handed_off):
+                try:
+                    await created.terminate()
+                except SandboxError as exc:
+                    typer.echo(
+                        f"warning: could not terminate {created.name or created.id}: {exc}",
+                        err=True,
+                    )
+            elif created is not None and keep and handed_off and keep_notice:
+                typer.echo(f"sandbox left running: {created.name or created.id}", err=True)
+        if retrying:
+            await asyncio.sleep(_CREATE_RETRY_BACKOFF_S)
+    # Every attempt hit a fast scheduling failure.
+    raise SandboxError(
+        f"sandbox could not be scheduled after {_CREATE_RETRY_ATTEMPTS} attempts — "
+        f"no hosts available (transient capacity, or the requested cpu/memory "
+        f"exceeds any host); last error: {last_exc}"
+    )
+
+
+def _shell_command(
+    sandbox_id: str | None = typer.Argument(
+        None,
+        metavar="[SANDBOX_ID]",
+        help="Existing sandbox (name or app id) to attach to. Omit to attach to a running one "
+        "(picker), or create a fresh ephemeral one if none are running. Use --new to always "
+        "create a fresh one.",
+    ),
+    cmd: str | None = typer.Option(
+        None,
+        "--cmd",
+        "--command",
+        metavar="COMMAND",
+        help="Command to run (streamed, exits with its code). Omit for an interactive "
+        "terminal. `--command` is an accepted alias (same flag on `run`).",
+    ),
+    image: str | None = typer.Option(
+        None,
+        "--image",
+        metavar="IMAGE",
+        help="Image for a fresh sandbox (default: the deployment's default image). Not valid with SANDBOX_ID.",
+    ),
+    memory: str | None = typer.Option(
+        None,
+        "--memory",
+        metavar="TIER",
+        help="Memory tier for a fresh sandbox: 1g/4g/8g/16g/32g/64g (default: 4g). Not valid with SANDBOX_ID.",
+    ),
+    cpu: float | None = typer.Option(
+        None,
+        "--cpu",
+        metavar="CORES",
+        help="CPU cores (0.25-8.0) for a fresh sandbox. Not valid with SANDBOX_ID.",
+    ),
+    envs: list[str] | None = typer.Option(
+        None,
+        "--env",
+        "-e",
+        metavar="KEY=VAL",
+        help=_ENV_HELP,
+    ),
+    eai: list[str] | None = typer.Option(
+        None,
+        "--eai",
+        metavar="NAME",
+        help=EAI_HELP,
+    ),
+    secret: list[str] | None = typer.Option(
+        None,
+        "--secret",
+        metavar="ENV_VAR=DB.SCHEMA.NAME@host",
+        help=SECRET_HELP,
+    ),
+    no_default_egress: bool = typer.Option(
+        False,
+        "--no-default-egress",
+        help=NO_DEFAULT_EGRESS_HELP,
+    ),
+    keep: bool = typer.Option(
+        False,
+        "--keep",
+        help="Leave a fresh sandbox running on exit (prints its name). Not valid with SANDBOX_ID.",
+    ),
+    new: bool = typer.Option(
+        False,
+        "--new",
+        help="Always create a fresh sandbox, even if some are already running "
+        "(skips the attach/picker path). Not valid with SANDBOX_ID.",
+    ),
+    cwd: str | None = typer.Option(
+        None,
+        "--cwd",
+        metavar="DIR",
+        help="Working directory for the shell/command.",
+    ),
+    no_sanitize: bool = typer.Option(
+        False,
+        "--no-sanitize",
+        help="Interactive only: pass the sandbox's output through byte-exact. See the trust note.",
+    ),
+    quiet: bool = typer.Option(
+        False,
+        "--quiet",
+        "-q",
+        help="Suppress the status lines on stderr (spinning up / ready / running). "
+        "The command's own output is unaffected, so `--quiet` is safe for scripting.",
+    ),
+    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:
+    from snowflake.sandbox.exceptions import SandboxError
+
+    # Opening the snow connection is a real login round-trip (a few seconds) and
+    # runs before anything else -- say so, so that gap does not look like a hang.
+    if use_snow_connection:
+        where = f" ({connection})" if connection else ""
+        _emit_status(f"Connecting to Snowflake{where}...", quiet=quiet)
+    _apply_connection(use_snow_connection, connection=connection)
+
+    # A reference that isn't a raw app id (cntr_...) is a vanity NAME — the identity
+    # `list` shows. Resolve to the app id up front so the attach and ssh-picker paths
+    # below work unchanged.
+    if sandbox_id:
+        sandbox_id = _resolve_sandbox_id(sandbox_id)
+
+    env_overrides: dict[str, str] = {}
+    for raw in envs or []:
+        k, v = _parse_env(raw)
+        env_overrides[k] = v
+
+    attaching = sandbox_id is not None
+    # "Am I being asked to make a NEW sandbox?" -- an explicit --new, a command to
+    # run, or any create-sizing option. Bare `shell`/`ssh` with none of these keeps
+    # the long-standing interactive behavior (attach by id, or pick from the running
+    # list -- and, when nothing is running, fall through to a fresh ephemeral box)
+    # rather than eagerly creating and billing one. Only these force the create path.
+    create_intent = _wants_fresh_sandbox(
+        new=new,
+        cmd=cmd,
+        image=image,
+        memory=memory,
+        cpu=cpu,
+        keep=keep,
+        env_overrides=env_overrides,
+    )
+
+    if attaching:
+        _reject_incompatible_attach_flags(
+            sandbox_id=sandbox_id,  # type: ignore[arg-type]  # attaching => not None
+            new=new,
+            image=image,
+            memory=memory,
+            cpu=cpu,
+            keep=keep,
+            env_overrides=env_overrides,
+            cmd=cmd,
+            eai=eai,
+            secret=secret,
+            no_default_egress=no_default_egress,
+        )
+
+    # Interactive with no create intent == the original behavior, preserved:
+    # attach to SANDBOX_ID, or (no id) pick from the running set. The one addition:
+    # when no id was given AND nothing is running, there is nothing to pick, so we
+    # fall through to create a fresh ephemeral box instead of dead-ending on an
+    # empty picker. The attach/picker itself is delegated to the SDK module command
+    # so the terminal wiring and exit codes stay identical to what `ssh` has always
+    # done and cannot drift.
+    # `and` short-circuits, so the picker only runs on the interactive no-create
+    # path; it returns True when it handled the session.
+    if (
+        cmd is None
+        and not create_intent
+        and _try_interactive_picker(
+            sandbox_id=sandbox_id, cwd=cwd, no_sanitize=no_sanitize, quiet=quiet
+        )
+    ):
+        return
+    # Otherwise nothing was running and no id was given (or a create was intended):
+    # fall through to the create path below and open a fresh ephemeral sandbox.
+
+    egress_cfg, secret_objs = build_egress(
+        eai=eai, secret=secret, no_default_egress=no_default_egress
+    )
+    plan = _ShellPlan(
+        attaching=attaching,
+        sandbox_id=sandbox_id,
+        cmd=cmd,
+        cwd=cwd,
+        no_sanitize=no_sanitize,
+        quiet=quiet,
+        keep=keep,
+        image=image,
+        egress=egress_cfg,
+        secrets=secret_objs,
+        memory=memory,
+        cpu=cpu,
+        env_overrides=env_overrides,
+    )
+    try:
+        code = asyncio.run(_run_shell(plan))
+    except KeyboardInterrupt:
+        raise typer.Exit(130) from None
+    except SandboxError as exc:
+        typer.echo(f"error: {exc}", err=True)
+        raise typer.Exit(1) from exc
+    raise typer.Exit(code)
+
+
+def _emit_status(msg: str, *, quiet: bool) -> None:
+    """Print a progress line to stderr, dimmed and marked with a leading ``»`` so
+    it reads as meta rather than the command's output. Click drops the styling
+    automatically when stderr is not a tty (piped), leaving plain ``» ...`` text.
+    ``quiet`` mutes it entirely."""
+    if quiet:
+        return
+    typer.secho(f"» {msg}", err=True, dim=True)
+
+
+async def _exec_streaming(sb: Any, cmd: str, cwd: str | None, env: dict[str, str] | None) -> int:
+    """Run `cmd` in `sb`, streaming output; return its exit code.
+
+    The command string is handed to `sh -c` so shell syntax (pipes, quoting,
+    `python -c '...'`) works exactly as typed.
+    """
+    kwargs: dict[str, Any] = {}
+    if cwd:
+        kwargs["working_dir"] = cwd
+    if env:
+        kwargs["env"] = env
+    stream = sb.exec_stream(["sh", "-c", cmd], **kwargs)
+    async for line in stream:
+        typer.echo(line.data, err=(line.stream == "stderr"))
+    # ExecStream sets exit_code once the process is terminal; None means we never
+    # saw a terminal frame (cancelled/torn down), which is a failure, not a 0.
+    return stream.exit_code if stream.exit_code is not None else 1
+
+
+async def _attach_pty(sb: Any, *, cwd: str | None, sanitize: bool, quiet: bool = False) -> int:
+    """Open an interactive PTY in `sb` and return the shell's exit code."""
+    from snowflake.sandbox.__main__ import _terminal_size
+
+    rows, cols = _terminal_size()
+    sh = await sb.shell(rows=rows, cols=cols, cwd=cwd)
+    _emit_status(
+        f"Connected to {sb.name or sb.id} (session {sh.session_id}); exit the shell to disconnect",
+        quiet=quiet,
+    )
+    try:
+        code = await sh.attach(sanitize=sanitize)
+    finally:
+        await sh.close()
+    return 0 if code == 0 else (code if isinstance(code, int) and code > 0 else 1)