MoviePilot/app/agent/middleware/output.py

"""
为当前会话保留有界工具结果，并按字符位置继续读取。
"""

import json
import threading
import time
import uuid
from collections import OrderedDict
from collections.abc import Awaitable, Callable
from dataclasses import dataclass
from typing import Annotated, Any

from langchain.agents.middleware.types import AgentMiddleware, ToolCallRequest
from langchain.tools import ToolRuntime
from langchain_core.messages import ToolMessage
from langchain_core.tools import StructuredTool, create_schema_from_function
from pydantic import Field

from app.agent.code.authority import RPC_RESULT_KEY
from app.agent.policy.contracts import ToolPolicyContext
from app.agent.tools.base import TOOL_RESULT_RECORDER, format_tool_result_for_agent
from app.agent.tools.result import inspect_tool_result

READ_TOOL_RESULT_NAME = "read_tool_result"
MAX_RESULT_BYTES = 1024 * 1024
MAX_TOTAL_BYTES = 4 * MAX_RESULT_BYTES
MAX_RESULTS = 8
RESULT_TTL_SECONDS = 15 * 60


@dataclass(frozen=True)
class _StoredResult:
    """
    内存中短期保留原始授权结果，不复制到日志、记忆或磁盘。

    :param thread_id: 允许读取该结果的会话线程标识
    :param text: 保存供后续分段读取的完整结果文本
    :param byte_size: 文本 UTF-8 编码后的字节数，用于限制存储容量
    :param tool_name: 产生结果的工具名称
    :param expires_at: 该缓存结果到期的单调时钟秒数
    :param requires_admin: 读取该结果是否要求管理员身份
    """

    thread_id: str
    text: str
    byte_size: int
    tool_name: str
    expires_at: float
    requires_admin: bool


class ToolOutputMiddleware(AgentMiddleware):
    """
    包装当前图的工具结果；读取按会话与当前管理员身份双重隔离。
    """

    def __init__(self, context: ToolPolicyContext) -> None:
        """
        注册内部续读工具，缓存容量和寿命均受固定上限约束。
        """
        self.context = context
        self._results: OrderedDict[str, _StoredResult] = OrderedDict()
        self._lock = threading.RLock()
        self.tools = [StructuredTool.from_function(
            name=READ_TOOL_RESULT_NAME,
            description=(
                "Read the next page of an archived tool result using its result_id and next_offset. "
                "Offsets count Unicode characters. Results expire after 15 minutes, on eviction, "
                "or when this conversation graph is rebuilt. Never guess result IDs."
            ),
            coroutine=self._read_result,
            args_schema=create_schema_from_function(
                "ReadToolResultInput", self._read_result, filter_args=["runtime"],
            ),
            tags=["agent_tool", "read"],
        )]
        object.__setattr__(self.tools[0], "_agent_tool_source", "middleware:output")

    @staticmethod
    def _thread_id(runtime: ToolRuntime) -> str:
        """
        使用图的真实线程身份，避免同一中间件被不同图线程误复用。
        """
        return str(runtime.config.get("configurable", {}).get("thread_id") or "")

    def _prune(self) -> None:
        """
        调用方持锁时删除到期结果，并按写入顺序执行 UTF-8 字节容量淘汰。
        """
        now = time.monotonic()
        for result_id, result in list(self._results.items()):
            if result.expires_at <= now:
                del self._results[result_id]
        size = sum(result.byte_size for result in self._results.values())
        while self._results and (len(self._results) > MAX_RESULTS or size > MAX_TOTAL_BYTES):
            _, result = self._results.popitem(last=False)
            size -= result.byte_size

    def _store(self, thread_id: str, tool_name: str, text: str, requires_admin: bool) -> dict[str, Any]:
        """
        只保存完整的有界结果，超过容量时明确说明无法续读。
        """
        if not thread_id:
            return {"result_unavailable": "thread_unavailable"}
        # 先用字符数下界拒绝超大文本，避免为明显无法归档的输出再分配编码副本。
        if len(text) > MAX_RESULT_BYTES or (byte_size := len(text.encode("utf-8"))) > MAX_RESULT_BYTES:
            return {"result_unavailable": "result_too_large", "result_limit_bytes": MAX_RESULT_BYTES}
        result_id = uuid.uuid4().hex
        with self._lock:
            self._results[result_id] = _StoredResult(
                thread_id=thread_id, text=text, byte_size=byte_size, tool_name=tool_name,
                expires_at=time.monotonic() + RESULT_TTL_SECONDS,
                requires_admin=requires_admin,
            )
            self._prune()
        return {
            "result_id": result_id,
            "read_tool": READ_TOOL_RESULT_NAME,
            "expires_in_seconds": RESULT_TTL_SECONDS,
            "offset_unit": "unicode_characters",
        }

    async def _read_result(
        self,
        result_id: Annotated[str, Field(min_length=32, max_length=32)],
        runtime: ToolRuntime,
        offset: Annotated[int, Field(ge=0)] = 0,
        limit: Annotated[int, Field(ge=1, le=16000)] = 4000,
    ) -> str:
        """
        按显式游标续读；过期、越权和未知 ID 使用同一不可用响应。
        """
        with self._lock:
            self._prune()
            result = self._results.get(result_id)
        if (
            not self._thread_id(runtime) or result is None or result.thread_id != self._thread_id(runtime)
            or (result.requires_admin and not self.context.agent_context.get("is_admin"))
        ):
            return json.dumps({"success": False, "error": "result_unavailable"})
        if offset > len(result.text):
            return json.dumps({"success": False, "error": "offset_out_of_range", "total_chars": len(result.text)})
        end = min(offset + limit, len(result.text))
        return json.dumps({
            "success": True, "result_id": result_id, "tool_name": result.tool_name,
            "offset": offset, "next_offset": end if end < len(result.text) else None,
            "total_chars": len(result.text), "content": result.text[offset:end],
        }, ensure_ascii=False)

    async def awrap_tool_call(
        self, request: ToolCallRequest, handler: Callable[[ToolCallRequest], Awaitable[Any]],
    ) -> Any:
        """
        让内置工具在截断前归档，并覆盖返回大文本的外部 MCP 工具。
        """
        if request.tool_call.get("name") == READ_TOOL_RESULT_NAME:
            return await handler(request)
        if request.tool_call.get('name') == 'execute_code' and getattr(request.tool, '_agent_tool_source', None) == 'builtin':
            return await handler(request)
        # 历史工具已经按消息/窗口预算裁剪；再次套通用预览会丢失锚点与书挡结构。
        if request.tool_call.get("name") == "session_search" and getattr(request.tool, "_agent_tool_source", None) == "middleware:recall":
            return await handler(request)
        # 学习工具已限制单文件、条目及批次；预览截断会破坏 JSON 和写前读取证据。
        if (request.tool_call.get('name') in {'memory', 'skill_manage', 'skill_view', 'skills_list'}
                and getattr(getattr(request, 'tool', None), '_agent_tool_source', None) == 'middleware:learning'):
            return await handler(request)
        thread_id = self._thread_id(request.runtime)
        requires_admin = bool(self.context.agent_context.get("is_admin"))
        def record(name: str, text: str) -> dict[str, Any]:
            """
            Python 聚合读取完整结果，同时保持普通工具的有界展示和归档引用。
            """
            state = getattr(request, 'state', None)
            raw = state.get(RPC_RESULT_KEY) if isinstance(state, dict) else None
            if isinstance(raw, dict):
                raw[name] = text
            return self._store(thread_id, name, text, requires_admin)

        token = TOOL_RESULT_RECORDER.set(record)
        try:
            result = await handler(request)
            if isinstance(result, ToolMessage) and isinstance(result.content, str):
                content = format_tool_result_for_agent(result.content, tool_name=result.name)
                if content != result.content:
                    payload = json.loads(content)
                    payload["execution_outcome"] = inspect_tool_result(result).value
                    content = json.dumps(payload, ensure_ascii=False, indent=2)
                return result.model_copy(update={"content": content})
            return result
        finally:
            TOOL_RESULT_RECORDER.reset(token)