forked from xiaohei/taiji-AI-PAD
chenchenandClaude Opus 4.7
1461051755
Two small follow-ups to the register hardening + Heicode P4 work:
1. heicode_client.list_user_models: the path `/api/user/{id}/models`
prescribed in §7.11.2 returns 404 `Invalid URL` on the live Heicode
NewAPI — that path is not registered on their router. Switched to
`/api/user/models` (no path segment), which Heicode binds to the
`New-Api-User: 26` admin header. End-to-end P4 smoke now 4/4 with
user 55@55.com (id=2 on Heicode): /balance /models /usage /logs.
Future: if Heicode ships an "admin-replaces-user" path, switch back
and pass the actual heicode_user_id.
2. routes/auth.register: previously line-744 SELECT only checked
req.username, but line 778 falls back to email.split("@")[0] when
blank — so two users registering with alice@foo.com and alice@bar.com
would both clear the predcheck, then the second would IntegrityError
on flush. Now predcheck uses `effective_username` matching what'll
actually be inserted.
Also append §7.15 to Heicode-对接进度与待办.md:
- 4-item agent-manager / Vault / Workload-Identity audit results
- §7.13 token rotation acknowledgement
- P4 end-to-end first-pass results
- This-session internal security hardening summary
Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
278 lines
11 KiB
Python
278 lines
11 KiB
Python
"""
|
||
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/models — 当前 token 持有人视角的可用模型列表。
|
||
|
||
修订历史：
|
||
- 原 `/api/models` 是 channel-dashboard 视角（key=channelId），
|
||
admin user 没渠道返回空，废弃。
|
||
- §7.11.2 曾改用 `/api/user/{id}/models`（admin 替指定用户查），
|
||
但 2026-05-12 实测 Heicode NewAPI 返回 404 `Invalid URL`，
|
||
该 path 在他们的 router 上根本没注册。
|
||
- 现切换到 `/api/user/models`（不带 user_id 段）。NewAPI 用
|
||
`New-Api-User` header（必须 == admin token 持有人 id，否则
|
||
CSRF check 拒绝）确定视角，返回 admin 视角的全量 model 列表。
|
||
后续 Heicode 那边如果上线"用户视角"接口，可以改回 path 段方案。
|
||
|
||
注意：当前实现下，所有 mcp-server 透传的用户拿到的是 admin（user 26）
|
||
视角的 models，即 NewAPI 全量。如果 Heicode 引入 group/channel 级
|
||
模型过滤，需要 Heicode 团队提供 `admin-replaces-user` 路由我方再切。
|
||
"""
|
||
# heicode_user_id 当前未直接使用 —— 仅做调试日志参考，将来如果 Heicode
|
||
# 上线 admin-replaces-user 路径会用到。保留参数避免上游 caller 改签名。
|
||
_ = heicode_user_id
|
||
body = await self._get(
|
||
"/api/user/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
|