badgerflow 0.1.0__py3-none-any.whl

agiel/__init__.py

@@ -0,0 +1,57 @@
+"""
+``agiel`` — deprecated import path for :mod:`badgerflow.client`.
+
+The ``agiel`` distribution moved into the ``badgerflow`` package as
+``badgerflow.client``. This shim keeps ``from agiel import AgiellClient`` (and
+the resource classes and exceptions) working for one release cycle so existing
+scripts, generated exports and docs do not break at import time.
+
+It warns once on import. New code should import from ``badgerflow.client``.
+"""
+
+import warnings
+
+from badgerflow import __version__
+from badgerflow.client import (
+    AgiellClient,
+    AgiellError,
+    AuthError,
+    ConflictError,
+    NotFoundError,
+    ServerError,
+    ValidationError,
+)
+from badgerflow.client.resources.agents import AgentClient
+from badgerflow.client.resources.auth import AuthClient
+from badgerflow.client.resources.guardrail_profiles import GuardrailProfilesClient
+from badgerflow.client.resources.guardrails import GuardrailsClient
+from badgerflow.client.resources.models import ModelRegistryClient
+from badgerflow.client.resources.namespaces import NamespaceClient
+from badgerflow.client.resources.rag import RAGClient
+from badgerflow.client.resources.webhooks import WebhookClient
+
+warnings.warn(
+    "the 'agiel' package is deprecated; import from 'badgerflow.client' instead "
+    "(e.g. 'from badgerflow.client import AgiellClient')",
+    DeprecationWarning,
+    stacklevel=2,
+)
+
+__all__ = [
+    "__version__",
+    "AgiellClient",
+    "AgiellError",
+    "AuthError",
+    "ConflictError",
+    "NotFoundError",
+    "ServerError",
+    "ValidationError",
+    "AgentClient",
+    "AuthClient",
+    "GuardrailProfilesClient",
+    "GuardrailsClient",
+    "ModelRegistryClient",
+    "NamespaceClient",
+    "RAGClient",
+    "WebhookClient",
+]

badgerflow/__init__.py

@@ -0,0 +1,38 @@
+"""
+badgerflow — the BadgerFlow Python SDK
+======================================
+
+One installable package with two halves:
+
+* ``badgerflow.client`` — the async REST client (``AgiellClient``) and the
+  ``agiel`` CLI that used to ship as the ``agiel`` distribution. The old
+  ``agiel`` import path still works through a deprecated shim.
+* ``badgerflow.procode`` and friends — the pro-code surface (a governed
+  ``@app.run`` server, ``Context``, the LangGraph bridge, the ``bf`` CLI).
+  These packages are skeletons filled in by later work packages; each one's
+  docstring names the WP that owns it.
+
+The package never imports the platform's ``libs/`` or ``services/``. The
+constants both sides must agree on are vendored under ``badgerflow.contract``
+and pinned to ``CONTRACT_VERSION``.
+
+This module stays import-light on purpose: ``import badgerflow`` must not
+pull in httpx, typer or any optional framework, so the version is the only
+thing defined eagerly. ``badgerflow.current()`` — the running invocation's
+``Context`` — is resolved lazily on first access so that the server package
+(fastapi, prometheus_client) loads only for code that actually calls it.
+"""
+
+from typing import Any
+
+__version__ = "0.1.0"
+
+__all__ = ["__version__", "current"]
+
+
+def __getattr__(name: str) -> Any:
+    if name == "current":
+        from .procode.current import current
+
+        return current
+    raise AttributeError(f"module {__name__!r} has no attribute {name!r}")

badgerflow/auth/__init__.py

@@ -0,0 +1,26 @@
+"""badgerflow.auth — run-token verification (JWKS) for the ``/invocations`` bearer.
+
+Vendors the platform's Ed25519/JWKS verification (``run_token.py``); never
+imports the platform's ``libs/``. Actor/exec token *forwarding* lives with the
+governed clients (S0.4) — the server here only carries them into ``Context``.
+"""
+
+from .run_token import (
+    RunTokenClaims,
+    RunTokenError,
+    RunTokenVerifier,
+    Verifier,
+    auth_disabled_from_env,
+    keys_from_jwks,
+    verifier_from_env,
+)
+
+__all__ = [
+    "RunTokenClaims",
+    "RunTokenError",
+    "RunTokenVerifier",
+    "Verifier",
+    "auth_disabled_from_env",
+    "keys_from_jwks",
+    "verifier_from_env",
+]

badgerflow/auth/run_token.py

@@ -0,0 +1,333 @@
+"""Verify the run token that arrives as the bearer of every ``/invocations``.
+
+The platform mints run tokens through ``libs/common/session_jwt.py`` (Ed25519,
+``alg=EdDSA``, a ``kid`` header) and publishes the public keys as a JWKS at
+``{BADGERFLOW_JWKS_URL}`` — ``http://auth-service:8090/v1/auth/jwks`` inside a
+cluster. This module *vendors* that verification: it never imports the
+platform, it reads the same JWKS any other verifier reads, and it checks the
+four facts the contract fixes (``constants.py``, "Run token"):
+
+* signature, against the key the ``kid`` names;
+* ``iss == RUN_TOKEN_ISSUER`` and ``tok == RUN_TOKEN_TYPE`` — a session or
+  actor token signed by the same key must not verify as a run token;
+* ``aud`` contains ``procode:<workflow_id>`` — a run token for a *different*
+  pro-code workflow, or a Tier-2 token whose audience is the platform's own
+  services, is refused;
+* ``exp`` in the future (60 s leeway, the platform's own tolerance is 30 s).
+
+Two more checks are the caller's: ``jti == body.run_id`` is done by the server
+because it needs the body, and ``wf``, the workflow claim the platform stamps on
+pro-code tokens, is checked *when present* — it was introduced on a separate
+platform branch and a token without it is still a valid token for the audience
+it names.
+
+Key handling copies ``libs/common/signing.py``: keys are looked up by ``kid``,
+cached for ``cache_ttl`` seconds, and a ``kid`` the cache does not know
+triggers exactly one refetch before the token is refused — that is what makes
+a key rotation invisible to a running service without letting a forged ``kid``
+turn every request into a JWKS fetch.
+"""
+
+from __future__ import annotations
+
+import asyncio
+import base64
+import logging
+import os
+import time
+from dataclasses import dataclass, field
+from datetime import datetime, timezone
+from typing import Any, Mapping, Protocol
+
+import httpx
+import jwt
+
+from badgerflow.contract.constants import RUN_TOKEN_ISSUER, RUN_TOKEN_TYPE, procode_audience
+
+log = logging.getLogger("badgerflow.auth")
+
+#: Leeway on ``exp``. Generous next to the platform's 30 s because a customer's
+#: cluster and the platform's are two clocks that nobody syncs on purpose.
+EXP_LEEWAY_SECONDS = 60
+
+#: Environment the server reads when the caller passes nothing explicit.
+ENV_JWKS_URL = "BADGERFLOW_JWKS_URL"
+ENV_WORKFLOW_ID = "BADGERFLOW_WORKFLOW_ID"
+ENV_AUTH_DISABLED = "BADGERFLOW_AUTH_DISABLED"
+ENV_ENV = "BADGERFLOW_ENV"
+
+
+class RunTokenError(Exception):
+    """The bearer is not a run token this service accepts. The message is safe
+    to log; the server answers 401 ``unauthorized`` without echoing it."""
+
+
+@dataclass(frozen=True)
+class RunTokenClaims:
+    """What a verified run token asserts. ``run_id`` is the ``jti`` — the binding
+    the server compares to ``body.run_id`` — and ``workflow_id`` is the id the
+    verifier was configured with (equal to the token's ``wf`` when it has one)."""
+
+    run_id: str
+    namespace: str
+    trust_tier: str
+    workflow_id: str
+    expires_at: datetime
+    subject: str = ""
+    audience: tuple[str, ...] = ()
+    claims: Mapping[str, Any] = field(default_factory=dict, repr=False)
+
+
+class Verifier(Protocol):
+    """What the server needs from a verifier; :class:`RunTokenVerifier` is the
+    real one, tests inject a stand-in with the same single method."""
+
+    async def verify(self, token: str) -> RunTokenClaims: ...
+
+
+def _b64url_decode(data: str) -> bytes:
+    return base64.urlsafe_b64decode(data + "=" * (-len(data) % 4))
+
+
+def _public_key_from_jwk(jwk: Mapping[str, Any]) -> Any:
+    """An ``Ed25519PublicKey`` from an OKP JWK; ``None`` for any other key type.
+
+    The platform's JWKS carries ``{"kty": "OKP", "crv": "Ed25519", "kid", "x"}``
+    (``libs/common/signing.get_public_key_jwk``). Keys of another type are
+    skipped rather than refused so a JWKS that one day also lists an RSA key
+    keeps working for the Ed25519 tokens this verifier is for.
+    """
+    if jwk.get("kty") != "OKP" or jwk.get("crv") != "Ed25519" or not jwk.get("x"):
+        return None
+    from cryptography.hazmat.primitives.asymmetric.ed25519 import Ed25519PublicKey
+
+    try:
+        return Ed25519PublicKey.from_public_bytes(_b64url_decode(str(jwk["x"])))
+    except (ValueError, TypeError):
+        return None
+
+
+def keys_from_jwks(jwks: Mapping[str, Any]) -> dict[str, list[Any]]:
+    """``kid -> [public keys]`` for every Ed25519 key in a JWKS document.
+
+    A list per ``kid`` because the platform may publish a rotation predecessor
+    under the same id for the length of a rollout (``_public_keys_for_kid``);
+    verification tries each.
+    """
+    out: dict[str, list[Any]] = {}
+    for entry in jwks.get("keys") or ():
+        if not isinstance(entry, Mapping):
+            continue
+        key = _public_key_from_jwk(entry)
+        if key is None:
+            continue
+        out.setdefault(str(entry.get("kid", "")), []).append(key)
+    return out
+
+
+class RunTokenVerifier:
+    """Verify run tokens for one pro-code workflow against the platform's JWKS.
+
+    ``jwks`` injects a static JWKS document (unit tests, air-gapped dev); with
+    it no HTTP fetch happens and a ``kid`` miss is simply a miss. Otherwise the
+    document is fetched from ``jwks_url`` on first use, re-fetched after
+    ``cache_ttl`` seconds, and re-fetched once, immediately, on a ``kid`` the
+    cache does not hold. Concurrent misses share one fetch.
+    """
+
+    def __init__(
+        self,
+        jwks_url: str | None,
+        workflow_id: str,
+        *,
+        cache_ttl: float = 300.0,
+        jwks: Mapping[str, Any] | None = None,
+        http_client: httpx.AsyncClient | None = None,
+        timeout: float = 5.0,
+    ):
+        if not workflow_id:
+            raise ValueError("RunTokenVerifier needs the workflow id the audience is checked against")
+        if jwks is None and not jwks_url:
+            raise ValueError("RunTokenVerifier needs a jwks_url (or an injected jwks document)")
+        self.jwks_url = jwks_url
+        self.workflow_id = workflow_id
+        self.audience = procode_audience(workflow_id)
+        self.cache_ttl = cache_ttl
+        self._static = jwks is not None
+        self._keys: dict[str, list[Any]] = keys_from_jwks(jwks) if jwks is not None else {}
+        self._fetched_at: float = time.monotonic() if jwks is not None else float("-inf")
+        self._client = http_client
+        self._timeout = timeout
+        self._lock = asyncio.Lock()
+        self.fetch_count = 0
+
+    # ── JWKS ─────────────────────────────────────────────────────────────────
+
+    async def _fetch(self) -> None:
+        assert self.jwks_url
+        client = self._client
+        try:
+            if client is None:
+                async with httpx.AsyncClient(timeout=self._timeout) as own:
+                    resp = await own.get(self.jwks_url)
+            else:
+                resp = await client.get(self.jwks_url)
+            resp.raise_for_status()
+            doc = resp.json()
+        except (httpx.HTTPError, ValueError) as exc:
+            raise RunTokenError(f"JWKS fetch from {self.jwks_url} failed: {exc}") from exc
+        if not isinstance(doc, dict):
+            raise RunTokenError("JWKS document is not an object")
+        self._keys = keys_from_jwks(doc)
+        self._fetched_at = time.monotonic()
+        self.fetch_count += 1
+        log.debug("badgerflow: loaded %d signing key ids from %s", len(self._keys), self.jwks_url)
+
+    async def _keys_for(self, kid: str) -> list[Any]:
+        """Keys for *kid*, refreshing the cache when it is stale or does not
+        know the id. Refresh happens at most once per call, so a token naming
+        an unknown ``kid`` costs one fetch, not a retry storm."""
+        if self._static:
+            return self._keys.get(kid, [])
+        async with self._lock:
+            stale = time.monotonic() - self._fetched_at > self.cache_ttl
+            if stale or kid not in self._keys:
+                await self._fetch()
+        return self._keys.get(kid, [])
+
+    # ── verification ─────────────────────────────────────────────────────────
+
+    async def verify(self, token: str) -> RunTokenClaims:
+        """Return the claims of a run token valid for this workflow, or raise
+        :class:`RunTokenError`. Every rejection reason is one exception type so
+        the server can map all of them to 401 without inspecting them."""
+        if not token:
+            raise RunTokenError("missing bearer token")
+        try:
+            header = jwt.get_unverified_header(token)
+        except jwt.PyJWTError as exc:
+            raise RunTokenError(f"malformed token: {exc}") from exc
+        if header.get("alg") not in ("EdDSA", "Ed25519"):
+            raise RunTokenError(f"unsupported alg {header.get('alg')!r}")
+        kid = str(header.get("kid") or "")
+        candidates = await self._keys_for(kid)
+        if not candidates:
+            raise RunTokenError(f"no signing key for kid={kid!r}")
+
+        payload: dict[str, Any] | None = None
+        last: Exception | None = None
+        for key in candidates:
+            try:
+                payload = jwt.decode(
+                    token,
+                    key=key,
+                    algorithms=["EdDSA"],
+                    audience=self.audience,
+                    issuer=RUN_TOKEN_ISSUER,
+                    leeway=EXP_LEEWAY_SECONDS,
+                    options={"require": ["exp", "iss", "aud", "jti"]},
+                )
+                break
+            except jwt.InvalidSignatureError as exc:
+                last = exc  # a rotation predecessor under the same kid may still match
+            except jwt.PyJWTError as exc:
+                raise RunTokenError(f"run token rejected: {exc}") from exc
+        if payload is None:
+            raise RunTokenError(f"run token rejected: {last}")
+
+        if payload.get("tok") != RUN_TOKEN_TYPE:
+            raise RunTokenError("not a run token (tok != run)")
+        wf = payload.get("wf")
+        if wf is not None and str(wf) != self.workflow_id:
+            raise RunTokenError(f"run token is for workflow {wf!r}, not {self.workflow_id!r}")
+        run_id = str(payload.get("jti") or "")
+        if not run_id:
+            raise RunTokenError("run token has an empty jti")
+
+        aud = payload.get("aud") or ()
+        return RunTokenClaims(
+            run_id=run_id,
+            namespace=str(payload.get("ns") or ""),
+            trust_tier=str(payload.get("trust_tier") or ""),
+            workflow_id=self.workflow_id,
+            expires_at=datetime.fromtimestamp(int(payload["exp"]), tz=timezone.utc),
+            subject=str(payload.get("sub") or ""),
+            audience=(aud,) if isinstance(aud, str) else tuple(str(a) for a in aud),
+            claims=dict(payload),
+        )
+
+
+def _truthy(value: str | None) -> bool:
+    return (value or "").strip().lower() in ("1", "true", "yes", "on")
+
+
+def auth_disabled_from_env(environ: Mapping[str, str] | None = None) -> bool:
+    """``True`` only when ``BADGERFLOW_AUTH_DISABLED`` is set *and*
+    ``BADGERFLOW_ENV=dev``. Set outside dev it raises rather than being ignored:
+    a deployment that asked for no auth and silently kept it would be as
+    surprising as one that silently dropped it.
+    """
+    env = os.environ if environ is None else environ
+    if not _truthy(env.get(ENV_AUTH_DISABLED)):
+        return False
+    if (env.get(ENV_ENV) or "").strip().lower() != "dev":
+        raise RuntimeError(
+            f"{ENV_AUTH_DISABLED} is only honoured with {ENV_ENV}=dev; "
+            "refusing to start without run-token verification"
+        )
+    log.warning(
+        "badgerflow: %s=true — run tokens are NOT verified; every /invocations caller is "
+        "trusted. This is a dev-only setting.",
+        ENV_AUTH_DISABLED,
+    )
+    return True
+
+
+def verifier_from_env(
+    workflow_id: str | None = None,
+    jwks_url: str | None = None,
+    *,
+    environ: Mapping[str, str] | None = None,
+    cache_ttl: float = 300.0,
+) -> RunTokenVerifier:
+    """A :class:`RunTokenVerifier` from explicit values or ``BADGERFLOW_JWKS_URL``
+    / ``BADGERFLOW_WORKFLOW_ID``. Raises ``RuntimeError`` naming the missing
+    variable — the server calls this at build so a misconfigured deployment
+    fails at startup, not on its first request."""
+    env = os.environ if environ is None else environ
+    wf = workflow_id or env.get(ENV_WORKFLOW_ID) or ""
+    url = jwks_url or env.get(ENV_JWKS_URL) or ""
+    if not wf:
+        raise RuntimeError(f"{ENV_WORKFLOW_ID} is not set (registration sets it)")
+    if not url:
+        raise RuntimeError(f"{ENV_JWKS_URL} is not set (e.g. http://auth-service:8090/v1/auth/jwks)")
+    return RunTokenVerifier(url, wf, cache_ttl=cache_ttl)
+
+
+def unverified_claims(token: str | None) -> dict[str, Any]:
+    """The payload of *token* WITHOUT verification — for the dev escape hatch
+    only, so a locally minted token still populates the context. Never
+    authentication."""
+    if not token:
+        return {}
+    try:
+        return dict(jwt.decode(token, options={"verify_signature": False}))
+    except jwt.PyJWTError:
+        return {}
+
+
+__all__ = [
+    "ENV_AUTH_DISABLED",
+    "ENV_ENV",
+    "ENV_JWKS_URL",
+    "ENV_WORKFLOW_ID",
+    "EXP_LEEWAY_SECONDS",
+    "RunTokenClaims",
+    "RunTokenError",
+    "RunTokenVerifier",
+    "Verifier",
+    "auth_disabled_from_env",
+    "keys_from_jwks",
+    "unverified_claims",
+    "verifier_from_env",
+]

badgerflow/cli/__init__.py

@@ -0,0 +1,5 @@
+"""badgerflow.cli — the ``bf`` CLI: ``init``, ``dev``, ``validate``, ``register``,
+``status``, ``version`` (S0.5). S2.1 adds ``release submit | status``, ``validate
+--against`` and ``compile --check``. ``project.py`` is the ``badgerflow.yaml`` +
+``agent.py`` loader, ``scaffold.py`` the ``bf init`` templates. The legacy ``agiel``
+CLI lives in ``badgerflow.client.cli``."""