hermes-agent/hermes_cli/middleware.py

"""Hermes middleware contract helpers.

Observer hooks report what happened. Middleware can change what happens by
rewriting a request or wrapping the actual execution callback. Keep the small
contract helpers here so agent-loop call sites and plugins share one vocabulary.
"""

from __future__ import annotations

import contextvars
import fnmatch
import json
import logging
import threading
from copy import deepcopy
from dataclasses import dataclass, field
from typing import Any, Callable, Dict, List, Optional

logger = logging.getLogger(__name__)

_TOOL_POLICY_PRIORITY = {"allow": 0, "ask": 1, "deny": 2}


class _RegistryDispatchCapability:
    """Opaque, single-use permission for one wrapper→registry handoff.

    Context copies share this object's consumed state, so only one dispatch can
    redeem it even if a caller copies the Context or reuses the same args.
    """

    __slots__ = ("_registry", "_tool_name", "_args", "_consumed", "_lock")

    def __init__(self, registry: Any, tool_name: str, args: Dict[str, Any]) -> None:
        self._registry = registry
        self._tool_name = tool_name
        self._args = args
        self._consumed = False
        self._lock = threading.Lock()

    def consume(self, registry: Any, tool_name: str, args: Dict[str, Any]) -> bool:
        if registry is not self._registry or tool_name != self._tool_name or args is not self._args:
            return False
        with self._lock:
            if self._consumed:
                return False
            self._consumed = True
            return True


_pending_registry_dispatch: contextvars.ContextVar[Optional[_RegistryDispatchCapability]] = (
    contextvars.ContextVar("pending_registry_dispatch", default=None)
)

OBSERVER_SCHEMA_VERSION = "hermes.observer.v1"
MIDDLEWARE_SCHEMA_VERSION = "hermes.middleware.v1"

TOOL_REQUEST_MIDDLEWARE = "tool_request"
TOOL_EXECUTION_MIDDLEWARE = "tool_execution"
LLM_REQUEST_MIDDLEWARE = "llm_request"
LLM_EXECUTION_MIDDLEWARE = "llm_execution"

# Back-compat aliases for older PoC branches that used API terminology.
API_REQUEST_MIDDLEWARE = LLM_REQUEST_MIDDLEWARE
API_EXECUTION_MIDDLEWARE = LLM_EXECUTION_MIDDLEWARE

VALID_MIDDLEWARE: set[str] = {
    TOOL_REQUEST_MIDDLEWARE,
    TOOL_EXECUTION_MIDDLEWARE,
    LLM_REQUEST_MIDDLEWARE,
    LLM_EXECUTION_MIDDLEWARE,
}


@dataclass
class RequestMiddlewareResult:
    """Result of applying request middleware to a mutable payload."""

    payload: Any
    original_payload: Any
    changed: bool = False
    trace: List[Dict[str, Any]] = field(default_factory=list)


def observer_payload(**kwargs: Any) -> Dict[str, Any]:
    kwargs.setdefault("telemetry_schema_version", OBSERVER_SCHEMA_VERSION)
    return kwargs


def middleware_payload(**kwargs: Any) -> Dict[str, Any]:
    kwargs.setdefault("telemetry_schema_version", OBSERVER_SCHEMA_VERSION)
    kwargs.setdefault("middleware_schema_version", MIDDLEWARE_SCHEMA_VERSION)
    return kwargs


def _safe_copy(payload: Any) -> Any:
    """Deep-copy a request payload, tolerating non-deepcopyable members.

    Request payloads are normally plain JSON-shaped dicts, but an LLM request
    can occasionally carry non-deepcopyable objects (clients, callbacks, file
    handles). A hard ``deepcopy`` failure there would otherwise abort the whole
    request-middleware pass. Fall back to a shallow ``dict`` copy so middleware
    still runs and the original nested objects are shared by reference rather
    than corrupting the live payload.
    """
    try:
        return deepcopy(payload)
    except Exception as exc:  # pragma: no cover - exercised via fallback test
        logger.debug("deepcopy failed for request payload (%s); using shallow copy", exc)
        if isinstance(payload, dict):
            return dict(payload)
        return payload


def apply_llm_request_middleware(
    request: Dict[str, Any],
    **context: Any,
) -> RequestMiddlewareResult:
    """Apply registered LLM request middleware.

    Middleware may return ``{"request": {...}}`` to replace the effective
    provider kwargs before Hermes sends them.
    """
    if not _has_middleware(LLM_REQUEST_MIDDLEWARE):
        return RequestMiddlewareResult(
            payload=request,
            original_payload=request,
            changed=False,
            trace=[],
        )

    original_request = _safe_copy(request)
    current_request = _safe_copy(original_request)
    trace: List[Dict[str, Any]] = []

    for result in _invoke_middleware(
        LLM_REQUEST_MIDDLEWARE,
        request=current_request,
        original_request=original_request,
        **context,
    ):
        if not isinstance(result, dict):
            continue
        next_request = result.get("request")
        if not isinstance(next_request, dict):
            continue
        current_request = _safe_copy(next_request)
        trace.append(_trace_entry(result))

    return RequestMiddlewareResult(
        payload=current_request,
        original_payload=original_request,
        changed=bool(trace),
        trace=trace,
    )


def apply_tool_request_middleware(
    tool_name: str,
    args: Dict[str, Any],
    **context: Any,
) -> RequestMiddlewareResult:
    """Apply registered tool request middleware.

    Middleware may return ``{"args": {...}}`` to replace the effective tool
    arguments before hooks, guardrails, approvals, and execution see them.
    """
    if not _has_middleware(TOOL_REQUEST_MIDDLEWARE):
        return RequestMiddlewareResult(
            payload=args,
            original_payload=args,
            changed=False,
            trace=[],
        )

    original_args = _safe_copy(args)
    current_args = _safe_copy(original_args)
    trace: List[Dict[str, Any]] = []

    for result in _invoke_middleware(
        TOOL_REQUEST_MIDDLEWARE,
        tool_name=tool_name,
        args=current_args,
        original_args=original_args,
        **context,
    ):
        if not isinstance(result, dict):
            continue
        next_args = result.get("args")
        if not isinstance(next_args, dict):
            continue
        current_args = _safe_copy(next_args)
        trace.append(_trace_entry(result))

    return RequestMiddlewareResult(
        payload=current_args,
        original_payload=original_args,
        changed=bool(trace),
        trace=trace,
    )


def apply_api_request_middleware(
    request: Dict[str, Any],
    **context: Any,
) -> RequestMiddlewareResult:
    """Compatibility wrapper for older ``api_request`` naming."""
    return apply_llm_request_middleware(request, **context)


def run_llm_execution_middleware(
    request: Dict[str, Any],
    next_call: Callable[[Dict[str, Any]], Any],
    **context: Any,
) -> Any:
    """Run provider execution through registered LLM execution middleware."""
    callbacks = _get_middleware_callbacks(LLM_EXECUTION_MIDDLEWARE)
    if not callbacks:
        return next_call(request)
    return _run_execution_chain(
        LLM_EXECUTION_MIDDLEWARE,
        callbacks,
        next_call,
        request=request,
        original_request=context.pop("original_request", request),
        **context,
    )


def run_tool_execution_middleware(
    tool_name: str,
    args: Dict[str, Any],
    next_call: Callable[[Dict[str, Any]], Any],
    **context: Any,
) -> Any:
    """Run a tool through the shared policy gate and plugin middleware chain.

    ``dispatch_registry`` marks a wrapper whose immediate ``next_call`` enters
    that registry. Only that exact wrapper→registry handoff is deduplicated;
    nested dispatches from inside the handler remain genuine executions.
    """
    dispatch_registry = context.pop("dispatch_registry", None)
    registry_entry = context.pop("registry_entry", None)
    registry_dispatch = context.pop("registry_dispatch", None)
    expected_dispatch = _pending_registry_dispatch.get()
    is_wrapped_registry_dispatch = bool(
        registry_dispatch
        and expected_dispatch is not None
        and expected_dispatch.consume(registry_dispatch, tool_name, args)
    )
    if is_wrapped_registry_dispatch:
        _pending_registry_dispatch.set(None)
        return next_call(args)

    policy = (
        resolve_tool_approval_policy(tool_name, entry=registry_entry)
        if registry_entry is not None
        else resolve_tool_approval_policy(tool_name)
    )
    blocked = _apply_tool_approval_policy(tool_name, policy)
    if blocked is not None:
        return blocked

    terminal_allow_token = None
    if tool_name == "terminal" and policy == "allow":
        from tools.approval import set_tool_policy_terminal_allow

        terminal_allow_token = set_tool_policy_terminal_allow()

    def call_next(next_args: Dict[str, Any]) -> Any:
        if dispatch_registry is None:
            return next_call(next_args)
        token = _pending_registry_dispatch.set(
            _RegistryDispatchCapability(dispatch_registry, tool_name, next_args)
        )
        try:
            return next_call(next_args)
        finally:
            _pending_registry_dispatch.reset(token)

    try:
        callbacks = _get_middleware_callbacks(TOOL_EXECUTION_MIDDLEWARE)
        if not callbacks:
            return call_next(args)
        return _run_execution_chain(
            TOOL_EXECUTION_MIDDLEWARE,
            callbacks,
            call_next,
            tool_name=tool_name,
            args=args,
            original_args=context.pop("original_args", args),
            **context,
        )
    finally:
        if terminal_allow_token is not None:
            from tools.approval import reset_tool_policy_terminal_allow

            reset_tool_policy_terminal_allow(terminal_allow_token)


def resolve_tool_approval_policy(
    tool_name: str, *, entry: Any = None
) -> Optional[str]:
    """Resolve glob policies for a tool name and its actual dispatch entry."""
    try:
        from hermes_cli.config import load_config

        config = load_config() or {}
        approvals = config.get("approvals", {}) or {}
        patterns = approvals.get("tool_policies", {}) or {}
    except Exception as exc:
        logger.warning("Failed to load tool approval policies: %s", exc)
        return None
    if not isinstance(patterns, dict):
        return None

    identifiers = [str(tool_name or "").strip().lower()]
    toolset = str(getattr(entry, "toolset", "") or "").strip().lower()
    if toolset:
        identifiers.append(toolset)

    matches: list[str] = []
    for pattern, raw_policy in patterns.items():
        if not isinstance(pattern, str) or not isinstance(raw_policy, str):
            continue
        normalized_pattern = pattern.strip().lower()
        policy = raw_policy.strip().lower()
        if not normalized_pattern or policy not in _TOOL_POLICY_PRIORITY:
            continue
        if any(
            fnmatch.fnmatchcase(identifier, normalized_pattern)
            for identifier in identifiers
        ):
            matches.append(policy)
    if not matches:
        return None
    return max(matches, key=_TOOL_POLICY_PRIORITY.__getitem__)


def _apply_tool_approval_policy(
    tool_name: str, policy: Optional[str]
) -> Optional[str]:
    """Return a synthetic blocked result, or None to continue execution."""
    if policy == "deny":
        return json.dumps(
            {
                "error": (
                    f"BLOCKED: Tool '{tool_name}' is denied by "
                    "approvals.tool_policies in config.yaml."
                ),
                "policy": "deny",
                "tool": tool_name,
            },
            ensure_ascii=False,
        )
    if policy != "ask":
        return None

    try:
        from tools.approval import request_tool_approval

        decision = request_tool_approval(
            tool_name,
            f"config.yaml requires approval for tool '{tool_name}'",
            rule_key=f"tool_policy:{tool_name}",
            enforce_under_yolo=True,
        )
    except Exception as exc:
        logger.warning("Tool approval policy check failed for %s: %s", tool_name, exc)
        decision = {
            "approved": False,
            "message": "BLOCKED: tool approval policy could not be evaluated.",
        }
    if decision.get("approved"):
        return None
    payload = dict(decision)
    payload["error"] = payload.pop("message", None) or (
        f"BLOCKED: Tool '{tool_name}' was not approved."
    )
    payload.setdefault("policy", "ask")
    payload.setdefault("tool", tool_name)
    return json.dumps(payload, ensure_ascii=False)


def run_api_execution_middleware(
    request: Dict[str, Any],
    next_call: Callable[[Dict[str, Any]], Any],
    **context: Any,
) -> Any:
    """Compatibility wrapper for older ``api_execution`` naming."""
    return run_llm_execution_middleware(request, next_call, **context)


def _invoke_middleware(kind: str, **kwargs: Any) -> List[Any]:
    from hermes_cli.plugins import invoke_middleware

    return invoke_middleware(kind, **middleware_payload(**kwargs))


def _has_middleware(kind: str) -> bool:
    from hermes_cli.plugins import has_middleware

    return has_middleware(kind)


def _get_middleware_callbacks(kind: str) -> List[Callable]:
    from hermes_cli.plugins import get_plugin_manager

    return list(get_plugin_manager()._middleware.get(kind, []))


def _run_execution_chain(
    kind: str,
    callbacks: List[Callable],
    terminal_call: Callable[[Any], Any],
    **kwargs: Any,
) -> Any:
    payload_key = "request" if "request" in kwargs else "args"

    class _DownstreamExecutionError(Exception):
        def __init__(self, original: BaseException) -> None:
            super().__init__(str(original))
            self.original = original

    def call_at(index: int, payload: Any) -> Any:
        if index >= len(callbacks):
            return terminal_call(payload)

        callback = callbacks[index]
        next_called = False
        next_succeeded = False
        next_result: Any = None

        def next_call(next_payload: Any = None) -> Any:
            nonlocal next_called, next_succeeded, next_result
            # ``next_call`` is single-use per middleware frame. Calling it more
            # than once would re-run the downstream provider/tool, so a second
            # invocation is a contract violation rather than a retry. Surface it
            # instead of silently executing the terminal call twice.
            if next_called:
                raise RuntimeError(
                    f"Middleware '{kind}' callback "
                    f"{getattr(callback, '__name__', repr(callback))} called "
                    "next_call() more than once; downstream execution is single-use"
                )
            next_called = True
            try:
                next_result = call_at(index + 1, payload if next_payload is None else next_payload)
                next_succeeded = True
                return next_result
            except Exception as exc:
                raise _DownstreamExecutionError(exc) from exc

        call_kwargs = middleware_payload(**kwargs)
        call_kwargs[payload_key] = payload
        call_kwargs["next_call"] = next_call
        try:
            return callback(**call_kwargs)
        except _DownstreamExecutionError as exc:
            raise exc.original
        except Exception as exc:
            logger.warning(
                "Middleware '%s' callback %s raised: %s",
                kind,
                getattr(callback, "__name__", repr(callback)),
                exc,
            )
            if next_succeeded:
                return next_result
            if next_called:
                raise
            return call_at(index + 1, payload)

    return call_at(0, kwargs[payload_key])


def _trace_entry(result: Dict[str, Any]) -> Dict[str, Any]:
    entry: Dict[str, Any] = {}
    for key in ("source", "reason", "name"):
        value = result.get(key)
        if isinstance(value, str) and value:
            entry[key] = value
    if not entry:
        entry["source"] = "plugin"
    return entry