agentenv-framework 0.9.1254__py3-none-any.whl

agent_env/store/image_store/google_credentials.py

@@ -0,0 +1,113 @@
+"""Artifact Registry credentials for OciRegistryImageStore (the ``gcp`` extra)."""
+
+from __future__ import annotations
+
+import copy
+import json
+import threading
+import time
+from datetime import UTC, datetime, timedelta
+
+import google.auth
+from google.auth import impersonated_credentials
+from google.api_core.exceptions import RetryError
+from google.auth.exceptions import DefaultCredentialsError, RefreshError, TransportError
+
+from agent_env.config.errors import ConfigError
+from agent_env.store import _google
+from agent_env.store.image_store.oci_registry_credentials import (
+    OciRegistryCredentials,
+    RegistryAuth,
+    normalize_registry_host,
+)
+
+# A token must outlast what it is handed to: an image build logs in with it to push at its end.
+_MIN_REMAINING = timedelta(minutes=45)
+_USERNAME = "oauth2accesstoken"
+_TRANSIENT_IAM_CODES = frozenset({429, 500, 502, 503, 504})
+# A mint that found IAM unreachable spent a whole retry window; callers that come just after it
+# share its failure rather than each spending another.
+_OUTAGE_REUSE_SECONDS = 10.0
+
+
+def _transient(error: Exception) -> bool:
+    """A dropped connection, or IAM declining for now. google-auth reports an IAM refusal as a
+    RefreshError carrying the response body: a 429 or 5xx, or a body that is no JSON at all, as a
+    front end's error page is. A token endpoint's refusal carries a parsed mapping, and is final."""
+    if _google.dropped_connection(error):
+        return True
+    if not isinstance(error, RefreshError) or len(error.args) < 2 or not isinstance(error.args[1], str):
+        return False
+    try:
+        body = json.loads(error.args[1])
+    except ValueError:
+        return True
+    status = body.get("error") if isinstance(body, dict) else None
+    return isinstance(status, dict) and status.get("code") in _TRANSIENT_IAM_CODES
+
+
+_MINT_RETRY = _google.RETRY.with_predicate(_transient)
+
+
+class GoogleAccessTokenCredentials(OciRegistryCredentials):
+    """Docker-login material for Artifact Registry: an access token of ``service_account``, which
+    Application Default Credentials impersonate. The token is handed to sandboxes, so the account
+    should hold nothing but write access to the registry's repository."""
+
+    def __init__(self, service_account: str) -> None:
+        if not service_account.strip():
+            raise ConfigError("GoogleAccessTokenCredentials needs a service_account")
+        self._service_account = service_account.strip()
+        self._credentials: impersonated_credentials.Credentials | None = None
+        self._request = None
+        self._lock = threading.Lock()
+        self._failure: tuple[float, Exception] | None = None
+
+    def mint(self, host: str) -> RegistryAuth:
+        arrived = time.monotonic()
+        with self._lock:
+            # A caller that queued behind a failed attempt, or arrives just after IAM proved
+            # unreachable, gets its failure rather than another retry window.
+            if self._failure is not None:
+                failed_at, failure = self._failure
+                outage = isinstance(failure, TransportError) and time.monotonic() - failed_at < _OUTAGE_REUSE_SECONDS
+                if failed_at > arrived or outage:
+                    raise copy.copy(failure).with_traceback(None) from failure
+            try:
+                token = self._token()
+            except Exception as e:
+                self._failure = (time.monotonic(), e)
+                raise
+            self._failure = None
+        return RegistryAuth(registry=normalize_registry_host(host), username=_USERNAME, password=token)
+
+    def _token(self) -> str:
+        if self._credentials is None:
+            self._credentials, self._request = self._impersonate()
+        if _remaining(self._credentials) < _MIN_REMAINING:
+            try:
+                _MINT_RETRY(self._credentials.refresh)(self._request)
+            except RefreshError as e:
+                # Look the credentials up again next time, in case they were replaced.
+                self._credentials = self._request = None
+                raise ConfigError(f"Cannot mint an access token as {self._service_account}: {e}") from e
+            except RetryError as e:
+                raise TransportError(f"IAM minted no access token as {self._service_account} within the retry window: {e.cause}") from e
+        return self._credentials.token
+
+    def _impersonate(self):
+        try:
+            source, _ = google.auth.default(scopes=_google.SCOPES)
+        except (DefaultCredentialsError, OSError) as e:
+            raise ConfigError(f"No Google credentials to impersonate {self._service_account} with: {e}") from e
+        credentials = impersonated_credentials.Credentials(
+            source_credentials=source, target_principal=self._service_account, target_scopes=_google.SCOPES
+        )
+        return credentials, _google.pooled_request()
+
+
+def _remaining(credentials) -> timedelta:
+    if credentials.token is None:
+        return timedelta(0)
+    # google-auth keeps expiry as naive UTC.
+    return credentials.expiry - datetime.now(UTC).replace(tzinfo=None)

agent_env/store/image_store/image_store.py

@@ -0,0 +1,113 @@
+"""Backend-agnostic container image (registry) store abstraction."""
+
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+from collections.abc import Mapping
+from typing import Any, Self
+
+from agent_env.store.image_store.oci_registry_credentials import (
+    OciRegistryCredentials,
+    RegistryAuth,
+    normalize_registry_host,
+    registry_host_from_ref,
+)
+
+
+class ImageStore(ABC):
+    """Image store interface. Images are addressed by a logical ``repository`` and ``tag``.
+
+    ``image_ref`` builds a fully-qualified, pullable ref; ``auth`` returns the
+    docker-login material a remote executor applies — the store never runs docker.
+    """
+
+    @classmethod
+    def from_config(cls, *args: Any, **config: Any) -> Self:
+        """Construct from a resolved config table; backends that build a client override this, taking
+        the table's keys as named parameters, which ``(*args, **config)`` lets a type checker accept."""
+        return cls(*args, **config)
+
+    @abstractmethod
+    def image_ref(self, repository: str, tag: str) -> str:
+        """Fully-qualified, pullable ref for ``(repository, tag)`` in this store."""
+
+    def ensure_repository(self, repository: str) -> None:
+        """Create the repository if the backend requires pre-creation; no-op otherwise."""
+
+    def owns(self, ref: str) -> bool:
+        """Whether ``ref`` is in this store, as ``auth`` decides it: a caller can tell a pull this
+        store authorizes from one it knows nothing about. The default owns nothing. Overriding
+        ``owns`` does not change which refs ``auth`` logs in for."""
+        return False
+
+    @abstractmethod
+    def auth(self, ref: str) -> RegistryAuth | None:
+        """Freshly-minted docker-login material for ``ref``, or None when none is needed
+        (anonymous/local registry, a public ref, or a ref not in this store)."""
+
+
+class OciRegistryImageStore(ImageStore):
+    """Generic OCI registry whose repository names and credentials are configurable."""
+
+    def __init__(
+        self,
+        registry_host: str,
+        repository_prefix: str = "",
+        credentials: OciRegistryCredentials | None = None,
+    ) -> None:
+        self._registry = normalize_registry_host(registry_host)
+        self._prefix = repository_prefix
+        self._credentials = credentials
+
+    @classmethod
+    def from_config(
+        cls,
+        *,
+        registry_host: str,
+        repository_prefix: str = "",
+        credentials: OciRegistryCredentials | Mapping[str, object] | None = None,
+    ) -> Self:
+        if isinstance(credentials, Mapping):
+            from agent_env.config import ConfigError, load_impl
+
+            credential_config = dict(credentials)
+            impl = credential_config.pop("impl", None)
+            if not impl:
+                raise ConfigError(
+                    "Registry credentials are missing an 'impl' dotted path"
+                )
+            credential_cls = load_impl(impl, OciRegistryCredentials)
+            credentials = credential_cls.from_config(**credential_config)
+        elif credentials is not None and not isinstance(
+            credentials, OciRegistryCredentials
+        ):
+            raise TypeError("credentials must implement OciRegistryCredentials")
+        return cls(
+            registry_host=registry_host,
+            repository_prefix=repository_prefix,
+            credentials=credentials,
+        )
+
+    @property
+    def registry_host(self) -> str:
+        return self._registry
+
+    @property
+    def credentials(self) -> OciRegistryCredentials | None:
+        return self._credentials
+
+    def image_ref(self, repository: str, tag: str) -> str:
+        return f"{self._registry}/{self._repo_name(repository)}:{tag}"
+
+    def owns(self, ref: str) -> bool:
+        """A ref on this store's registry, whatever its repository: a registry login covers the
+        whole host."""
+        return registry_host_from_ref(ref) == self._registry
+
+    def auth(self, ref: str) -> RegistryAuth | None:
+        if registry_host_from_ref(ref) != self._registry or self._credentials is None:
+            return None
+        return self._credentials.mint(self._registry)
+
+    def _repo_name(self, repository: str) -> str:
+        return f"{self._prefix}/{repository}" if self._prefix else repository

agent_env/store/image_store/local_registry_image_store.py

@@ -0,0 +1,106 @@
+"""Local OCI registry implementation of the ImageStore interface (no cloud services)."""
+
+from __future__ import annotations
+
+import subprocess
+import time
+from urllib.parse import urlsplit
+
+from agent_env.store.image_store.image_store import OciRegistryImageStore
+from agent_env.store.image_store.oci_registry_credentials import OciRegistryCredentials
+
+_REGISTRY_CONTAINER = "agentenv-registry"
+
+
+class LocalRegistryImageStore(OciRegistryImageStore):
+    """ImageStore backed by an anonymous local OCI registry.
+
+    Repositories are created automatically on push.
+    """
+
+    def __init__(
+        self,
+        registry_host: str = "localhost:5000",
+        repository_prefix: str = "",
+        credentials: OciRegistryCredentials | None = None,
+    ) -> None:
+        super().__init__(
+            registry_host=registry_host,
+            repository_prefix=repository_prefix,
+            credentials=credentials,
+        )
+
+    def ensure_repository(self, repository: str) -> None:
+        """A local registry creates repositories on push, so this instead lazily brings up the
+        backing ``registry:2`` — idempotently, and safe to call from many concurrent pushes. A
+        no-op when something already serves the registry at ``registry_host`` (a registry you
+        run yourself, or one a concurrent push already started); it never removes a container
+        that might be serving. Called from the image-push path, so no caller manages it."""
+        port = urlsplit(f"//{self.registry_host}").port or 5000
+        if self._registry_listening(port):
+            return
+        run = self._run_registry(port)
+        if run.returncode != 0:
+            # `docker run` creates the named container atomically, so a failure is either a name
+            # conflict (our container already exists) or something unrelated — disambiguate on
+            # the existing container instead of blindly retrying, so a real error isn't masked.
+            existing_port = self._container_configured_port(_REGISTRY_CONTAINER)
+            if existing_port is None:
+                # No such container: the failure is real — the port is taken by another process,
+                # the image pull failed, or the daemon is down. Surface it now, don't wait 30s.
+                raise RuntimeError(f"could not start the local registry: {run.stderr.strip()}")
+            if existing_port != port:
+                # A stale container of ours on a different port (registry_host changed since it
+                # was created). It isn't serving our port — the check above proved that — so
+                # replacing it is safe.
+                subprocess.run(["docker", "rm", "-f", _REGISTRY_CONTAINER], capture_output=True, text=True)
+                run = self._run_registry(port)
+                if run.returncode != 0:
+                    raise RuntimeError(f"could not start the local registry: {run.stderr.strip()}")
+            else:
+                # A concurrent push won the race, or a stopped container of ours is on our port:
+                # (re)start it (a no-op if it's already running) and surface a genuine failure.
+                start = subprocess.run(["docker", "start", _REGISTRY_CONTAINER], capture_output=True, text=True)
+                if start.returncode != 0:
+                    raise RuntimeError(
+                        f"local registry container exists but could not be started: {start.stderr.strip()}"
+                    )
+        deadline = time.time() + 30
+        while time.time() < deadline:
+            if self._registry_listening(port):
+                return
+            time.sleep(0.5)
+        raise RuntimeError(f"the local registry started but did not answer /v2/ on port {port} within 30s")
+
+    @staticmethod
+    def _run_registry(port: int) -> subprocess.CompletedProcess:
+        return subprocess.run(
+            ["docker", "run", "-d", "--restart", "unless-stopped", "--name", _REGISTRY_CONTAINER,
+             "-p", f"127.0.0.1:{port}:5000", "registry:2"],
+            capture_output=True, text=True,
+        )
+
+    @staticmethod
+    def _container_configured_port(name: str) -> int | None:
+        """The host port ``name`` maps container-port 5000 to (works for stopped containers too),
+        or None when no such container exists — i.e. the ``docker run`` failure was not a name
+        clash but a real error (port taken elsewhere / pull / daemon)."""
+        r = subprocess.run(
+            ["docker", "inspect", "-f",
+             '{{ (index (index .HostConfig.PortBindings "5000/tcp") 0).HostPort }}', name],
+            capture_output=True, text=True,
+        )
+        if r.returncode != 0:
+            return None
+        out = r.stdout.strip()
+        return int(out) if out.isdigit() else None
+
+    @staticmethod
+    def _registry_listening(port: int) -> bool:
+        """True if an OCI registry answers the ``/v2/`` ping on ``port`` (ours or one you run)."""
+        import httpx
+
+        try:
+            return httpx.get(f"http://localhost:{port}/v2/", timeout=2).status_code in (200, 401)
+        except Exception:
+            return False

agent_env/store/image_store/oci_registry_credentials.py

@@ -0,0 +1,196 @@
+"""Credential providers for OCI registries."""
+
+from __future__ import annotations
+
+import base64
+import binascii
+import json
+import logging
+from abc import ABC, abstractmethod
+from collections.abc import Mapping
+from dataclasses import dataclass, field
+from urllib.parse import urlsplit
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass(frozen=True)
+class RegistryAuth:
+    """docker-login material for a registry (the password may be a short-lived token)."""
+
+    registry: str
+    username: str
+    password: str = field(repr=False)
+
+
+class OciRegistryCredentials(ABC):
+    """Produces docker-login material for a registry.
+
+    Implementations resolve secret or identity inputs inside ``mint`` rather than
+    ``__init__`` so a process-cached image store can observe credential rotation.
+    """
+
+    @classmethod
+    def from_config(cls, **config) -> OciRegistryCredentials:
+        return cls(**config)
+
+    @abstractmethod
+    def mint(self, host: str) -> RegistryAuth | None:
+        """Return credentials for ``host``, or ``None`` when none are available."""
+
+
+# Docker Hub is special: unqualified image references resolve there, Docker config
+# commonly keys its credentials by ``index.docker.io``, and pulls may use
+# ``registry-1.docker.io``. Canonicalizing these aliases lets one credential entry
+# cover all equivalent forms; every other registry still requires an exact host match.
+_DOCKER_HUB_HOSTS = {
+    "docker.io",
+    "index.docker.io",
+    "registry-1.docker.io",
+}
+
+
+def normalize_registry_host(host: str) -> str:
+    """Normalize a registry authority without weakening exact host matching."""
+    normalized = host.strip().lower()
+    if not normalized:
+        raise ValueError("registry_host must not be empty")
+    if "://" in normalized or "/" in normalized:
+        raise ValueError("registry_host must be an authority without a scheme or path")
+    return "docker.io" if normalized in _DOCKER_HUB_HOSTS else normalized
+
+
+def registry_host_from_ref(ref: str) -> str | None:
+    """Return the normalized registry authority from a Docker/OCI image reference.
+
+    Docker treats a first path component as a registry only when it is ``localhost``
+    or contains a dot or port. Unqualified references therefore resolve to Docker Hub.
+    Schemed values are not valid image references and deliberately do not match.
+    """
+    if not ref or "://" in ref:
+        return None
+    first, separator, _ = ref.partition("/")
+    if not first:
+        return None
+    if not separator or not (
+        first.lower() == "localhost" or "." in first or ":" in first
+    ):
+        return "docker.io"
+    try:
+        return normalize_registry_host(first)
+    except ValueError:
+        return None
+
+
+class SecretStoreCredentials(OciRegistryCredentials):
+    """Read inline Docker-config credentials from the configured secret store."""
+
+    # The default is a secret-store lookup key, not credential material.
+    def __init__(self, secret_key: str = "registry_auths") -> None:  # nosec B107
+        self._secret_key = secret_key
+
+    def mint(self, host: str) -> RegistryAuth | None:
+        from agent_env.config import ConfigError, get_config
+
+        raw = get_config().get_secret_store().get(self._secret_key)
+        if raw is None:
+            return None
+        try:
+            document = json.loads(raw)
+        except (TypeError, json.JSONDecodeError):
+            raise ConfigError(
+                f"Registry credential secret {self._secret_key!r} must be a JSON string "
+                "containing a Docker config 'auths' mapping; nested YAML mappings are not supported"
+            ) from None
+        if not isinstance(document, Mapping):
+            raise ConfigError(
+                f"Registry credential secret {self._secret_key!r} must contain a JSON object"
+            )
+
+        if "credsStore" in document or "credHelpers" in document:
+            logger.debug(
+                "Registry credential secret %r contains Docker credential-helper settings; "
+                "agent-env only reads inline auth entries",
+                self._secret_key,
+            )
+
+        auths = document.get("auths", document)
+        if not isinstance(auths, Mapping):
+            raise ConfigError(
+                f"Registry credential secret {self._secret_key!r} has a non-object 'auths' value"
+            )
+
+        normalized_host = normalize_registry_host(host)
+        entry = next(
+            (
+                candidate
+                for key, candidate in auths.items()
+                if _docker_config_registry_host(str(key)) == normalized_host
+            ),
+            None,
+        )
+        if entry is None:
+            return None
+        if not isinstance(entry, Mapping):
+            raise ConfigError(
+                f"Registry credential secret {self._secret_key!r} has invalid credentials "
+                f"for registry {normalized_host!r}"
+            )
+
+        if entry.get("auth") is not None:
+            username, password = _decode_auth(
+                entry["auth"], self._secret_key, normalized_host
+            )
+        else:
+            username = entry.get("username")
+            password = entry.get("password")
+            if not isinstance(username, str) or not isinstance(password, str):
+                raise ConfigError(
+                    f"Registry credential secret {self._secret_key!r} must provide either 'auth' "
+                    f"or string 'username'/'password' fields for registry {normalized_host!r}"
+                )
+        return RegistryAuth(
+            registry=normalized_host, username=username, password=password
+        )
+
+
+def _docker_config_registry_host(key: str) -> str | None:
+    candidate = key.strip()
+    if not candidate:
+        return None
+    if "://" in candidate:
+        parsed = urlsplit(candidate)
+        if not parsed.hostname:
+            return None
+        try:
+            candidate = parsed.hostname
+            if parsed.port is not None:
+                candidate = f"{candidate}:{parsed.port}"
+        except ValueError:
+            return None
+    else:
+        candidate = candidate.split("/", 1)[0]
+    try:
+        return normalize_registry_host(candidate)
+    except ValueError:
+        return None
+
+
+def _decode_auth(value: object, secret_key: str, host: str) -> tuple[str, str]:
+    if not isinstance(value, str):
+        raise _invalid_auth(secret_key, host)
+    try:
+        decoded = base64.b64decode(value, validate=True).decode()
+        username, password = decoded.split(":", 1)
+    except (binascii.Error, UnicodeDecodeError, ValueError):
+        raise _invalid_auth(secret_key, host) from None
+    return username, password
+
+
+def _invalid_auth(secret_key: str, host: str):
+    from agent_env.config import ConfigError
+
+    return ConfigError(
+        f"Registry credential secret {secret_key!r} has invalid base64 'auth' "
+        f"for registry {host!r}"
+    )

agent_env/store/local_state.py

@@ -0,0 +1,27 @@
+"""On-disk state directories for the local (no-infra) stores."""
+
+from __future__ import annotations
+
+from pathlib import Path
+
+from agent_env.config.errors import ConfigError
+
+
+def ensure_state_dir(path: Path) -> None:
+    """Create ``path`` if missing. A directory this creates is private to the user (``0o700``)
+    and gets a ``.gitignore`` of ``*`` so generated state is never committed; a directory that
+    already existed is left as is. A directory that cannot be created is a ``ConfigError``."""
+    try:
+        path.parent.mkdir(parents=True, exist_ok=True)
+        try:
+            path.mkdir(mode=0o700)
+        except FileExistsError:
+            if path.is_dir():
+                return
+            raise
+        (path / ".gitignore").write_text("*\n")
+    except OSError as e:
+        raise ConfigError(
+            f"cannot create the local store directory {path} ({e}); set XDG_STATE_HOME to a "
+            "writable directory, or point the store somewhere writable in config.toml"
+        ) from e

agent_env/store/object_store/__init__.py

@@ -0,0 +1,24 @@
+"""Backend-agnostic object store abstraction + its implementations.
+
+``object_store`` defines the abstraction (ObjectStore); ``s3_object_store``,
+``local_object_store`` and ``gcs_object_store`` are implementations. The last needs
+the ``gcp`` extra, so it is not re-exported here: an impl pointer names its module.
+"""
+
+from agent_env.store.object_store.local_object_store import LocalFilesystemObjectStore
+from agent_env.store.object_store.object_store import (
+    DEFAULT_CONTENT_TYPE,
+    ObjectMetadata,
+    ObjectStore,
+    UploadPolicy,
+)
+from agent_env.store.object_store.s3_object_store import S3ObjectStore
+
+__all__ = [
+    "ObjectStore",
+    "ObjectMetadata",
+    "UploadPolicy",
+    "DEFAULT_CONTENT_TYPE",
+    "S3ObjectStore",
+    "LocalFilesystemObjectStore",
+]