Plugin Architecture for Tools in a Custom Runtime

from __future__ import annotations
 
import json
from dataclasses import dataclass, field
from typing import Any, Callable, Protocol
 
 
@dataclass(frozen=True)
class AuthContext:
    user_id: str
    roles: frozenset[str]
    tenant_id: str
 
 
@dataclass
class ToolResult:
    ok: bool
    data: Any = None
    error_type: str | None = None
    message: str | None = None
 
 
class ToolPlugin(Protocol):
    name: str
    description: str
    parameters: dict[str, Any]  # JSON Schema object
    risk: str  # "read" | "write" | "dangerous"
 
    def handle(self, args: dict[str, Any], ctx: AuthContext) -> ToolResult: ...
 
 
@dataclass
class FunctionPlugin:
    name: str
    description: str
    parameters: dict[str, Any]
    risk: str
    fn: Callable[[dict[str, Any], AuthContext], Any]
 
    def handle(self, args: dict[str, Any], ctx: AuthContext) -> ToolResult:
        try:
            return ToolResult(ok=True, data=self.fn(args, ctx))
        except Exception as exc:
            return ToolResult(ok=False, error_type="tool_error", message=str(exc))
 
 
def _get_order_status(args: dict[str, Any], ctx: AuthContext) -> dict[str, Any]:
    # demo: ignore tenant; production: scope by ctx.tenant_id
    order_id = args["order_id"]
    return {"order_id": order_id, "status": "shipped", "tenant": ctx.tenant_id}
 
 
def _refund_order(args: dict[str, Any], ctx: AuthContext) -> dict[str, Any]:
    if "refunds" not in ctx.roles:
        raise PermissionError("role refunds required")
    return {"order_id": args["order_id"], "refunded": True}
 
 
class ToolRegistry:
    def __init__(self) -> None:
        self._tools: dict[str, ToolPlugin] = {}
 
    def register(self, plugin: ToolPlugin) -> None:
        if plugin.name in self._tools:
            raise ValueError(f"duplicate tool: {plugin.name}")
        self._tools[plugin.name] = plugin
 
    def specs(self, allow: set[str] | None = None) -> list[dict[str, Any]]:
        names = allow if allow is not None else set(self._tools)
        out: list[dict[str, Any]] = []
        for name in sorted(names):
            p = self._tools[name]
            out.append(
                {
                    "type": "function",
                    "function": {
                        "name": p.name,
                        "description": p.description,
                        "parameters": p.parameters,
                    },
                }
            )
        return out
 
    def dispatch(
        self,
        name: str,
        args: dict[str, Any],
        ctx: AuthContext,
        allow: set[str] | None = None,
    ) -> ToolResult:
        if allow is not None and name not in allow:
            return ToolResult(
                ok=False, error_type="tool_policy", message=f"tool not allowed: {name}"
            )
        plugin = self._tools.get(name)
        if plugin is None:
            return ToolResult(
                ok=False, error_type="unknown_tool", message=f"unknown tool: {name}"
            )
        # optional: jsonschema.validate(args, plugin.parameters)
        return plugin.handle(args, ctx)
 
 
def build_default_registry() -> ToolRegistry:
    reg = ToolRegistry()
    reg.register(
        FunctionPlugin(
            name="get_order_status",
            description="Look up an order status by id.",
            parameters={
                "type": "object",
                "properties": {"order_id": {"type": "string"}},
                "required": ["order_id"],
            },
            risk="read",
            fn=_get_order_status,
        )
    )
    reg.register(
        FunctionPlugin(
            name="refund_order",
            description="Refund an order. Requires refunds role.",
            parameters={
                "type": "object",
                "properties": {"order_id": {"type": "string"}},
                "required": ["order_id"],
            },
            risk="write",
            fn=_refund_order,
        )
    )
    return reg
 
 
# --- use inside the loop ---
registry = build_default_registry()
SUPPORT_ALLOW = {"get_order_status"}  # support agent cannot refund
 
 
def run_tool_for_support(name: str, raw_args: str, ctx: AuthContext) -> str:
    args = json.loads(raw_args or "{}")
    result = registry.dispatch(name, args, ctx, allow=SUPPORT_ALLOW)
    return json.dumps(result.__dict__)

MCP client list_tools -> ToolPlugin adapters -> registry
MCP call_tool        <- dispatch

PROFILES = {
    "support": {"get_order_status", "search_kb"},
    "billing": {"get_order_status", "refund_order"},
    "admin_read": {"get_order_status", "list_audit_events"},
}

domain/orders.py          # refund_order_domain(order_id, tenant_id)
runtime/plugins/orders.py # FunctionPlugin wrapping domain
runtime/registry.py
runtime/loop.py