agentenv-framework 0.9.1254__py3-none-any.whl

agent_env/store/document_store/sqlite_document_store.py

@@ -0,0 +1,314 @@
+"""SQLite implementation of the DocumentStore interface (stdlib, no infra).
+
+Stores each document as a JSON row in a per-collection table; SQLite provides
+uniqueness (via unique indexes on ``json_extract``), atomicity (transactions),
+and durability. ``Filter`` / ``UpdateSpec`` / ``Sort`` semantics are evaluated in
+Python, mirroring ``mongo_document_store`` as the oracle. Datetimes are
+serialized as ISO strings.
+
+SQL safety note: every user-controlled *value* is bound as a ``?`` placeholder.
+The only interpolated tokens are SQL *identifiers* (table / index / column
+names), which cannot be passed as bound parameters. Those are derived from a
+strict allowlist (see ``_safe_ident`` / ``_safe_path``) before interpolation, so
+the ``execute(f"...")`` sites carry ``# nosemgrep`` for the raw-query rule.
+"""
+
+from __future__ import annotations
+
+import json
+import logging
+import re
+import sqlite3
+import threading
+import time
+from contextlib import contextmanager
+from datetime import datetime
+from pathlib import Path
+from typing import Optional
+
+logger = logging.getLogger(__name__)
+
+from agent_env.store.document_store.document_store import (
+    DocumentStore,
+    DuplicateKeyError,
+    Filter,
+    Sort,
+    UpdateSpec,
+)
+from agent_env.store.document_store import evaluation
+from agent_env.store.ids import aliases_local_encoding
+from agent_env.store.local_state import ensure_state_dir
+
+# SQL identifiers (table/index names) must be a bare word; JSON-path field
+# references additionally allow dots for nested paths (e.g. ``metadata.rev``).
+_IDENT_RE = re.compile(r"[A-Za-z0-9_]+")
+_PATH_RE = re.compile(r"[A-Za-z0-9_.]+")
+
+_BUSY_TIMEOUT_SECONDS = 5.0
+_WAL_SWITCH_RETRY_DELAY_SECONDS = 0.01
+
+
+def _enable_wal(conn: sqlite3.Connection) -> None:
+    """Switch the database to WAL. Processes opening a new database at the same time all need
+    the exclusive lock this takes, and SQLite fails the losers at once instead of waiting, so
+    retry until the winner is done."""
+    deadline = time.monotonic() + _BUSY_TIMEOUT_SECONDS
+    while True:
+        try:
+            conn.execute("PRAGMA journal_mode=WAL")
+            return
+        except sqlite3.OperationalError as e:
+            if e.sqlite_errorcode != sqlite3.SQLITE_BUSY or time.monotonic() >= deadline:
+                raise
+            time.sleep(_WAL_SWITCH_RETRY_DELAY_SECONDS)
+
+
+class LocalSqliteDocumentStore(DocumentStore):
+    """DocumentStore backed by a single stdlib-``sqlite3`` database file."""
+
+    def __init__(self, path: str) -> None:
+        # Nothing touches disk until the first operation: constructing a store creates nothing.
+        self._path = Path(path)
+        self._connection: Optional[sqlite3.Connection] = None
+        self._lock = threading.RLock()
+        self._tables: set[str] = set()
+
+    @property
+    def path(self) -> Path:
+        return self._path
+
+    def check_id(self, entity_id: str) -> None:
+        super().check_id(entity_id)
+        if aliases_local_encoding(entity_id):
+            raise ValueError(
+                f"id {entity_id!r} is spelled like an encoded @local id, so it would share that id's objects, "
+                "images and files; a bare id in a local store can't start with 'local/' or 'local-<name>-<12 hex>'"
+            )
+
+    @property
+    def _conn(self) -> sqlite3.Connection:
+        with self._lock:
+            if self._connection is None:
+                ensure_state_dir(self._path.parent)
+                conn = sqlite3.connect(
+                    str(self._path), timeout=_BUSY_TIMEOUT_SECONDS, check_same_thread=False, isolation_level=None
+                )
+                _enable_wal(conn)
+                self._connection = conn
+            return self._connection
+
+    def find_one(
+        self, collection: str, filter: Filter, sort: Optional[Sort] = None
+    ) -> Optional[dict]:
+        with self._lock:
+            matches = [doc for _, doc in self._load(collection) if evaluation.matches(doc, filter)]
+            matches = evaluation.sort_docs(matches, sort)
+            return matches[0] if matches else None
+
+    def query(
+        self,
+        collection: str,
+        filter: Filter,
+        sort: Optional[Sort] = None,
+        limit: Optional[int] = None,
+        offset: Optional[int] = None,
+    ) -> list[dict]:
+        with self._lock:
+            docs = [doc for _, doc in self._load(collection) if evaluation.matches(doc, filter)]
+            docs = evaluation.sort_docs(docs, sort)
+            if offset:
+                docs = docs[offset:]
+            if limit:
+                docs = docs[:limit]
+            return docs
+
+    def count(self, collection: str, filter: Filter) -> int:
+        with self._lock:
+            return sum(1 for _, doc in self._load(collection) if evaluation.matches(doc, filter))
+
+    def insert(self, collection: str, doc: dict) -> None:
+        with self._lock:
+            tbl = self._ensure_table(collection)
+            self._insert_raw(tbl, dict(doc))
+
+    def update(
+        self, collection: str, filter: Filter, update: UpdateSpec, upsert: bool = False
+    ) -> int:
+        update.validate()
+        with self._lock, self._immediate():
+            match = self._first_match(collection, filter)
+            if match is not None:
+                rowid, doc = match
+                evaluation.apply_update(doc, update)
+                self._conn.execute(
+                    # nosemgrep: sqlalchemy-execute-raw-query -- table is a validated identifier; doc/rowid are bound
+                    f'UPDATE "{self._table(collection)}" SET doc=? WHERE rowid=?',
+                    (self._dumps(doc), rowid),
+                )
+                return 1
+            if upsert:
+                self._insert_raw(self._ensure_table(collection), evaluation.synthesize(filter, update))
+                return 1
+            return 0
+
+    def update_one_and_get(
+        self,
+        collection: str,
+        filter: Filter,
+        update: UpdateSpec,
+        return_after: bool = True,
+        upsert: bool = False,
+    ) -> Optional[dict]:
+        update.validate()
+        with self._lock, self._immediate():
+            match = self._first_match(collection, filter)
+            if match is not None:
+                rowid, doc = match
+                before = json.loads(self._dumps(doc))
+                evaluation.apply_update(doc, update)
+                self._conn.execute(
+                    # nosemgrep: sqlalchemy-execute-raw-query -- table is a validated identifier; doc/rowid are bound
+                    f'UPDATE "{self._table(collection)}" SET doc=? WHERE rowid=?',
+                    (self._dumps(doc), rowid),
+                )
+                return doc if return_after else before
+            if upsert:
+                new_doc = evaluation.synthesize(filter, update)
+                self._insert_raw(self._ensure_table(collection), new_doc)
+                return new_doc if return_after else None
+            return None
+
+    def replace(
+        self, collection: str, filter: Filter, doc: dict, upsert: bool = False
+    ) -> int:
+        with self._lock, self._immediate():
+            match = self._first_match(collection, filter)
+            if match is not None:
+                rowid, _ = match
+                self._conn.execute(
+                    # nosemgrep: sqlalchemy-execute-raw-query -- table is a validated identifier; doc/rowid are bound
+                    f'UPDATE "{self._table(collection)}" SET doc=? WHERE rowid=?',
+                    (self._dumps(dict(doc)), rowid),
+                )
+                return 1
+            if upsert:
+                self._insert_raw(self._ensure_table(collection), dict(doc))
+                return 1
+            return 0
+
+    def delete(self, collection: str, filter: Filter) -> int:
+        with self._lock, self._immediate():
+            match = self._first_match(collection, filter)
+            if match is None:
+                return 0
+            rowid, _ = match
+            # nosemgrep: sqlalchemy-execute-raw-query -- table is a validated identifier; rowid is bound
+            self._conn.execute(f'DELETE FROM "{self._table(collection)}" WHERE rowid=?', (rowid,))
+            return 1
+
+    def ensure_index(
+        self,
+        collection: str,
+        fields: list[str],
+        unique: bool = False,
+        ttl_seconds: Optional[int] = None,
+    ) -> None:
+        if ttl_seconds is not None:
+            logger.warning(
+                "LocalSqliteDocumentStore has no TTL support; ttl_seconds=%s on %s%s ignored — rows will not auto-expire.",
+                ttl_seconds, collection, fields,
+            )
+        with self._lock:
+            tbl = self._ensure_table(collection)
+            raw_name = tbl + "_" + "_".join(fields) + ("_unique" if unique else "_index")
+            name = self._safe_ident(re.sub(r"\W", "_", raw_name))
+            cols = ", ".join(f"json_extract(doc, '$.{self._safe_path(f)}')" for f in fields)
+            uniq = "UNIQUE " if unique else ""
+            self._conn.execute(
+                # nosemgrep: sqlalchemy-execute-raw-query -- name/tbl/cols are validated identifiers
+                f'CREATE {uniq}INDEX IF NOT EXISTS "{name}" ON "{tbl}" ({cols})'
+            )
+
+    # --- sqlite plumbing ---
+
+    @staticmethod
+    def _safe_ident(name: str) -> str:
+        """Validate a SQL identifier (table/index name) against a strict allowlist.
+
+        Identifiers cannot be passed as bound parameters, so we assert they are a
+        bare ``[A-Za-z0-9_]`` word before interpolating them into SQL. User input
+        only ever reaches SQL as a bound ``?`` value, never as an identifier.
+        """
+        if not _IDENT_RE.fullmatch(name):
+            raise ValueError(f"unsafe SQL identifier: {name!r}")
+        return name
+
+    @staticmethod
+    def _safe_path(field: str) -> str:
+        """Validate a JSON-path field reference used inside ``json_extract('$.…')``."""
+        if not _PATH_RE.fullmatch(field):
+            raise ValueError(f"unsafe JSON field path: {field!r}")
+        return field
+
+    def _table(self, collection: str) -> str:
+        # The table name is an identifier (not bindable): sanitize the collection
+        # to a bare word, then assert the allowlist as defense-in-depth.
+        return self._safe_ident("docs_" + re.sub(r"\W", "_", collection))
+
+    def _ensure_table(self, collection: str) -> str:
+        tbl = self._table(collection)
+        if tbl not in self._tables:
+            # nosemgrep: sqlalchemy-execute-raw-query -- tbl is a validated identifier
+            self._conn.execute(f'CREATE TABLE IF NOT EXISTS "{tbl}" (doc TEXT NOT NULL)')
+            self._tables.add(tbl)
+        return tbl
+
+    def _adopt_if_created(self, tbl: str) -> bool:
+        """Adopt a table that exists on disk but is not yet in ``self._tables``."""
+        if self._conn.execute(
+            "SELECT 1 FROM sqlite_master WHERE type='table' AND name=?", (tbl,)
+        ).fetchone():
+            self._tables.add(tbl)
+            return True
+        return False
+
+    def _load(self, collection: str) -> list[tuple[int, dict]]:
+        tbl = self._table(collection)
+        if tbl not in self._tables and not self._adopt_if_created(tbl):
+            return []
+        return [
+            (row[0], json.loads(row[1]))
+            # nosemgrep: sqlalchemy-execute-raw-query -- tbl is a validated identifier
+            for row in self._conn.execute(f'SELECT rowid, doc FROM "{tbl}"')
+        ]
+
+    def _first_match(self, collection: str, filter: Filter) -> Optional[tuple[int, dict]]:
+        for rowid, doc in self._load(collection):
+            if evaluation.matches(doc, filter):
+                return rowid, doc
+        return None
+
+    @contextmanager
+    def _immediate(self):
+        self._conn.execute("BEGIN IMMEDIATE")
+        try:
+            yield
+            self._conn.execute("COMMIT")
+        except BaseException:
+            self._conn.execute("ROLLBACK")
+            raise
+
+    def _insert_raw(self, tbl: str, doc: dict) -> None:
+        try:
+            # nosemgrep: sqlalchemy-execute-raw-query -- tbl is a validated identifier; doc is bound
+            self._conn.execute(f'INSERT INTO "{tbl}"(doc) VALUES(?)', (self._dumps(doc),))
+        except sqlite3.IntegrityError as e:
+            raise DuplicateKeyError(str(e)) from e
+
+    def _json_default(self, o):
+        if isinstance(o, datetime):
+            return o.isoformat()
+        raise TypeError(f"Object of type {type(o).__name__} is not JSON-serializable")
+
+    def _dumps(self, doc: dict) -> str:
+        return json.dumps(doc, default=self._json_default)

agent_env/store/ids.py

@@ -0,0 +1,123 @@
+"""Namespaced entity ids, and the object keys, image repositories and filenames derived from them.
+
+An id opening ``@<namespace>/`` names the store it lives in; ``@local/…`` is the only namespace
+today. Documents keep an id verbatim. Everything derived from an ``@local`` id uses one bounded
+segment, ``local/<slug>-<sha256(id)[:12]>``, and every other id passes through byte-identical.
+The encoders never raise: they expect ids that ``validate_local_id`` accepts, and still return a
+safe segment for any other input. An id may contain spaces, ``&`` and parentheses, so quote one
+before it reaches a shell or a URL, and derive object keys, image repositories and filenames with
+the helpers here rather than from the raw id.
+"""
+
+from __future__ import annotations
+
+import hashlib
+import re
+import unicodedata
+from typing import Optional
+
+LOCAL_NAMESPACE = "local"
+LOCAL_PREFIX = f"@{LOCAL_NAMESPACE}/"
+MAX_LOCAL_ID_BYTES = 4096
+
+_SLUG_LENGTH = 48
+_HASH_LENGTH = 12
+_NON_SLUG_RUN = re.compile(r"[^a-z0-9]+")
+_ENCODED_LOCAL_ID = re.compile(
+    rf"{LOCAL_NAMESPACE}(?:/|-(?:[a-z0-9]+(?:-[a-z0-9]+)*-)?[0-9a-f]{{{_HASH_LENGTH}}}(?:/|$))"
+)
+_ID_PUNCTUATION = frozenset(" ._-~/@()+,&")
+_ID_CATEGORIES = frozenset("LNM")
+_INVISIBLE = re.compile(
+    "[\u034f\u115f\u1160\u17b4\u17b5\u180b-\u180d\u180f\u3164\ufe00-\ufe0f\uffa0\U000e0100-\U000e01ef]"
+)
+
+
+def parse_namespace(entity_id: str) -> Optional[str]:
+    """The namespace an ``@<namespace>/…`` id names, or None for a bare id."""
+    if not entity_id.startswith("@"):
+        return None
+    return entity_id[1:].split("/", 1)[0]
+
+
+def is_local_id(entity_id: str) -> bool:
+    return entity_id.startswith(LOCAL_PREFIX)
+
+
+def validate_local_id(entity_id: str) -> None:
+    """Raise ValueError unless ``entity_id`` is ``@local/`` followed by non-empty path segments,
+    none of them ``.`` or ``..`` or opening or closing with a space, made only of letters, digits,
+    spaces and ``. _ - ~ / @ ( ) + , &``, and at most 4096 bytes of UTF-8 in all (Linux's PATH_MAX).
+    Only the leading ``@`` names a namespace, so one inside a segment is an ordinary character.
+
+    Letters include every script's, with its combining marks, as long as each mark follows a
+    letter, digit or mark. The invisible (default-ignorable) letters and marks are refused."""
+    if not entity_id.startswith(LOCAL_PREFIX):
+        raise ValueError(f"local id {entity_id!r} must start with {LOCAL_PREFIX!r}")
+    size = len(entity_id.encode("utf-8", "surrogatepass"))
+    if size > MAX_LOCAL_ID_BYTES:
+        raise ValueError(f"local id is {size} bytes; the limit is {MAX_LOCAL_ID_BYTES}: {entity_id[:64]!r}…")
+    disallowed = sorted({c for c in entity_id[len(LOCAL_PREFIX):] if not _allowed_in_id(c)})
+    if disallowed:
+        raise ValueError(
+            f"local id {entity_id!r} contains {ascii(''.join(disallowed))}; "
+            "a local id may contain only letters, digits, spaces and . _ - ~ / @ ( ) + , &"
+        )
+    segments = entity_id[len(LOCAL_PREFIX):].split("/")
+    if any(segment in ("", ".", "..") for segment in segments):
+        raise ValueError(f"local id {entity_id!r} has an empty, '.' or '..' path segment")
+    if any(segment != segment.strip(" ") for segment in segments):
+        raise ValueError(f"local id {entity_id!r} has a path segment that opens or closes with a space")
+    if any(_detached_mark(segment) for segment in segments):
+        raise ValueError(f"local id {entity_id!r} has a combining mark that does not follow a letter or digit")
+
+
+def _allowed_in_id(character: str) -> bool:
+    if _INVISIBLE.match(character):
+        return False
+    return character in _ID_PUNCTUATION or unicodedata.category(character)[0] in _ID_CATEGORIES
+
+
+def _detached_mark(segment: str) -> bool:
+    categories = [unicodedata.category(character)[0] for character in segment]
+    return any(this == "M" and before not in _ID_CATEGORIES for before, this in zip(["Z", *categories], categories))
+
+
+def _local_segment(entity_id: str) -> str:
+    digest = hashlib.sha256(entity_id.encode("utf-8", "surrogatepass")).hexdigest()[:_HASH_LENGTH]
+    decomposed = unicodedata.normalize("NFKD", entity_id.removeprefix(LOCAL_PREFIX))
+    folded = "".join(c for c in decomposed if not unicodedata.combining(c)).casefold().replace("~", "")
+    slug = _NON_SLUG_RUN.sub("-", folded).strip("-")[:_SLUG_LENGTH].rstrip("-")
+    return f"{LOCAL_NAMESPACE}/{slug}-{digest}" if slug else f"{LOCAL_NAMESPACE}/{digest}"
+
+
+def aliases_local_encoding(entity_id: str) -> bool:
+    """Whether a bare id is spelled like an ``@local`` id's encoding: ``local/…``, or the filename form
+    ``local-<slug>-<hash>`` alone or opening a path. Its objects, images and files would take that id's."""
+    return _ENCODED_LOCAL_ID.match(entity_id) is not None
+
+
+def key_segment(entity_id: str) -> str:
+    """The object-key segment for ``entity_id`` (``artifacts/<type>/<segment>/<version>/…``)."""
+    return _local_segment(entity_id) if is_local_id(entity_id) else entity_id
+
+
+def image_repository(entity_id: str) -> str:
+    """The image repository for ``entity_id``; an ``@local`` id yields a valid OCI repository path."""
+    return key_segment(entity_id)
+
+
+def fs_safe(entity_id: str) -> str:
+    """A filename for an ``@local`` id: its key segment with no ``/``. Any other id is returned
+    unchanged, so it is only as filename-safe as it already was."""
+    return _local_segment(entity_id).replace("/", "-", 1) if is_local_id(entity_id) else entity_id
+
+
+# The longest id derive_id makes is a skill's or CLI's per-file artifact, ``{id}__files__<16 hex>``,
+# so an authored id leaves room for those 25 bytes.
+MAX_AUTHORED_LOCAL_ID_BYTES = MAX_LOCAL_ID_BYTES - len("__files__") - 16
+
+
+def derive_id(base: str, suffix: str) -> str:
+    """The id of an entity derived from ``base``: a suffix, so the base's namespace carries over."""
+    return f"{base}__{suffix}"

agent_env/store/image_store/__init__.py

@@ -0,0 +1,34 @@
+"""Backend-agnostic image (registry) store abstraction + its implementations.
+
+``image_store`` defines the abstractions and generic OCI implementation;
+``ecr_image_store`` and ``local_registry_image_store`` retain compatibility names.
+"""
+
+from agent_env.store.image_store.ecr_image_store import EcrCredentials, EcrImageStore
+from agent_env.store.image_store.image_store import (
+    ImageStore,
+    OciRegistryImageStore,
+)
+from agent_env.store.image_store.local_registry_image_store import (
+    LocalRegistryImageStore,
+)
+from agent_env.store.image_store.oci_registry_credentials import (
+    OciRegistryCredentials,
+    RegistryAuth,
+    SecretStoreCredentials,
+    normalize_registry_host,
+    registry_host_from_ref,
+)
+
+__all__ = [
+    "ImageStore",
+    "OciRegistryImageStore",
+    "RegistryAuth",
+    "OciRegistryCredentials",
+    "normalize_registry_host",
+    "registry_host_from_ref",
+    "EcrCredentials",
+    "SecretStoreCredentials",
+    "EcrImageStore",
+    "LocalRegistryImageStore",
+]

agent_env/store/image_store/ecr_image_store.py

@@ -0,0 +1,131 @@
+"""ECR implementation of the ImageStore interface."""
+
+from __future__ import annotations
+
+import base64
+import binascii
+import logging
+from typing import Any
+
+import boto3
+from botocore.exceptions import ClientError
+
+from agent_env.config.errors import ConfigError
+from agent_env.store.image_store.image_store import OciRegistryImageStore
+from agent_env.store.image_store.oci_registry_credentials import (
+    OciRegistryCredentials,
+    RegistryAuth,
+    normalize_registry_host,
+)
+
+logger = logging.getLogger(__name__)
+
+
+class EcrCredentials(OciRegistryCredentials):
+    """Mint short-lived Docker credentials through the AWS ECR API."""
+
+    def __init__(
+        self,
+        region: str | None = None,
+        access_key: str | None = None,
+        secret_key: str | None = None,
+        *,
+        client: Any = None,
+    ) -> None:
+        if (access_key is None) != (secret_key is None):
+            raise ValueError("access_key and secret_key must be provided together")
+        self._region = region
+        self._access_key = access_key
+        self._secret_key = secret_key
+        self._client = client
+
+    @classmethod
+    def from_config(
+        cls,
+        *,
+        region: str | None = None,
+        access_key: str | None = None,
+        secret_key: str | None = None,
+    ) -> EcrCredentials:
+        """Build explicitly configured ECR credentials.
+
+        The keyless constructor (the legacy ``EcrImageStore`` default) resolves
+        the AWS default credential chain; generic OCI configuration must name its
+        credential inputs.
+        """
+        if not region or not access_key or not secret_key:
+            raise ValueError("region, access_key, and secret_key must be configured")
+        return cls(region=region, access_key=access_key, secret_key=secret_key)
+
+    @property
+    def client(self):
+        if self._client is None:
+            if self._region is None:
+                raise ConfigError(
+                    "EcrCredentials has no region: configure it on the store's "
+                    "credentials (e.g. [stores.image.config] region)."
+                )
+            if self._access_key is None:
+                self._client = boto3.client("ecr", region_name=self._region)
+            else:
+                self._client = boto3.client(
+                    "ecr",
+                    region_name=self._region,
+                    aws_access_key_id=self._access_key,
+                    aws_secret_access_key=self._secret_key,
+                )
+        return self._client
+
+    def mint(self, host: str) -> RegistryAuth:
+        token = self.client.get_authorization_token()["authorizationData"][0][
+            "authorizationToken"
+        ]
+        try:
+            username, password = (
+                base64.b64decode(token, validate=True).decode().split(":", 1)
+            )
+        except (binascii.Error, UnicodeDecodeError, ValueError):
+            raise ValueError("ECR returned an invalid authorization token") from None
+        return RegistryAuth(
+            registry=normalize_registry_host(host), username=username, password=password
+        )
+
+
+class EcrImageStore(OciRegistryImageStore):
+    """ImageStore backed by AWS ECR; the client is built lazily on the first ECR op."""
+
+    def __init__(
+        self,
+        registry_host: str,
+        repository_prefix: str = "",
+        *,
+        credentials: EcrCredentials | None = None,
+    ) -> None:
+        if credentials is not None and not isinstance(credentials, EcrCredentials):
+            raise TypeError(
+                "EcrImageStore credentials must be an EcrCredentials instance"
+            )
+        self._ecr_credentials = credentials if credentials is not None else EcrCredentials()
+        super().__init__(
+            registry_host=registry_host,
+            repository_prefix=repository_prefix,
+            credentials=self._ecr_credentials,
+        )
+
+    @property
+    def _ecr(self):
+        return self._ecr_credentials.client
+
+    def ensure_repository(self, repository: str) -> None:
+        name = self._repo_name(repository)
+        try:
+            self._ecr.create_repository(repositoryName=name)
+            logger.info(f"Created ECR repository {name}")
+        except self._ecr.exceptions.RepositoryAlreadyExistsException:
+            pass
+        except ClientError as e:
+            if e.response["Error"]["Code"] != "AccessDeniedException":
+                raise
+            logger.info(
+                f"No CreateRepository permission for {name}; assuming it exists"
+            )