evotraders/backend/agents/base/tool_guard.py

# -*- coding: utf-8 -*-
"""ToolGuardMixin - Security interception for dangerous tool calls.

Provides ``_acting`` and ``_reasoning`` overrides that intercept
sensitive tool calls before execution, implementing the deny /
guard / approve flow.

Based on CoPaw's tool_guard_mixin.py design.
"""
from __future__ import annotations

import asyncio
import json
import logging
from dataclasses import dataclass, field
from datetime import datetime
from enum import Enum

from typing import Any, Callable, Dict, Iterable, List, Optional, Set

from agentscope.message import Msg
from backend.runtime.manager import get_global_runtime_manager
logger = logging.getLogger(__name__)

class SeverityLevel(str, Enum):
    """Risk severity level."""

    LOW = "low"
    MEDIUM = "medium"
    HIGH = "high"
    CRITICAL = "critical"


class ApprovalStatus(str, Enum):
    """Approval lifecycle state."""

    PENDING = "pending"
    APPROVED = "approved"
    DENIED = "denied"
    EXPIRED = "expired"


class ToolFindingRecord:
    """Internal representation of a guard finding."""

    def __init__(self, severity: SeverityLevel, message: str, field: Optional[str] = None) -> None:
        self.severity = severity
        self.message = message
        self.field = field

    def to_dict(self) -> Dict[str, Any]:
        return {
            "severity": self.severity.value,
            "message": self.message,
            "field": self.field,
        }


class ApprovalRecord:
    """Stores the state of an approval request."""

    def __init__(
        self,
        approval_id: str,
        tool_name: str,
        tool_input: Dict[str, Any],
        agent_id: str,
        workspace_id: str,
        session_id: Optional[str] = None,
        findings: Optional[List[ToolFindingRecord]] = None,
    ) -> None:
        self.approval_id = approval_id
        self.tool_name = tool_name
        self.tool_input = tool_input
        self.agent_id = agent_id
        self.workspace_id = workspace_id
        self.session_id = session_id
        self.status = ApprovalStatus.PENDING
        self.findings = findings or []
        self.created_at = datetime.utcnow()
        self.resolved_at: Optional[datetime] = None
        self.resolved_by: Optional[str] = None
        self.metadata: Dict[str, Any] = {}
        self.pending_request: "ToolApprovalRequest" | None = None

    def to_dict(self) -> Dict[str, Any]:
        return {
            "approval_id": self.approval_id,
            "status": self.status.value,
            "tool_name": self.tool_name,
            "tool_input": self.tool_input,
            "agent_id": self.agent_id,
            "workspace_id": self.workspace_id,
            "session_id": self.session_id,
            "findings": [f.to_dict() for f in self.findings],
            "created_at": self.created_at.isoformat(),
            "resolved_at": self.resolved_at.isoformat() if self.resolved_at else None,
            "resolved_by": self.resolved_by,
        }


class ToolGuardStore:
    """Simple in-memory approval store for development/testing."""

    def __init__(self) -> None:
        self._records: Dict[str, ApprovalRecord] = {}
        self._counter = 0

    def next_id(self) -> str:
        self._counter += 1
        return f"approval_{self._counter:06d}"

    def list(
        self,
        status: ApprovalStatus | None = None,
        workspace_id: Optional[str] = None,
        agent_id: Optional[str] = None,
    ) -> Iterable[ApprovalRecord]:
        for record in self._records.values():
            if status and record.status != status:
                continue
            if workspace_id and record.workspace_id != workspace_id:
                continue
            if agent_id and record.agent_id != agent_id:
                continue
            yield record

    def get(self, approval_id: str) -> Optional[ApprovalRecord]:
        return self._records.get(approval_id)

    def create_pending(
        self,
        tool_name: str,
        tool_input: Dict[str, Any],
        agent_id: str,
        workspace_id: str,
        session_id: Optional[str] = None,
        findings: Optional[List[ToolFindingRecord]] = None,
    ) -> ApprovalRecord:
        record = ApprovalRecord(
            approval_id=self.next_id(),
            tool_name=tool_name,
            tool_input=tool_input,
            agent_id=agent_id,
            workspace_id=workspace_id,
            session_id=session_id,
            findings=findings,
        )
        self._records[record.approval_id] = record
        return record

    def set_status(
        self,
        approval_id: str,
        status: ApprovalStatus,
        resolved_by: Optional[str] = None,
        notify_request: bool = True,
    ) -> ApprovalRecord:
        record = self._records[approval_id]
        if record.status == status:
            return record

        record.status = status
        record.resolved_at = datetime.utcnow()
        record.resolved_by = resolved_by
        if notify_request and record.pending_request:
            if status == ApprovalStatus.APPROVED:
                record.pending_request.approve()
            elif status == ApprovalStatus.DENIED:
                record.pending_request.deny()
        return record

    def cancel(self, approval_id: str) -> None:
        self._records.pop(approval_id, None)


TOOL_GUARD_STORE = ToolGuardStore()


def get_tool_guard_store() -> ToolGuardStore:
    return TOOL_GUARD_STORE


# Default tools that require approval
DEFAULT_GUARDED_TOOLS: Set[str] = {
    "execute_shell_command",
    "write_file",
    "edit_file",
    "place_order",
    "modify_position",
    "delete_file",
}

# Default denied tools (cannot be approved)
DEFAULT_DENIED_TOOLS: Set[str] = {
    "execute_shell_command",  # Shell execution is dangerous
}

# Mark for tool guard denied messages
TOOL_GUARD_DENIED_MARK = "tool_guard_denied"


def default_findings_for_tool(tool_name: str) -> List[ToolFindingRecord]:
    findings: List[ToolFindingRecord] = []
    if tool_name in {"execute_trade", "modify_portfolio"}:
        findings.append(
            ToolFindingRecord(
                severity=SeverityLevel.HIGH,
                message=f"Tool '{tool_name}' touches portfolio state",
            )
        )
    return findings


class ToolApprovalRequest:
    """Represents a pending tool approval request."""

    def __init__(
        self,
        approval_id: str,
        tool_name: str,
        tool_input: Dict[str, Any],
        tool_call_id: str,
        session_id: Optional[str] = None,
    ):
        self.approval_id = approval_id
        self.tool_name = tool_name
        self.tool_input = tool_input
        self.tool_call_id = tool_call_id
        self.session_id = session_id
        self.approved: Optional[bool] = None
        self._event = asyncio.Event()

    async def wait_for_approval(self, timeout: Optional[float] = None) -> bool:
        """Wait for approval decision.

        Args:
            timeout: Maximum time to wait in seconds

        Returns:
            True if approved, False otherwise
        """
        try:
            await asyncio.wait_for(self._event.wait(), timeout=timeout)
        except asyncio.TimeoutError:
            return False
        return self.approved is True

    def approve(self) -> None:
        """Approve this request."""
        self.approved = True
        self._event.set()

    def deny(self) -> None:
        """Deny this request."""
        self.approved = False
        self._event.set()


class ToolGuardMixin:
    """Mixin that adds tool-guard interception to a ReActAgent.

    At runtime this class is combined with ReActAgent via MRO,
    so ``super()._acting`` and ``super()._reasoning`` resolve to
    the concrete agent methods.

    Usage:
        class MyAgent(ToolGuardMixin, ReActAgent):
            def __init__(self, ...):
                super().__init__(...)
                self._init_tool_guard()
    """

    def _init_tool_guard(
        self,
        guarded_tools: Optional[Set[str]] = None,
        denied_tools: Optional[Set[str]] = None,
        approval_timeout: float = 300.0,
    ) -> None:
        """Initialize tool guard.

        Args:
            guarded_tools: Set of tool names requiring approval
            denied_tools: Set of tool names that are always denied
            approval_timeout: Timeout for approval requests in seconds
        """
        self._guarded_tools = guarded_tools or DEFAULT_GUARDED_TOOLS.copy()
        self._denied_tools = denied_tools or DEFAULT_DENIED_TOOLS.copy()
        self._approval_timeout = approval_timeout
        self._pending_approval: Optional[ToolApprovalRequest] = None
        self._approval_callback: Optional[Callable[[ToolApprovalRequest], None]] = None

    def set_approval_callback(
        self,
        callback: Callable[[ToolApprovalRequest], None],
    ) -> None:
        """Set callback for approval requests.

        Args:
            callback: Function called when approval is needed
        """
        self._approval_callback = callback

    def _is_tool_guarded(self, tool_name: str) -> bool:
        """Check if a tool requires approval.

        Args:
            tool_name: Name of the tool

        Returns:
            True if tool requires approval
        """
        return tool_name in self._guarded_tools

    def _is_tool_denied(self, tool_name: str) -> bool:
        """Check if a tool is always denied.

        Args:
            tool_name: Name of the tool

        Returns:
            True if tool is denied
        """
        return tool_name in self._denied_tools

    def _last_tool_response_is_denied(self) -> bool:
        """Check if the last message is a guard-denied tool result."""
        if not hasattr(self, "memory") or not self.memory.content:
            return False

        msg, marks = self.memory.content[-1]
        return TOOL_GUARD_DENIED_MARK in marks and msg.role == "system"

    async def _cleanup_tool_guard_denied_messages(
        self,
        include_denial_response: bool = True,
    ) -> None:
        """Remove tool-guard denied messages from memory.

        Args:
            include_denial_response: Also remove the assistant's denial explanation
        """
        if not hasattr(self, "memory"):
            return

        ids_to_delete: list[str] = []
        last_marked_idx = -1

        for i, (msg, marks) in enumerate(self.memory.content):
            if TOOL_GUARD_DENIED_MARK in marks:
                ids_to_delete.append(msg.id)
                last_marked_idx = i

        if (
            include_denial_response
            and last_marked_idx >= 0
            and last_marked_idx + 1 < len(self.memory.content)
        ):
            next_msg, _ = self.memory.content[last_marked_idx + 1]
            if next_msg.role == "assistant":
                ids_to_delete.append(next_msg.id)

        if ids_to_delete:
            removed = await self.memory.delete(ids_to_delete)
            logger.info("Tool guard: cleaned up %d denied message(s)", removed)

    async def _request_guard_approval(
        self,
        tool_name: str,
        tool_input: Dict[str, Any],
        tool_call_id: str,
    ) -> bool:
        """Request approval for a guarded tool call.

        This method creates a ToolApprovalRequest and waits for
        external approval via approve_guard_call() or deny_guard_call().

        Args:
            tool_name: Name of the tool
            tool_input: Tool input parameters
            tool_call_id: ID of the tool call

        Returns:
            True if approved, False otherwise
        """
        record = TOOL_GUARD_STORE.create_pending(
            tool_name=tool_name,
            tool_input=tool_input,
            agent_id=getattr(self, "agent_id", "unknown"),
            workspace_id=getattr(self, "workspace_id", "default"),
            session_id=getattr(self, "session_id", None),
            findings=default_findings_for_tool(tool_name),
        )

        manager = get_global_runtime_manager()
        if manager:
            manager.register_pending_approval(
                record.approval_id,
                {
                    "tool_name": record.tool_name,
                    "agent_id": record.agent_id,
                    "workspace_id": record.workspace_id,
                    "session_id": record.session_id,
                    "tool_input": record.tool_input,
                },
            )

        self._pending_approval = ToolApprovalRequest(
            approval_id=record.approval_id,
            tool_name=tool_name,
            tool_input=tool_input,
            tool_call_id=tool_call_id,
            session_id=getattr(self, "session_id", None),
        )
        record.pending_request = self._pending_approval

        # Notify via callback if set
        if self._approval_callback:
            self._approval_callback(self._pending_approval)

        # Wait for approval
        approval_request = self._pending_approval
        approved = await approval_request.wait_for_approval(
            timeout=self._approval_timeout
        )

        if approval_request:
            status = (
                ApprovalStatus.APPROVED
                if approval_request.approved is True
                else ApprovalStatus.DENIED
                if approval_request.approved is False
                else ApprovalStatus.EXPIRED
            )
            TOOL_GUARD_STORE.set_status(
                approval_request.approval_id,
                status,
                resolved_by="agent",
                notify_request=False,
            )
            manager = get_global_runtime_manager()
            if manager:
                manager.resolve_pending_approval(
                    approval_request.approval_id,
                    resolved_by="agent",
                    status=status.value,
                )

        self._pending_approval = None
        return approved

    def approve_guard_call(self, request_id: Optional[str] = None) -> bool:
        """Approve a pending guard request.

        This method is called externally to approve a tool call
        that is waiting for approval.

        Args:
            request_id: Optional request ID to verify (not yet implemented)

        Returns:
            True if a request was approved, False if no pending request
        """
        if self._pending_approval is None:
            logger.warning("No pending approval request to approve")
            return False

        TOOL_GUARD_STORE.set_status(
            self._pending_approval.approval_id,
            ApprovalStatus.APPROVED,
            resolved_by="agent",
            notify_request=False,
        )
        manager = get_global_runtime_manager()
        if manager:
            manager.resolve_pending_approval(
                self._pending_approval.approval_id,
                resolved_by="agent",
                status=ApprovalStatus.APPROVED.value,
            )
        self._pending_approval.approve()
        logger.info("Approved tool call: %s", self._pending_approval.tool_name)
        return True

    def deny_guard_call(self, request_id: Optional[str] = None) -> bool:
        """Deny a pending guard request.

        This method is called externally to deny a tool call
        that is waiting for approval.

        Args:
            request_id: Optional request ID to verify (not yet implemented)

        Returns:
            True if a request was denied, False if no pending request
        """
        if self._pending_approval is None:
            logger.warning("No pending approval request to deny")
            return False

        TOOL_GUARD_STORE.set_status(
            self._pending_approval.approval_id,
            ApprovalStatus.DENIED,
            resolved_by="agent",
            notify_request=False,
        )
        manager = get_global_runtime_manager()
        if manager:
            manager.resolve_pending_approval(
                self._pending_approval.approval_id,
                resolved_by="agent",
                status=ApprovalStatus.DENIED.value,
            )
        self._pending_approval.deny()
        logger.info("Denied tool call: %s", self._pending_approval.tool_name)
        return True

    async def _acting(self, tool_call) -> dict | None:
        """Intercept sensitive tool calls before execution.

        1. If tool is in denied_tools, auto-deny unconditionally.
        2. Check for a one-shot pre-approval.
        3. If tool is in the guarded scope, request approval.
        4. Otherwise, delegate to parent _acting.

        Args:
            tool_call: Tool call from the model

        Returns:
            Tool result dict or None
        """
        tool_name: str = tool_call.get("name", "")
        tool_input: dict = tool_call.get("input", {})
        tool_call_id: str = tool_call.get("id", "")

        # Check if tool is denied
        if tool_name and self._is_tool_denied(tool_name):
            logger.warning("Tool '%s' is in the denied set, auto-denying", tool_name)
            return await self._acting_auto_denied(tool_call, tool_name)

        # Check if tool is guarded
        if tool_name and self._is_tool_guarded(tool_name):
            approved = await self._request_guard_approval(
                tool_name=tool_name,
                tool_input=tool_input,
                tool_call_id=tool_call_id,
            )

            if not approved:
                return await self._acting_with_denial(tool_call, tool_name)

        # Call parent _acting
        return await super()._acting(tool_call)  # type: ignore[misc]

    async def _acting_auto_denied(
        self,
        tool_call: Dict[str, Any],
        tool_name: str,
    ) -> dict | None:
        """Auto-deny a tool call without offering approval.

        Args:
            tool_call: Tool call from the model
            tool_name: Name of the denied tool

        Returns:
            Denial result
        """
        from agentscope.message import ToolResultBlock

        denied_text = (
            f"⛔ **Tool Blocked / 工具已拦截**\n\n"
            f"- Tool / 工具: `{tool_name}`\n"
            f"- Reason / 原因: This tool is blocked for security reasons\n\n"
            f"This tool is blocked and cannot be approved.\n"
            f"该工具已被禁止，无法批准执行。"
        )

        tool_res_msg = Msg(
            "system",
            [
                ToolResultBlock(
                    type="tool_result",
                    id=tool_call.get("id", ""),
                    name=tool_name,
                    output=[{"type": "text", "text": denied_text}],
                ),
            ],
            "system",
        )

        await self.print(tool_res_msg, True)
        await self.memory.add(tool_res_msg)
        return None

    async def _acting_with_denial(
        self,
        tool_call: Dict[str, Any],
        tool_name: str,
    ) -> dict | None:
        """Deny the tool call after approval was rejected.

        Args:
            tool_call: Tool call from the model
            tool_name: Name of the tool

        Returns:
            Denial result
        """
        from agentscope.message import ToolResultBlock

        params_text = json.dumps(
            tool_call.get("input", {}),
            ensure_ascii=False,
            indent=2,
        )

        denied_text = (
            f"⚠️ **Tool Call Denied / 工具调用被拒绝**\n\n"
            f"- Tool / 工具: `{tool_name}`\n"
            f"- Parameters / 参数:\n"
            f"```json\n{params_text}\n```\n\n"
            f"The tool call was denied by the user or timed out.\n"
            f"工具调用被用户拒绝或已超时。"
        )

        tool_res_msg = Msg(
            "system",
            [
                ToolResultBlock(
                    type="tool_result",
                    id=tool_call.get("id", ""),
                    name=tool_name,
                    output=[{"type": "text", "text": denied_text}],
                ),
            ],
            "system",
        )

        await self.print(tool_res_msg, True)
        await self.memory.add(tool_res_msg, marks=TOOL_GUARD_DENIED_MARK)
        return None

    async def _reasoning(self, **kwargs) -> Msg:
        """Short-circuit reasoning when awaiting guard approval.

        If the last message was a guard denial, return a waiting message
        instead of continuing reasoning.

        Returns:
            Response message
        """
        if self._last_tool_response_is_denied():
            msg = Msg(
                self.name,
                "⏳ Waiting for approval / 等待审批...\n\n"
                "Type `/approve` to approve, or send any message to deny.\n"
                "输入 `/approve` 批准执行，或发送任意消息拒绝。",
                "assistant",
            )
            await self.print(msg, True)
            await self.memory.add(msg)
            return msg

        return await super()._reasoning(**kwargs)  # type: ignore[misc]


__all__ = [
    "ToolGuardMixin",
    "ToolApprovalRequest",
    "DEFAULT_GUARDED_TOOLS",
    "DEFAULT_DENIED_TOOLS",
    "TOOL_GUARD_DENIED_MARK",
]