@qingflow-tech/qingflow-cli 1.1.61-beta.1 → 1.1.61-beta.3

package/src/qingflow_mcp/tool_policy.py

@@ -0,0 +1,197 @@
+from __future__ import annotations
+
+import json
+import os
+from dataclasses import dataclass
+from typing import Any, Literal, Mapping, Sequence
+from urllib.parse import parse_qs
+
+from mcp.server.fastmcp import FastMCP
+from mcp.server.fastmcp.exceptions import ToolError
+
+from .cloud_context import current_http_tool_policy, current_tool_policy_header
+from .tool_catalog import TOOL_ID_PATTERN, require_registered_tool_id, tool_is_exempt
+
+
+TOOL_POLICY_ENV = "QINGFLOW_TOOL_POLICY"
+TOOL_POLICY_HEADER = "X-Qingflow-Tool-Policy"
+TOOL_POLICY_INVALID = "TOOL_POLICY_INVALID"
+TOOL_POLICY_DENIED = "TOOL_POLICY_DENIED"
+
+
+@dataclass(frozen=True, slots=True)
+class ToolPolicySnapshot:
+    mode: Literal["unrestricted", "restricted", "invalid"]
+    allowed_tools: tuple[str, ...] = ()
+    error: str | None = None
+
+    def violation_for(self, stable_id: str) -> "ToolPolicyViolation | None":
+        # 基础恢复能力在非法策略下也必须可用。
+        if tool_is_exempt(stable_id):
+            return None
+        if self.mode == "unrestricted":
+            return None
+        if self.mode == "invalid":
+            return ToolPolicyViolation(TOOL_POLICY_INVALID, stable_id, self.error or "工具策略无效")
+        if stable_id in self.allowed_tools:
+            return None
+        return ToolPolicyViolation(TOOL_POLICY_DENIED, stable_id, "当前策略未允许该工具")
+
+
+class ToolPolicyViolation(Exception):
+    def __init__(self, code: str, tool_id_value: str, message: str) -> None:
+        super().__init__(message)
+        self.code = code
+        self.tool_id = tool_id_value
+        self.message = message
+
+    def payload(self, *, camel_case: bool = False) -> dict[str, Any]:
+        if camel_case:
+            return {
+                "status": "failed",
+                "errorCode": self.code,
+                "message": self.message,
+                "toolId": self.tool_id,
+            }
+        return {
+            "status": "failed",
+            "error_code": self.code,
+            "message": self.message,
+            "toolId": self.tool_id,
+        }
+
+
+def parse_tool_policy(value: str | Mapping[str, Any] | None) -> ToolPolicySnapshot:
+    if value is None:
+        return ToolPolicySnapshot("unrestricted")
+    payload: Any = value
+    if isinstance(value, str):
+        if value == "":
+            return ToolPolicySnapshot("unrestricted")
+        try:
+            payload = json.loads(value)
+        except json.JSONDecodeError:
+            return ToolPolicySnapshot("invalid", error="工具策略不是合法 JSON")
+    if not isinstance(payload, Mapping):
+        return ToolPolicySnapshot("invalid", error="工具策略必须是对象")
+    if set(payload) != {"version", "allowed_tools"}:
+        return ToolPolicySnapshot("invalid", error="工具策略字段必须为 version 和 allowed_tools")
+    if type(payload.get("version")) is not int or payload.get("version") != 1:  # noqa: E721
+        return ToolPolicySnapshot("invalid", error="工具策略 version 必须为 1")
+    allowed = payload.get("allowed_tools")
+    if not isinstance(allowed, list):
+        return ToolPolicySnapshot("invalid", error="allowed_tools 必须是数组")
+    normalized: set[str] = set()
+    for stable_id in allowed:
+        if not isinstance(stable_id, str) or TOOL_ID_PATTERN.fullmatch(stable_id) is None:
+            return ToolPolicySnapshot("invalid", error="allowed_tools 包含非法工具 ID")
+        normalized.add(stable_id)
+    return ToolPolicySnapshot("restricted", tuple(sorted(normalized)))
+
+
+def parse_http_tool_policy(header: str | None, query_string: bytes) -> ToolPolicySnapshot:
+    header_policy = parse_tool_policy(header)
+    # 保留空值并在 URL 解码后检查重复，tools= 与缺省不能合并为同一种状态。
+    values = parse_qs(query_string.decode("latin1"), keep_blank_values=True).get("tools", [])
+    if len(values) > 1:
+        return ToolPolicySnapshot("invalid", error="tools 查询参数不能重复")
+    url_policy = parse_tool_policy(
+        {"version": 1, "allowed_tools": values[0].split(",") if values[0] else []}
+        if values else None
+    )
+    if header_policy.mode == "invalid":
+        return header_policy
+    if url_policy.mode == "invalid":
+        return ToolPolicySnapshot("invalid", error=f"tools 查询参数无效: {url_policy.error}")
+    if header_policy.mode == "unrestricted":
+        return url_policy
+    if url_policy.mode == "unrestricted":
+        return header_policy
+    return ToolPolicySnapshot(
+        "restricted", tuple(sorted(set(header_policy.allowed_tools).intersection(url_policy.allowed_tools)))
+    )
+
+
+class PolicyAwareFastMCP(FastMCP):
+    """让 MCP 列表与调用在 FastMCP handler 边界共用策略快照。"""
+
+    def __init__(self, *args: Any, **kwargs: Any) -> None:
+        self._tool_policy_bindings: dict[str, str] | None = None
+        self._tool_policy_source: Literal["environment", "request"] = "environment"
+        self._tool_policy_environment_snapshot = ToolPolicySnapshot("unrestricted")
+        super().__init__(*args, **kwargs)
+
+    def configure_tool_policy(
+        self,
+        bindings: Mapping[str, str],
+        *,
+        source: Literal["environment", "request"],
+    ) -> None:
+        registered = set(self._tool_manager._tools)
+        bound = set(bindings)
+        if registered != bound:
+            missing = sorted(registered - bound)
+            extra = sorted(bound - registered)
+            raise ValueError(f"工具策略绑定不完整: missing={missing}, extra={extra}")
+        for stable_id in bindings.values():
+            # 登记校验只依赖提供方稳定 ID，不依赖 Wingent-X 的产品分类。
+            require_registered_tool_id(stable_id)
+        self._tool_policy_bindings = dict(bindings)
+        self._tool_policy_source = source
+        if source == "environment":
+            self._tool_policy_environment_snapshot = parse_tool_policy(os.getenv(TOOL_POLICY_ENV))
+
+    @property
+    def tool_policy_bindings(self) -> dict[str, str]:
+        return dict(self._tool_policy_bindings or {})
+
+    async def list_tools(self):  # type: ignore[no-untyped-def]
+        tools = await super().list_tools()
+        bindings = self._require_bindings()
+        snapshot = self._current_snapshot()
+        return [
+            tool
+            for tool in tools
+            if snapshot.violation_for(bindings[tool.name]) is None
+        ]
+
+    async def call_tool(self, name: str, arguments: dict[str, Any]) -> Sequence[Any] | dict[str, Any]:
+        bindings = self._require_bindings()
+        stable_id = bindings.get(name)
+        if stable_id is None:
+            raise ToolError(json.dumps(
+                {
+                    "status": "failed",
+                    "errorCode": TOOL_POLICY_INVALID,
+                    "message": "工具未登记稳定 ID",
+                },
+                ensure_ascii=False,
+            ))
+        violation = self._current_snapshot().violation_for(stable_id)
+        if violation is not None:
+            raise ToolError(json.dumps(violation.payload(camel_case=True), ensure_ascii=False))
+        return await super().call_tool(name, arguments)
+
+    def _require_bindings(self) -> dict[str, str]:
+        if self._tool_policy_bindings is None:
+            raise RuntimeError("MCP 工具策略绑定尚未配置")
+        return self._tool_policy_bindings
+
+    def _current_snapshot(self) -> ToolPolicySnapshot:
+        if self._tool_policy_source == "request":
+            snapshot = current_http_tool_policy.get()
+            if snapshot is not None:
+                return snapshot
+            # 保留未经过 HTTP 入口的请求头上下文用法，绝不回退共享 env。
+            return parse_tool_policy(current_tool_policy_header.get())
+        return self._tool_policy_environment_snapshot
+
+
+def configure_tool_policy(
+    server: PolicyAwareFastMCP,
+    bindings: Mapping[str, str],
+    *,
+    source: Literal["environment", "request"],
+) -> PolicyAwareFastMCP:
+    server.configure_tool_policy(bindings, source=source)
+    return server