sandbox.py

"""
Modal sandbox (https://modal.com) implementation.

Run `python -m modal setup` to configure Modal locally.

This module provides a Modal-backed sandbox client/session implementation backed by
`modal.Sandbox`.

Note: The `modal` dependency is intended to be optional (installed via an extra),
so package-level exports should guard imports of this module. Within this module,
we import Modal normally so IDEs can resolve and navigate Modal types.
"""

from __future__ import annotations

import asyncio
import functools
import io
import json
import logging
import math
import os
import shlex
import time
import uuid
from collections.abc import AsyncIterator, Awaitable, Callable
from contextlib import asynccontextmanager
from dataclasses import dataclass, field
from pathlib import Path
from typing import Any, Literal, TypeVar, cast

import modal
from modal.config import config as modal_config
from modal.container_process import ContainerProcess

from ....sandbox.config import DEFAULT_PYTHON_SANDBOX_IMAGE
from ....sandbox.entries import Mount
from ....sandbox.errors import (
    ExecTimeoutError,
    ExecTransportError,
    ExposedPortUnavailableError,
    MountConfigError,
    WorkspaceArchiveReadError,
    WorkspaceArchiveWriteError,
    WorkspaceReadNotFoundError,
    WorkspaceStartError,
    WorkspaceStopError,
    WorkspaceWriteTypeError,
)
from ....sandbox.manifest import Manifest
from ....sandbox.session import SandboxSession, SandboxSessionState
from ....sandbox.session.base_sandbox_session import BaseSandboxSession
from ....sandbox.session.dependencies import Dependencies
from ....sandbox.session.manager import Instrumentation
from ....sandbox.session.pty_types import (
    PTY_PROCESSES_MAX,
    PTY_PROCESSES_WARNING,
    PtyExecUpdate,
    allocate_pty_process_id,
    clamp_pty_yield_time_ms,
    process_id_to_prune_from_meta,
    resolve_pty_write_yield_time_ms,
    truncate_text_by_tokens,
)
from ....sandbox.session.runtime_helpers import RESOLVE_WORKSPACE_PATH_HELPER, RuntimeHelperScript
from ....sandbox.session.sandbox_client import BaseSandboxClient, BaseSandboxClientOptions
from ....sandbox.snapshot import SnapshotBase, SnapshotSpec, resolve_snapshot
from ....sandbox.types import ExecResult, ExposedPortEndpoint, User
from ....sandbox.util.retry import (
    TRANSIENT_HTTP_STATUS_CODES,
    exception_chain_contains_type,
    exception_chain_has_status_code,
    retry_async,
)
from ....sandbox.util.tar_utils import UnsafeTarMemberError, validate_tar_bytes
from ....sandbox.workspace_paths import (
    coerce_posix_path,
    posix_path_as_path,
    posix_path_for_error,
    sandbox_path_str,
)
from .mounts import ModalCloudBucketMountStrategy

_DEFAULT_TIMEOUT_S = 30.0
_DEFAULT_IMAGE_TAG = DEFAULT_PYTHON_SANDBOX_IMAGE
_DEFAULT_IMAGE_BUILDER_VERSION = "2025.06"
_DEFAULT_SNAPSHOT_FILESYSTEM_TIMEOUT_S = 60.0
_MODAL_STDIN_CHUNK_SIZE = 8 * 1024 * 1024
_PTY_POLL_INTERVAL_S = 0.05

WorkspacePersistenceMode = Literal["tar", "snapshot_filesystem", "snapshot_directory"]

_WORKSPACE_PERSISTENCE_TAR: WorkspacePersistenceMode = "tar"
_WORKSPACE_PERSISTENCE_SNAPSHOT_FILESYSTEM: WorkspacePersistenceMode = "snapshot_filesystem"
_WORKSPACE_PERSISTENCE_SNAPSHOT_DIRECTORY: WorkspacePersistenceMode = "snapshot_directory"

# Magic prefixes for snapshot payloads that cannot be represented as tar bytes.
_MODAL_SANDBOX_FS_SNAPSHOT_MAGIC = b"MODAL_SANDBOX_FS_SNAPSHOT_V1\n"
_MODAL_SANDBOX_DIR_SNAPSHOT_MAGIC = b"MODAL_SANDBOX_DIR_SNAPSHOT_V1\n"

logger = logging.getLogger(__name__)
R = TypeVar("R")


@asynccontextmanager
async def _override_modal_image_builder_version(
    image_builder_version: str | None,
) -> AsyncIterator[None]:
    """Apply a process-local Modal image builder version for the duration of a build."""

    if image_builder_version is None:
        yield
        return

    previous_value = os.environ.get("MODAL_IMAGE_BUILDER_VERSION")
    modal_config.override_locally("image_builder_version", image_builder_version)
    try:
        yield
    finally:
        if previous_value is None:
            os.environ.pop("MODAL_IMAGE_BUILDER_VERSION", None)
        else:
            os.environ["MODAL_IMAGE_BUILDER_VERSION"] = previous_value


def _maybe_set_sandbox_cmd(
    image: modal.Image,
    *,
    use_sleep_cmd: bool,
) -> modal.Image:
    if not use_sleep_cmd:
        return image
    return image.cmd(["sleep", "infinity"])


async def _write_process_stdin(proc: ContainerProcess[bytes], data: bytes | bytearray) -> None:
    """
    Stream stdin to Modal in bounded chunks so command-router backed writers do not overflow.
    """

    view = memoryview(data)
    for start in range(0, len(view), _MODAL_STDIN_CHUNK_SIZE):
        proc.stdin.write(view[start : start + _MODAL_STDIN_CHUNK_SIZE])
        await proc.stdin.drain.aio()
    proc.stdin.write_eof()
    await proc.stdin.drain.aio()


class ModalSandboxClientOptions(BaseSandboxClientOptions):
    type: Literal["modal"] = "modal"
    app_name: str
    sandbox_create_timeout_s: float | None = None
    workspace_persistence: WorkspacePersistenceMode = _WORKSPACE_PERSISTENCE_TAR
    snapshot_filesystem_timeout_s: float | None = None
    snapshot_filesystem_restore_timeout_s: float | None = None
    exposed_ports: tuple[int, ...] = ()
    gpu: str | None = None  # Modal GPU type, e.g. "A100" or "H100:8"
    timeout: int = 300  # Lifetime of a sandbox from creation in seconds, defaults to 5 minutes
    use_sleep_cmd: bool = True
    image_builder_version: str | None = _DEFAULT_IMAGE_BUILDER_VERSION
    idle_timeout: int | None = None

    def __init__(
        self,
        app_name: str,
        sandbox_create_timeout_s: float | None = None,
        workspace_persistence: WorkspacePersistenceMode = _WORKSPACE_PERSISTENCE_TAR,
        snapshot_filesystem_timeout_s: float | None = None,
        snapshot_filesystem_restore_timeout_s: float | None = None,
        exposed_ports: tuple[int, ...] = (),
        gpu: str | None = None,
        timeout: int = 300,  # 5 minutes
        use_sleep_cmd: bool = True,
        image_builder_version: str | None = _DEFAULT_IMAGE_BUILDER_VERSION,
        idle_timeout: int | None = None,
        *,
        type: Literal["modal"] = "modal",
    ) -> None:
        super().__init__(
            type=type,
            app_name=app_name,
            sandbox_create_timeout_s=sandbox_create_timeout_s,
            workspace_persistence=workspace_persistence,
            snapshot_filesystem_timeout_s=snapshot_filesystem_timeout_s,
            snapshot_filesystem_restore_timeout_s=snapshot_filesystem_restore_timeout_s,
            exposed_ports=exposed_ports,
            gpu=gpu,
            timeout=timeout,
            use_sleep_cmd=use_sleep_cmd,
            image_builder_version=image_builder_version,
            idle_timeout=idle_timeout,
        )


def _encode_modal_snapshot_ref(
    *,
    snapshot_id: str,
    workspace_persistence: WorkspacePersistenceMode,
) -> bytes:
    # Small JSON envelope so we can round-trip a non-tar snapshot reference
    # through Snapshot.persist().
    body = json.dumps(
        {"snapshot_id": snapshot_id, "workspace_persistence": workspace_persistence},
        separators=(",", ":"),
        sort_keys=True,
    ).encode("utf-8")
    if workspace_persistence == _WORKSPACE_PERSISTENCE_SNAPSHOT_DIRECTORY:
        return _MODAL_SANDBOX_DIR_SNAPSHOT_MAGIC + body
    return _MODAL_SANDBOX_FS_SNAPSHOT_MAGIC + body


def _encode_snapshot_filesystem_ref(*, snapshot_id: str) -> bytes:
    return _encode_modal_snapshot_ref(
        snapshot_id=snapshot_id,
        workspace_persistence=_WORKSPACE_PERSISTENCE_SNAPSHOT_FILESYSTEM,
    )


def _encode_snapshot_directory_ref(*, snapshot_id: str) -> bytes:
    return _encode_modal_snapshot_ref(
        snapshot_id=snapshot_id,
        workspace_persistence=_WORKSPACE_PERSISTENCE_SNAPSHOT_DIRECTORY,
    )


def _decode_modal_snapshot_ref(raw: bytes) -> tuple[WorkspacePersistenceMode, str] | None:
    if raw.startswith(_MODAL_SANDBOX_DIR_SNAPSHOT_MAGIC):
        prefix = _MODAL_SANDBOX_DIR_SNAPSHOT_MAGIC
        default_persistence = _WORKSPACE_PERSISTENCE_SNAPSHOT_DIRECTORY
    elif raw.startswith(_MODAL_SANDBOX_FS_SNAPSHOT_MAGIC):
        prefix = _MODAL_SANDBOX_FS_SNAPSHOT_MAGIC
        default_persistence = _WORKSPACE_PERSISTENCE_SNAPSHOT_FILESYSTEM
    else:
        return None
    body = raw[len(prefix) :]
    try:
        obj = json.loads(body.decode("utf-8"))
    except Exception:
        return None
    snapshot_id = obj.get("snapshot_id")
    workspace_persistence = obj.get("workspace_persistence", default_persistence)
    if workspace_persistence not in (
        _WORKSPACE_PERSISTENCE_SNAPSHOT_FILESYSTEM,
        _WORKSPACE_PERSISTENCE_SNAPSHOT_DIRECTORY,
    ):
        return None
    if not isinstance(snapshot_id, str) or not snapshot_id:
        return None
    return cast(WorkspacePersistenceMode, workspace_persistence), snapshot_id


@dataclass(frozen=True)
class ModalImageSelector:
    """
    A single "image selector" type to avoid juggling image/image_id/image_tag separately.
    """

    kind: Literal["image", "id", "tag"]
    value: modal.Image | str

    @classmethod
    def from_image(cls, image: modal.Image) -> ModalImageSelector:
        return cls(kind="image", value=image)

    @classmethod
    def from_id(cls, image_id: str) -> ModalImageSelector:
        return cls(kind="id", value=image_id)

    @classmethod
    def from_tag(cls, image_tag: str) -> ModalImageSelector:
        return cls(kind="tag", value=image_tag)


@dataclass(frozen=True)
class ModalSandboxSelector:
    """
    A single "sandbox selector" type to avoid juggling sandbox/sandbox_id separately.
    """

    kind: Literal["sandbox", "id"]
    value: modal.Sandbox | str

    @classmethod
    def from_sandbox(cls, sandbox: modal.Sandbox) -> ModalSandboxSelector:
        return cls(kind="sandbox", value=sandbox)

    @classmethod
    def from_id(cls, sandbox_id: str) -> ModalSandboxSelector:
        return cls(kind="id", value=sandbox_id)


class ModalSandboxSessionState(SandboxSessionState):
    """
    Serializable state for a Modal-backed session.

    We store only values that can be safely persisted and later used by `resume()`.
    """

    type: Literal["modal"] = "modal"
    app_name: str
    # Optional Modal image object id (enables reconstructing a custom image via Image.from_id()).
    image_id: str | None = None
    # Registry image tag (e.g. "debian:bookworm" or "ghcr.io/org/img:tag").
    # Used when `image_id` isn't available and no in-memory image override was provided.
    image_tag: str | None = None
    # Timeout for creating a sandbox (Modal calls are synchronous from the user's perspective
    # and can block; we wrap them in a thread with asyncio timeout).
    sandbox_create_timeout_s: float = _DEFAULT_TIMEOUT_S
    sandbox_id: str | None = None
    # Workspace persistence mode:
    # - "tar": create a tar stream in the sandbox via `tar cf - ...` and pull bytes back via stdout.
    # - "snapshot_filesystem": use Modal's `Sandbox.snapshot_filesystem()`
    #   (if available) and persist a snapshot reference.
    # - "snapshot_directory": use Modal's `Sandbox.snapshot_directory()` on the workspace root
    #   and reattach it during resume.
    workspace_persistence: WorkspacePersistenceMode = _WORKSPACE_PERSISTENCE_TAR
    # Async timeouts for snapshot_filesystem-based persistence and restore.
    snapshot_filesystem_timeout_s: float = _DEFAULT_SNAPSHOT_FILESYSTEM_TIMEOUT_S
    snapshot_filesystem_restore_timeout_s: float = _DEFAULT_SNAPSHOT_FILESYSTEM_TIMEOUT_S
    gpu: str | None = None  # Modal GPU type, e.g. "A100" or "H100:8"
    # Maximum lifetime of the sandbox in seconds
    timeout: int = 300  # 5 minutes
    use_sleep_cmd: bool = True
    image_builder_version: str | None = _DEFAULT_IMAGE_BUILDER_VERSION
    idle_timeout: int | None = None


@dataclass
class _ModalPtyProcessEntry:
    process: ContainerProcess[bytes]
    tty: bool
    last_used: float = field(default_factory=time.monotonic)
    stdout_iter: AsyncIterator[object] | None = None
    stderr_iter: AsyncIterator[object] | None = None
    stdout_read_task: asyncio.Task[object] | None = None
    stderr_read_task: asyncio.Task[object] | None = None


class ModalSandboxSession(BaseSandboxSession):
    """
    SandboxSession implementation backed by a Modal Sandbox.
    """

    state: ModalSandboxSessionState

    _sandbox: modal.Sandbox | None
    _image: modal.Image | None
    _running: bool
    _pty_lock: asyncio.Lock
    _pty_processes: dict[int, _ModalPtyProcessEntry]
    _reserved_pty_process_ids: set[int]
    _modal_snapshot_ephemeral_backup: bytes | None
    _modal_snapshot_ephemeral_backup_path: Path | None

    def __init__(
        self,
        *,
        state: ModalSandboxSessionState,
        # Optional in-memory handles. These are not guaranteed to be resumable; state holds ids.
        image: modal.Image | None = None,
        sandbox: modal.Sandbox | None = None,
    ) -> None:
        self.state = state
        self._image = None
        if image is not None:
            self._image = _maybe_set_sandbox_cmd(
                image,
                use_sleep_cmd=self.state.use_sleep_cmd,
            )
        self._sandbox = sandbox
        if self._image is not None:
            self.state.image_id = getattr(self._image, "object_id", self.state.image_id)
        if sandbox is not None:
            self.state.sandbox_id = getattr(sandbox, "object_id", self.state.sandbox_id)
        self._running = False
        self._pty_lock = asyncio.Lock()
        self._pty_processes = {}
        self._reserved_pty_process_ids = set()
        self._modal_snapshot_ephemeral_backup = None
        self._modal_snapshot_ephemeral_backup_path = None

    async def _validate_path_access(self, path: Path | str, *, for_write: bool = False) -> Path:
        return await self._validate_remote_path_access(path, for_write=for_write)

    def _runtime_helpers(self) -> tuple[RuntimeHelperScript, ...]:
        return (RESOLVE_WORKSPACE_PATH_HELPER,)

    def _current_runtime_helper_cache_key(self) -> object | None:
        return self.state.sandbox_id

    @classmethod
    def from_state(
        cls,
        state: ModalSandboxSessionState,
        *,
        image: modal.Image | None = None,
        sandbox: modal.Sandbox | None = None,
    ) -> ModalSandboxSession:
        return cls(state=state, image=image, sandbox=sandbox)

    async def _call_modal(
        self,
        fn: Callable[..., R],
        *args: object,
        call_timeout: float | None = None,
        **kwargs: object,
    ) -> R:
        """
        Prefer Modal's async interface (`fn.aio(...)`) when available.

        Falls back to running the blocking call in a thread to preserve compatibility
        with SDK surfaces that do not expose `.aio`.
        """

        aio_fn = getattr(fn, "aio", None)
        if callable(aio_fn):
            coro = cast(Awaitable[R], aio_fn(*args, **kwargs))
        else:
            loop = asyncio.get_running_loop()
            bound = functools.partial(fn, *args, **kwargs)
            coro = loop.run_in_executor(None, bound)
        if call_timeout is None:
            return await coro
        return await asyncio.wait_for(coro, timeout=call_timeout)

    async def _ensure_backend_started(self) -> None:
        await self._ensure_sandbox()

    async def _prepare_backend_workspace(self) -> None:
        # Ensure workspace root exists before the base workspace flow needs it.
        root = self._workspace_path_policy().sandbox_root().as_posix()
        await self.exec("mkdir", "-p", "--", root, shell=False)

    async def _after_start(self) -> None:
        self._running = True

    async def _after_start_failed(self) -> None:
        self._running = False

    def _wrap_start_error(self, error: Exception) -> Exception:
        if isinstance(error, WorkspaceStartError):
            return error
        return WorkspaceStartError(path=self._workspace_root_path(), cause=error)

    async def _resolve_exposed_port(self, port: int) -> ExposedPortEndpoint:
        await self._ensure_sandbox()
        assert self._sandbox is not None

        try:
            tunnels = await asyncio.wait_for(self._sandbox.tunnels.aio(), timeout=10.0)
        except Exception as e:
            raise ExposedPortUnavailableError(
                port=port,
                exposed_ports=self.state.exposed_ports,
                reason="backend_unavailable",
                context={"backend": "modal", "detail": "tunnels_lookup_failed"},
                cause=e,
            ) from e

        if not isinstance(tunnels, dict):
            raise ExposedPortUnavailableError(
                port=port,
                exposed_ports=self.state.exposed_ports,
                reason="backend_unavailable",
                context={"backend": "modal", "detail": "invalid_tunnels_response"},
            )

        tunnel = tunnels.get(port)
        host = getattr(tunnel, "host", None)
        host_port = getattr(tunnel, "port", None)
        if not isinstance(host, str) or not host or not isinstance(host_port, int):
            raise ExposedPortUnavailableError(
                port=port,
                exposed_ports=self.state.exposed_ports,
                reason="backend_unavailable",
                context={"backend": "modal", "detail": "port_not_exposed"},
            )
        return ExposedPortEndpoint(host=host, port=host_port, tls=True)

    def _wrap_stop_error(self, error: Exception) -> Exception:
        if isinstance(error, WorkspaceStopError):
            return error
        return WorkspaceStopError(path=self._workspace_root_path(), cause=error)

    async def _shutdown_backend(self) -> None:
        try:
            sandbox = self._sandbox
            if sandbox is not None:
                await self._call_modal(
                    sandbox.terminate,
                    call_timeout=_DEFAULT_TIMEOUT_S,
                )
            elif self.state.sandbox_id:
                sid = self.state.sandbox_id
                assert sid is not None
                sb = await self._call_modal(
                    modal.Sandbox.from_id,
                    sid,
                    call_timeout=_DEFAULT_TIMEOUT_S,
                )
                await self._call_modal(
                    sb.terminate,
                    call_timeout=_DEFAULT_TIMEOUT_S,
                )
        except Exception:
            pass
        finally:
            self.state.sandbox_id = None
            self.state.workspace_root_ready = False
            self._sandbox = None
            self._running = False

    async def _ensure_sandbox(self) -> bool:
        if self._sandbox is not None:
            return False

        # If resuming, try to rehydrate the sandbox handle from the persisted id.
        sid = self.state.sandbox_id
        if sid:
            try:
                sb = await self._call_modal(
                    modal.Sandbox.from_id,
                    sid,
                    call_timeout=self.state.sandbox_create_timeout_s,
                )

                # `poll()` returns an exit code when the sandbox is terminated, else None.
                poll_result = await self._call_modal(sb.poll, call_timeout=_DEFAULT_TIMEOUT_S)
                is_running = poll_result is None
                if is_running:
                    self._sandbox = sb
                    self._running = True
                    return True
            except Exception:
                pass

            # Resumed sandbox handle is dead or invalid; clear and create a fresh one.
            self._sandbox = None
            self.state.sandbox_id = None

        app = await self._call_modal(
            modal.App.lookup,
            self.state.app_name,
            create_if_missing=True,
            call_timeout=10.0,
        )
        if not self._image:
            image_id = self.state.image_id
            if image_id:
                self._image = modal.Image.from_id(image_id)
            else:
                tag = self.state.image_tag
                if not isinstance(tag, str) or not tag:
                    tag = _DEFAULT_IMAGE_TAG
                    # Record the default for better debuggability/resume.
                    self.state.image_tag = tag
                self._image = await self._call_modal(
                    modal.Image.from_registry,
                    tag,
                    call_timeout=_DEFAULT_TIMEOUT_S,
                )
            self._image = _maybe_set_sandbox_cmd(
                self._image,
                use_sleep_cmd=self.state.use_sleep_cmd,
            )

        manifest_envs = cast(dict[str, str | None], await self.state.manifest.environment.resolve())
        volumes = self._modal_cloud_bucket_mounts_for_manifest()
        create_coro = modal.Sandbox.create.aio(
            app=app,
            image=self._image,
            workdir=self.state.manifest.root,
            env=manifest_envs,
            encrypted_ports=self.state.exposed_ports,
            volumes=volumes,
            gpu=self.state.gpu,
            timeout=self.state.timeout,
            idle_timeout=self.state.idle_timeout,
        )
        async with _override_modal_image_builder_version(self.state.image_builder_version):
            if self.state.sandbox_create_timeout_s is None:
                self._sandbox = await create_coro
            else:
                self._sandbox = await asyncio.wait_for(
                    create_coro, timeout=self.state.sandbox_create_timeout_s
                )

        # Persist sandbox id for future resume.
        assert self._sandbox is not None
        self.state.sandbox_id = self._sandbox.object_id
        self.state.workspace_root_ready = False

        assert self._image is not None
        self.state.image_id = self._image.object_id
        return False

    async def snapshot_filesystem(self) -> str:
        """Snapshot the current sandbox filesystem and return the resulting Modal image ID.

        The returned ID can be passed as ``image_id`` when creating a new sandbox to boot
        from this filesystem state.  The image ID is also stored in ``state.image_id`` for future
        resume.
        """
        await self._ensure_sandbox()
        assert self._sandbox is not None
        snap_coro = self._sandbox.snapshot_filesystem.aio()
        if self.state.snapshot_filesystem_timeout_s is None:
            snap = await snap_coro
        else:
            snap = await asyncio.wait_for(
                snap_coro, timeout=self.state.snapshot_filesystem_timeout_s
            )
        image_id: str | None
        if isinstance(snap, str):
            image_id = snap
        else:
            image_id = getattr(snap, "object_id", None) or getattr(snap, "id", None)
        if not isinstance(image_id, str) or not image_id:
            raise RuntimeError(
                f"snapshot_filesystem returned unexpected type: {type(snap).__name__}"
            )
        self.state.image_id = image_id
        self._image = modal.Image.from_id(image_id)
        return image_id

    async def _exec_internal(
        self, *command: str | Path, timeout: float | None = None
    ) -> ExecResult:
        await self._ensure_sandbox()
        assert self._sandbox is not None

        modal_timeout: int | None = None
        if timeout is not None:
            # Modal's Sandbox.exec timeout is integer seconds; use ceil so the command
            # is guaranteed to be terminated server-side at or before our timeout window
            # (modulo 1s granularity).
            modal_timeout = int(max(_DEFAULT_TIMEOUT_S, math.ceil(timeout)))

        async def _run_async() -> ExecResult:
            assert self._sandbox is not None
            argv: tuple[str, ...] = tuple(str(part) for part in command)
            proc = await self._sandbox.exec.aio(*argv, text=False, timeout=modal_timeout)
            # Drain full output; Modal buffers process output server-side.
            stdout = await proc.stdout.read.aio()
            stderr = await proc.stderr.read.aio()
            exit_code = await proc.wait.aio()
            return ExecResult(stdout=stdout or b"", stderr=stderr or b"", exit_code=exit_code or 0)

        try:
            run_coro = _run_async()
            if timeout is None:
                return await run_coro
            return await asyncio.wait_for(run_coro, timeout=timeout)
        except asyncio.TimeoutError as e:
            sandbox = self._sandbox
            if sandbox is not None:
                try:
                    await self._call_modal(sandbox.terminate, call_timeout=_DEFAULT_TIMEOUT_S)
                except Exception:
                    pass
            self._sandbox = None
            self.state.sandbox_id = None
            self._running = False
            raise ExecTimeoutError(command=command, timeout_s=timeout, cause=e) from e
        except ExecTimeoutError:
            raise
        except Exception as e:
            raise ExecTransportError(command=command, cause=e) from e

    def supports_pty(self) -> bool:
        return True

    async def pty_exec_start(
        self,
        *command: str | Path,
        timeout: float | None = None,
        shell: bool | list[str] = True,
        user: str | User | None = None,
        tty: bool = False,
        yield_time_s: float | None = None,
        max_output_tokens: int | None = None,
    ) -> PtyExecUpdate:
        await self._ensure_sandbox()
        assert self._sandbox is not None

        sanitized_command = self._prepare_exec_command(*command, shell=shell, user=user)
        argv: tuple[str, ...] = tuple(str(part) for part in sanitized_command)
        modal_timeout: int | None = None
        if timeout is not None:
            modal_timeout = int(max(_DEFAULT_TIMEOUT_S, math.ceil(timeout)))

        entry: _ModalPtyProcessEntry | None = None
        registered = False
        pruned_entry: _ModalPtyProcessEntry | None = None
        process_id = 0
        process_count = 0
        try:
            process = cast(
                Any,
                await self._call_modal(
                    self._sandbox.exec,
                    *argv,
                    text=False,
                    timeout=modal_timeout,
                    pty=tty,
                ),
            )
            entry = _ModalPtyProcessEntry(process=process, tty=tty)

            async with self._pty_lock:
                process_id = allocate_pty_process_id(self._reserved_pty_process_ids)
                self._reserved_pty_process_ids.add(process_id)
                pruned_entry = await self._prune_pty_processes_if_needed()
                self._pty_processes[process_id] = entry
                registered = True
                process_count = len(self._pty_processes)
        except asyncio.TimeoutError as e:
            if entry is not None and not registered:
                await self._terminate_pty_entry(entry)
            raise ExecTimeoutError(command=command, timeout_s=timeout, cause=e) from e
        except asyncio.CancelledError:
            if entry is not None and not registered:
                await self._terminate_pty_entry(entry)
            raise
        except Exception as e:
            if entry is not None and not registered:
                await self._terminate_pty_entry(entry)
            raise ExecTransportError(command=command, cause=e) from e

        if pruned_entry is not None:
            await self._terminate_pty_entry(pruned_entry)

        if process_count >= PTY_PROCESSES_WARNING:
            logger.warning(
                "PTY process count reached warning threshold: %s active sessions",
                process_count,
            )

        yield_time_ms = 10_000 if yield_time_s is None else int(yield_time_s * 1000)
        output, original_token_count = await self._collect_pty_output(
            entry=entry,
            yield_time_ms=clamp_pty_yield_time_ms(yield_time_ms),
            max_output_tokens=max_output_tokens,
        )
        return await self._finalize_pty_update(
            process_id=process_id,
            entry=entry,
            output=output,
            original_token_count=original_token_count,
        )

    async def pty_write_stdin(
        self,
        *,
        session_id: int,
        chars: str,
        yield_time_s: float | None = None,
        max_output_tokens: int | None = None,
    ) -> PtyExecUpdate:
        async with self._pty_lock:
            entry = self._resolve_pty_session_entry(
                pty_processes=self._pty_processes,
                session_id=session_id,
            )

        if chars:
            if not entry.tty:
                raise RuntimeError("stdin is not available for this process")
            await self._write_pty_stdin(entry.process, chars.encode("utf-8"))
            await asyncio.sleep(0.1)

        yield_time_ms = 250 if yield_time_s is None else int(yield_time_s * 1000)
        output, original_token_count = await self._collect_pty_output(
            entry=entry,
            yield_time_ms=resolve_pty_write_yield_time_ms(
                yield_time_ms=yield_time_ms, input_empty=chars == ""
            ),
            max_output_tokens=max_output_tokens,
        )
        entry.last_used = time.monotonic()
        return await self._finalize_pty_update(
            process_id=session_id,
            entry=entry,
            output=output,
            original_token_count=original_token_count,
        )

    async def pty_terminate_all(self) -> None:
        async with self._pty_lock:
            entries = list(self._pty_processes.values())
            self._pty_processes.clear()
            self._reserved_pty_process_ids.clear()

        for entry in entries:
            await self._terminate_pty_entry(entry)

    async def _write_pty_stdin(self, process: ContainerProcess[bytes], payload: bytes) -> None:
        stdin = process.stdin
        write = getattr(stdin, "write", None)
        if not callable(write):
            raise RuntimeError("stdin is not writable for this process")
        await self._call_modal(write, payload, call_timeout=5.0)

        drain = getattr(stdin, "drain", None)
        if callable(drain):
            await self._call_modal(drain, call_timeout=5.0)

    async def _collect_pty_output(
        self,
        *,
        entry: _ModalPtyProcessEntry,
        yield_time_ms: int,
        max_output_tokens: int | None,
    ) -> tuple[bytes, int | None]:
        deadline = time.monotonic() + (yield_time_ms / 1000)
        chunks = bytearray()

        while True:
            stdout_chunk = await self._read_modal_stream(entry=entry, stream_name="stdout")
            stderr_chunk = await self._read_modal_stream(entry=entry, stream_name="stderr")
            if stdout_chunk:
                chunks.extend(stdout_chunk)
            if stderr_chunk:
                chunks.extend(stderr_chunk)

            if time.monotonic() >= deadline:
                break

            exit_code = await self._peek_exit_code(entry.process)
            if exit_code is not None:
                stdout_chunks = await self._drain_modal_stream(entry=entry, stream_name="stdout")
                stderr_chunks = await self._drain_modal_stream(entry=entry, stream_name="stderr")
                chunks.extend(stdout_chunks)
                chunks.extend(stderr_chunks)
                break

            if not stdout_chunk and not stderr_chunk:
                remaining_s = deadline - time.monotonic()
                if remaining_s <= 0:
                    break
                await asyncio.sleep(min(_PTY_POLL_INTERVAL_S, remaining_s))

        text = chunks.decode("utf-8", errors="replace")
        truncated_text, original_token_count = truncate_text_by_tokens(text, max_output_tokens)
        return truncated_text.encode("utf-8", errors="replace"), original_token_count

    async def _drain_modal_stream(
        self,
        *,
        entry: _ModalPtyProcessEntry,
        stream_name: Literal["stdout", "stderr"],
    ) -> bytes:
        chunks = bytearray()
        while True:
            chunk = await self._read_modal_stream(
                entry=entry,
                stream_name=stream_name,
                await_pending=True,
            )
            if not chunk:
                break
            chunks.extend(chunk)
        return bytes(chunks)

    async def _read_modal_stream(
        self,
        *,
        entry: _ModalPtyProcessEntry,
        stream_name: Literal["stdout", "stderr"],
        await_pending: bool = False,
    ) -> bytes:
        stream = entry.process.stdout if stream_name == "stdout" else entry.process.stderr
        if stream is None:
            return b""

        iter_attr = "stdout_iter" if stream_name == "stdout" else "stderr_iter"
        task_attr = "stdout_read_task" if stream_name == "stdout" else "stderr_read_task"
        stream_iter = getattr(entry, iter_attr)
        if stream_iter is None:
            aiter_method = getattr(stream, "__aiter__", None)
            if callable(aiter_method):
                try:
                    stream_iter = aiter_method()
                except Exception:
                    stream_iter = None
                else:
                    setattr(entry, iter_attr, stream_iter)

        task = getattr(entry, task_attr)
        if task is None and stream_iter is not None:
            task = asyncio.create_task(stream_iter.__anext__())
            setattr(entry, task_attr, task)

        if task is not None:
            wait_timeout = 0.2 if await_pending else 0
            done, _pending = await asyncio.wait({task}, timeout=wait_timeout)
            if not done:
                return b""

            setattr(entry, task_attr, None)
            try:
                value = task.result()
            except StopAsyncIteration:
                setattr(entry, iter_attr, None)
                return b""
            except Exception:
                setattr(entry, iter_attr, None)
                return b""

            return self._coerce_modal_stream_chunk(value)

        read = getattr(stream, "read", None)
        if not callable(read):
            return b""

        try:
            value = await self._call_modal(read, 16_384, call_timeout=0.2)
        except TypeError:
            return b""
        except Exception:
            return b""

        return self._coerce_modal_stream_chunk(value)

    def _coerce_modal_stream_chunk(self, value: object) -> bytes:
        if value is None:
            return b""
        if isinstance(value, bytes):
            return value
        if isinstance(value, bytearray):
            return bytes(value)
        if isinstance(value, str):
            return value.encode("utf-8", errors="replace")
        return str(value).encode("utf-8", errors="replace")

    async def _finalize_pty_update(
        self,
        *,
        process_id: int,
        entry: _ModalPtyProcessEntry,
        output: bytes,
        original_token_count: int | None,
    ) -> PtyExecUpdate:
        exit_code = await self._peek_exit_code(entry.process)
        live_process_id: int | None = process_id
        if exit_code is not None:
            async with self._pty_lock:
                removed = self._pty_processes.pop(process_id, None)
                self._reserved_pty_process_ids.discard(process_id)
            if removed is not None:
                await self._terminate_pty_entry(removed)
            live_process_id = None

        return PtyExecUpdate(
            process_id=live_process_id,
            output=output,
            exit_code=exit_code,
            original_token_count=original_token_count,
        )

    async def _prune_pty_processes_if_needed(self) -> _ModalPtyProcessEntry | None:
        if len(self._pty_processes) < PTY_PROCESSES_MAX:
            return None

        meta: list[tuple[int, float, bool]] = []
        for process_id, entry in self._pty_processes.items():
            exit_code = await self._peek_exit_code(entry.process)
            meta.append((process_id, entry.last_used, exit_code is not None))
        process_id_to_prune = process_id_to_prune_from_meta(meta)
        if process_id_to_prune is None:
            return None

        self._reserved_pty_process_ids.discard(process_id_to_prune)
        return self._pty_processes.pop(process_id_to_prune, None)

    async def _peek_exit_code(self, process: ContainerProcess[bytes]) -> int | None:
        try:
            value = await self._call_modal(process.poll, call_timeout=0.2)
        except Exception:
            return None

        if value is None:
            return None
        if isinstance(value, int):
            return value
        try:
            return int(value)
        except (TypeError, ValueError):
            return None

    async def _terminate_pty_entry(self, entry: _ModalPtyProcessEntry) -> None:
        process = entry.process
        for task in (entry.stdout_read_task, entry.stderr_read_task):
            if task is not None and not task.done():
                task.cancel()

        try:
            terminated = False
            terminate = getattr(process, "terminate", None)
            if callable(terminate):
                await self._call_modal(terminate, call_timeout=5.0)
                terminated = True

            if not terminated:
                stdin = getattr(process, "stdin", None)
            else:
                stdin = None
            if stdin is not None:
                write_eof = getattr(stdin, "write_eof", None)
                if callable(write_eof):
                    await self._call_modal(write_eof, call_timeout=5.0)
        except Exception:
            pass
        finally:
            await asyncio.gather(
                *(
                    task
                    for task in (entry.stdout_read_task, entry.stderr_read_task)
                    if task is not None
                ),
                return_exceptions=True,
            )

    async def read(self, path: Path, *, user: str | User | None = None) -> io.IOBase:
        if user is not None:
            await self._check_read_with_exec(path, user=user)

        # Read by `cat` so the payload is returned as bytes.
        workspace_path = await self._validate_path_access(path)
        cmd = ["sh", "-lc", f"cat -- {shlex.quote(sandbox_path_str(workspace_path))}"]
        try:
            out = await self.exec(*cmd, shell=False)
        except ExecTimeoutError as e:
            raise WorkspaceArchiveReadError(path=workspace_path, cause=e) from e
        except ExecTransportError as e:
            raise WorkspaceArchiveReadError(path=workspace_path, cause=e) from e

        if not out.ok():
            raise WorkspaceReadNotFoundError(
                path=path, context={"stderr": out.stderr.decode("utf-8", "replace")}
            )

        return io.BytesIO(out.stdout)

    async def write(
        self,
        path: Path,
        data: io.IOBase,
        *,
        user: str | User | None = None,
    ) -> None:
        if user is not None:
            await self._check_write_with_exec(path, user=user)

        payload = data.read()
        if isinstance(payload, str):
            payload = payload.encode("utf-8")
        if not isinstance(payload, bytes | bytearray):
            raise WorkspaceWriteTypeError(path=path, actual_type=type(payload).__name__)

        await self._ensure_sandbox()
        assert self._sandbox is not None

        workspace_path = await self._validate_path_access(path, for_write=True)

        async def _run_write() -> None:
            assert self._sandbox is not None
            # Ensure parent directory exists.
            parent = sandbox_path_str(workspace_path.parent)
            mkdir_proc = await self._sandbox.exec.aio("mkdir", "-p", "--", parent, text=False)
            await mkdir_proc.wait.aio()

            # Stream bytes into `cat > file` to avoid quoting/binary issues.
            cmd = ["sh", "-lc", f"cat > {shlex.quote(sandbox_path_str(workspace_path))}"]
            proc = await self._sandbox.exec.aio(*cmd, text=False)
            await _write_process_stdin(proc, payload)
            exit_code = await proc.wait.aio()
            if exit_code != 0:
                stderr = await proc.stderr.read.aio()
                raise WorkspaceArchiveWriteError(
                    path=workspace_path,
                    context={
                        "reason": "write_nonzero_exit",
                        "exit_code": exit_code,
                        "stderr": stderr.decode("utf-8", "replace"),
                    },
                )

        try:
            await asyncio.wait_for(_run_write(), timeout=30.0)
        except WorkspaceArchiveWriteError:
            raise
        except Exception as e:
            raise WorkspaceArchiveWriteError(path=workspace_path, cause=e) from e

    async def running(self) -> bool:
        if not self._running or self._sandbox is None:
            return False

        try:
            assert self._sandbox is not None
            poll_result = await asyncio.wait_for(self._sandbox.poll.aio(), timeout=5.0)
            return poll_result is None
        except Exception:
            return False

    async def persist_workspace(self) -> io.IOBase:
        if self.state.workspace_persistence == _WORKSPACE_PERSISTENCE_SNAPSHOT_FILESYSTEM:
            return await self._persist_workspace_via_snapshot_filesystem()
        if self.state.workspace_persistence == _WORKSPACE_PERSISTENCE_SNAPSHOT_DIRECTORY:
            return await self._persist_workspace_via_snapshot_directory()
        return await self._persist_workspace_via_tar()

    async def hydrate_workspace(self, data: io.IOBase) -> None:
        if self.state.workspace_persistence == _WORKSPACE_PERSISTENCE_SNAPSHOT_FILESYSTEM:
            return await self._hydrate_workspace_via_snapshot_filesystem(data)
        if self.state.workspace_persistence == _WORKSPACE_PERSISTENCE_SNAPSHOT_DIRECTORY:
            return await self._hydrate_workspace_via_snapshot_directory(data)
        return await self._hydrate_workspace_via_tar(data)

    async def _persist_workspace_via_snapshot_filesystem(self) -> io.IOBase:
        """
        Persist the workspace using Modal's snapshot_filesystem API when available.

        Modal's snapshot_filesystem is expected to return a snapshot reference
        (a Modal Image handle). We serialize a small reference envelope that
        `_hydrate_workspace_via_snapshot_filesystem` can interpret.
        """

        await self._ensure_sandbox()
        assert self._sandbox is not None
        if not hasattr(self._sandbox, "snapshot_filesystem"):
            return await self._persist_workspace_via_tar()
        if self._native_snapshot_requires_tar_fallback():
            return await self._persist_workspace_via_tar()
        root = self._workspace_root_path()
        error_root = posix_path_for_error(root)
        plain_skip = self._modal_snapshot_plain_skip_relpaths(root)
        skip_abs = [root / rel for rel in sorted(plain_skip, key=lambda p: p.as_posix())]
        self._modal_snapshot_ephemeral_backup = None
        self._modal_snapshot_ephemeral_backup_path = None

        async def restore_ephemeral_paths() -> WorkspaceArchiveReadError | None:
            backup = self._modal_snapshot_ephemeral_backup
            if not backup:
                return None

            try:
                assert self._sandbox is not None
                proc = await self._sandbox.exec.aio(
                    "tar", "xf", "-", "-C", root.as_posix(), text=False
                )
                await _write_process_stdin(proc, bytes(backup))
                exit_code = await proc.wait.aio()
                if exit_code != 0:
                    stderr = await proc.stderr.read.aio()
                    return WorkspaceArchiveReadError(
                        path=error_root,
                        context={
                            "reason": "snapshot_filesystem_ephemeral_restore_failed",
                            "exit_code": exit_code,
                            "stderr": stderr.decode("utf-8", "replace"),
                        },
                    )
            except Exception as exc:
                if isinstance(exc, WorkspaceArchiveReadError):
                    return exc
                return WorkspaceArchiveReadError(
                    path=error_root,
                    context={"reason": "snapshot_filesystem_ephemeral_restore_failed"},
                    cause=exc,
                )
            return None

        if skip_abs:
            rel_args = " ".join(shlex.quote(p.relative_to(root).as_posix()) for p in skip_abs)
            cmd = (
                f"cd -- {shlex.quote(root.as_posix())} && "
                f"(tar cf - -- {rel_args} 2>/dev/null || true)"
            )
            out = await self.exec("sh", "-lc", cmd, shell=False)
            self._modal_snapshot_ephemeral_backup = out.stdout or b""

            rm_cmd = ["rm", "-rf", "--", *[p.as_posix() for p in skip_abs]]
            rm_out = await self.exec(*rm_cmd, shell=False)
            if not rm_out.ok():
                cleanup_restore_error = await restore_ephemeral_paths()
                if cleanup_restore_error is not None:
                    logger.warning(
                        "Failed to restore Modal ephemeral paths after cleanup failure: %s",
                        cleanup_restore_error,
                    )
                raise WorkspaceArchiveReadError(
                    path=error_root,
                    context={
                        "reason": "snapshot_filesystem_ephemeral_remove_failed",
                        "exit_code": rm_out.exit_code,
                        "stderr": rm_out.stderr.decode("utf-8", "replace"),
                    },
                )

        try:
            snapshot_sandbox = await self._refresh_sandbox_handle_for_snapshot()
            snap_coro = snapshot_sandbox.snapshot_filesystem.aio()
            if self.state.snapshot_filesystem_timeout_s is None:
                snap = await snap_coro
            else:
                snap = await asyncio.wait_for(
                    snap_coro, timeout=self.state.snapshot_filesystem_timeout_s
                )
        except Exception as e:
            restore_error = await restore_ephemeral_paths()
            if restore_error is not None:
                logger.warning(
                    "Failed to restore Modal ephemeral paths after snapshot failure: %s",
                    restore_error,
                )
            raise WorkspaceArchiveReadError(
                path=error_root, context={"reason": "snapshot_filesystem_failed"}, cause=e
            ) from e

        snapshot_id, snapshot_error = self._extract_modal_snapshot_id(
            snap=snap, root=root, snapshot_kind="snapshot_filesystem"
        )

        restore_error = await restore_ephemeral_paths()
        if restore_error is not None:
            raise restore_error

        if snapshot_error is not None:
            raise snapshot_error

        assert snapshot_id is not None
        return io.BytesIO(_encode_snapshot_filesystem_ref(snapshot_id=snapshot_id))

    async def _persist_workspace_via_snapshot_directory(self) -> io.IOBase:
        """
        Persist the workspace using Modal's snapshot_directory API when available.
        """

        root = self._workspace_root_path()
        error_root = posix_path_for_error(root)
        await self._ensure_sandbox()
        assert self._sandbox is not None
        if not hasattr(self._sandbox, "snapshot_directory"):
            return await self._persist_workspace_via_tar()
        if self._native_snapshot_requires_tar_fallback():
            return await self._persist_workspace_via_tar()
        plain_skip = self._modal_snapshot_plain_skip_relpaths(root)
        skip_abs = [root / rel for rel in sorted(plain_skip, key=lambda p: p.as_posix())]
        self._modal_snapshot_ephemeral_backup = None
        self._modal_snapshot_ephemeral_backup_path = None
        detached_mounts: list[tuple[Mount, Path]] = []

        async def restore_ephemeral_paths() -> WorkspaceArchiveReadError | None:
            backup_path = self._modal_snapshot_ephemeral_backup_path
            if backup_path is None:
                return None

            restore_cmd = (
                f"if [ ! -f {shlex.quote(backup_path.as_posix())} ]; then "
                f"echo missing ephemeral backup archive >&2; "
                f"exit 1; "
                f"fi; "
                f"tar xf {shlex.quote(backup_path.as_posix())} -C "
                f"{shlex.quote(root.as_posix())} && "
                f"rm -f -- {shlex.quote(backup_path.as_posix())}"
            )
            out = await self.exec("sh", "-lc", restore_cmd, shell=False)
            if not out.ok():
                return WorkspaceArchiveReadError(
                    path=error_root,
                    context={
                        "reason": "snapshot_directory_ephemeral_restore_failed",
                        "exit_code": out.exit_code,
                        "stderr": out.stderr.decode("utf-8", "replace"),
                    },
                )
            return None

        async def restore_detached_mounts() -> WorkspaceArchiveReadError | None:
            remount_error: WorkspaceArchiveReadError | None = None
            for mount_entry, mount_path in reversed(detached_mounts):
                try:
                    await mount_entry.mount_strategy.restore_after_snapshot(
                        mount_entry,
                        self,
                        mount_path,
                    )
                except Exception as e:
                    current_error = WorkspaceArchiveReadError(path=error_root, cause=e)
                    if remount_error is None:
                        remount_error = current_error
                    else:
                        additional_remount_errors = remount_error.context.setdefault(
                            "additional_remount_errors", []
                        )
                        assert isinstance(additional_remount_errors, list)
                        additional_remount_errors.append(
                            {
                                "message": current_error.message,
                                "cause_type": type(e).__name__,
                                "cause": str(e),
                            }
                        )
            return remount_error

        snapshot_error: WorkspaceArchiveReadError | None = None
        snapshot_id: str | None = None
        try:
            if skip_abs:
                backup_path = posix_path_as_path(
                    coerce_posix_path(
                        "/tmp/openai-agents/session-state"
                        f"/{self.state.session_id.hex}/modal-snapshot-directory-ephemeral.tar"
                    )
                )
                rel_args = " ".join(shlex.quote(p.relative_to(root).as_posix()) for p in skip_abs)
                backup_cmd = (
                    f"mkdir -p -- {shlex.quote(backup_path.parent.as_posix())} && "
                    f"cd -- {shlex.quote(root.as_posix())} && "
                    "{ "
                    f"for rel in {rel_args}; do "
                    'if [ -e "$rel" ]; then printf \'%s\\n\' "$rel"; fi; '
                    "done; "
                    "} | "
                    f"tar cf {shlex.quote(backup_path.as_posix())} -T - 2>/dev/null && "
                    f"test -f {shlex.quote(backup_path.as_posix())}"
                )
                backup_out = await self.exec("sh", "-lc", backup_cmd, shell=False)
                if not backup_out.ok():
                    raise WorkspaceArchiveReadError(
                        path=error_root,
                        context={
                            "reason": "snapshot_directory_ephemeral_backup_failed",
                            "exit_code": backup_out.exit_code,
                            "stderr": backup_out.stderr.decode("utf-8", "replace"),
                        },
                    )
                self._modal_snapshot_ephemeral_backup_path = backup_path

                rm_cmd = ["rm", "-rf", "--", *[sandbox_path_str(p) for p in skip_abs]]
                rm_out = await self.exec(*rm_cmd, shell=False)
                if not rm_out.ok():
                    raise WorkspaceArchiveReadError(
                        path=error_root,
                        context={
                            "reason": "snapshot_directory_ephemeral_remove_failed",
                            "exit_code": rm_out.exit_code,
                            "stderr": rm_out.stderr.decode("utf-8", "replace"),
                        },
                    )

            for mount_entry, mount_path in self._snapshot_directory_mount_targets_to_restore(root):
                await mount_entry.mount_strategy.teardown_for_snapshot(
                    mount_entry,
                    self,
                    mount_path,
                )
                detached_mounts.append((mount_entry, mount_path))

            snapshot_sandbox = await self._refresh_sandbox_handle_for_snapshot()
            snap_coro = snapshot_sandbox.snapshot_directory.aio(root.as_posix())
            if self.state.snapshot_filesystem_timeout_s is None:
                snap = await snap_coro
            else:
                snap = await asyncio.wait_for(
                    snap_coro, timeout=self.state.snapshot_filesystem_timeout_s
                )
            snapshot_id, snapshot_error = self._extract_modal_snapshot_id(
                snap=snap, root=root, snapshot_kind="snapshot_directory"
            )
        except WorkspaceArchiveReadError as e:
            snapshot_error = e
        except Exception as e:
            snapshot_error = WorkspaceArchiveReadError(
                path=error_root, context={"reason": "snapshot_directory_failed"}, cause=e
            )
        finally:
            remount_error = await restore_detached_mounts()
            restore_error = await restore_ephemeral_paths()
            cleanup_error = remount_error
            if restore_error is not None:
                if cleanup_error is None:
                    cleanup_error = restore_error
                else:
                    additional_restore_errors = cleanup_error.context.setdefault(
                        "additional_restore_errors", []
                    )
                    assert isinstance(additional_restore_errors, list)
                    additional_restore_errors.append(
                        {
                            "message": restore_error.message,
                            "cause_type": (
                                type(restore_error.cause).__name__
                                if restore_error.cause is not None
                                else None
                            ),
                            "cause": str(restore_error.cause) if restore_error.cause else None,
                        }
                    )

            if cleanup_error is not None:
                if snapshot_error is not None:
                    cleanup_error.context["snapshot_error_before_restore_corruption"] = {
                        "message": snapshot_error.message
                    }
                raise cleanup_error

        if snapshot_error is not None:
            raise snapshot_error

        assert snapshot_id is not None
        return io.BytesIO(_encode_snapshot_directory_ref(snapshot_id=snapshot_id))

    def _extract_modal_snapshot_id(
        self,
        *,
        snap: object,
        root: Path,
        snapshot_kind: Literal["snapshot_filesystem", "snapshot_directory"],
    ) -> tuple[str | None, WorkspaceArchiveReadError | None]:
        if isinstance(snap, bytes | bytearray):
            return None, WorkspaceArchiveReadError(
                path=posix_path_for_error(root),
                context={
                    "reason": f"{snapshot_kind}_unexpected_bytes",
                    "type": type(snap).__name__,
                },
            )
        if not hasattr(snap, "object_id") and not isinstance(snap, str):
            return None, WorkspaceArchiveReadError(
                path=posix_path_for_error(root),
                context={
                    "reason": f"{snapshot_kind}_unexpected_return",
                    "type": type(snap).__name__,
                },
            )
        if isinstance(snap, str):
            return snap, None
        snapshot_id = getattr(snap, "object_id", None)
        if snapshot_id is not None and not isinstance(snapshot_id, str):
            snapshot_id = None
        if not snapshot_id:
            return None, WorkspaceArchiveReadError(
                path=posix_path_for_error(root),
                context={
                    "reason": f"{snapshot_kind}_unexpected_return",
                    "type": type(snap).__name__,
                },
            )
        return snapshot_id, None

    async def _refresh_sandbox_handle_for_snapshot(self) -> modal.Sandbox:
        await self._ensure_sandbox()
        assert self._sandbox is not None

        sandbox_module = type(self._sandbox).__module__
        if not sandbox_module.startswith("modal"):
            return self._sandbox

        sandbox_id = self.state.sandbox_id or getattr(self._sandbox, "object_id", None)
        if not sandbox_id:
            return self._sandbox

        try:
            refreshed = await self._call_modal(
                modal.Sandbox.from_id,
                sandbox_id,
                call_timeout=_DEFAULT_TIMEOUT_S,
            )
        except Exception:
            return self._sandbox

        self._sandbox = refreshed
        return refreshed

    def _modal_snapshot_plain_skip_relpaths(self, root: Path) -> set[Path]:
        plain_skip = set(self.state.manifest.ephemeral_entry_paths())
        if self._runtime_persist_workspace_skip_relpaths:
            plain_skip.update(self._runtime_persist_workspace_skip_relpaths)

        mount_skip_rel_paths: set[Path] = set()
        for rel_path, artifact in self.state.manifest.iter_entries():
            if isinstance(artifact, Mount) and artifact.ephemeral:
                mount_skip_rel_paths.add(rel_path)
        for _mount_entry, mount_path in self.state.manifest.ephemeral_mount_targets():
            try:
                mount_skip_rel_paths.add(mount_path.relative_to(root))
            except ValueError:
                continue
        return plain_skip - mount_skip_rel_paths

    def _modal_tar_skip_relpaths(self, root: Path) -> set[Path]:
        """Return Modal tar-capture skip paths, including resolved mount targets."""

        skip = self._persist_workspace_skip_relpaths()
        for _mount_entry, mount_path in self.state.manifest.mount_targets():
            try:
                skip.add(mount_path.relative_to(root))
            except ValueError:
                continue
        return skip

    @retry_async(
        retry_if=lambda exc, self: (
            exception_chain_contains_type(exc, (ExecTransportError,))
            or exception_chain_has_status_code(exc, TRANSIENT_HTTP_STATUS_CODES)
        )
    )
    async def _persist_workspace_via_tar(self) -> io.IOBase:
        # Existing tar implementation extracted so snapshot_filesystem mode can fall back cleanly.
        root = self._workspace_root_path()
        error_root = posix_path_for_error(root)
        skip = self._modal_tar_skip_relpaths(root)

        excludes: list[str] = []
        for rel in sorted(skip, key=lambda p: p.as_posix()):
            excludes.extend(["--exclude", f"./{rel.as_posix().lstrip('./')}"])

        cmd: list[str] = [
            "tar",
            "cf",
            "-",
            *excludes,
            "-C",
            root.as_posix(),
            ".",
        ]

        try:
            out = await self.exec(*cmd, shell=False)
            if not out.ok():
                raise WorkspaceArchiveReadError(
                    path=error_root,
                    context={
                        "reason": "tar_nonzero_exit",
                        "exit_code": out.exit_code,
                        "stderr": out.stderr.decode("utf-8", "replace"),
                    },
                )
            return io.BytesIO(out.stdout)
        except WorkspaceArchiveReadError:
            raise
        except Exception as e:
            raise WorkspaceArchiveReadError(path=error_root, cause=e) from e

    async def _hydrate_workspace_via_snapshot_filesystem(self, data: io.IOBase) -> None:
        """
        Hydrate using Modal's snapshot_filesystem restore API when the
        persisted payload is a snapshot ref. Otherwise, fall back to tar
        extraction (to support SDKs that return tar bytes).
        """
        root = self._workspace_root_path()
        raw, snapshot_id = self._read_modal_snapshot_id_from_archive(
            data=data.read(),
            expected_persistence=_WORKSPACE_PERSISTENCE_SNAPSHOT_FILESYSTEM,
            invalid_reason="snapshot_filesystem_invalid_snapshot_id",
        )
        if snapshot_id is None:
            return await self._hydrate_workspace_via_tar(io.BytesIO(raw))
        await self._restore_snapshot_filesystem_image(snapshot_id=snapshot_id, root=root)

    async def _hydrate_workspace_via_snapshot_directory(self, data: io.IOBase) -> None:
        """
        Hydrate using Modal's snapshot_directory restore API when the
        persisted payload is a snapshot ref. Otherwise, fall back to tar extraction.
        """

        root = self._workspace_root_path()
        raw, snapshot_id = self._read_modal_snapshot_id_from_archive(
            data=data.read(),
            expected_persistence=_WORKSPACE_PERSISTENCE_SNAPSHOT_DIRECTORY,
            invalid_reason="snapshot_directory_invalid_snapshot_id",
        )
        if snapshot_id is None:
            return await self._hydrate_workspace_via_tar(io.BytesIO(raw))
        await self._restore_snapshot_directory_image(snapshot_id=snapshot_id, root=root)

    def _read_modal_snapshot_id_from_archive(
        self,
        *,
        data: object,
        expected_persistence: WorkspacePersistenceMode,
        invalid_reason: str,
    ) -> tuple[bytes, str | None]:
        root = self._workspace_root_path()
        raw = data
        if isinstance(raw, str):
            raw = raw.encode("utf-8")
        if not isinstance(raw, bytes | bytearray):
            raise WorkspaceArchiveWriteError(path=root, context={"reason": "non_bytes_payload"})
        raw_bytes = bytes(raw)

        snapshot_ref = _decode_modal_snapshot_ref(raw_bytes)
        if snapshot_ref is None:
            return raw_bytes, None
        workspace_persistence, snapshot_id = snapshot_ref
        if workspace_persistence != expected_persistence:
            raise WorkspaceArchiveWriteError(
                path=root,
                context={"reason": invalid_reason, "workspace_persistence": workspace_persistence},
            )
        if not snapshot_id:
            raise WorkspaceArchiveWriteError(path=root, context={"reason": invalid_reason})
        return raw_bytes, snapshot_id

    async def _restore_snapshot_filesystem_image(self, *, snapshot_id: str, root: Path) -> None:
        prior = self._sandbox
        if prior is not None:
            try:
                await self._call_modal(prior.terminate, call_timeout=_DEFAULT_TIMEOUT_S)
            except Exception:
                pass
            finally:
                self._sandbox = None
                self.state.sandbox_id = None

        manifest_envs = cast(dict[str, str | None], await self.state.manifest.environment.resolve())

        async def _run_restore() -> None:
            image = modal.Image.from_id(snapshot_id)
            app = await modal.App.lookup.aio(self.state.app_name, create_if_missing=True)
            sb = await modal.Sandbox.create.aio(
                app=app,
                image=image,
                workdir=self.state.manifest.root,
                env=manifest_envs,
                encrypted_ports=self.state.exposed_ports,
                volumes=self._modal_cloud_bucket_mounts_for_manifest(),
                gpu=self.state.gpu,
                timeout=self.state.timeout,
                idle_timeout=self.state.idle_timeout,
            )
            try:
                mkdir_proc = await sb.exec.aio("mkdir", "-p", "--", root.as_posix(), text=False)
                await mkdir_proc.wait.aio()
            except Exception:
                pass
            self._image = image
            self.state.image_id = snapshot_id
            self._sandbox = sb
            self.state.sandbox_id = sb.object_id

        try:
            await asyncio.wait_for(
                _run_restore(), timeout=self.state.snapshot_filesystem_restore_timeout_s
            )
        except Exception as e:
            raise WorkspaceArchiveWriteError(
                path=root,
                context={
                    "reason": "snapshot_filesystem_restore_failed",
                    "snapshot_id": snapshot_id,
                },
                cause=e,
            ) from e

    async def _restore_snapshot_directory_image(self, *, snapshot_id: str, root: Path) -> None:
        await self._ensure_sandbox()
        assert self._sandbox is not None
        sandbox = self._sandbox

        async def _run_restore() -> None:
            image = modal.Image.from_id(snapshot_id)
            await self._call_modal(
                sandbox.mount_image,
                root.as_posix(),
                image,
                call_timeout=self.state.snapshot_filesystem_restore_timeout_s,
            )
            for mount_entry, mount_path in reversed(
                self._snapshot_directory_mount_targets_to_restore(root)
            ):
                await mount_entry.mount_strategy.restore_after_snapshot(
                    mount_entry,
                    self,
                    mount_path,
                )

        try:
            await asyncio.wait_for(
                _run_restore(), timeout=self.state.snapshot_filesystem_restore_timeout_s
            )
        except Exception as e:
            raise WorkspaceArchiveWriteError(
                path=root,
                context={
                    "reason": "snapshot_directory_restore_failed",
                    "snapshot_id": snapshot_id,
                },
                cause=e,
            ) from e

    def _snapshot_directory_mount_targets_to_restore(self, root: Path) -> list[tuple[Mount, Path]]:
        mount_targets: list[tuple[Mount, Path]] = []
        for mount_entry, mount_path in self.state.manifest.mount_targets():
            if mount_entry.ephemeral:
                continue
            if isinstance(mount_entry.mount_strategy, ModalCloudBucketMountStrategy):
                continue
            if mount_path != root and root not in mount_path.parents:
                continue
            mount_targets.append((mount_entry, mount_path))
        return mount_targets

    async def _hydrate_workspace_via_tar(self, data: io.IOBase) -> None:
        root = self._workspace_root_path()

        raw = data.read()
        if isinstance(raw, str):
            raw = raw.encode("utf-8")
        if not isinstance(raw, bytes | bytearray):
            raise WorkspaceArchiveWriteError(path=root, context={"reason": "non_bytes_tar_payload"})

        try:
            validate_tar_bytes(
                bytes(raw),
                skip_rel_paths=self.state.manifest.ephemeral_persistence_paths(),
            )
        except UnsafeTarMemberError as e:
            raise WorkspaceArchiveWriteError(
                path=root, context={"reason": e.reason, "member": e.member}, cause=e
            ) from e

        await self._ensure_sandbox()
        assert self._sandbox is not None

        async def _run_extract() -> None:
            assert self._sandbox is not None
            mkdir_proc = await self._sandbox.exec.aio(
                "mkdir", "-p", "--", root.as_posix(), text=False
            )
            await mkdir_proc.wait.aio()
            proc = await self._sandbox.exec.aio("tar", "xf", "-", "-C", root.as_posix(), text=False)
            await _write_process_stdin(proc, raw)
            exit_code = await proc.wait.aio()
            if exit_code != 0:
                stderr = await proc.stderr.read.aio()
                raise WorkspaceArchiveWriteError(
                    path=root,
                    context={
                        "reason": "tar_extract_nonzero_exit",
                        "exit_code": exit_code,
                        "stderr": stderr.decode("utf-8", "replace"),
                    },
                )

        try:
            await asyncio.wait_for(_run_extract(), timeout=60.0)
        except WorkspaceArchiveWriteError:
            raise
        except Exception as e:
            raise WorkspaceArchiveWriteError(path=root, cause=e) from e

    def _modal_cloud_bucket_mounts_for_manifest(
        self,
    ) -> dict[str | os.PathLike[Any], modal.Volume | modal.CloudBucketMount]:
        volumes: dict[str | os.PathLike[Any], modal.Volume | modal.CloudBucketMount] = {}
        for mount_entry, mount_path in self.state.manifest.mount_targets():
            strategy = mount_entry.mount_strategy
            if not isinstance(strategy, ModalCloudBucketMountStrategy):
                continue
            config = strategy._build_modal_cloud_bucket_mount_config(mount_entry)
            secret = None
            if config.secret_name is not None:
                secret = modal.Secret.from_name(
                    config.secret_name,
                    environment_name=config.secret_environment_name,
                )
            elif config.credentials is not None:
                secret = modal.Secret.from_dict(cast(dict[str, str | None], config.credentials))
            volumes[mount_path.as_posix()] = modal.CloudBucketMount(
                bucket_name=config.bucket_name,
                bucket_endpoint_url=config.bucket_endpoint_url,
                key_prefix=config.key_prefix,
                secret=secret,
                read_only=config.read_only,
            )
        return volumes


class ModalSandboxClient(BaseSandboxClient[ModalSandboxClientOptions]):
    backend_id = "modal"
    _default_image: ModalImageSelector | None
    _default_sandbox: ModalSandboxSelector | None
    _instrumentation: Instrumentation

    def __init__(
        self,
        *,
        image: ModalImageSelector | None = None,
        sandbox: ModalSandboxSelector | None = None,
        instrumentation: Instrumentation | None = None,
        dependencies: Dependencies | None = None,
    ) -> None:
        self._default_image = image
        self._default_sandbox = sandbox
        self._instrumentation = instrumentation or Instrumentation()
        self._dependencies = dependencies

    def _validate_manifest_for_workspace_persistence(
        self,
        *,
        manifest: Manifest,
        workspace_persistence: WorkspacePersistenceMode,
    ) -> None:
        if workspace_persistence != _WORKSPACE_PERSISTENCE_SNAPSHOT_DIRECTORY:
            return

        root = posix_path_as_path(coerce_posix_path(manifest.root))
        for mount_entry, mount_path in manifest.mount_targets():
            if not isinstance(mount_entry.mount_strategy, ModalCloudBucketMountStrategy):
                continue
            if mount_path == root or root in mount_path.parents:
                raise MountConfigError(
                    message=(
                        "snapshot_directory is not supported when a Modal cloud bucket mount "
                        "lives at or under the workspace root"
                    ),
                    context={
                        "workspace_root": root.as_posix(),
                        "mount_path": mount_path.as_posix(),
                        "workspace_persistence": workspace_persistence,
                    },
                )

    async def create(
        self,
        *,
        snapshot: SnapshotSpec | SnapshotBase | None = None,
        manifest: Manifest | None = None,
        options: ModalSandboxClientOptions,
    ) -> SandboxSession:
        """
        Create a new Modal-backed session.

        Expected options:
        - app_name: str (required)
        - sandbox_create_timeout_s: float | None (async timeout for sandbox creation call)
        - workspace_persistence: Literal["tar", "snapshot_filesystem", "snapshot_directory"]
          (optional)
        - snapshot_filesystem_timeout_s: float | None
          (async timeout for snapshot_filesystem call)
        - snapshot_filesystem_restore_timeout_s: float | None
          (async timeout for snapshot restore call)
        - timeout: int (maximum sandbox lifetime in seconds, default 300)
        - idle_timeout: int | None (maximum sandbox inactivity in seconds, default None)
        - image_builder_version: str | None (Modal image builder version, default "2025.06")
        """

        if options is None:
            raise ValueError("ModalSandboxClient.create requires options with app_name")
        manifest = manifest or Manifest()
        app_name = options.app_name
        if not app_name:
            raise ValueError("ModalSandboxClient.create requires a valid app_name")

        image_sel = self._default_image

        sandbox_sel = self._default_sandbox

        sandbox_create_timeout_s = options.sandbox_create_timeout_s
        if sandbox_create_timeout_s is not None and not isinstance(
            sandbox_create_timeout_s, int | float
        ):
            raise ValueError(
                "ModalSandboxClient.create requires sandbox_create_timeout_s to be a number"
            )

        workspace_persistence = options.workspace_persistence
        if workspace_persistence not in (
            _WORKSPACE_PERSISTENCE_TAR,
            _WORKSPACE_PERSISTENCE_SNAPSHOT_FILESYSTEM,
            _WORKSPACE_PERSISTENCE_SNAPSHOT_DIRECTORY,
        ):
            raise ValueError(
                "ModalSandboxClient.create requires workspace_persistence to be one of "
                f"{_WORKSPACE_PERSISTENCE_TAR!r}, "
                f"{_WORKSPACE_PERSISTENCE_SNAPSHOT_FILESYSTEM!r}, or "
                f"{_WORKSPACE_PERSISTENCE_SNAPSHOT_DIRECTORY!r}"
            )
        snapshot_filesystem_timeout_s = options.snapshot_filesystem_timeout_s
        if snapshot_filesystem_timeout_s is not None and not isinstance(
            snapshot_filesystem_timeout_s, int | float
        ):
            raise ValueError(
                "ModalSandboxClient.create requires snapshot_filesystem_timeout_s to be a number"
            )

        snapshot_filesystem_restore_timeout_s = options.snapshot_filesystem_restore_timeout_s
        if snapshot_filesystem_restore_timeout_s is not None and not isinstance(
            snapshot_filesystem_restore_timeout_s, int | float
        ):
            raise ValueError(
                "ModalSandboxClient.create requires "
                "snapshot_filesystem_restore_timeout_s to be a number"
            )
        image_builder_version = options.image_builder_version
        if "image_builder_version" not in options.model_fields_set or image_builder_version == "":
            image_builder_version = _DEFAULT_IMAGE_BUILDER_VERSION
        elif image_builder_version is not None and not isinstance(image_builder_version, str):
            raise ValueError(
                "ModalSandboxClient.create requires image_builder_version to be a string or None"
            )

        self._validate_manifest_for_workspace_persistence(
            manifest=manifest,
            workspace_persistence=workspace_persistence,
        )

        session_id = uuid.uuid4()
        state_image_id: str | None = None
        state_image_tag: str | None = None
        session_image: modal.Image | None = None
        if image_sel is not None:
            if image_sel.kind == "image":
                if not isinstance(image_sel.value, modal.Image):
                    raise ValueError(
                        "ModalSandboxClient.__init__ requires image to be a modal.Image"
                    )
                session_image = image_sel.value
                state_image_id = getattr(session_image, "object_id", None)
            elif image_sel.kind == "id":
                if not isinstance(image_sel.value, str) or not image_sel.value:
                    raise ValueError(
                        "ModalSandboxClient.__init__ requires image_id to be a non-empty string"
                    )
                state_image_id = image_sel.value
            else:
                if not isinstance(image_sel.value, str) or not image_sel.value:
                    raise ValueError(
                        "ModalSandboxClient.__init__ requires image_tag to be a non-empty string"
                    )
                state_image_tag = image_sel.value

        state_sandbox_id: str | None = None
        session_sandbox: modal.Sandbox | None = None
        if sandbox_sel is not None:
            if sandbox_sel.kind == "sandbox":
                if not isinstance(sandbox_sel.value, modal.Sandbox):
                    raise ValueError(
                        "ModalSandboxClient.__init__ requires sandbox to be a modal.Sandbox"
                    )
                session_sandbox = sandbox_sel.value
                state_sandbox_id = getattr(session_sandbox, "object_id", None)
            else:
                if not isinstance(sandbox_sel.value, str) or not sandbox_sel.value:
                    raise ValueError(
                        "ModalSandboxClient.__init__ requires sandbox_id to be a non-empty string"
                    )
                state_sandbox_id = sandbox_sel.value

        snapshot_id = str(session_id)
        snapshot_instance = resolve_snapshot(snapshot, snapshot_id)
        state = ModalSandboxSessionState(
            session_id=session_id,
            manifest=manifest,
            snapshot=snapshot_instance,
            app_name=app_name,
            image_tag=state_image_tag,
            image_id=state_image_id,
            sandbox_id=state_sandbox_id,
            workspace_persistence=workspace_persistence,
            exposed_ports=options.exposed_ports,
            gpu=options.gpu,
            timeout=options.timeout,
            use_sleep_cmd=options.use_sleep_cmd,
            image_builder_version=image_builder_version,
            idle_timeout=options.idle_timeout,
        )
        if sandbox_create_timeout_s is not None:
            state.sandbox_create_timeout_s = float(sandbox_create_timeout_s)
        if snapshot_filesystem_timeout_s is not None:
            state.snapshot_filesystem_timeout_s = float(snapshot_filesystem_timeout_s)
        if snapshot_filesystem_restore_timeout_s is not None:
            state.snapshot_filesystem_restore_timeout_s = float(
                snapshot_filesystem_restore_timeout_s
            )

        # Pass the in-memory handles through to the session (they may not be resumable).
        inner = ModalSandboxSession.from_state(
            state,
            image=session_image,
            sandbox=session_sandbox,
        )
        await inner._ensure_sandbox()
        return self._wrap_session(inner, instrumentation=self._instrumentation)

    async def delete(self, session: SandboxSession) -> SandboxSession:
        """
        Best-effort cleanup of Modal sandbox resources.
        """

        inner = session._inner
        if not isinstance(inner, ModalSandboxSession):
            raise TypeError("ModalSandboxClient.delete expects a ModalSandboxSession")

        # Prefer the live handle if present.
        sandbox = getattr(inner, "_sandbox", None)
        try:
            if sandbox is not None:
                await inner._call_modal(sandbox.terminate, call_timeout=_DEFAULT_TIMEOUT_S)
                return session
        except Exception:
            return session

        # Otherwise, best-effort terminate via sandbox_id.
        sid = inner.state.sandbox_id
        if sid:
            try:
                sb = await inner._call_modal(
                    modal.Sandbox.from_id,
                    sid,
                    call_timeout=_DEFAULT_TIMEOUT_S,
                )
                await inner._call_modal(sb.terminate, call_timeout=_DEFAULT_TIMEOUT_S)
            except Exception:
                pass

        return session

    async def resume(
        self,
        state: SandboxSessionState,
    ) -> SandboxSession:
        if not isinstance(state, ModalSandboxSessionState):
            raise TypeError("ModalSandboxClient.resume expects a ModalSandboxSessionState")
        inner = ModalSandboxSession.from_state(state)
        reconnected = await inner._ensure_sandbox()
        if reconnected:
            inner._set_start_state_preserved(True)
        return self._wrap_session(inner, instrumentation=self._instrumentation)

    def deserialize_session_state(self, payload: dict[str, object]) -> SandboxSessionState:
        return ModalSandboxSessionState.model_validate(payload)