agentrust-py 0.0.3__py3-none-any.whl

agentrust_sdk/adapters/mcp.py

@@ -0,0 +1,193 @@
+"""
+MCP (Model Context Protocol) adapter — AgentTrust tool-call governance.
+
+Architecture doc §3: "MCP wrapper" as a key Adapter Layer capability.
+
+Wraps an MCP server so every tool execution is validated by AgentTrust
+before the result is returned to the calling model.
+
+Two usage patterns:
+
+1. Tool guard decorator (wrap individual tools):
+
+    from agentrust_sdk.adapters.mcp import mcp_tool_guard
+
+    @mcp_tool_guard(agent_id="my-mcp-server")
+    async def read_file(path: str) -> str:
+        return open(path).read()
+
+2. MCPServerGuard (wrap a FastMCP / mcp.Server instance):
+
+    from mcp.server.fastmcp import FastMCP
+    from agentrust_sdk.adapters.mcp import MCPServerGuard
+
+    mcp = FastMCP("my-server")
+    guard = MCPServerGuard(mcp, agent_id="my-mcp-server")
+    # All @mcp.tool() handlers are now governed by AgentTrust.
+
+Requires Developer tier (tier check at construction time).
+"""
+from __future__ import annotations
+
+import functools
+import inspect
+import logging
+import time
+from typing import Any, Callable
+
+logger = logging.getLogger(__name__)
+
+
+def _check_tier(api_key: str | None = None) -> None:
+    from agentrust_sdk.auth import resolve_key
+    from agentrust_sdk.tiers import Capability, is_allowed, UPGRADE_MESSAGES
+    info = resolve_key(api_key)
+    cap = Capability.MCP_ADAPTER if hasattr(Capability, "MCP_ADAPTER") else None
+    if cap and not is_allowed(cap, info.tier):
+        msg = UPGRADE_MESSAGES.get(cap, "MCP adapter requires Developer tier or higher.")
+        raise RuntimeError(f"[AgentTrust] {msg}")
+
+
+def mcp_tool_guard(
+    agent_id: str,
+    *,
+    base_url: str = "http://localhost:8000",
+    api_key: str | None = None,
+    user: str = "mcp",
+    block_on_block: bool = True,
+    framework: str = "MCP",
+) -> Callable:
+    """Decorator that validates every MCP tool call via AgentTrust.
+
+    Works with both sync and async tool functions.
+    """
+    def decorator(fn: Callable) -> Callable:
+        is_async = inspect.iscoroutinefunction(fn)
+
+        async def _validate(tool_name: str, args: dict, result: Any, latency_ms: float) -> None:
+            from agentrust_sdk.client import AgentTrustClient
+            from agentrust_sdk.decorator import _check_decision, BlockedError
+
+            output = result if isinstance(result, dict) else {"result": str(result)[:2000]}
+            try:
+                with AgentTrustClient(base_url=base_url, api_key=api_key) as client:
+                    resp = client.validate(
+                        agent_id=agent_id,
+                        user=user,
+                        input=tool_name,
+                        output={**output, "_tool_args": args},
+                        framework=framework,
+                        tools_called=[{"name": tool_name, "args": args, "result": result}],
+                        latency_ms=latency_ms,
+                    )
+                _check_decision(resp, block_on_block, False)
+            except BlockedError:
+                raise
+            except Exception as exc:
+                logger.warning("[AgentTrust] MCP tool guard failed (non-fatal): %s", exc)
+
+        if is_async:
+            @functools.wraps(fn)
+            async def async_wrapper(*args: Any, **kwargs: Any) -> Any:
+                t0 = time.perf_counter()
+                result = await fn(*args, **kwargs)
+                latency_ms = (time.perf_counter() - t0) * 1000
+                call_args = inspect.signature(fn).bind(*args, **kwargs).arguments
+                await _validate(fn.__name__, dict(call_args), result, latency_ms)
+                return result
+            return async_wrapper
+        else:
+            @functools.wraps(fn)
+            def sync_wrapper(*args: Any, **kwargs: Any) -> Any:
+                import asyncio
+                t0 = time.perf_counter()
+                result = fn(*args, **kwargs)
+                latency_ms = (time.perf_counter() - t0) * 1000
+                call_args = inspect.signature(fn).bind(*args, **kwargs).arguments
+                try:
+                    loop = asyncio.get_event_loop()
+                    if loop.is_running():
+                        loop.create_task(_validate(fn.__name__, dict(call_args), result, latency_ms))
+                    else:
+                        loop.run_until_complete(_validate(fn.__name__, dict(call_args), result, latency_ms))
+                except Exception as exc:
+                    logger.warning("[AgentTrust] MCP sync validate error: %s", exc)
+                return result
+            return sync_wrapper
+
+    return decorator
+
+
+class MCPServerGuard:
+    """Wraps a FastMCP / mcp.Server instance and governs all registered tools.
+
+    Usage::
+
+        from mcp.server.fastmcp import FastMCP
+        from agentrust_sdk.adapters.mcp import MCPServerGuard
+
+        mcp = FastMCP("payments-server")
+
+        @mcp.tool()
+        async def process_payment(amount: float, account: str) -> dict:
+            ...
+
+        guard = MCPServerGuard(mcp, agent_id="payments-mcp")
+        # process_payment is now governed — AgentTrust validates every call.
+    """
+
+    def __init__(
+        self,
+        server: Any,
+        agent_id: str,
+        *,
+        base_url: str = "http://localhost:8000",
+        api_key: str | None = None,
+        user: str = "mcp",
+        block_on_block: bool = True,
+    ) -> None:
+        _check_tier(api_key)
+        self._server = server
+        self._agent_id = agent_id
+        self._base_url = base_url
+        self._api_key = api_key
+        self._user = user
+        self._block_on_block = block_on_block
+        self._wrap_tools()
+
+    def _wrap_tools(self) -> None:
+        """Monkey-patch all registered tool handlers on the server."""
+        tools_attr = None
+        for attr in ("_tool_handlers", "tools", "_tools", "_handlers"):
+            if hasattr(self._server, attr):
+                tools_attr = getattr(self._server, attr)
+                break
+
+        if tools_attr is None:
+            logger.warning(
+                "[AgentTrust] MCPServerGuard: could not locate tool registry on %s. "
+                "Use @mcp_tool_guard decorator instead.",
+                type(self._server).__name__,
+            )
+            return
+
+        guard = mcp_tool_guard(
+            self._agent_id,
+            base_url=self._base_url,
+            api_key=self._api_key,
+            user=self._user,
+            block_on_block=self._block_on_block,
+        )
+
+        if isinstance(tools_attr, dict):
+            for name, fn in list(tools_attr.items()):
+                tools_attr[name] = guard(fn)
+        elif isinstance(tools_attr, list):
+            for i, fn in enumerate(tools_attr):
+                tools_attr[i] = guard(fn)
+
+        logger.info(
+            "[AgentTrust] MCPServerGuard: wrapped %d tool(s) on agent '%s'",
+            len(tools_attr) if hasattr(tools_attr, "__len__") else -1,
+            self._agent_id,
+        )

agentrust_sdk/adapters/openai_agents.py

@@ -0,0 +1,263 @@
+"""
+OpenAI Agents SDK adapter — AgentTrust governance for OpenAI Agents (Swarm / Agents SDK).
+
+Architecture doc (two-pager): "Connects to … OpenAI Agents."
+
+Supports three patterns:
+
+1. @agent_guard decorator — wrap a single Agent run() call:
+
+    from agents import Agent, Runner
+    from agentrust_sdk.adapters.openai_agents import agent_guard
+
+    @agent_guard(agent_id="support-bot")
+    async def run_support(query: str) -> str:
+        result = await Runner.run(my_agent, query)
+        return result.final_output
+
+2. OpenAIAgentsGuard — wrap a Runner:
+
+    from agentrust_sdk.adapters.openai_agents import OpenAIAgentsGuard
+
+    guard = OpenAIAgentsGuard(agent_id="support-bot")
+    result = await guard.run(agent, "process this query")
+
+3. Swarm (legacy openai-swarm) wrapper:
+
+    from agentrust_sdk.adapters.openai_agents import SwarmGuard
+
+    guard = SwarmGuard(agent_id="swarm-agent")
+    response = guard.run(swarm_client, agent, messages)
+
+Requires Team tier.
+"""
+from __future__ import annotations
+
+import functools
+import inspect
+import logging
+import time
+from typing import Any
+
+logger = logging.getLogger(__name__)
+
+
+def _check_tier(api_key: str | None = None) -> None:
+    from agentrust_sdk.auth import resolve_key
+    from agentrust_sdk.tiers import Capability, is_allowed
+    info = resolve_key(api_key)
+    if not is_allowed(Capability.CREWAI_ADAPTER, info.tier):
+        raise RuntimeError("[AgentTrust] OpenAI Agents adapter requires Team tier or higher.")
+
+
+def _validate(
+    agent_id: str,
+    base_url: str,
+    api_key: str | None,
+    user: str,
+    input_text: str,
+    output: dict,
+    framework: str,
+    block_on_block: bool,
+    latency_ms: float = 0.0,
+    tool_calls: list | None = None,
+) -> None:
+    from agentrust_sdk.client import AgentTrustClient
+    from agentrust_sdk.decorator import _check_decision, BlockedError
+    try:
+        with AgentTrustClient(base_url=base_url, api_key=api_key) as client:
+            resp = client.validate(
+                agent_id=agent_id,
+                user=user,
+                input=input_text[:1000],
+                output=output,
+                framework=framework,
+                tools_called=tool_calls or [],
+                latency_ms=latency_ms,
+            )
+        _check_decision(resp, block_on_block, False)
+    except BlockedError:
+        raise
+    except Exception as exc:
+        logger.warning("[AgentTrust] OpenAI Agents validate failed (non-fatal): %s", exc)
+
+
+def _extract_run_output(result: Any) -> tuple[str, list]:
+    """Extract final_output and tool calls from an openai-agents RunResult."""
+    final_output = ""
+    tool_calls: list[dict] = []
+
+    if hasattr(result, "final_output"):
+        final_output = str(result.final_output or "")
+    elif isinstance(result, str):
+        final_output = result
+
+    # Extract tool calls from new_items (RunResult in openai-agents SDK)
+    if hasattr(result, "new_items"):
+        for item in result.new_items or []:
+            item_type = getattr(item, "type", None)
+            if item_type == "tool_call_item" or hasattr(item, "tool_call"):
+                call = getattr(item, "tool_call", item)
+                tool_calls.append({
+                    "name": getattr(call, "name", ""),
+                    "args": getattr(call, "arguments", {}),
+                })
+
+    return final_output, tool_calls
+
+
+def agent_guard(
+    agent_id: str,
+    *,
+    base_url: str = "http://localhost:8000",
+    api_key: str | None = None,
+    user: str = "openai-agent",
+    block_on_block: bool = True,
+    framework: str = "OpenAI-Agents",
+) -> Any:
+    """Decorator for async agent runner functions."""
+    _check_tier(api_key)
+
+    def decorator(fn: Any) -> Any:
+        is_async = inspect.iscoroutinefunction(fn)
+
+        if is_async:
+            @functools.wraps(fn)
+            async def async_wrapper(*args: Any, **kwargs: Any) -> Any:
+                t0 = time.perf_counter()
+                result = await fn(*args, **kwargs)
+                latency_ms = (time.perf_counter() - t0) * 1000
+                input_text = str(args[0]) if args else ""
+                final_output, tool_calls = _extract_run_output(result)
+                _validate(
+                    agent_id, base_url, api_key, user, input_text,
+                    {"final_output": final_output[:2000]},
+                    framework, block_on_block, latency_ms, tool_calls,
+                )
+                return result
+            return async_wrapper
+        else:
+            @functools.wraps(fn)
+            def sync_wrapper(*args: Any, **kwargs: Any) -> Any:
+                t0 = time.perf_counter()
+                result = fn(*args, **kwargs)
+                latency_ms = (time.perf_counter() - t0) * 1000
+                input_text = str(args[0]) if args else ""
+                final_output, tool_calls = _extract_run_output(result)
+                _validate(
+                    agent_id, base_url, api_key, user, input_text,
+                    {"final_output": final_output[:2000]},
+                    framework, block_on_block, latency_ms, tool_calls,
+                )
+                return result
+            return sync_wrapper
+
+    return decorator
+
+
+class OpenAIAgentsGuard:
+    """
+    Wraps openai-agents Runner.run() with AgentTrust governance.
+
+    Usage::
+
+        from agents import Agent
+        from agentrust_sdk.adapters.openai_agents import OpenAIAgentsGuard
+
+        guard = OpenAIAgentsGuard(agent_id="invoice-agent")
+        result = await guard.run(my_agent, "process invoice #42")
+    """
+
+    def __init__(
+        self,
+        agent_id: str,
+        *,
+        base_url: str = "http://localhost:8000",
+        api_key: str | None = None,
+        user: str = "openai-agent",
+        block_on_block: bool = True,
+        framework: str = "OpenAI-Agents",
+    ) -> None:
+        _check_tier(api_key)
+        self._agent_id = agent_id
+        self._base_url = base_url
+        self._api_key = api_key
+        self._user = user
+        self._block_on_block = block_on_block
+        self._framework = framework
+
+    async def run(self, agent: Any, input_text: str, **kwargs: Any) -> Any:
+        """Run *agent* via openai-agents Runner and validate the result."""
+        try:
+            from agents import Runner
+        except ImportError:
+            raise RuntimeError(
+                "openai-agents package not installed. pip install openai-agents"
+            )
+        t0 = time.perf_counter()
+        result = await Runner.run(agent, input_text, **kwargs)
+        latency_ms = (time.perf_counter() - t0) * 1000
+
+        final_output, tool_calls = _extract_run_output(result)
+        _validate(
+            self._agent_id, self._base_url, self._api_key, self._user,
+            input_text, {"final_output": final_output[:2000]},
+            self._framework, self._block_on_block, latency_ms, tool_calls,
+        )
+        return result
+
+    def run_sync(self, agent: Any, input_text: str, **kwargs: Any) -> Any:
+        """Synchronous version for non-async code."""
+        import asyncio
+        return asyncio.run(self.run(agent, input_text, **kwargs))
+
+
+class SwarmGuard:
+    """
+    Wraps legacy openai-swarm client.run() with AgentTrust governance.
+
+    Usage::
+
+        from swarm import Swarm
+        from agentrust_sdk.adapters.openai_agents import SwarmGuard
+
+        client = Swarm()
+        guard = SwarmGuard(agent_id="swarm-agent")
+        response = guard.run(client, agent, [{"role": "user", "content": query}])
+    """
+
+    def __init__(
+        self,
+        agent_id: str,
+        *,
+        base_url: str = "http://localhost:8000",
+        api_key: str | None = None,
+        user: str = "swarm",
+        block_on_block: bool = True,
+        framework: str = "OpenAI-Swarm",
+    ) -> None:
+        _check_tier(api_key)
+        self._agent_id = agent_id
+        self._base_url = base_url
+        self._api_key = api_key
+        self._user = user
+        self._block_on_block = block_on_block
+        self._framework = framework
+
+    def run(self, swarm_client: Any, agent: Any, messages: list, **kwargs: Any) -> Any:
+        t0 = time.perf_counter()
+        response = swarm_client.run(agent, messages, **kwargs)
+        latency_ms = (time.perf_counter() - t0) * 1000
+
+        last_user = next(
+            (m.get("content", "") for m in reversed(messages) if m.get("role") == "user"), ""
+        )
+        final_messages = getattr(response, "messages", []) or []
+        last_content = final_messages[-1].get("content", "") if final_messages else ""
+
+        _validate(
+            self._agent_id, self._base_url, self._api_key, self._user,
+            str(last_user)[:500], {"content": last_content[:2000]},
+            self._framework, self._block_on_block, latency_ms,
+        )
+        return response

agentrust_sdk/auth.py

@@ -0,0 +1,192 @@
+"""
+API key resolution and tier detection.
+
+Key resolution order (first found wins):
+  1. AGENTRUST_KEY  env var
+  2. .agentrust/config.yaml  (project-level)
+  3. ~/.agentrust/config.yaml  (global)
+  4. Passed explicitly to client/decorator
+
+No key found → Tier.OSS (schema validation only, no telemetry).
+
+API keys are signed JWTs (HS256) containing:
+  { "sub": "<org_id>", "tier": "team", "iat": ..., "exp": ... }
+
+For local/dev keys issued by `agentrust init --local`, the key is a
+plain opaque string and tier defaults to FREE.
+"""
+from __future__ import annotations
+
+import base64
+import json
+import logging
+import os
+from pathlib import Path
+from typing import Any
+
+import yaml
+
+from .tiers import Tier
+
+logger = logging.getLogger(__name__)
+
+_ENV_VAR = "AGENTRUST_KEY"
+_CONFIG_FILENAME = "config.yaml"
+_CONFIG_DIR = ".agentrust"
+_GLOBAL_CONFIG = Path.home() / _CONFIG_DIR / _CONFIG_FILENAME
+
+
+# ---------------------------------------------------------------------------
+# Config file helpers
+# ---------------------------------------------------------------------------
+
+def _project_config() -> Path | None:
+    """Walk up from cwd looking for .agentrust/config.yaml"""
+    current = Path.cwd()
+    for parent in [current, *current.parents]:
+        candidate = parent / _CONFIG_DIR / _CONFIG_FILENAME
+        if candidate.exists():
+            return candidate
+    return None
+
+
+def _read_config(path: Path) -> dict[str, Any]:
+    try:
+        return yaml.safe_load(path.read_text()) or {}
+    except Exception:
+        return {}
+
+
+def _write_config(path: Path, data: dict[str, Any]) -> None:
+    path.parent.mkdir(parents=True, exist_ok=True)
+    path.write_text(yaml.dump(data, default_flow_style=False))
+    # Restrict to owner-only: API key must not be world-readable.
+    path.chmod(0o600)
+
+
+# ---------------------------------------------------------------------------
+# JWT decode (no signature verification in SDK — server verifies)
+# ---------------------------------------------------------------------------
+
+def _decode_jwt_payload(token: str) -> dict[str, Any]:
+    """
+    Decode JWT payload without verifying signature.
+    The control plane verifies; the SDK only reads the tier claim.
+    """
+    try:
+        parts = token.split(".")
+        if len(parts) != 3:
+            return {}
+        # Add padding if needed
+        payload = parts[1]
+        payload += "=" * (4 - len(payload) % 4)
+        decoded = base64.urlsafe_b64decode(payload)
+        return json.loads(decoded)
+    except Exception:
+        return {}
+
+
+# ---------------------------------------------------------------------------
+# Public API
+# ---------------------------------------------------------------------------
+
+class KeyInfo:
+    __slots__ = ("key", "tier", "org_id", "source")
+
+    def __init__(self, key: str | None, tier: Tier, org_id: str, source: str) -> None:
+        self.key = key
+        self.tier = tier
+        self.org_id = org_id
+        self.source = source  # "env", "project_config", "global_config", "explicit", "none"
+
+    def __repr__(self) -> str:
+        return f"KeyInfo(tier={self.tier.value}, org={self.org_id}, source={self.source})"
+
+
+_OSS_KEY_INFO = KeyInfo(key=None, tier=Tier.OSS, org_id="anonymous", source="none")
+
+
+def resolve_key(explicit_key: str | None = None) -> KeyInfo:
+    """
+    Resolve API key and decode tier. Never raises — falls back to OSS on any error.
+    """
+    candidates: list[tuple[str, str]] = []  # (key, source)
+
+    if explicit_key:
+        candidates.append((explicit_key, "explicit"))
+
+    env_key = os.environ.get(_ENV_VAR, "").strip()
+    if env_key:
+        candidates.append((env_key, "env"))
+
+    project_cfg_path = _project_config()
+    if project_cfg_path:
+        cfg = _read_config(project_cfg_path)
+        if cfg.get("api_key"):
+            candidates.append((cfg["api_key"], "project_config"))
+
+    if _GLOBAL_CONFIG.exists():
+        cfg = _read_config(_GLOBAL_CONFIG)
+        if cfg.get("api_key"):
+            candidates.append((cfg["api_key"], "global_config"))
+
+    for key, source in candidates:
+        info = _parse_key(key, source)
+        if info is not None:
+            return info
+
+    return _OSS_KEY_INFO
+
+
+def _parse_key(key: str, source: str) -> KeyInfo | None:
+    if not key:
+        return None
+
+    # JWT keys have 2 dots.
+    # The payload tier is used for client-side capability hints (UX only — early errors,
+    # upgrade prompts). The gateway verifies the JWT signature server-side before
+    # granting any access, so a forged tier in the payload is harmless here.
+    if key.count(".") == 2:
+        payload = _decode_jwt_payload(key)
+        org_id = payload.get("sub", "unknown")
+        tier_str = payload.get("tier", "free")
+        try:
+            tier = Tier(tier_str)
+        except ValueError:
+            tier = Tier.FREE
+        return KeyInfo(key=key, tier=tier, org_id=org_id, source=source)
+
+    # Plain opaque key → FREE tier (issued by agentrust init --local)
+    return KeyInfo(key=key, tier=Tier.FREE, org_id="local", source=source)
+
+
+# ---------------------------------------------------------------------------
+# Config file management (used by CLI)
+# ---------------------------------------------------------------------------
+
+def save_key_to_config(api_key: str, scope: str = "global") -> Path:
+    """Write API key to config file. scope='global' or 'project'."""
+    if scope == "project":
+        path = Path.cwd() / _CONFIG_DIR / _CONFIG_FILENAME
+    else:
+        path = _GLOBAL_CONFIG
+
+    existing = _read_config(path) if path.exists() else {}
+    existing["api_key"] = api_key
+
+    # Decode and store tier for convenience (humans reading the file)
+    info = _parse_key(api_key, "save")
+    if info:
+        existing["tier"] = info.tier.value
+        existing["org_id"] = info.org_id
+
+    _write_config(path, existing)
+    logger.info("API key saved to %s", path)
+    return path
+
+
+def read_config(scope: str = "global") -> dict[str, Any]:
+    if scope == "project":
+        p = _project_config()
+        return _read_config(p) if p else {}
+    return _read_config(_GLOBAL_CONFIG) if _GLOBAL_CONFIG.exists() else {}