agentenv-framework 0.9.1254__py3-none-any.whl

agent_env/env/legacy_protocol.py

@@ -0,0 +1,146 @@
+"""Legacy (pre-/v1) data-plane wire protocol for MCP server / website envs.
+
+Standalone async helpers for the current gateway endpoints: ``POST /api/reset``
+(with an MCP-tool fallback for older servers), ``POST /api/add``, and
+``GET /export-state``. Mirrors the module-level function style of
+``agent_env.a2a_agent.protocol`` — callers pass URLs + params, not ``self``.
+"""
+from __future__ import annotations
+
+import asyncio
+import logging
+from typing import TYPE_CHECKING, Optional
+
+import httpx
+from agentenv_protocol import RPC_PATH, client as protocol_v1
+
+if TYPE_CHECKING:
+    from agent_env.env.env import DeployedEnv
+
+logger = logging.getLogger(__name__)
+
+DEFAULT_MCP_MAX_RETRIES = 5
+
+
+def environment_base_url(gateway_url: Optional[str], environment_name: str, mcp: bool = True) -> str:
+    """Gateway path for a service: ``/svc/mcp-{name}`` (MCP) or ``/svc/{name}`` (website)."""
+    if not gateway_url:  # an env deployed without a gateway has no gateway paths
+        from agent_env.env.env import EnvNeedsGateway
+
+        raise EnvNeedsGateway(f"Reaching '{environment_name}' here")
+    prefix = "mcp-" if mcp else ""
+    return f"{gateway_url}/svc/{prefix}{environment_name}"
+
+
+async def v1_base_url(deployed: Optional[DeployedEnv], gateway_url: Optional[str], environment_name: str, mcp: bool = True) -> Optional[str]:
+    """A child env's v1 data-plane base from the stored env card, or None for legacy; without a stored card, the live probe decides."""
+    if deployed is None or not deployed.environment_card:
+        base_url = environment_base_url(gateway_url, environment_name, mcp=mcp)
+        return base_url if await protocol_v1.supports_v1(base_url) else None
+    child = deployed.get_child_env_card(environment_name)
+    return deployed.environment_url + child.get("url", RPC_PATH).removesuffix(RPC_PATH) if child else None
+
+
+
+async def child_env_card(deployed: Optional[DeployedEnv], gateway_url: Optional[str], environment_name: str, timeout: int = 10) -> tuple[str, Optional[dict]]:
+    """An MCP child env's card and the base its endpoints resolve against: from the stored env card, or without one a live read (None on 404)."""
+    if deployed is None or not deployed.environment_card:
+        base_url = environment_base_url((gateway_url or "").rstrip("/"), environment_name)
+        try:
+            return base_url, await protocol_v1.get_card(base_url, timeout=timeout)
+        except httpx.HTTPStatusError as e:
+            if e.response.status_code == 404:
+                return base_url, None
+            raise
+    return deployed.environment_url, deployed.get_child_env_card(environment_name)
+
+async def reset_via_rest(
+    base_url: str,
+    mock_data_path: Optional[str] = None,
+    timeout: int = 30,
+    verify: bool = True,
+) -> dict:
+    """POST ``{base_url}/api/reset``. Sends the seed path when given. Raises on non-2xx."""
+    body = {"mock_data_path": mock_data_path} if mock_data_path is not None else None
+    async with httpx.AsyncClient(verify=verify) as client:
+        response = await client.post(f"{base_url}/api/reset", json=body, timeout=timeout)
+        response.raise_for_status()
+        result = response.json()
+        logger.info(f"load_environment_artifact: REST reset succeeded: {result}")
+        return result
+
+
+async def add_via_rest(base_url: str, file_path: str, timeout: int = 120, verify: bool = True) -> dict:
+    """POST ``{base_url}/api/add`` with a backend file path. Raises on non-2xx."""
+    async with httpx.AsyncClient(verify=verify) as client:
+        response = await client.post(f"{base_url}/api/add", json={"file_path": file_path}, timeout=timeout)
+        response.raise_for_status()
+        result = response.json()
+        logger.info(f"load_environment_artifact: add response: {result}")
+        return result
+
+
+async def export_state(gateway_url: str, environment_name: str, timeout: int = 60, verify: bool = True) -> dict:
+    """GET ``/svc/mcp-{name}/export-state`` — the legacy state snapshot for one service."""
+    base_url = environment_base_url(gateway_url, environment_name, mcp=True)
+    async with httpx.AsyncClient(verify=verify) as client:
+        response = await client.get(f"{base_url}/export-state", timeout=timeout)
+        response.raise_for_status()
+        return response.json()
+
+
+async def reset_via_mcp_tool(
+    mcp_url: str,
+    environment_name: str,
+    mock_data_path: str,
+    max_retries: int = DEFAULT_MCP_MAX_RETRIES,
+) -> None:
+    """Call the ``{environment_name}_reset`` MCP tool, retrying on transient failures."""
+    from mcp import ClientSession
+    from mcp.client.streamable_http import streamable_http_client
+
+    for attempt in range(max_retries):
+        reset_done = False
+        try:
+            async with streamable_http_client(mcp_url) as (read_stream, write_stream, _):
+                async with ClientSession(read_stream, write_stream) as session:
+                    await session.initialize()
+                    reset_tool_name = f"{environment_name}_reset"
+                    result = await session.call_tool(reset_tool_name, {"mock_data_path": mock_data_path})
+                    result_text = "".join(c.text for c in result.content if hasattr(c, "text"))
+                    logger.info(f"load_environment_artifact: {reset_tool_name} returned: {result_text}")
+                    reset_done = True
+        except BaseExceptionGroup as eg:
+            if reset_done:
+                logger.info(f"load_environment_artifact: succeeded (ignoring cleanup ExceptionGroup)")
+                break
+            logger.info(f"load_environment_artifact attempt {attempt + 1}/{max_retries} failed: {eg}")
+            if attempt < max_retries - 1:
+                await asyncio.sleep(5)
+            else:
+                raise
+        except Exception as e:
+            logger.info(f"load_environment_artifact attempt {attempt + 1}/{max_retries} failed: {type(e).__name__}")
+            if attempt < max_retries - 1:
+                await asyncio.sleep(5)
+            else:
+                raise
+        else:
+            break
+
+
+async def reset(
+    gateway_url: str,
+    environment_name: str,
+    mock_data_path: str,
+    max_retries: int = DEFAULT_MCP_MAX_RETRIES,
+) -> Optional[dict]:
+    """Reset an MCP service to a seed: try REST ``/api/reset``, fall back to the MCP tool."""
+    base_url = environment_base_url(gateway_url, environment_name, mcp=True)
+    try:
+        return await reset_via_rest(base_url, mock_data_path)
+    except Exception as rest_err:
+        # TODO remove MCP tool fallback once all usages of MCP servers are on REST
+        logger.info(f"load_environment_artifact: REST reset failed ({type(rest_err).__name__}), falling back to MCP tool")
+        await reset_via_mcp_tool(f"{gateway_url}/mcp", environment_name, mock_data_path, max_retries=max_retries)
+        return None

agent_env/env/registry.py

@@ -0,0 +1,93 @@
+"""Env registry for type-based deserialization."""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+from agent_env.plugins import _registration
+
+if TYPE_CHECKING:
+    from agent_env.env.env import Env
+    from agent_env.config.runtime import Config
+
+# Env defines this only to raise, so a registered class must override it. Not ``deploy``: a
+# component env such as ``service_db`` is deployed by the env that contains it.
+_MUST_IMPLEMENT = ("from_dict",)
+
+
+def _get_type(cls: type["Env"]) -> str:
+    return cls.type
+
+
+def _builtin_registry() -> dict[str, type["Env"]]:
+    """The built-in env types, by type."""
+    from agent_env.env.envs.gateway_server import GatewayEnv
+    from agent_env.env.envs.mcp_server import MCPServerEnv
+    from agent_env.env.envs.multi_env import MultiEnv
+    from agent_env.env.envs.service_db import ServiceDBEnv
+    from agent_env.env.envs.website import WebsiteEnv
+
+    return {
+        _get_type(GatewayEnv): GatewayEnv,
+        _get_type(MCPServerEnv): MCPServerEnv,
+        _get_type(MultiEnv): MultiEnv,
+        _get_type(ServiceDBEnv): ServiceDBEnv,
+        _get_type(WebsiteEnv): WebsiteEnv,
+    }
+
+
+def _build_registry(source: Config | None = None) -> dict[str, type["Env"]]:
+    """The built-in env types, then ``agent_env.envs`` plugins, then ``[envs] impls`` in ``source``."""
+    from agent_env.env.env import Env
+
+    registry = _builtin_registry()
+    validate = _registration.typed_validator(Env, builtins=frozenset(registry), required=_MUST_IMPLEMENT)
+    from_plugins = _registration.merge(registry, _registration.ENVS, validate, source=source)
+    _merge_config_toml_envs(registry, source=source, from_plugins=from_plugins)
+    return registry
+
+
+def get_env_registry() -> dict[str, type["Env"]]:
+    from agent_env.config import runtime
+
+    return runtime.get_config().env_registry()
+
+
+def _merge_config_toml_envs(
+    registry: dict[str, type["Env"]],
+    *,
+    source: Config | None = None,
+    from_plugins: _registration.Registrations | None = None,
+) -> None:
+    from agent_env.config import runtime
+    from agent_env.config import ConfigError, load_impl
+    from agent_env.env.env import Env
+
+    section = (source or runtime.get_config()).section("envs")
+    impls = section.get("impls", [])
+    if not isinstance(impls, list):
+        raise ConfigError(
+            f"[envs] impls must be a list of 'module:Class' strings, got {type(impls).__name__}"
+        )
+    from_plugins = from_plugins if from_plugins is not None else _registration.Registrations.empty()
+    for impl in impls:
+        if not isinstance(impl, str):
+            raise ConfigError(
+                f"[envs] impl must be a 'module:Class' string, got {type(impl).__name__}: {impl!r}"
+            )
+        cls = load_impl(impl, Env)
+        if cls.type == Env.type:
+            raise ConfigError(
+                f"[envs] impl {impl!r} does not define its own 'type' "
+                f"(inherits the base default {Env.type!r}); set a unique 'type' ClassVar"
+            )
+        if problem := _registration.unimplemented(cls, Env, _MUST_IMPLEMENT):
+            raise ConfigError(f"[envs] impl {impl!r}: {problem}")
+        if from_plugins.refuse(cls.type, f"[envs] impl {impl!r}"):
+            continue
+        if not from_plugins.release(cls.type, "envs", impl, cls) and cls.type in registry:
+            raise ConfigError(
+                f"[envs] impl {impl!r} type {cls.type!r} is already registered "
+                f"(conflicts with a built-in or another custom env)"
+            )
+        registry[cls.type] = cls

agent_env/env/snapshot_store.py

@@ -0,0 +1,372 @@
+"""Env snapshot store for caching pre-loaded servicedb images."""
+
+from __future__ import annotations
+
+import asyncio
+import hashlib
+import logging
+from dataclasses import dataclass
+from datetime import datetime, timezone
+from typing import Callable, Optional
+
+from agent_env.providers.sandbox_providers.sandbox_provider import SANDBOX_MODE_CONTAINER
+from agent_env.config import get_config
+from agent_env.store.document_store import AbsentOrNull, Filter, Sort
+from agent_env.store.routing import refuse_local_derivation
+
+logger = logging.getLogger(__name__)
+
+ENV_SNAPSHOTS_COLLECTION = "env_snapshots"
+
+ProgressCallback = Callable[[str, str, int], None]
+
+
+def compute_env_fingerprint(env) -> str:
+    """A stable digest of the service images a MultiEnv is built from.
+
+    Snapshots carry pre-ingested PGDATA, and the schema in that PGDATA is produced by
+    the *service images* — not by the env id. So the env id is the wrong cache key in
+    both directions: re-registering the same images under a new id needlessly orphans a
+    valid snapshot, while rebuilding a service under the same id would silently serve a
+    snapshot whose schema no longer matches.
+
+    Keying on the images fixes both. Sorted by environment name so dict/list ordering
+    can't change the digest, and truncated because this is a cache key, not a security
+    claim.
+    """
+    parts: list[str] = []
+    for mcp_env in getattr(env, "mcp_server_envs", None) or []:
+        image = getattr(getattr(mcp_env, "docker_image_artifact", None), "image_name", "")
+        parts.append(f"mcp:{mcp_env.environment_name}={image}")
+    for site_env in getattr(env, "website_envs", None) or []:
+        backend = getattr(getattr(site_env, "backend_docker_image_artifact", None), "image_name", "")
+        frontend = getattr(getattr(site_env, "frontend_docker_image_artifact", None), "image_name", "")
+        parts.append(f"web:{site_env.environment_name}={backend}+{frontend}")
+    digest = hashlib.sha256("\n".join(sorted(parts)).encode()).hexdigest()
+    return digest[:32]
+
+
+@dataclass
+class EnvSnapshot:
+    env_id: str
+    environment_universe_id: str
+    environment_universe_version: int
+    instance_id: str
+    is_clean: bool
+    db_image_artifact_id: str
+    db_image_artifact_version: int
+    created_at_utc: datetime | None = None
+    # Absent on rows written before image-fingerprint keying; those match on env_id only.
+    env_fingerprint: str | None = None
+
+    @classmethod
+    def from_dict(cls, data: dict) -> EnvSnapshot:
+        return cls(
+            env_id=data["env_id"],
+            environment_universe_id=data["service_universe_id"],
+            environment_universe_version=data["service_universe_version"],
+            instance_id=data["instance_id"],
+            is_clean=data["is_clean"],
+            db_image_artifact_id=data["db_image_artifact_id"],
+            db_image_artifact_version=data["db_image_artifact_version"],
+            created_at_utc=data.get("created_at_utc"),
+            env_fingerprint=data.get("env_fingerprint"),
+        )
+
+    @classmethod
+    async def create(cls, instance_id: str, on_progress: ProgressCallback | None = None) -> EnvSnapshot:
+        """Create a snapshot of a deployed MultiEnv's servicedb.
+
+        Reconnects to the sandbox, checks changelog status, exports the
+        PostgreSQL data directory, builds a Docker image with the data
+        pre-loaded, stores it as a DockerImageArtifact, and persists an
+        EnvSnapshot record.
+        """
+        from agent_env.artifact import DockerImageArtifact
+        from agent_env.env.env import DeployedSandboxEnv, Env
+        from agent_env.env.envs.multi_env import MultiEnv, _gateway_state_provider
+        from agent_env.env.store import get_env_instance_store
+
+        log = on_progress or (lambda step, msg, pct: None)
+
+        log("validate", "Fetching env instance...", 5)
+        deployed_env = get_env_instance_store().get(instance_id)
+        refuse_local_derivation(deployed_env.env_id, "env", "snapshotting")
+        env = Env.get(deployed_env.env_id, deployed_env.env_version)
+        if not isinstance(env, MultiEnv):
+            raise ValueError(
+                f"Env '{deployed_env.env_id}' is type '{env.type}', "
+                "only MultiEnv is supported"
+            )
+
+        env_id = deployed_env.env_id
+        without_our_gateway = (f"Snapshot capture needs our gateway and its local Postgres store; env '{env_id}' was deployed by "
+                               f"env_provider_type '{deployed_env.env_provider_type}'")
+        if not isinstance(deployed_env, DeployedSandboxEnv):
+            raise NotImplementedError(without_our_gateway)
+
+        environment_universe = get_env_instance_store().get_environment_universe(instance_id)
+        if environment_universe is None:
+            raise ValueError(f"Instance '{instance_id}' has no service_universe loaded")
+        universe_id = environment_universe["id"]
+        universe_version = environment_universe["version"]
+        log("reconnect", f"Reconnecting to sandbox {deployed_env.sandbox_id}...", 10)
+        multi_env = await MultiEnv.from_deployed_env(deployed_env)
+        if multi_env._sandbox is None:
+            raise NotImplementedError(without_our_gateway)
+        if multi_env._sandbox.mode == SANDBOX_MODE_CONTAINER:
+            raise NotImplementedError("Snapshot creation not supported on container mode")
+
+        # Snapshot capture bakes the local servicedb container into an image (docker cp PGDATA →
+        # image), so it's only meaningful for a backend that can also restore from that image
+        state_provider = _gateway_state_provider(multi_env)
+        if state_provider is None or not state_provider.supports_restore_from_snapshot:
+            raise NotImplementedError(
+                "Snapshot capture is only supported for the local Postgres backend "
+                f"(this env's state backend does not support it: {type(state_provider).__name__ if state_provider else None}). "
+                "Remote-backed deploys use the normal per-service load."
+            )
+
+        log("check_changelog", "Checking changelog...", 20)
+        is_clean = await _check_changelog_empty(multi_env._sandbox)
+        logger.info(f"Changelog empty: {is_clean} ({'clean' if is_clean else 'dirty'} snapshot)")
+        log("check_changelog", f"Changelog: {'clean' if is_clean else 'dirty'}", 30)
+
+        log("snapshot_db", "Exporting servicedb data...", 40)
+        tar_gz_s3_url = await _snapshot_servicedb(multi_env._sandbox, env_id, universe_id, log)
+
+        log("store_artifact", "Storing Docker image artifact...", 70)
+        logger.info("Storing snapshot as DockerImageArtifact...")
+        image_tag = f"env-snapshot-{env_id}-{universe_id}"
+        artifact = DockerImageArtifact.put_tar(
+            id=f"env-snapshot-{env_id}",
+            description=(
+                f"Snapshot of servicedb for env={env_id} "
+                f"universe={universe_id} v{universe_version}"
+            ),
+            image_name=image_tag,
+            tar_gz_s3_url=tar_gz_s3_url,
+        )
+        logger.info(f"Artifact: id={artifact.id} version={artifact.version}")
+        log("store_artifact", f"Artifact: id={artifact.id} version={artifact.version}", 80)
+
+        log("store_snapshot", "Persisting snapshot record...", 90)
+        snapshot = get_env_snapshot_store().put(
+            env_id=env_id,
+            environment_universe_id=universe_id,
+            environment_universe_version=universe_version,
+            instance_id=instance_id,
+            is_clean=is_clean,
+            db_image_artifact_id=artifact.id,
+            db_image_artifact_version=artifact.version,
+            # Recorded from the env whose images produced this PGDATA, so a later
+            # image-identical env (any id) can match it — and a rebuilt one cannot.
+            env_fingerprint=compute_env_fingerprint(env),
+        )
+        log("done", "Snapshot complete", 100)
+        return snapshot
+
+
+async def _check_changelog_empty(sandbox) -> bool:
+    """Return True if the _changelog table is empty (no modifications since data load)."""
+    from agent_env.env.envs.service_db import DB_NAME, DB_USER
+    from agent_env.env.envs.service_db import DATABASE_SERVICE_NAME
+    from agent_env.providers.env_providers.constants import DOCKER_COMPOSE_PATH
+
+    cmd = (
+        f"docker compose -f {DOCKER_COMPOSE_PATH} exec -T {DATABASE_SERVICE_NAME} "
+        f'psql -U {DB_USER} -d {DB_NAME} -tAc "SELECT count(*) FROM public._changelog"'
+    )
+    output = await sandbox.exec_script(cmd)
+    count = int(output.strip())
+    return count == 0
+
+
+async def _snapshot_servicedb(sandbox, env_id: str, environment_universe_id: str, log: ProgressCallback) -> str:
+    """Export pgdata from sandbox servicedb, build Docker image on sandbox, upload to S3.
+
+    Returns the S3 URL of the uploaded tar.gz.
+    """
+    from agent_env.env.envs.service_db import DATABASE_SERVICE_NAME
+    from agent_env.providers.env_providers.constants import DOCKER_COMPOSE_PATH
+
+    container_id_output = await sandbox.exec_script(
+        f"docker compose -f {DOCKER_COMPOSE_PATH} ps -a -q {DATABASE_SERVICE_NAME}"
+    )
+    container_id = container_id_output.strip()
+    if not container_id:
+        raise RuntimeError("servicedb container not found on sandbox")
+
+    log("snapshot_db", "Copying pgdata from container...", 45)
+    logger.info("Exporting pgdata from servicedb container...")
+    await sandbox.exec_script(f"docker cp {container_id}:/var/lib/postgresql/data /tmp/pgdata")
+
+    image_tag = f"env-snapshot-{env_id}-{environment_universe_id}"
+    log("snapshot_db", f"Building Docker image on sandbox: {image_tag}...", 50)
+    logger.info(f"Building Docker image on sandbox: {image_tag}...")
+    await sandbox.exec_script(
+        "mkdir -p /tmp/snapshot-build && "
+        "mv /tmp/pgdata /tmp/snapshot-build/pgdata && "
+        "cat > /tmp/snapshot-build/Dockerfile << 'DEOF'\n"
+        "FROM public.ecr.aws/docker/library/postgres:16-alpine\n"
+        "COPY pgdata /var/lib/postgresql/data\n"
+        "RUN chown -R postgres:postgres /var/lib/postgresql/data\n"
+        "DEOF"
+    )
+    await sandbox.exec_script(
+        f"docker build --platform linux/amd64 -t {image_tag} /tmp/snapshot-build"
+    )
+
+    log("snapshot_db", "Saving Docker image...", 55)
+    logger.info("Saving Docker image to tar.gz on sandbox...")
+    await sandbox.exec_script(f"docker save {image_tag} | gzip > /tmp/snapshot-image.tar.gz")
+
+    config = get_config()
+    object_store = config.get_object_store()
+    object_url = object_store.object_url(
+        f"{config.get_artifact_key_prefix()}env-snapshots/{env_id}/{environment_universe_id}/{image_tag}.tar.gz"
+    )
+    put_url = await asyncio.to_thread(object_store.signed_put_url, object_url)
+    if put_url is None:
+        raise RuntimeError(
+            f"{type(object_store).__name__} can't presign uploads; env snapshots need a signable object store or local execution."
+        )
+    log("snapshot_db", "Uploading Docker image...", 60)
+    logger.info(f"Uploading Docker image to {object_url}...")
+    await sandbox.exec_script(
+        f'curl -fsSL -X PUT --upload-file /tmp/snapshot-image.tar.gz "{put_url}"'
+    )
+
+    await sandbox.exec_script(f"rm -rf /tmp/snapshot-build /tmp/snapshot-image.tar.gz && docker rmi {image_tag}")
+
+    return object_url
+
+
+class EnvSnapshotStore:
+    def __init__(self) -> None:
+        self._indexed = None
+
+    @property
+    def _doc_store(self):
+        # Resolved per call: a cached store outlives reset_config(), so a process that
+        # re-pointed would read the new config and write the old backend.
+        store = get_config().get_document_store()
+        if self._indexed is not store:
+            store.ensure_index(
+                ENV_SNAPSHOTS_COLLECTION,
+                ["env_id", "service_universe_id", "instance_id"],
+                unique=True,
+            )
+            # The get_clean lookup shape — hit on every universe load, so it should
+            # not table-scan as snapshots accumulate.
+            store.ensure_index(
+                ENV_SNAPSHOTS_COLLECTION,
+                ["env_fingerprint", "service_universe_id", "service_universe_version", "is_clean"],
+            )
+            self._indexed = store
+        return store
+
+    def put(
+        self,
+        env_id: str,
+        environment_universe_id: str,
+        environment_universe_version: int,
+        instance_id: str,
+        is_clean: bool,
+        db_image_artifact_id: str,
+        db_image_artifact_version: int,
+        env_fingerprint: str | None = None,
+    ) -> EnvSnapshot:
+        doc = {
+            "env_id": env_id,
+            "service_universe_id": environment_universe_id,
+            "service_universe_version": environment_universe_version,
+            "instance_id": instance_id,
+            "is_clean": is_clean,
+            "db_image_artifact_id": db_image_artifact_id,
+            "db_image_artifact_version": db_image_artifact_version,
+            "created_at_utc": datetime.now(timezone.utc),
+            "env_fingerprint": env_fingerprint,
+        }
+        self._doc_store.replace(
+            ENV_SNAPSHOTS_COLLECTION,
+            # `Filter.of` keywords are stored document keys, not Python names: this one keeps
+            # the on-disk spelling while the parameter follows the rename.
+            Filter.of(env_id=env_id, service_universe_id=environment_universe_id, instance_id=instance_id),
+            doc,
+            upsert=True,
+        )
+        return EnvSnapshot.from_dict(doc)
+
+    def get(self, env_id: str, environment_universe_id: str, instance_id: str) -> EnvSnapshot | None:
+        doc = self._doc_store.find_one(
+            ENV_SNAPSHOTS_COLLECTION,
+            Filter.of(env_id=env_id, service_universe_id=environment_universe_id, instance_id=instance_id),
+        )
+        if not doc:
+            return None
+        return EnvSnapshot.from_dict(doc)
+
+    def get_clean(
+        self,
+        env_id: str,
+        environment_universe_id: str,
+        environment_universe_version: int | None = None,
+        env_fingerprint: str | None = None,
+    ) -> EnvSnapshot | None:
+        """Newest clean snapshot that can be restored into this env, or None.
+
+        With an ``env_fingerprint``, an image-identical env matches regardless of its
+        id — so re-registering the same services under a new env id reuses the existing
+        snapshot instead of paying for a fresh full load. The env_id lookup is kept as a
+        fallback for rows written before fingerprinting, which carry no digest to match.
+
+        Note the fingerprint match is deliberately NOT widened to a "close enough" env
+        id comparison: an env whose service images differ has a different schema, and
+        restoring a stale snapshot into it would corrupt the env rather than slow it down.
+        """
+        base = {"service_universe_id": environment_universe_id, "is_clean": True}
+        if environment_universe_version is not None:
+            base["service_universe_version"] = environment_universe_version
+
+        candidates: list[Filter] = []
+        # Fingerprint first: it is the correct key, and it lets a renamed env hit.
+        if env_fingerprint:
+            candidates.append(Filter.of(**base, env_fingerprint=env_fingerprint))
+        # Legacy rows predate fingerprinting, so the exact env id is the only provenance
+        # they carry. Restricted to rows with NO fingerprint on purpose: a row that has
+        # one and did not match above was built from different service images, and
+        # matching it on env_id alone would restore a stale schema into this env.
+        candidates.append(
+            Filter.of(**base, env_id=env_id).where("env_fingerprint", AbsentOrNull())
+        )
+
+        for query in candidates:
+            doc = self._doc_store.find_one(
+                ENV_SNAPSHOTS_COLLECTION,
+                query,
+                sort=Sort.by("created_at_utc", descending=True),
+            )
+            if doc:
+                return EnvSnapshot.from_dict(doc)
+        return None
+
+
+_env_snapshot_store: Optional[EnvSnapshotStore] = None
+
+
+def get_env_snapshot_store() -> EnvSnapshotStore:
+    global _env_snapshot_store
+    if _env_snapshot_store is None:
+        _env_snapshot_store = EnvSnapshotStore()
+    return _env_snapshot_store
+
+
+def set_env_snapshot_store(store: EnvSnapshotStore) -> None:
+    global _env_snapshot_store
+    _env_snapshot_store = store
+
+
+def reset_env_snapshot_store() -> None:
+    global _env_snapshot_store
+    _env_snapshot_store = None