toolbound 0.2.0__py3-none-any.whl

toolbound/core/signing.py

@@ -0,0 +1,137 @@
+"""Expiring tokens for cursors, approval requests and approval grants.
+
+* :meth:`Signer.sign` makes integrity-protected tokens (HMAC-SHA256). Their
+  payload is **readable** by anyone holding the token. Used for grants, which
+  hold only ids and digests.
+* :meth:`Signer.seal` makes authenticated **encrypted** tokens (AES-256-GCM).
+  Used for cursors (which reach the model) and approval requests (which hold
+  arguments and trusted bound values).
+
+Both kinds expire and are bound to a purpose prefix, so a token minted for one
+purpose is never accepted for another.
+"""
+
+from __future__ import annotations
+
+import base64
+import hashlib
+import hmac
+import json
+import secrets
+import time
+from typing import Any
+
+from cryptography.exceptions import InvalidTag
+from cryptography.hazmat.primitives.ciphers.aead import AESGCM
+
+from .errors import ConfigurationError
+
+_MIN_SECRET_BYTES = 32
+_DEFAULT_MAX_TOKEN = 8192
+_NONCE_BYTES = 12
+
+
+def _b64(data: bytes) -> str:
+    return base64.urlsafe_b64encode(data).rstrip(b"=").decode("ascii")
+
+
+def _unb64(text: str) -> bytes:
+    return base64.urlsafe_b64decode(text + "=" * (-len(text) % 4))
+
+
+def canonical_digest(value: Any) -> str:
+    """SHA-256 over canonical JSON (sorted keys, no whitespace)."""
+    raw = json.dumps(value, sort_keys=True, separators=(",", ":"), ensure_ascii=False, default=str)
+    return hashlib.sha256(raw.encode()).hexdigest()
+
+
+class Signer:
+    """Signs small JSON payloads. Tokens look like ``<prefix>.<payload>.<mac>``.
+
+    Tokens are integrity-protected, not encrypted: never put secrets in them.
+    """
+
+    def __init__(self, secret: bytes | str | None) -> None:
+        if secret is None:
+            key = secrets.token_bytes(32)
+            self.ephemeral = True
+        else:
+            key = secret.encode() if isinstance(secret, str) else bytes(secret)
+            self.ephemeral = False
+        if len(key) < _MIN_SECRET_BYTES:
+            raise ConfigurationError("Gateway secret must be at least 32 bytes")
+        self._key = key
+        # Separate keys per primitive, derived from the configured secret.
+        self._mac_key = hmac.new(key, b"toolbound/v1/mac", hashlib.sha256).digest()
+        self._aead = AESGCM(hmac.new(key, b"toolbound/v1/seal", hashlib.sha256).digest())
+
+    def seal(self, prefix: str, payload: dict[str, Any], ttl_seconds: float) -> str:
+        """Encrypt and authenticate ``payload``; the prefix is bound as associated data."""
+        body = dict(payload)
+        body["exp"] = int(time.time() + ttl_seconds)
+        plaintext = json.dumps(body, sort_keys=True, separators=(",", ":")).encode()
+        nonce = secrets.token_bytes(_NONCE_BYTES)
+        ciphertext = self._aead.encrypt(nonce, plaintext, prefix.encode())
+        return f"{prefix}.{_b64(nonce + ciphertext)}"
+
+    def unseal(
+        self, prefix: str, token: Any, *, max_length: int = _DEFAULT_MAX_TOKEN
+    ) -> dict[str, Any] | None:
+        """Return the payload, or None if malformed, forged, for another purpose, or expired."""
+        if not isinstance(token, str) or len(token) > max_length:
+            return None
+        head, sep, data = token.partition(".")
+        if not sep or head != prefix or "." in data:
+            return None
+        try:
+            raw = _unb64(data)
+        except ValueError:
+            return None
+        if len(raw) <= _NONCE_BYTES:
+            return None
+        try:
+            plaintext = self._aead.decrypt(raw[:_NONCE_BYTES], raw[_NONCE_BYTES:], prefix.encode())
+            payload = json.loads(plaintext)
+        except (InvalidTag, ValueError, UnicodeDecodeError):
+            return None
+        return _unexpired(payload)
+
+    def sign(self, prefix: str, payload: dict[str, Any], ttl_seconds: float) -> str:
+        body = dict(payload)
+        body["exp"] = int(time.time() + ttl_seconds)
+        encoded = _b64(json.dumps(body, sort_keys=True, separators=(",", ":")).encode())
+        mac = self._mac(prefix, encoded)
+        return f"{prefix}.{encoded}.{mac}"
+
+    def verify(
+        self, prefix: str, token: Any, *, max_length: int = _DEFAULT_MAX_TOKEN
+    ) -> dict[str, Any] | None:
+        """Return the payload, or None if the token is malformed, forged or expired."""
+        if not isinstance(token, str) or len(token) > max_length:
+            return None
+        parts = token.split(".")
+        if len(parts) != 3 or parts[0] != prefix:
+            return None
+        expected = self._mac(prefix, parts[1])
+        if not hmac.compare_digest(expected, parts[2]):
+            return None
+        try:
+            payload = json.loads(_unb64(parts[1]))
+        except (ValueError, UnicodeDecodeError):
+            return None
+        return _unexpired(payload)
+
+    def _mac(self, prefix: str, encoded: str) -> str:
+        digest = hmac.new(self._mac_key, f"{prefix}.{encoded}".encode(), hashlib.sha256).digest()
+        return _b64(digest[:24])
+
+    def __repr__(self) -> str:
+        return f"Signer(ephemeral={self.ephemeral})"
+
+
+def _unexpired(payload: Any) -> dict[str, Any] | None:
+    if not isinstance(payload, dict) or not isinstance(payload.get("exp"), int):
+        return None
+    if payload["exp"] < time.time():
+        return None
+    return payload

toolbound/core/state.py

@@ -0,0 +1,102 @@
+"""Default in-memory :class:`~toolbound.core.ports.StateStore`, and claim helpers."""
+
+from __future__ import annotations
+
+import heapq
+import logging
+import threading
+import time
+
+from .errors import ErrorCode, StateStoreError, ToolFailure
+from .ports import StateStore
+
+logger = logging.getLogger("toolbound")
+
+
+class MemoryStateStore:
+    """Per-process TTL set. Suitable for a single worker or for tests.
+
+    It never forgets a key before its TTL expires: when ``max_entries`` live
+    keys exist, new claims raise :class:`StateStoreError` and Toolbound refuses
+    the call (fail closed) instead of evicting a live key, which would let an
+    approval or call id be used twice.
+
+    Keys are partitioned by their first two ``:``-separated segments. Toolbound's
+    keys are ``call:<actor+tenant>:…`` and ``approval:<actor+tenant>:…``, so one
+    actor filling its partition (``max_per_partition``) cannot block anyone else.
+    ``max_entries`` bounds total memory.
+
+    It is thread-safe and not bound to an event loop. With several processes,
+    replay protection needs a shared store; see ``docs/security.md`` for a Redis
+    recipe.
+    """
+
+    def __init__(self, *, max_entries: int = 200_000, max_per_partition: int = 10_000) -> None:
+        if max_entries < 1 or max_per_partition < 1:
+            raise ValueError("max_entries and max_per_partition must be at least 1")
+        self._entries: dict[str, float] = {}
+        self._expiry: list[tuple[float, str]] = []
+        self._partitions: dict[str, int] = {}
+        self._max = max_entries
+        self._max_partition = max_per_partition
+        self._lock = threading.Lock()
+
+    async def claim(self, key: str, ttl_seconds: float) -> bool:
+        with self._lock:
+            now = time.monotonic()
+            self._evict(now)
+            expiry = self._entries.get(key)
+            if expiry is not None and expiry > now:
+                return False
+            partition = _partition(key)
+            if expiry is None and self._partitions.get(partition, 0) >= self._max_partition:
+                raise StateStoreError(
+                    f"MemoryStateStore partition is full ({self._max_partition} live keys); "
+                    "refusing to evict unexpired replay-protection entries"
+                )
+            if expiry is None and len(self._entries) >= self._max:
+                raise StateStoreError(
+                    f"MemoryStateStore is full ({self._max} live keys); refusing to evict "
+                    "unexpired replay-protection entries"
+                )
+            deadline = now + ttl_seconds
+            if expiry is None:
+                self._partitions[partition] = self._partitions.get(partition, 0) + 1
+            self._entries[key] = deadline
+            heapq.heappush(self._expiry, (deadline, key))
+            return True
+
+    def __len__(self) -> int:
+        return len(self._entries)
+
+    def _evict(self, now: float) -> None:
+        while self._expiry and self._expiry[0][0] <= now:
+            deadline, key = heapq.heappop(self._expiry)
+            if self._entries.get(key) == deadline:
+                del self._entries[key]
+                partition = _partition(key)
+                left = self._partitions.get(partition, 1) - 1
+                if left:
+                    self._partitions[partition] = left
+                else:
+                    self._partitions.pop(partition, None)
+
+
+def _partition(key: str) -> str:
+    return ":".join(key.split(":", 2)[:2])
+
+
+async def claim_or_fail(store: StateStore, key: str, ttl_seconds: float) -> bool:
+    """``store.claim`` that fails closed: any store error stops the call."""
+    try:
+        return bool(await store.claim(key, ttl_seconds))
+    except Exception as exc:
+        logger.error("state store claim failed: %s", type(exc).__name__)
+        full = isinstance(exc, StateStoreError)
+        raise ToolFailure(
+            ErrorCode.INTERNAL_ERROR,
+            "replay protection is unavailable, so the call was not executed"
+            + (" (too many recent calls for this user)" if full else ""),
+            retryable=not full,
+            request_sent=False,
+        ) from None

toolbound/estimators.py

@@ -0,0 +1,24 @@
+"""Optional tokenizer-backed estimators (install ``toolbound[tiktoken]``)."""
+
+from __future__ import annotations
+
+from typing import Any
+
+
+class TiktokenEstimator:
+    """Counts tokens with a tiktoken encoding (default ``o200k_base``).
+
+    tiktoken implements OpenAI tokenizers. For other model families it is still
+    a closer approximation than byte counting, but it is not exact.
+    """
+
+    def __init__(self, encoding: str = "o200k_base") -> None:
+        try:
+            import tiktoken
+        except ImportError as exc:  # pragma: no cover - environment dependent
+            raise ImportError("pip install 'toolbound[tiktoken]'") from exc
+        self._encoding: Any = tiktoken.get_encoding(encoding)
+        self.name = encoding
+
+    def count(self, text: str) -> int:
+        return len(self._encoding.encode(text, disallowed_special=()))

toolbound/loaders/__init__.py

@@ -0,0 +1,202 @@
+"""Load declarative contracts (YAML or JSON) into :class:`Operation` objects.
+
+Files are validated against :data:`CONTRACT_SCHEMA` and then against the same
+checks as Python-defined contracts. Files with ``status: draft`` (for example
+generated by an importer) are refused until a reviewer changes the status.
+Custom Python transforms cannot be referenced from files except by name via
+``Gateway(transforms={...})``.
+"""
+
+from __future__ import annotations
+
+import json
+from pathlib import Path
+from typing import Any, cast
+
+from jsonschema import Draft202012Validator
+
+from toolbound.core.contracts import (
+    ApiKeyAuth,
+    Auth,
+    BasicAuth,
+    BearerAuth,
+    Effect,
+    Field,
+    FieldSpec,
+    Limits,
+    NetworkPolicy,
+    NoAuth,
+    Operation,
+    Output,
+    Pagination,
+    RateLimit,
+    RetryPolicy,
+    Service,
+)
+from toolbound.core.errors import ConfigurationError
+from toolbound.core.paths import parse_path
+
+from .meta_schema import CONTRACT_SCHEMA
+
+_VALIDATOR = Draft202012Validator(CONTRACT_SCHEMA)
+_SUFFIXES = (".yaml", ".yml", ".json")
+
+__all__ = ["CONTRACT_SCHEMA", "load_contract_file", "load_contracts", "parse_contract"]
+
+
+def load_contracts(path: str | Path) -> list[Operation]:
+    """Load one file, or every ``*.yaml|*.yml|*.json`` file in a directory (sorted)."""
+    root = Path(path)
+    if root.is_dir():
+        files = sorted(p for p in root.iterdir() if p.suffix in _SUFFIXES and p.is_file())
+        if not files:
+            raise ConfigurationError(f"no contract files found in {root}")
+        operations: list[Operation] = []
+        for file in files:
+            operations.extend(load_contract_file(file))
+        return operations
+    return load_contract_file(root)
+
+
+def load_contract_file(path: str | Path) -> list[Operation]:
+    file = Path(path)
+    text = file.read_text(encoding="utf-8")
+    if file.suffix == ".json":
+        document = json.loads(text)
+    elif file.suffix in (".yaml", ".yml"):
+        try:
+            import yaml
+        except ImportError as exc:  # pragma: no cover - depends on environment
+            raise ConfigurationError(
+                "YAML contracts need PyYAML: pip install 'toolbound[yaml]'"
+            ) from exc
+        document = yaml.safe_load(text)
+    else:
+        raise ConfigurationError(f"unsupported contract file type: {file.name}")
+    try:
+        return parse_contract(document)
+    except ConfigurationError as exc:
+        raise ConfigurationError(f"{file}: {exc}") from exc
+
+
+def parse_contract(document: Any) -> list[Operation]:
+    """Validate a parsed contract document and build its operations."""
+    errors = sorted(_VALIDATOR.iter_errors(document), key=lambda e: list(e.absolute_path))
+    if errors:
+        first = errors[0]
+        where = "/".join(str(p) for p in first.absolute_path) or "(root)"
+        raise ConfigurationError(f"invalid contract at {where}: {first.message}")
+    if document.get("status", "active") == "draft":
+        raise ConfigurationError(
+            "contract is marked 'status: draft'; review it and set 'status: active' to load it"
+        )
+    service = _service(document["service"])
+    return [_operation(op, service) for op in document["operations"]]
+
+
+def _service(raw: dict[str, Any]) -> Service:
+    return Service(
+        name=raw["name"],
+        base_url=raw["base_url"],
+        auth=_auth(raw.get("auth")),
+        headers=raw.get("headers", {}),
+        network=_network(raw.get("network", {})),
+        limits=Limits(**raw.get("limits", {})),
+        retry=_retry(raw.get("retry", {})),
+    )
+
+
+def _auth(raw: dict[str, Any] | None) -> Auth:
+    if raw is None or raw["type"] == "none":
+        return NoAuth()
+    credential = raw.get("credential")
+    if not credential:
+        raise ConfigurationError(f"auth type {raw['type']!r} needs 'credential'")
+    if raw["type"] == "bearer":
+        return BearerAuth(credential)
+    if raw["type"] == "basic":
+        return BasicAuth(credential)
+    if "query" in raw:
+        return ApiKeyAuth(credential, header=None, query=raw["query"])
+    return ApiKeyAuth(credential, header=raw.get("header", "X-API-Key"))
+
+
+def _network(raw: dict[str, Any]) -> NetworkPolicy:
+    return NetworkPolicy(
+        allow_hosts=tuple(raw.get("allow_hosts", ())),
+        allow_ports=tuple(raw.get("allow_ports", ())),
+        allow_http=raw.get("allow_http", False),
+        allow_networks=tuple(raw.get("allow_networks", ())),
+        follow_redirects=raw.get("follow_redirects", False),
+        max_redirects=raw.get("max_redirects", 3),
+    )
+
+
+def _retry(raw: dict[str, Any]) -> RetryPolicy:
+    data = dict(raw)
+    if "retry_statuses" in data:
+        data["retry_statuses"] = frozenset(data["retry_statuses"])
+    return RetryPolicy(**data)
+
+
+def _shape(raw: dict[str, Any] | None) -> dict[str, FieldSpec] | None:
+    if raw is None:
+        return None
+    shape: dict[str, FieldSpec] = {}
+    for key, value in raw.items():
+        parse_path(value if isinstance(value, str) else value["from"])  # fail early
+        if isinstance(value, str):
+            shape[key] = value
+        else:
+            kwargs: dict[str, Any] = {
+                "source": value["from"],
+                "max_chars": value.get("max_chars"),
+                "max_items": value.get("max_items"),
+                "transform": value.get("transform"),
+                "shape": _shape(value.get("shape")),
+            }
+            if "default" in value:
+                kwargs["default"] = value["default"]
+            shape[key] = Field(**kwargs)
+    return shape
+
+
+def _output(raw: dict[str, Any] | None) -> Output:
+    if not raw:
+        return Output()
+    data = dict(raw)
+    parse_path(raw.get("root", "$"))
+    data["shape"] = _shape(raw.get("shape"))
+    data["error_shape"] = _shape(raw.get("error_shape"))
+    if "required" in data:
+        data["required"] = tuple(data["required"])
+    if "redact_keys" in data:
+        data["redact_keys"] = tuple(data["redact_keys"])
+    return Output(**data)
+
+
+def _operation(raw: dict[str, Any], service: Service) -> Operation:
+    rate = raw.get("rate_limit")
+    pagination = raw.get("pagination")
+    return Operation(
+        name=raw["name"],
+        description=raw["description"],
+        service=service,
+        method=raw.get("method", "GET"),
+        path=raw["path"],
+        effect=cast(Effect, raw.get("effect")),  # None: inferred from the method
+        input=raw.get("input", {}),
+        bind=raw.get("bind", {}),
+        required_scopes=tuple(raw.get("required_scopes", ())),
+        tags=tuple(raw.get("tags", ())),
+        group=raw.get("group"),
+        requires_approval=raw.get("requires_approval"),
+        rate_limit=RateLimit(**rate) if rate else None,
+        query=tuple(raw["query"]) if "query" in raw else None,
+        body=tuple(raw["body"]) if "body" in raw else None,
+        body_format=raw.get("body_format", "json"),
+        headers=raw.get("headers", {}),
+        idempotency_header=raw.get("idempotency_header"),
+        pagination=Pagination(**pagination) if pagination else None,
+        output=_output(raw.get("output")),
+    )

toolbound/loaders/meta_schema.py

@@ -0,0 +1,162 @@
+"""JSON Schema for declarative contract files (format version 1)."""
+
+from __future__ import annotations
+
+from typing import Any
+
+_FIELD: dict[str, Any] = {
+    "oneOf": [
+        {"type": "string"},
+        {
+            "type": "object",
+            "additionalProperties": False,
+            "required": ["from"],
+            "properties": {
+                "from": {"type": "string"},
+                "max_chars": {"type": "integer", "minimum": 1},
+                "max_items": {"type": "integer", "minimum": 1},
+                "transform": {"type": "string"},
+                "default": {},
+                "shape": {"$ref": "#/$defs/shape"},
+            },
+        },
+    ]
+}
+
+CONTRACT_SCHEMA: dict[str, Any] = {
+    "$schema": "https://json-schema.org/draft/2020-12/schema",
+    "title": "Toolbound contract file",
+    "type": "object",
+    "additionalProperties": False,
+    "required": ["version", "service", "operations"],
+    "properties": {
+        "version": {"const": 1},
+        "status": {"enum": ["active", "draft"]},
+        "service": {
+            "type": "object",
+            "additionalProperties": False,
+            "required": ["name", "base_url"],
+            "properties": {
+                "name": {"type": "string"},
+                "base_url": {"type": "string"},
+                "auth": {
+                    "type": "object",
+                    "additionalProperties": False,
+                    "required": ["type"],
+                    "properties": {
+                        "type": {"enum": ["none", "bearer", "api_key", "basic"]},
+                        "credential": {"type": "string"},
+                        "header": {"type": "string"},
+                        "query": {"type": "string"},
+                    },
+                },
+                "headers": {"type": "object", "additionalProperties": {"type": "string"}},
+                "network": {
+                    "type": "object",
+                    "additionalProperties": False,
+                    "properties": {
+                        "allow_hosts": {"type": "array", "items": {"type": "string"}},
+                        "allow_ports": {"type": "array", "items": {"type": "integer"}},
+                        "allow_http": {"type": "boolean"},
+                        "allow_networks": {"type": "array", "items": {"type": "string"}},
+                        "follow_redirects": {"type": "boolean"},
+                        "max_redirects": {"type": "integer", "minimum": 0},
+                    },
+                },
+                "limits": {
+                    "type": "object",
+                    "additionalProperties": False,
+                    "properties": {
+                        "timeout": {"type": "number", "exclusiveMinimum": 0},
+                        "connect_timeout": {"type": "number", "exclusiveMinimum": 0},
+                        "max_request_bytes": {"type": "integer", "minimum": 1},
+                        "max_response_bytes": {"type": "integer", "minimum": 1},
+                        "max_concurrency": {"type": "integer", "minimum": 1},
+                        "total_timeout": {"type": "number", "exclusiveMinimum": 0},
+                    },
+                },
+                "retry": {
+                    "type": "object",
+                    "additionalProperties": False,
+                    "properties": {
+                        "max_attempts": {"type": "integer", "minimum": 1},
+                        "backoff_base": {"type": "number", "minimum": 0},
+                        "backoff_max": {"type": "number", "minimum": 0},
+                        "retry_statuses": {"type": "array", "items": {"type": "integer"}},
+                        "max_retry_after": {"type": "number", "minimum": 0},
+                    },
+                },
+            },
+        },
+        "operations": {"type": "array", "minItems": 1, "items": {"$ref": "#/$defs/operation"}},
+    },
+    "$defs": {
+        "shape": {"type": "object", "additionalProperties": _FIELD},
+        "operation": {
+            "type": "object",
+            "additionalProperties": False,
+            "required": ["name", "description", "path"],
+            "properties": {
+                "name": {"type": "string"},
+                "description": {"type": "string"},
+                "method": {"type": "string"},
+                "path": {"type": "string"},
+                "effect": {"enum": ["read", "write", "destructive"]},
+                "input": {"type": "object"},
+                "bind": {"type": "object"},
+                "required_scopes": {"type": "array", "items": {"type": "string"}},
+                "tags": {"type": "array", "items": {"type": "string"}},
+                "group": {"type": "string"},
+                "requires_approval": {"type": "boolean"},
+                "rate_limit": {
+                    "type": "object",
+                    "additionalProperties": False,
+                    "required": ["calls", "per_seconds"],
+                    "properties": {
+                        "calls": {"type": "integer", "minimum": 1},
+                        "per_seconds": {"type": "number", "exclusiveMinimum": 0},
+                        "scope": {"enum": ["actor", "tenant", "global"]},
+                    },
+                },
+                "query": {"type": "array", "items": {"type": "string"}},
+                "body": {"type": "array", "items": {"type": "string"}},
+                "body_format": {"enum": ["json", "form"]},
+                "headers": {"type": "object", "additionalProperties": {"type": "string"}},
+                "idempotency_header": {"type": "string"},
+                "pagination": {
+                    "type": "object",
+                    "additionalProperties": False,
+                    "required": ["type"],
+                    "properties": {
+                        "type": {"enum": ["cursor", "link", "page"]},
+                        "param": {"type": "string"},
+                        "next_path": {"type": "string"},
+                        "start": {"type": "integer"},
+                        "page_size": {"type": "integer", "minimum": 1},
+                    },
+                },
+                "output": {
+                    "type": "object",
+                    "additionalProperties": False,
+                    "properties": {
+                        "root": {"type": "string"},
+                        "shape": {"$ref": "#/$defs/shape"},
+                        "required": {"type": "array", "items": {"type": "string"}},
+                        "max_items": {"type": "integer", "minimum": 1},
+                        "max_chars": {"type": "integer", "minimum": 1},
+                        "max_tokens": {"type": ["integer", "null"], "minimum": 16},
+                        "overflow": {"enum": ["truncate", "reject"]},
+                        "selectable": {"type": "boolean"},
+                        "redact_keys": {"type": "array", "items": {"type": "string"}},
+                        "redact_secrets": {"type": "boolean"},
+                        "schema": {"type": "object"},
+                        "schema_mode": {"enum": ["strict", "warn"]},
+                        "error_shape": {"$ref": "#/$defs/shape"},
+                        "accept_text": {"type": "boolean"},
+                        "continuation": {"type": "boolean"},
+                    },
+                },
+            },
+        },
+    },
+}