taiji-AI-PAD/services/mcp-server/app/heicode_client.py

"""
Heicode NewAPI 出站客户端（P4）

mcp-server 通过 admin service token 调用 Heicode NewAPI（`code.xinghanlab.com`），
拉取用户视角的元数据（余额 / 模型 / 用量 / 调用日志），用于 Manager 控制台展示。

关键约束（NewAPI middleware/auth.go 强制）：
- `Authorization: Bearer <admin_access_token>` — admin 身份
- `New-Api-User: <admin_user_id>` — **必须等于** access token 对应的用户 id（CSRF 检查）

→ 这意味着 admin token **不能"切身份"调 user-self 路由**（`/api/user/self*`），
  必须用 admin 路由 + `?user_id=X` query 参数定位目标用户。

→ "把当前 mcp-server 用户的 email 解析成 heicode 本地 user_id" 通过 admin search 实现：
  `GET /api/user/search?keyword=<email>` → cache 结果。

依据：
- Docs/Heicode-对接进度与待办.md §2.3 + §2.3.1（Heicode 团队 2026-05-07 决策 B 方案）
- Docs/Heicode-完整调用流程图.md §2.5 4 路径架构
"""
from __future__ import annotations

import asyncio
from typing import Any, Dict, List, Optional, Tuple

import httpx
import structlog

from config import settings


logger = structlog.get_logger(__name__)


class HeicodeNewAPIError(Exception):
    """Heicode NewAPI 调用失败"""
    def __init__(self, message: str, status_code: int = 0, body: Any = None):
        super().__init__(message)
        self.status_code = status_code
        self.body = body


class HeicodeNewAPIClient:
    """Heicode NewAPI HTTP 客户端（admin token-based）"""

    def __init__(
        self,
        base_url: Optional[str] = None,
        admin_token: Optional[str] = None,
        admin_user_id: Optional[str] = None,
        timeout: Optional[int] = None,
    ):
        self.base_url = (base_url or settings.heicode_newapi_base_url).rstrip("/")
        self.admin_token = admin_token or settings.heicode_newapi_admin_token
        self.admin_user_id = admin_user_id or settings.heicode_newapi_admin_user_id
        self.timeout = timeout or settings.heicode_newapi_timeout

        # 简单内存缓存：email → heicode local user_id（int）
        # TTL 30 分钟，命中即用
        self._user_id_cache: Dict[str, Tuple[int, float]] = {}
        self._cache_ttl_sec = 1800
        self._cache_lock = asyncio.Lock()

    def is_configured(self) -> bool:
        """判断是否已配好可用的 admin token + user_id"""
        return bool(self.admin_token and self.admin_user_id)

    def _admin_headers(self, request_id: Optional[str] = None) -> Dict[str, str]:
        """构造 admin 调用的 headers（含 Authorization + New-Api-User）"""
        if not self.is_configured():
            raise HeicodeNewAPIError(
                "Heicode NewAPI admin token / admin_user_id 未配置 "
                "(设置 HEICODE_NEWAPI_SERVICE_TOKEN + HEICODE_NEWAPI_ADMIN_USER_ID)",
                status_code=503,
            )
        h = {
            "Authorization": f"Bearer {self.admin_token}",
            "New-Api-User": str(self.admin_user_id),
            "Content-Type": "application/json",
            "Accept": "application/json",
        }
        if request_id:
            h["X-Request-Id"] = request_id
        return h

    async def _get(
        self,
        path: str,
        params: Optional[Dict[str, Any]] = None,
        request_id: Optional[str] = None,
    ) -> Dict[str, Any]:
        """统一 GET 调用"""
        url = f"{self.base_url}{path}"
        headers = self._admin_headers(request_id)
        async with httpx.AsyncClient(timeout=self.timeout) as client:
            try:
                resp = await client.get(url, headers=headers, params=params)
            except httpx.HTTPError as e:
                logger.error(
                    "heicode_newapi_request_failed",
                    method="GET", url=url, error=str(e),
                )
                raise HeicodeNewAPIError(f"请求 Heicode NewAPI 失败: {e}", status_code=502)

        try:
            body = resp.json()
        except Exception:
            body = {"raw": resp.text[:500]}

        if resp.status_code >= 400:
            logger.warning(
                "heicode_newapi_non_2xx",
                method="GET", url=url, status=resp.status_code,
                body_preview=str(body)[:200],
            )
            raise HeicodeNewAPIError(
                f"Heicode NewAPI HTTP {resp.status_code}: {body}",
                status_code=resp.status_code,
                body=body,
            )

        # NewAPI 响应惯例：{"success": bool, "message": str, "data": ...}
        if isinstance(body, dict) and body.get("success") is False:
            raise HeicodeNewAPIError(
                f"Heicode NewAPI business error: {body.get('message', '')}",
                status_code=resp.status_code,
                body=body,
            )

        return body if isinstance(body, dict) else {"raw": body}

    # ==================== email → heicode user_id 解析 ====================

    async def resolve_user_id_by_email(self, email: str) -> Optional[int]:
        """通过 admin search 把 email 解析成 heicode 本地 user_id（int）

        命中缓存 → 直接返回；否则调 admin /api/user/search?keyword=<email> 查询。
        """
        import time
        if not email:
            return None
        email_key = email.strip().lower()
        now = time.time()

        async with self._cache_lock:
            cached = self._user_id_cache.get(email_key)
            if cached and (now - cached[1]) < self._cache_ttl_sec:
                return cached[0]

        try:
            body = await self._get(
                "/api/user/search",
                params={"keyword": email_key, "group": ""},
            )
        except HeicodeNewAPIError as e:
            logger.warning(
                "heicode_user_resolve_failed",
                email=email_key, error=str(e),
            )
            return None

        # NewAPI 返回结构通常为 {"data": {"items": [...]}} 或 {"data": [...]}
        data = body.get("data") if isinstance(body, dict) else None
        items: List[Dict[str, Any]] = []
        if isinstance(data, list):
            items = data
        elif isinstance(data, dict):
            items = data.get("items") or data.get("users") or []

        # 精确匹配 email
        for u in items:
            if not isinstance(u, dict):
                continue
            if (u.get("email") or "").strip().lower() == email_key:
                uid = u.get("id")
                if isinstance(uid, int) and uid > 0:
                    async with self._cache_lock:
                        self._user_id_cache[email_key] = (uid, now)
                    return uid

        # 找不到精确匹配
        return None

    # ==================== 用户视角元数据查询（admin 路由 + user_id 参数） ====================

    async def get_user_info(self, heicode_user_id: int, request_id: Optional[str] = None) -> Dict[str, Any]:
        """GET /api/user/{id} — 用户详情（含 quota/group/status/...）"""
        body = await self._get(f"/api/user/{heicode_user_id}", request_id=request_id)
        return body.get("data") if isinstance(body, dict) else body

    async def list_user_models(
        self, heicode_user_id: int,
        request_id: Optional[str] = None,
    ) -> List[Any]:
        """GET /api/user/{id}/models — 按用户 id 列出该用户可用模型清单。

        修订（2026-05-08，按 Heicode §7.11.2）：原 `/api/models` 是渠道仪表盘
        视角（key 是 channelId），mcp-server-service 没渠道，自然空。改用
        `/api/user/{id}/models`（admin 替指定用户查），返回 string[]。
        """
        body = await self._get(
            f"/api/user/{heicode_user_id}/models",
            request_id=request_id,
        )
        data = body.get("data") if isinstance(body, dict) else body
        if isinstance(data, list):
            return data
        if isinstance(data, dict):
            return data.get("items") or data.get("models") or []
        return []

    async def get_user_quota_dates(
        self, heicode_user_id: int,
        days: int = 30,
        request_id: Optional[str] = None,
    ) -> List[Dict[str, Any]]:
        """GET /api/data/users — 按用户 id 拉取最近 N 天的用量"""
        body = await self._get(
            "/api/data/users",
            params={"user_id": heicode_user_id, "default_time": str(days)},
            request_id=request_id,
        )
        data = body.get("data") if isinstance(body, dict) else None
        if isinstance(data, list):
            return data
        if isinstance(data, dict):
            return data.get("items") or []
        return []

    async def get_user_logs(
        self, heicode_user_id: int,
        page: int = 1, page_size: int = 50,
        request_id: Optional[str] = None,
    ) -> List[Dict[str, Any]]:
        """GET /api/log/?user_id=X — 用户调用日志（admin 路由）"""
        body = await self._get(
            "/api/log/",
            params={
                "p": page,
                "page_size": page_size,
                "user_id": heicode_user_id,
                "type": 0,  # 0=全部
            },
            request_id=request_id,
        )
        data = body.get("data") if isinstance(body, dict) else None
        if isinstance(data, list):
            return data
        if isinstance(data, dict):
            return data.get("items") or []
        return []


# ==================== 单例 ====================

_default_client: Optional[HeicodeNewAPIClient] = None


def get_heicode_client() -> HeicodeNewAPIClient:
    global _default_client
    if _default_client is None:
        _default_client = HeicodeNewAPIClient()
    return _default_client