openagent-control 0.1.0__py3-none-any.whl

openagent_control/adapters/policy/opa.py

@@ -0,0 +1,46 @@
+"""OPA/Rego policy engine adapter. See docs/adr/0002-opa-rego-as-the-v1-policy-engine.md."""
+
+from __future__ import annotations
+
+import httpx
+
+from openagent_control.domain.errors import PolicyEngineUnavailableError
+from openagent_control.domain.models import Decision, PolicyDecision, ToolCallRequest
+
+
+class OPAPolicyEngine:
+    """Evaluates a tool call against OPA's HTTP API."""
+
+    def __init__(self, opa_url: str, client: httpx.AsyncClient | None = None) -> None:
+        self._opa_url = opa_url
+        self._client = client or httpx.AsyncClient(timeout=5.0)
+
+    async def evaluate(self, request: ToolCallRequest) -> PolicyDecision:
+        opa_input = {
+            "input": {
+                "method": request.method,
+                "spiffe_id": request.agent.spiffe_id,
+                # Registry facts (ADR-0008): policy logic evaluates against these
+                # instead of data hardcoded in the Rego file.
+                "agent": (
+                    request.registration.model_dump(mode="json") if request.registration else None
+                ),
+                "params": {
+                    "name": request.tool_name,
+                    "arguments": request.arguments,
+                },
+            }
+        }
+        try:
+            response = await self._client.post(self._opa_url, json=opa_input)
+            response.raise_for_status()
+        except httpx.HTTPError as exc:
+            raise PolicyEngineUnavailableError(str(exc)) from exc
+        result = response.json().get("result", {})
+
+        decision = Decision.ALLOW if result.get("allow", False) else Decision.DENY
+        reason = result.get("reason", "" if decision is Decision.ALLOW else "Denied by policy")
+        return PolicyDecision(decision=decision, reason=reason)
+
+    async def aclose(self) -> None:
+        await self._client.aclose()

openagent_control/adapters/registry/caching.py

@@ -0,0 +1,34 @@
+"""Redis-caching decorator over any AgentRegistry. See docs/adr/0009.
+
+Wraps — does not replace — the file or Postgres registry, so caching composes
+with either backend unchanged. Short TTL (default 30s) bounds how stale a status
+read can be; there is no invalidate-on-write yet, which is a deliberate,
+documented trade-off (see the ADR) until an admin/kill-switch surface exists.
+"""
+
+from __future__ import annotations
+
+from redis.asyncio import Redis
+
+from openagent_control.domain.models import RegisteredAgent
+from openagent_control.domain.ports import AgentRegistry
+
+_KEY_PREFIX = "oac:registry:"
+
+
+class CachingAgentRegistry:
+    def __init__(self, inner: AgentRegistry, redis: Redis, ttl_seconds: int) -> None:
+        self._inner = inner
+        self._redis = redis
+        self._ttl_seconds = ttl_seconds
+
+    async def lookup(self, spiffe_id: str) -> RegisteredAgent | None:
+        key = _KEY_PREFIX + spiffe_id
+        cached = await self._redis.get(key)
+        if cached is not None:
+            # Empty string is the cached-negative marker (agent confirmed absent).
+            return RegisteredAgent.model_validate_json(cached) if cached else None
+
+        agent = await self._inner.lookup(spiffe_id)
+        await self._redis.set(key, agent.model_dump_json() if agent else "", ex=self._ttl_seconds)
+        return agent

openagent_control/adapters/registry/file.py

@@ -0,0 +1,38 @@
+"""File-backed agent registry. See docs/adr/0008-agent-registry-as-source-of-truth.md.
+
+Config-as-code: agents are registered by adding a record to a YAML file reviewed
+through normal git workflow. A database- or IdP-backed registry is a later
+adapter behind the same `AgentRegistry` port.
+
+The parsed file is cached but re-read when its mtime changes. Caching it for the
+process lifetime would mean suspending an agent had no effect until the gateway
+restarted — revocation is a control-plane guarantee (ADR-0008), so it cannot
+depend on a redeploy. The cost is one `stat` per lookup.
+"""
+
+from __future__ import annotations
+
+from pathlib import Path
+
+import yaml
+
+from openagent_control.domain.models import RegisteredAgent
+
+
+class FileAgentRegistry:
+    def __init__(self, path: str | Path) -> None:
+        self._path = Path(path)
+        self._agents: dict[str, RegisteredAgent] | None = None
+        self._loaded_mtime: float | None = None
+
+    def _load(self) -> dict[str, RegisteredAgent]:
+        mtime = self._path.stat().st_mtime
+        if self._agents is None or mtime != self._loaded_mtime:
+            raw = yaml.safe_load(self._path.read_text()) or {}
+            agents = [RegisteredAgent.model_validate(entry) for entry in raw.get("agents", [])]
+            self._agents = {agent.spiffe_id: agent for agent in agents}
+            self._loaded_mtime = mtime
+        return self._agents
+
+    async def lookup(self, spiffe_id: str) -> RegisteredAgent | None:
+        return self._load().get(spiffe_id)

openagent_control/adapters/registry/postgres.py

@@ -0,0 +1,42 @@
+"""Postgres-backed Agent Registry. See docs/adr/0008 and docs/adr/0009.
+
+The production registry adapter: agent facts live as queryable rows in `oac.agents`
+rather than a git-reviewed file, so status is inventoried and (with the caching
+layer's short TTL, see adapters/registry/caching.py) close to instantly readable
+after a change — the capability an admin kill-switch feature would build on.
+"""
+
+from __future__ import annotations
+
+from sqlalchemy import select
+from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker
+
+from openagent_control.adapters.db.tables import AgentRow
+from openagent_control.domain.models import AgentStatus, RegisteredAgent, RiskTier
+
+
+def _to_domain(row: AgentRow) -> RegisteredAgent:
+    return RegisteredAgent(
+        spiffe_id=row.spiffe_id,
+        display_name=row.display_name,
+        purpose=row.purpose,
+        owner=row.owner,
+        risk_tier=RiskTier(row.risk_tier),
+        status=AgentStatus(row.status),
+        granted_tools=list(row.granted_tools),
+        created_at=row.created_at,
+        updated_at=row.updated_at,
+        status_changed_at=row.status_changed_at,
+    )
+
+
+class PostgresAgentRegistry:
+    def __init__(self, session_factory: async_sessionmaker[AsyncSession]) -> None:
+        self._session_factory = session_factory
+
+    async def lookup(self, spiffe_id: str) -> RegisteredAgent | None:
+        async with self._session_factory() as session:
+            row = (
+                await session.execute(select(AgentRow).where(AgentRow.spiffe_id == spiffe_id))
+            ).scalar_one_or_none()
+            return _to_domain(row) if row is not None else None

openagent_control/adapters/token_exchange/caching.py

@@ -0,0 +1,73 @@
+"""Redis-caching decorator over any TokenExchange. See docs/adr/0009.
+
+Cache TTL is derived from the exchanged token's own `exp` claim (peeked via an
+unverified JWT decode — we already trust this token because it came straight
+from the IdP's response over TLS; we are not re-validating its signature, only
+reading a cache-lifetime hint), minus a safety margin, capped at a configured
+maximum. Opaque (non-JWT) tokens fall back to the safety-margin value as a short,
+conservative TTL. Never caches past a token's real expiry.
+"""
+
+from __future__ import annotations
+
+import hashlib
+import time
+
+import jwt
+from redis.asyncio import Redis
+
+from openagent_control.domain.ports import TokenExchange
+
+_KEY_PREFIX = "oac:token-exchange:"
+
+
+def _cache_key(subject_token: str, audience: str) -> str:
+    digest = hashlib.sha256(f"{subject_token}:{audience}".encode()).hexdigest()
+    return _KEY_PREFIX + digest
+
+
+def _ttl_from_token(token: str, max_ttl: int, safety_margin: int) -> int:
+    try:
+        claims = jwt.decode(token, options={"verify_signature": False})
+        exp = int(claims["exp"])
+    except (jwt.InvalidTokenError, KeyError, TypeError, ValueError):
+        return safety_margin
+
+    remaining = exp - int(time.time()) - safety_margin
+    return max(0, min(remaining, max_ttl))
+
+
+class CachingTokenExchange:
+    def __init__(
+        self,
+        inner: TokenExchange,
+        redis: Redis,
+        max_ttl_seconds: int,
+        safety_margin_seconds: int,
+    ) -> None:
+        self._inner = inner
+        self._redis = redis
+        self._max_ttl_seconds = max_ttl_seconds
+        self._safety_margin_seconds = safety_margin_seconds
+
+    async def exchange(self, subject_token: str, audience: str) -> str:
+        key = _cache_key(subject_token, audience)
+        cached = await self._redis.get(key)
+        if cached is not None:
+            # A client without decode_responses=True hands back bytes; str() on
+            # bytes would corrupt the credential into "b'...'", so decode
+            # explicitly instead of coercing.
+            return cached.decode() if isinstance(cached, bytes) else str(cached)
+
+        token = await self._inner.exchange(subject_token, audience)
+        ttl = _ttl_from_token(token, self._max_ttl_seconds, self._safety_margin_seconds)
+        if ttl > 0:
+            await self._redis.set(key, token, ex=ttl)
+        return token
+
+    async def aclose(self) -> None:
+        """Forwards to the wrapped adapter; the decorator owns no resources of
+        its own (the Redis client is shared and closed once by the Container)."""
+        closer = getattr(self._inner, "aclose", None)
+        if closer is not None:
+            await closer()

openagent_control/adapters/token_exchange/entra_obo.py

@@ -0,0 +1,52 @@
+"""Microsoft Entra ID On-Behalf-Of token exchange adapter.
+
+Entra's OBO flow predates RFC 8693 and uses the jwt-bearer grant with
+`requested_token_use=on_behalf_of`. The `audience` argument of the port maps to
+Entra's `scope` parameter (callers typically pass `api://<app-id>/.default`).
+See ADR-0004.
+"""
+
+from __future__ import annotations
+
+import httpx
+
+from openagent_control.domain.errors import TokenExchangeError
+
+_GRANT_TYPE = "urn:ietf:params:oauth:grant-type:jwt-bearer"
+
+
+class EntraOnBehalfOfTokenExchange:
+    def __init__(
+        self,
+        token_url: str,
+        client_id: str,
+        client_secret: str,
+        client: httpx.AsyncClient | None = None,
+    ) -> None:
+        self._token_url = token_url
+        self._client_id = client_id
+        self._client_secret = client_secret
+        self._client = client or httpx.AsyncClient(timeout=10.0)
+
+    async def exchange(self, subject_token: str, audience: str) -> str:
+        data = {
+            "grant_type": _GRANT_TYPE,
+            "client_id": self._client_id,
+            "client_secret": self._client_secret,
+            "assertion": subject_token,
+            "scope": audience,
+            "requested_token_use": "on_behalf_of",
+        }
+        try:
+            response = await self._client.post(self._token_url, data=data)
+            response.raise_for_status()
+        except httpx.HTTPError as exc:
+            raise TokenExchangeError(str(exc)) from exc
+
+        access_token = response.json().get("access_token")
+        if not access_token:
+            raise TokenExchangeError("Entra returned no access_token")
+        return str(access_token)
+
+    async def aclose(self) -> None:
+        await self._client.aclose()

openagent_control/adapters/token_exchange/rfc8693.py

@@ -0,0 +1,50 @@
+"""RFC 8693 OAuth 2.0 Token Exchange adapter (Okta-compatible).
+
+Swaps the human sponsor's subject token for a short-lived, audience-scoped access
+token. Works against any authorization server implementing RFC 8693 (Okta org
+authorization servers, Keycloak, etc.). See ADR-0004.
+"""
+
+from __future__ import annotations
+
+import httpx
+
+from openagent_control.domain.errors import TokenExchangeError
+
+_GRANT_TYPE = "urn:ietf:params:oauth:grant-type:token-exchange"
+_ACCESS_TOKEN_TYPE = "urn:ietf:params:oauth:token-type:access_token"
+
+
+class Rfc8693TokenExchange:
+    def __init__(
+        self,
+        token_url: str,
+        client_id: str,
+        client_secret: str,
+        client: httpx.AsyncClient | None = None,
+    ) -> None:
+        self._token_url = token_url
+        self._auth = (client_id, client_secret)
+        self._client = client or httpx.AsyncClient(timeout=10.0)
+
+    async def exchange(self, subject_token: str, audience: str) -> str:
+        data = {
+            "grant_type": _GRANT_TYPE,
+            "subject_token": subject_token,
+            "subject_token_type": _ACCESS_TOKEN_TYPE,
+            "requested_token_type": _ACCESS_TOKEN_TYPE,
+            "audience": audience,
+        }
+        try:
+            response = await self._client.post(self._token_url, data=data, auth=self._auth)
+            response.raise_for_status()
+        except httpx.HTTPError as exc:
+            raise TokenExchangeError(str(exc)) from exc
+
+        access_token = response.json().get("access_token")
+        if not access_token:
+            raise TokenExchangeError("authorization server returned no access_token")
+        return str(access_token)
+
+    async def aclose(self) -> None:
+        await self._client.aclose()

openagent_control/adapters/token_exchange/stub.py

@@ -0,0 +1,12 @@
+"""Stub RFC 8693 token exchange. See docs/adr/0004-mcp-as-the-v1-protocol-surface.md.
+
+Real adapters (Okta, Microsoft Entra ID) implement the same `TokenExchange` port —
+see docs/adr/0006-hexagonal-architecture-for-the-control-plane.md.
+"""
+
+from __future__ import annotations
+
+
+class StubTokenExchange:
+    async def exchange(self, subject_token: str, audience: str) -> str:
+        return f"stub-obo-token::aud={audience}::subj={subject_token[:8]}"

openagent_control/alembic.ini

@@ -0,0 +1,36 @@
+[alembic]
+script_location = migrations
+prepend_sys_path = .
+path_separator = os
+
+[loggers]
+keys = root,sqlalchemy,alembic
+
+[handlers]
+keys = console
+
+[formatters]
+keys = generic
+
+[logger_root]
+level = WARNING
+handlers = console
+
+[logger_sqlalchemy]
+level = WARNING
+handlers =
+qualname = sqlalchemy.engine
+
+[logger_alembic]
+level = INFO
+handlers =
+qualname = alembic
+
+[handler_console]
+class = StreamHandler
+args = (sys.stderr,)
+level = NOTSET
+formatter = generic
+
+[formatter_generic]
+format = %(levelname)-5.5s [%(name)s] %(message)s

openagent_control/application/governed_execution.py

@@ -0,0 +1,177 @@
+"""The governed-execution use case: the transport-agnostic core of the gateway.
+
+Every deployment pattern from ADR-0001 (egress gateway today; sidecar and native SDK
+later) drives this same service rather than re-implementing the sequence:
+
+    identify -> evaluate policy -> record + export receipt
+             -> (deny: semantic error) | (allow: broker credential, forward upstream)
+
+Failure posture:
+- Policy engine unreachable  => fail closed: an explicit DENY is recorded and
+  exported, so the outage itself leaves an audit trail (never a silent 500).
+- Upstream failure after ALLOW => JSON-RPC error with a stop instruction, so an LLM
+  agent halts instead of retry-looping; the ALLOW receipt has already been recorded.
+- Identity / missing subject token => typed errors for the transport layer to map
+  (HTTP 401 in the FastAPI adapter).
+"""
+
+from __future__ import annotations
+
+from typing import Any
+
+from openagent_control.domain.errors import (
+    MissingSubjectTokenError,
+    PolicyEngineUnavailableError,
+    TokenExchangeError,
+    UpstreamError,
+)
+from openagent_control.domain.models import (
+    AgentIdentity,
+    AgentStatus,
+    Decision,
+    PolicyDecision,
+    ToolCallRequest,
+)
+from openagent_control.domain.ports import (
+    AgentRegistry,
+    AuditExporter,
+    IdentityProvider,
+    Ledger,
+    MCPUpstream,
+    PolicyEngine,
+    TokenExchange,
+)
+
+_SUBJECT_TOKEN_HEADER = "x-subject-token"
+
+_STOP_INSTRUCTION = "Stop execution and request user approval."
+_FAIL_CLOSED_REASON = "Policy engine unavailable; denied (fail-closed)"
+
+
+class GovernedExecutionService:
+    def __init__(
+        self,
+        *,
+        identity_provider: IdentityProvider,
+        agent_registry: AgentRegistry,
+        policy_engine: PolicyEngine,
+        ledger: Ledger,
+        audit_exporter: AuditExporter,
+        token_exchange: TokenExchange,
+        mcp_upstream: MCPUpstream,
+        delegated_audience: str,
+    ) -> None:
+        self._identity_provider = identity_provider
+        self._agent_registry = agent_registry
+        self._policy_engine = policy_engine
+        self._ledger = ledger
+        self._audit_exporter = audit_exporter
+        self._token_exchange = token_exchange
+        self._mcp_upstream = mcp_upstream
+        self._delegated_audience = delegated_audience
+
+    async def execute(self, headers: dict[str, str], payload: dict[str, Any]) -> dict[str, Any]:
+        """Runs one governed tool call; returns a JSON-RPC response object.
+
+        Raises IdentityError or MissingSubjectTokenError for the transport layer
+        to map to its own auth failure shape.
+        """
+        agent = await self._identity_provider.identify(headers)
+        registration = await self._agent_registry.lookup(agent.spiffe_id)
+
+        params = payload.get("params") or {}
+        call = ToolCallRequest(
+            method=payload.get("method", ""),
+            tool_name=params.get("name"),
+            arguments=params.get("arguments", {}),
+            agent=agent,
+            registration=registration,
+            request_id=payload.get("id"),
+        )
+
+        # Registry gate (ADR-0008): orphaned or suspended agents never reach the
+        # policy engine, but the attempt is still receipted below.
+        if registration is None:
+            decision = PolicyDecision(
+                decision=Decision.DENY,
+                reason="Agent not registered in the Agent Registry (orphaned agents are refused)",
+            )
+        elif registration.status is not AgentStatus.ACTIVE:
+            decision = PolicyDecision(
+                decision=Decision.DENY,
+                reason=f"Agent is {registration.status.value} in the Agent Registry",
+            )
+        else:
+            try:
+                decision = await self._policy_engine.evaluate(call)
+            except PolicyEngineUnavailableError:
+                decision = PolicyDecision(decision=Decision.DENY, reason=_FAIL_CLOSED_REASON)
+
+        receipt = await self._ledger.record(agent, call, decision)
+        await self._audit_exporter.export(receipt)
+
+        if decision.decision is not Decision.ALLOW:
+            return _jsonrpc_error(
+                call.request_id,
+                code=-32000,
+                message=f"Policy violation: {decision.reason}",
+                instruction=_STOP_INSTRUCTION,
+            )
+
+        try:
+            credential = await self._broker_credential(agent, headers)
+        except TokenExchangeError as exc:
+            return _jsonrpc_error(
+                call.request_id,
+                code=-32004,
+                message=f"Credential brokering failed: {exc}",
+                instruction="Do not retry automatically; report the failure to the user.",
+            )
+
+        try:
+            return await self._mcp_upstream.forward(call, credential)
+        except UpstreamError as exc:
+            return _jsonrpc_error(
+                call.request_id,
+                code=-32002,
+                message=f"Upstream execution failed: {exc}",
+                instruction="Do not retry automatically; report the failure to the user.",
+            )
+
+    async def _broker_credential(self, agent: AgentIdentity, headers: dict[str, str]) -> str:
+        """Obtains the short-lived, audience-scoped credential for the upstream.
+
+        Delegated calls exchange the human sponsor's subject token. Autonomous
+        calls exchange the agent's own access token — RFC 8693 permits either,
+        and a real resource server validates the result's audience and scope, so
+        the agent must never be handed a credential it could reuse elsewhere.
+
+        The placeholder below is reachable only when no bearer token exists at
+        all, which means identity_mode="header" — the dev stub of ADR-0005,
+        where there is no real credential to broker in the first place.
+        """
+        normalized = {k.lower(): v for k, v in headers.items()}
+
+        if agent.human_sponsor:
+            subject_token = normalized.get(_SUBJECT_TOKEN_HEADER)
+            if not subject_token:
+                raise MissingSubjectTokenError(
+                    f"delegated call for sponsor '{agent.human_sponsor}' "
+                    f"requires a '{_SUBJECT_TOKEN_HEADER}' header"
+                )
+            return await self._token_exchange.exchange(subject_token, self._delegated_audience)
+
+        scheme, _, agent_token = normalized.get("authorization", "").partition(" ")
+        if scheme.lower() != "bearer" or not agent_token:
+            return f"autonomous::{agent.spiffe_id}"
+        return await self._token_exchange.exchange(agent_token, self._delegated_audience)
+
+
+def _jsonrpc_error(
+    request_id: str | int | None, *, code: int, message: str, instruction: str
+) -> dict[str, Any]:
+    return {
+        "jsonrpc": "2.0",
+        "id": request_id,
+        "error": {"code": code, "message": message, "data": {"instruction": instruction}},
+    }