agentenv-framework 0.9.1254__py3-none-any.whl

agent_env/providers/env_providers/env_provider.py

@@ -0,0 +1,237 @@
+"""Environment providers: the contract a provider implements, the built-ins' shared base, and finding a provider by type."""
+
+from __future__ import annotations
+
+import asyncio
+import logging
+import time
+from abc import ABC, abstractmethod
+from typing import TYPE_CHECKING, Any, Awaitable, Callable, ClassVar
+
+import httpx
+from mcp import ClientSession
+from mcp.client.streamable_http import streamable_http_client
+from agent_env.env.env import DeployedEnv
+from agent_env.env.gateway.constants import WELL_KNOWN_PATH
+from agent_env.plugins import _registration
+from agent_env.providers.sandbox_providers.sandbox import Sandbox
+from agent_env.providers.sandbox_providers.sandbox_provider import SANDBOX_MODE_CONTAINER
+
+if TYPE_CHECKING:
+    from agent_env.config.runtime import Config
+    from agent_env.env.env import Env
+    from agent_env.providers.sandbox_providers.sandbox_provider import SandboxProvider
+
+logger = logging.getLogger(__name__)
+
+
+class EnvironmentProvider(ABC):
+    """Deploys an env and hands back its record; tears down what it created."""
+
+    type: ClassVar[str]  # the name an env declares it by, and that its records carry as env_provider_type
+
+    @abstractmethod
+    async def deploy(self, env: Env, sandbox_provider: SandboxProvider) -> DeployedEnv:
+        """Deploy the env and return its record, unregistered: the env registers it. Take ``**options``, which receives every option
+        the env's deploy was given; a deploy that names its options instead is refused any other one set away from its default."""
+
+    @abstractmethod
+    async def close(self) -> None:
+        """Tear down everything deploy() created."""
+
+
+def record_class_for(env_provider_type: str | None) -> type[DeployedEnv] | None:
+    """The record class a built-in provider type writes; None for any other type, whose records load by shape."""
+    provider = _builtin_env_providers().get(env_provider_type)
+    return provider.record_class if provider is not None else None
+
+
+def build_env_provider(env_provider_type: str) -> EnvironmentProvider:
+    """A fresh provider of the type an env declares, built in or from a plugin; raises for a type this process can't find."""
+    return _env_provider_class(env_provider_type)()
+
+
+def _env_provider_class(env_provider_type: str) -> type[EnvironmentProvider]:
+    """The registered provider class of that type; raises, naming the registered types, for one this process can't find."""
+    from agent_env.config import runtime
+
+    providers = runtime.get_config().env_provider_registry()
+    if env_provider_type not in providers:
+        raise ValueError(f"Unknown env_provider_type: {env_provider_type!r} (expected one of {sorted(providers)})"
+                         f"{_registration.failure_note(_registration.ENV_PROVIDERS, env_provider_type)}")
+    return providers[env_provider_type]
+
+
+class _SandboxEnvironmentProvider(EnvironmentProvider):
+    """The built-in providers' base: tracks the containers it creates in our sandboxes, closes them, reattaches to the records it wrote,
+    and retries across a chained spec."""
+
+    record_class: ClassVar[type[DeployedEnv]] = DeployedEnv  # the class its records load back as
+
+    def __init__(self):
+        self._container_sandboxes: list[Sandbox] = []
+        self._environment_sandboxes: dict[str, Sandbox] = {}
+
+    def environment_sandbox(self, environment_name: str) -> Sandbox | None:
+        return self._environment_sandboxes.get(environment_name)
+
+    async def close(self) -> None:
+        for sb in self._container_sandboxes:
+            try:
+                await sb.terminate()
+            except Exception as e:
+                logger.warning(f"Failed to terminate {type(sb).__name__} {sb.sandbox_id}: {e}")
+        self._container_sandboxes.clear()
+        self._environment_sandboxes.clear()
+
+    async def _reattach(self, env: Env, deployed: DeployedEnv) -> Sandbox:
+        """A record this provider wrote, reattached: its primary sandbox, or in container mode the server's container, with every recorded one for close()."""
+        from agent_env.env.envs.mcp_server import MCPServerEnv
+        from agent_env.providers import build_sandbox_provider, get_env_sandbox_provider
+        provider = build_sandbox_provider(deployed.sandbox_type) if deployed.sandbox_type else get_env_sandbox_provider()
+        sandbox = await provider.get_sandbox(deployed.sandbox_id)
+        if env.environment_name in deployed.sandbox_ids.get(MCPServerEnv.type, {}):
+            # Container mode.
+            sandbox = await self._reattach_containers(provider, sandbox, deployed.sandbox_ids, env.environment_name)
+        return sandbox
+
+    async def _reattach_containers(self, provider: SandboxProvider, gateway: Sandbox, sandbox_ids: dict, environment_name: str) -> Sandbox:
+        """Reattach the recorded containers for close(), skipping a dead servicedb or sidecar; returns the server's container, raising if it's dead."""
+        from agent_env.env.envs.mcp_server import MCPServerEnv
+        from agent_env.env.envs.service_db import ServiceDBEnv
+        server_id = sandbox_ids[MCPServerEnv.type][environment_name]
+        db_ids = list(sandbox_ids.get(ServiceDBEnv.type, {}).values())
+        server, *dbs = await asyncio.gather(*(provider.get_sandbox(i) for i in [server_id, *db_ids]), return_exceptions=True)
+        if isinstance(server, BaseException):
+            raise RuntimeError(f"[{environment_name}] can't reattach the server's container {server_id}: {type(server).__name__}: {server}") from server
+        server.mode = SANDBOX_MODE_CONTAINER  # recorded under mcp_server, so create_container made it
+        for db_id, db in zip(db_ids, dbs):
+            if isinstance(db, BaseException):
+                logger.warning(f"[{environment_name}] skipping recorded container {db_id}: {type(db).__name__}: {db}")
+        self._container_sandboxes = [server, *(db for db in dbs if not isinstance(db, BaseException)), gateway]
+        return server
+
+    async def _wait_for_tunnel(self, url: str, timeout: int = 60) -> dict | None:
+        """Wait until the deployed env serves its card to this process; returns the card.
+
+        Probes the well-known document rather than ``/mcp``, whose healthy endpoint holds the
+        connection open. Only a 200 carrying a card counts, and that card is what the deploy
+        records. Returns None if none arrives within ``timeout`` attempts.
+        """
+        probe_url = f"{url}{WELL_KNOWN_PATH}"
+        logger.info(f"Waiting for tunnel to be ready (timeout={timeout}s) via {probe_url}...")
+        last = "no response"
+        async with httpx.AsyncClient() as client:
+            for i in range(timeout):
+                try:
+                    response = await client.get(probe_url, timeout=5)
+                    card = _card_from(response)
+                    if card is not None:
+                        logger.info(f"  Tunnel ready (HTTP {response.status_code}) after {i+1}s")
+                        return card
+                    last = f"HTTP {response.status_code} {response.headers.get('content-type', '')}".rstrip()
+                except Exception as e:
+                    last = f"error: {type(e).__name__}"
+                if i % 10 == 0:
+                    logger.info(f"  Still waiting... ({last})")
+                await asyncio.sleep(1)
+        logger.info(f"  Tunnel not ready after {timeout}s (last: {last})")
+        return None
+
+    def _sandbox_providers(self, sandbox_provider: SandboxProvider, env_id: str | None) -> list[SandboxProvider]:
+        """The sandbox providers to try, in order: a chained one's members, or itself."""
+        from agent_env.providers.sandbox_providers.chained_sandbox_provider import ChainedSandboxProvider
+
+        if isinstance(sandbox_provider, ChainedSandboxProvider):
+            return list(sandbox_provider._providers)
+        return [sandbox_provider]
+
+    async def _run(self, sandbox_provider: SandboxProvider, env_id: str | None, attempt: Callable[[SandboxProvider], Awaitable[Any]]) -> Any:
+        """Run ``attempt`` on the sandbox provider, or on each chained member until one succeeds; a failed deploy is closed."""
+        try:
+            sandbox_providers = self._sandbox_providers(sandbox_provider, env_id)
+            if len(sandbox_providers) == 1:
+                return await self._attempt(sandbox_providers[0], env_id, attempt)
+            errors: list[tuple[str, Exception]] = []
+            for p in sandbox_providers:
+                try:
+                    return await self._attempt(p, env_id, attempt)
+                except Exception as e:
+                    errors.append((type(p).__name__, e))
+                    await self.close()
+            detail = "; ".join(f"{n}: {e!r}" for n, e in errors)
+            raise RuntimeError(f"All {len(sandbox_providers)} chained providers failed to deploy: {detail}")
+        except Exception:
+            await self.close()
+            raise
+
+    async def _attempt(self, sandbox_provider: SandboxProvider, env_id: str | None, attempt: Callable[[SandboxProvider], Awaitable[Any]]) -> Any:
+        """One deploy on one sandbox provider, logged as an ``env_deploy_provider_attempt``."""
+        start = time.monotonic()
+        provider_name = type(sandbox_provider).__name__
+        try:
+            result = await attempt(sandbox_provider)
+            logger.info(f"env_deploy_provider_attempt env_id={env_id or 'unknown'} provider={provider_name} status=success duration_s={time.monotonic() - start:.1f}")
+        except Exception as e:
+            logger.warning(f"env_deploy_provider_attempt env_id={env_id or 'unknown'} provider={provider_name} status=failure duration_s={time.monotonic() - start:.1f} error_type={type(e).__name__} error={e!r}")
+            raise
+        return result
+
+
+def _size_kwargs(cpu: float | None, memory_mb: int | None) -> dict:
+    """Forward only what the caller set, so the provider's own default stands otherwise.
+
+    Passing None through would not reach that default: the providers hand cpu/memory
+    straight to the backend, so an explicit None means "no reservation at all" rather than
+    "your floor".
+    """
+    kwargs: dict = {}
+    if cpu is not None:
+        kwargs["cpu"] = cpu
+    if memory_mb is not None:
+        kwargs["memory"] = memory_mb
+    return kwargs
+
+
+def _card_from(response: httpx.Response) -> dict | None:
+    """The env card in a readiness response: a 200 whose JSON body is an object with a name."""
+    if response.status_code != 200:
+        return None
+    try:
+        card = response.json()
+    except ValueError:
+        return None
+    return card if isinstance(card, dict) and isinstance(card.get("name"), str) and card["name"] else None
+
+
+async def _tool_names(mcp_url: str) -> list[str]:
+    """The server's tool names as it lists them, duplicates included."""
+    async with streamable_http_client(mcp_url) as (read_stream, write_stream, _):
+        async with ClientSession(read_stream, write_stream) as session:
+            await session.initialize()
+            return [tool.name for tool in (await session.list_tools()).tools]
+
+
+def _builtin_env_providers() -> dict[str, type[_SandboxEnvironmentProvider]]:
+    """The built-in providers, by type."""
+    from agent_env.providers.env_providers.env_gateway_provider import EnvironmentGatewayProvider
+    from agent_env.providers.env_providers.env_server_provider import EnvironmentServerProvider
+
+    return {p.type: p for p in (EnvironmentGatewayProvider, EnvironmentServerProvider)}
+
+
+def _build_registry(source: Config | None = None) -> dict[str, type[EnvironmentProvider]]:
+    """The built-in providers, then ``agent_env.env_providers`` plugins."""
+    registry: dict[str, type[EnvironmentProvider]] = dict(_builtin_env_providers())
+    _registration.merge(registry, _registration.ENV_PROVIDERS, _validate_plugin, source=source)
+    return registry
+
+
+def _validate_plugin(name: str, loaded: Any) -> type[EnvironmentProvider]:
+    cls = _registration.require_subclass(loaded, EnvironmentProvider)
+    if getattr(cls, "type", None) != name:
+        raise TypeError(f"{cls.__qualname__} has type {getattr(cls, 'type', None)!r}; it must equal the entry-point "
+                        "name, which its records carry as env_provider_type")
+    if problem := _registration.unimplemented(cls, EnvironmentProvider):
+        raise TypeError(problem)
+    return cls

agent_env/providers/env_providers/env_server_provider.py

@@ -0,0 +1,133 @@
+"""The server environment provider: one MCP server in its own container on the configured sandbox provider, with no gateway in front and no state store."""
+
+from __future__ import annotations
+
+import asyncio
+import functools
+from datetime import datetime, timezone
+from typing import TYPE_CHECKING, ClassVar, Optional
+
+import anyio
+
+from agent_env.attribution import Attribution
+from agent_env.env.env import DeployedSandboxEnv
+from agent_env.env.gateway import AGENT_ENV_GATEWAY_MCP_PORT
+from agent_env.env.gateway.constants import WELL_KNOWN_PATH
+from agent_env.providers.env_providers.env_provider import _SandboxEnvironmentProvider, _size_kwargs, _tool_names
+from agent_env.providers.sandbox_providers.sandbox_provider import SandboxProvider
+
+if TYPE_CHECKING:
+    from agent_env.env.env import Env
+
+
+class EnvironmentServerProvider(_SandboxEnvironmentProvider):
+    """Deploys one MCP server with no gateway in front; the server's own card is the env card."""
+
+    type: ClassVar[str] = "server"
+    record_class = DeployedSandboxEnv
+
+    async def deploy(
+        self,
+        env: Env,
+        sandbox_provider: SandboxProvider,
+        *,
+        ttl_seconds: int = 10800,
+        disk_size_gb: float = 10,
+        cpu: float | None = None,
+        memory_mb: int | None = None,
+        priority: Optional[int] = None,
+        attribution: Optional[Attribution] = None,
+    ) -> DeployedSandboxEnv:
+        """The env's MCP server in its own container; returns the record once it serves its card and its tools."""
+        from agent_env.env.envs.multi_env import MultiEnv
+        from agent_env.env.envs.website import WebsiteEnv
+
+        if isinstance(env, (MultiEnv, WebsiteEnv)):
+            raise TypeError(f"env_provider_type '{self.type}' deploys one MCP server, not a {env.type} env")
+        attempt = functools.partial(
+            self._deploy_server, env, ttl_seconds=ttl_seconds, disk_size_gb=disk_size_gb,
+            cpu=cpu, memory_mb=memory_mb, priority=priority, attribution=attribution,
+        )
+        return await self._run(sandbox_provider, env.id, attempt)
+
+    async def install_changelog_triggers(self, environment_name: str) -> None:
+        """No state store, so no changelog to install."""
+
+    async def _deploy_server(
+        self,
+        env: Env,
+        sandbox_provider: SandboxProvider,
+        *,
+        ttl_seconds: int,
+        disk_size_gb: float,
+        cpu: float | None,
+        memory_mb: int | None,
+        priority: Optional[int],
+        attribution: Optional[Attribution],
+    ) -> DeployedSandboxEnv:
+        """One attempt: the server's container, then its card and tools; returns its record."""
+        from agent_env.env.envs.mcp_server import MCPServerEnv
+
+        name = env.environment_name
+        server = await sandbox_provider.create_container(
+            image_name=env.docker_image_artifact.image_name,
+            port=AGENT_ENV_GATEWAY_MCP_PORT,
+            env={
+                "ENVIRONMENT_NAME": name,
+                "MCP_HOST": "0.0.0.0",  # IPv4 too: an IPv6-only "::" bind starves tunnels and published ports
+                "DATABASE_URL": f"sqlite:////tmp/{name}.sqlite",
+            },
+            **_size_kwargs(cpu, memory_mb),
+            disk_size_gb=disk_size_gb, timeout=ttl_seconds,
+            attribution=dict(attribution or {}),
+            priority=priority,
+        )
+        self._container_sandboxes.append(server)
+        self._environment_sandboxes[name] = server
+        server_url = server.tunnel_urls[AGENT_ENV_GATEWAY_MCP_PORT]
+        try:
+            async with asyncio.timeout(_READINESS_TIMEOUT_S):
+                card = await self._wait_for_tunnel(server_url, timeout=int(_READINESS_TIMEOUT_S))
+        except TimeoutError:
+            card = None
+        if card is None:
+            raise RuntimeError(f"'{name}' did not serve its env card at {server_url}{WELL_KNOWN_PATH} in time; "
+                               "an env deployed without a gateway must serve its own card")
+        card_read_at_utc = datetime.now(timezone.utc).isoformat()
+        if card["name"] != name:
+            raise RuntimeError(f"'{name}' serves a card named '{card['name']}'; "
+                               "an env deployed without a gateway must serve its card under its environment_name")
+        record = DeployedSandboxEnv(
+            env_id=env.id,
+            env_version=env.version,
+            env_provider_type=self.type,
+            environment_card_url=f"{server_url}{WELL_KNOWN_PATH}",
+            environment_card=card,
+            environment_card_read_at_utc=card_read_at_utc,
+            sandbox_id=server.sandbox_id,
+            sandbox_type=server.type,
+            sandbox_ids={MCPServerEnv.type: {name: server.sandbox_id}},
+        )
+        await _require_tools(name, record.mcp_url)
+        return record
+
+
+# Wall-clock cap on the server's readiness: attempts alone don't bound a tunnel that hangs each request.
+_READINESS_TIMEOUT_S = 180.0
+
+_TOOLS_GATE_TIMEOUT_S = 60.0
+
+
+async def _require_tools(environment_name: str, mcp_url: str) -> None:
+    """Fail unless the server at ``mcp_url`` lists at least one MCP tool and no tool name twice."""
+    try:
+        # anyio re-cancels in cleanup; asyncio.timeout would wait out the MCP client's session-closing DELETE (up to 300 s).
+        with anyio.fail_after(_TOOLS_GATE_TIMEOUT_S):
+            tools = await _tool_names(mcp_url)
+    except Exception as e:
+        raise RuntimeError(f"'{environment_name}' did not list its MCP tools at {mcp_url}: {e!r}") from e
+    if not tools:
+        raise RuntimeError(f"'{environment_name}' lists no MCP tools at {mcp_url}; an env deployed without a gateway must serve at least one")
+    duplicates = sorted({t for t in tools if tools.count(t) > 1})
+    if duplicates:
+        raise RuntimeError(f"'{environment_name}' lists duplicate MCP tool names at {mcp_url}: {', '.join(duplicates)}")

agent_env/providers/env_state/__init__.py

@@ -0,0 +1,49 @@
+"""Env-state providers — abstract WHERE an environment's data/state lives"""
+
+from agent_env.providers.env_state.env_state_provider import (
+    DatabaseStateProvider,
+    DEFAULT_STATE_TTL_SECONDS,
+    EnvStateInstance,
+    EnvStateProvider,
+    StateContext,
+    acquire_state_for_deploy,
+    build_state_provider,
+    LOCAL_POSTGRES_STATE_TYPE,
+)
+from agent_env.providers.env_state.local_postgres import (
+    LocalPostgresStateContext,
+    LocalPostgresStateProvider,
+    LocalPostgresStoreSpec,
+)
+from agent_env.providers.env_state.store import (
+    ENV_STATE_INSTANCES_COLLECTION,
+    EnvStateInstanceStore,
+    get_env_state_instance_store,
+    register_env_state_instance,
+    reset_env_state_instance_store,
+    retire_env_state_instance,
+    set_env_state_instance_store,
+    update_env_state_instance_metadata,
+)
+
+__all__ = [
+    "DatabaseStateProvider",
+    "DEFAULT_STATE_TTL_SECONDS",
+    "EnvStateInstance",
+    "EnvStateProvider",
+    "StateContext",
+    "acquire_state_for_deploy",
+    "build_state_provider",
+    "LOCAL_POSTGRES_STATE_TYPE",
+    "LocalPostgresStateContext",
+    "LocalPostgresStateProvider",
+    "LocalPostgresStoreSpec",
+    "ENV_STATE_INSTANCES_COLLECTION",
+    "EnvStateInstanceStore",
+    "get_env_state_instance_store",
+    "set_env_state_instance_store",
+    "reset_env_state_instance_store",
+    "register_env_state_instance",
+    "retire_env_state_instance",
+    "update_env_state_instance_metadata",
+]