deer-flow/backend/packages/harness/deerflow/guardrails/middleware.py

"""GuardrailMiddleware - evaluates tool calls against a GuardrailProvider before execution."""

import logging
from collections.abc import Awaitable, Callable
from datetime import UTC, datetime
from typing import override

from langchain.agents import AgentState
from langchain.agents.middleware import AgentMiddleware
from langchain_core.messages import ToolMessage
from langgraph.errors import GraphBubbleUp
from langgraph.prebuilt.tool_node import ToolCallRequest
from langgraph.types import Command

from deerflow.authz.principal import normalize_authz_attributes
from deerflow.guardrails.provider import GuardrailDecision, GuardrailProvider, GuardrailReason, GuardrailRequest
from deerflow.runtime.events.catalog import MIDDLEWARE_GUARDRAIL_TAG

logger = logging.getLogger(__name__)

_REASON_MESSAGE_LIMIT = 500


class GuardrailMiddleware(AgentMiddleware[AgentState]):
    """Evaluate tool calls against a GuardrailProvider before execution.

    Denied calls return an error ToolMessage so the agent can adapt.
    If the provider raises, behavior depends on fail_closed:
      - True (default): block the call
      - False: allow it through with a warning
    """

    def __init__(self, provider: GuardrailProvider, *, fail_closed: bool = True, passport: str | None = None):
        self.provider = provider
        self.fail_closed = fail_closed
        self.passport = passport

    @staticmethod
    def _resolve_context(request: ToolCallRequest) -> dict:
        runtime = getattr(request, "runtime", None)
        context = getattr(runtime, "context", None) if runtime is not None else None
        return context if isinstance(context, dict) else {}

    def _build_request(self, request: ToolCallRequest, context: dict) -> GuardrailRequest:
        return GuardrailRequest(
            tool_name=str(request.tool_call.get("name", "")),
            tool_input=request.tool_call.get("args", {}),
            agent_id=self.passport,
            thread_id=context.get("thread_id"),
            is_subagent=bool(context.get("is_subagent")),
            timestamp=datetime.now(UTC).isoformat(),
            user_id=context.get("user_id"),
            user_role=context.get("user_role"),
            oauth_provider=context.get("oauth_provider"),
            oauth_id=context.get("oauth_id"),
            run_id=context.get("run_id"),
            tool_call_id=request.tool_call.get("id"),
            channel_user_id=context.get("channel_user_id"),
            is_internal=context.get("is_internal") is True,
            authz_attributes=normalize_authz_attributes(context.get("authz_attributes")),
        )

    def _build_denied_message(self, request: ToolCallRequest, decision: GuardrailDecision) -> ToolMessage:
        tool_name = str(request.tool_call.get("name", "unknown_tool"))
        tool_call_id = str(request.tool_call.get("id", "missing_id"))
        reason_text = decision.reasons[0].message if decision.reasons else "blocked by guardrail policy"
        reason_code = decision.reasons[0].code if decision.reasons else "oap.denied"
        return ToolMessage(
            content=f"Guardrail denied: tool '{tool_name}' was blocked ({reason_code}). Reason: {reason_text}. Choose an alternative approach.",
            tool_call_id=tool_call_id,
            name=tool_name,
            status="error",
        )

    def _record_guardrail_event(
        self,
        context: dict,
        guardrail_request: GuardrailRequest,
        decision: GuardrailDecision,
        *,
        action: str,
        provider_error: bool,
    ) -> None:
        """Persist a security-relevant guardrail decision to RunJournal.

        This follows the optional-Journal pattern used by existing middleware:
        audit persistence is best-effort and must never change tool execution
        behavior. Runtimes without ``__run_journal`` (including embedded and
        subagent execution) skip persistence.
        """
        journal = context.get("__run_journal")
        if journal is None:
            return

        reason_codes = [reason.code for reason in decision.reasons if reason.code]
        reason_messages = [reason.message[:_REASON_MESSAGE_LIMIT] for reason in decision.reasons if reason.message]

        changes = {
            "tool_name": guardrail_request.tool_name,
            "tool_call_id": guardrail_request.tool_call_id,
            "agent_id": guardrail_request.agent_id,
            # Native subagents do not currently inherit __run_journal; custom
            # runtimes may still provide one with subagent attribution.
            "is_subagent": guardrail_request.is_subagent,
            "user_role": guardrail_request.user_role,
            "allow": decision.allow,
            "policy_id": decision.policy_id,
            "reason_codes": reason_codes,
            "reason_messages": reason_messages,
            "fail_closed": self.fail_closed,
            "provider_error": provider_error,
        }

        try:
            journal.record_middleware(
                tag=MIDDLEWARE_GUARDRAIL_TAG,
                name=type(self).__name__,
                hook="wrap_tool_call",
                action=action,
                changes=changes,
            )
        except Exception:  # noqa: BLE001
            logger.warning("Failed to record middleware:guardrail event", exc_info=True)

    @override
    def wrap_tool_call(
        self,
        request: ToolCallRequest,
        handler: Callable[[ToolCallRequest], ToolMessage | Command],
    ) -> ToolMessage | Command:
        context = self._resolve_context(request)
        gr = self._build_request(request, context)
        try:
            decision = self.provider.evaluate(gr)
        except GraphBubbleUp:
            # Preserve LangGraph control-flow signals (interrupt/pause/resume).
            raise
        except Exception:
            logger.exception("Guardrail provider error (sync)")
            if self.fail_closed:
                decision = GuardrailDecision(allow=False, reasons=[GuardrailReason(code="oap.evaluator_error", message="guardrail provider error (fail-closed)")])
                self._record_guardrail_event(
                    context,
                    gr,
                    decision,
                    action="deny_tool_call",
                    provider_error=True,
                )
                return self._build_denied_message(request, decision)
            else:
                decision = GuardrailDecision(allow=True, reasons=[GuardrailReason(code="oap.evaluator_error", message="guardrail provider error (fail-open)")])
                self._record_guardrail_event(
                    context,
                    gr,
                    decision,
                    action="allow_tool_call_after_provider_error",
                    provider_error=True,
                )
                return handler(request)
        if not decision.allow:
            logger.warning("Guardrail denied: tool=%s policy=%s code=%s", gr.tool_name, decision.policy_id, decision.reasons[0].code if decision.reasons else "unknown")
            self._record_guardrail_event(
                context,
                gr,
                decision,
                action="deny_tool_call",
                provider_error=False,
            )
            return self._build_denied_message(request, decision)
        return handler(request)

    @override
    async def awrap_tool_call(
        self,
        request: ToolCallRequest,
        handler: Callable[[ToolCallRequest], Awaitable[ToolMessage | Command]],
    ) -> ToolMessage | Command:
        context = self._resolve_context(request)
        gr = self._build_request(request, context)
        try:
            decision = await self.provider.aevaluate(gr)
        except GraphBubbleUp:
            # Preserve LangGraph control-flow signals (interrupt/pause/resume).
            raise
        except Exception:
            logger.exception("Guardrail provider error (async)")
            if self.fail_closed:
                decision = GuardrailDecision(allow=False, reasons=[GuardrailReason(code="oap.evaluator_error", message="guardrail provider error (fail-closed)")])
                self._record_guardrail_event(
                    context,
                    gr,
                    decision,
                    action="deny_tool_call",
                    provider_error=True,
                )
                return self._build_denied_message(request, decision)
            else:
                decision = GuardrailDecision(allow=True, reasons=[GuardrailReason(code="oap.evaluator_error", message="guardrail provider error (fail-open)")])
                self._record_guardrail_event(
                    context,
                    gr,
                    decision,
                    action="allow_tool_call_after_provider_error",
                    provider_error=True,
                )
                return await handler(request)
        if not decision.allow:
            logger.warning("Guardrail denied: tool=%s policy=%s code=%s", gr.tool_name, decision.policy_id, decision.reasons[0].code if decision.reasons else "unknown")
            self._record_guardrail_event(
                context,
                gr,
                decision,
                action="deny_tool_call",
                provider_error=False,
            )
            return self._build_denied_message(request, decision)
        return await handler(request)