yuxuanhui
128 lines
8.5 KiB
Python
128 lines
8.5 KiB
Python
"""MCP transport over shared research operations, with minimal durable audit evidence."""
|
|
|
|
import asyncio
|
|
import json
|
|
import time
|
|
from uuid import uuid4
|
|
|
|
import anyio
|
|
from fastapi import HTTPException
|
|
from mcp import types
|
|
from mcp.server.lowlevel import Server
|
|
from mcp.server.transport_security import TransportSecuritySettings
|
|
from pydantic import ValidationError
|
|
|
|
from ..alphas import sanitize
|
|
from ..backtests.contracts import fingerprint
|
|
from ..models import MCPAudit, now
|
|
from ..research.serialization import encode_snapshot
|
|
from ..research_access import contracts as c
|
|
from ..research_access.service import ResearchAccess, ResearchError
|
|
|
|
# Name, schema, business method, required scope, description. No generic arbitrary HTTP tool.
|
|
TOOLS = {
|
|
"get_research_capabilities": (c.Empty, "capabilities", "research:read", "读取直接研究能力、完整设置 schema 和调度阻塞，不代表平台剩余额度。"),
|
|
"search_catalog": (c.CatalogSearch, "catalog", "research:read", "分页查询指定范围的数据集或字段元数据；无缓存不等于无数据，不隐式刷新。"),
|
|
"get_research_metadata": (c.Metadata, "metadata", "research:read", "读取范围、设置快照、算子定义或字段可用性；未知不认定通过。"),
|
|
"refresh_research_data": (c.Refresh, "refresh", "research:refresh", "显式刷新目录、算子、设置、字段可用性或 PnL；不会创建模拟。任务返回 job_id。"),
|
|
"get_refresh_job": (c.JobReference, "refresh_job", "research:read", "查询研究刷新任务的状态和产物引用。"),
|
|
"search_backtests": (c.History, "history", "research:read", "分页查历史候选与固定设置；candidates 按完整输入精确匹配，不推断数学等价。"),
|
|
"submit_backtests": (c.Submit, "submit", "backtests:execute", "执行用户已授权的固定批次，自动留痕并立即返回运行 ID。每项必须完整设置；重复默认拒绝，rerun 明确重跑。不需要研究资产。"),
|
|
"get_backtest": (c.RunReference, "run", "research:read", "读取真实运行进度、提交数量和可选增量事件；受理不等于成功。"),
|
|
"get_backtest_results": (c.Results, "results", "research:read", "分页读取固定快照指标、全部非通过检查及三层状态；缺失指标不补零。"),
|
|
"get_backtest_artifact": (c.Artifact, "artifact", "research:read", "分页读取候选脱敏快照的顶层键值或独立采集的 PnL；缺缓存不自动刷新。"),
|
|
"control_backtest": (c.Control, "control", "backtests:control", "对已授权运行暂停、继续、停止或恢复采集；不远程取消、不重提未知模拟。需要版本和幂等键。"),
|
|
}
|
|
|
|
|
|
def tool_result(data, error=False):
|
|
data = encode_snapshot(data)
|
|
return types.CallToolResult(content=[types.TextContent(type="text", text=json.dumps(data, ensure_ascii=False))],
|
|
structuredContent=data, isError=error)
|
|
|
|
|
|
class MCPResearchServer:
|
|
def __init__(self, sessions, runner, settings):
|
|
self.sessions, self.runner, self.settings = sessions, runner, settings
|
|
# The existing deployment has one owner; this also gives SQLite test transactions a fair queue.
|
|
self.mutation_lock = asyncio.Lock()
|
|
self.server = Server("wq-alpha-research", version="1.0.0", on_list_tools=self.list_tools,
|
|
on_call_tool=self.call_tool,
|
|
instructions="自由探索，直接固定候选回测，无需先建研究资产。工具不安排定时研究；结果按运行 ID 查询。")
|
|
from urllib.parse import urlsplit
|
|
|
|
host = urlsplit(settings.public_origin).netloc
|
|
self.app = self.server.streamable_http_app(
|
|
streamable_http_path="/", stateless_http=True, json_response=True,
|
|
transport_security=TransportSecuritySettings(enable_dns_rebinding_protection=True,
|
|
allowed_hosts=[host], allowed_origins=[settings.public_origin.rstrip("/")]),
|
|
)
|
|
|
|
async def list_tools(self, ctx, params):
|
|
principal = ctx.request.state.mcp_principal
|
|
return types.ListToolsResult(tools=[types.Tool(name=name, description=description,
|
|
inputSchema=schema.model_json_schema(), annotations=types.ToolAnnotations(
|
|
readOnlyHint=scope == "research:read", destructiveHint=method == "control",
|
|
idempotentHint=method in {"submit", "control"} or scope == "research:read",
|
|
openWorldHint=method in {"refresh", "submit", "metadata"}))
|
|
for name, (schema, method, scope, description) in TOOLS.items()
|
|
if scope in principal.scopes and "research:read" in principal.scopes])
|
|
|
|
async def call_tool(self, ctx, params):
|
|
principal = ctx.request.state.mcp_principal
|
|
return await self.invoke(principal, params.name, params.arguments or {}, str(ctx.request_id or uuid4()))
|
|
|
|
async def invoke(self, principal, name, arguments, request_id=None):
|
|
"""Invoke with a server-authenticated principal; atomic success audit and post-commit wake."""
|
|
started = time.monotonic()
|
|
request_id = request_id or str(uuid4())
|
|
entry = TOOLS.get(name)
|
|
if not entry:
|
|
return tool_result({"error": ResearchError("UNKNOWN_TOOL", "工具不存在").data}, True)
|
|
schema, method, scope, _ = entry
|
|
if "research:read" not in principal.scopes or scope not in principal.scopes:
|
|
raise HTTPException(403, "MCP 令牌缺少所需权限")
|
|
digest = fingerprint(arguments)
|
|
async with self.mutation_lock:
|
|
# Disconnect does not roll back an already accepted operation or lose its wake-up.
|
|
with anyio.CancelScope(shield=True):
|
|
async with self.sessions.begin() as db:
|
|
access = ResearchAccess(db, principal, self.runner.client, self.settings.public_origin)
|
|
code, error = "OK", False
|
|
try:
|
|
async with db.begin_nested():
|
|
args = schema.model_validate(arguments)
|
|
async with asyncio.timeout(30 if method in {"refresh", "metadata"} else None):
|
|
data = encode_snapshot(await getattr(access, method)(args))
|
|
data.setdefault("_meta", {"schema_version": 1, "observed_at": now().isoformat(),
|
|
"nulls": "null 表示来源未提供，不等于零", "source": "system"})
|
|
except ValidationError as exc:
|
|
code, error = "INVALID_INPUT", True
|
|
data = {"error": ResearchError(code, "; ".join(
|
|
f"{'.'.join(map(str, e['loc']))}: {e['msg']}" for e in exc.errors())).data}
|
|
except TimeoutError:
|
|
code, error = "UPSTREAM_TIMEOUT", True
|
|
data = {"error": ResearchError(code, "元数据读取或刷新超时，未发布新快照", retryable=True).data}
|
|
except ResearchError as exc:
|
|
code, error, data = exc.data["code"], True, {"error": exc.data}
|
|
except HTTPException as exc:
|
|
code = {404: "NOT_FOUND", 409: "CONFLICT", 422: "INVALID_INPUT", 429: "RATE_LIMITED", 502: "UPSTREAM_ERROR"}.get(exc.status_code, "REQUEST_FAILED")
|
|
error = True
|
|
data = {"error": ResearchError(code, str(sanitize(exc.detail)),
|
|
retryable=exc.status_code in {429, 502, 503},
|
|
retry_after=(exc.headers or {}).get("Retry-After")).data}
|
|
except Exception:
|
|
# Never expose SQL parameters, exception reprs or credentials in unexpected errors.
|
|
code, error = "INTERNAL_ERROR", True
|
|
data = {"error": ResearchError(code, "研究操作失败；可使用原幂等键重试或查询历史", retryable=True).data}
|
|
db.add(MCPAudit(id=str(uuid4()), token_id=principal.token_id, tool=name,
|
|
request_id=fingerprint({"request_id": request_id}), input_digest=digest,
|
|
business_id=data.get("backtest_run_id", data.get("job_id")),
|
|
result_code=code, elapsed_ms=int((time.monotonic()-started)*1000)))
|
|
if not error:
|
|
if access.wake == "backtests":
|
|
self.runner.backtests.wake.set()
|
|
elif access.wake == "jobs":
|
|
self.runner.wake.set()
|
|
return tool_result(data, error)
|