feat: add MCP research access and browser key management

This commit is contained in:
yuxuanhui
2026-09-09 16:18:27 +08:00
parent 4debca7dbd
commit 45238280e3
47 changed files with 2642 additions and 44 deletions
+50 -6
View File
@@ -4,11 +4,11 @@ import argparse
import asyncio
import getpass
from sqlalchemy import delete
from sqlalchemy import delete, update
from .config import Settings
from .db import create_database
from .models import Admin, LoginSession
from .models import Admin, LoginSession, MCPToken, now
from .security import password_hasher
@@ -21,13 +21,57 @@ async def reset_password():
admin = await db.get(Admin, 1)
admin.password_hash = password_hasher.hash(password)
await db.execute(delete(LoginSession))
await db.execute(update(MCPToken).where(MCPToken.revoked_at.is_(None)).values(revoked_at=now()))
await db.commit()
await engine.dispose()
print("Admin password updated; all system sessions revoked.")
print("Admin password updated; all system sessions and MCP tokens revoked.")
async def token_command(args):
import json
from sqlalchemy import select
from .mcp_api.auth import create_token
from .models import MCPToken, now
from .research.serialization import encode_snapshot
engine, sessions = create_database(Settings().database_url)
try:
async with sessions.begin() as db:
if args.command == "mcp-token-create":
row, secret = await create_token(db, args.name, args.scope, args.days)
result = {"id": row.id, "name": row.name, "scopes": row.scopes,
"expires_at": row.expires_at, "token": secret}
elif args.command == "mcp-token-revoke":
row = await db.get(MCPToken, args.token_id)
if not row:
raise ValueError("令牌不存在")
row.revoked_at = row.revoked_at or now()
result = {"id": row.id, "revoked": True}
else:
rows = list(await db.scalars(select(MCPToken).order_by(MCPToken.created_at.desc())))
result = [{k: getattr(row, k) for k in
("id", "name", "scopes", "created_at", "expires_at", "revoked_at")} for row in rows]
# Reveal only after the transaction has committed successfully.
print(json.dumps(encode_snapshot(result), ensure_ascii=False, indent=2))
finally:
await engine.dispose()
if __name__ == "__main__":
parser = argparse.ArgumentParser()
parser.add_argument("command", choices=["reset-password"])
parser.parse_args()
asyncio.run(reset_password())
commands = parser.add_subparsers(dest="command", required=True)
commands.add_parser("reset-password")
create = commands.add_parser("mcp-token-create")
create.add_argument("--name", required=True)
create.add_argument("--scope", action="append", default=None)
create.add_argument("--days", type=int, default=90)
commands.add_parser("mcp-token-list")
revoke = commands.add_parser("mcp-token-revoke")
revoke.add_argument("token_id")
args = parser.parse_args()
try:
asyncio.run(reset_password() if args.command == "reset-password" else token_command(args))
except ValueError as exc:
parser.error(str(exc))
+1
View File
@@ -25,6 +25,7 @@ class Settings(BaseSettings):
request_timeout: float = 30
retry_attempts: int = Field(default=4, ge=1, le=8)
enable_runner: bool = True
mcp_enabled: bool = False
ai_request_limit: int = Field(default=12, ge=1, le=30)
ai_tool_limit: int = Field(default=12, ge=1, le=100)
ai_output_tokens: int = Field(default=4096, ge=128, le=32768)
+64 -3
View File
@@ -5,7 +5,7 @@ import csv
import io
import time
from collections import defaultdict
from contextlib import asynccontextmanager
from contextlib import AsyncExitStack, asynccontextmanager
from typing import Annotated
from fastapi import APIRouter, Depends, FastAPI, HTTPException, Query, Request, Response
@@ -24,6 +24,7 @@ from .catalog.routes import router as catalog_router
from .config import Settings
from .db import create_database
from .jobs import AUTH_KINDS, Runner, create_job
from .mcp_api.token_routes import router as mcp_token_router
from .models import Account, Admin, BacktestConfig, Job, JobItem, LoginSession
from .research.routes import router as research_router
from .research.runtime import ResearchRuntime
@@ -93,6 +94,12 @@ def create_app(settings=None, wq_client=None, ai_model_factory=None):
ai_runtime = AIRuntime(sessions, settings, runner, ai_model_factory)
research_runtime = ResearchRuntime(sessions, ai_runtime, runner)
mcp_runtime = None
if settings.mcp_enabled:
from .mcp_api.server import MCPResearchServer
mcp_runtime = MCPResearchServer(sessions, runner, settings)
@asynccontextmanager
async def lifespan(app):
async with sessions() as db:
@@ -104,7 +111,10 @@ def create_app(settings=None, wq_client=None, ai_model_factory=None):
if settings.enable_runner:
await runner.start()
await research_runtime.start()
yield
async with AsyncExitStack() as stack:
if mcp_runtime:
await stack.enter_async_context(mcp_runtime.server.session_manager.run())
yield
if settings.enable_runner:
await research_runtime.stop()
await ai_runtime.stop()
@@ -124,6 +134,7 @@ def create_app(settings=None, wq_client=None, ai_model_factory=None):
app.state.settings = settings
app.state.ai = ai_runtime
app.state.research = research_runtime
app.state.mcp = mcp_runtime
login_failures = defaultdict(list)
@app.exception_handler(RequestValidationError)
@@ -139,7 +150,54 @@ def create_app(settings=None, wq_client=None, ai_model_factory=None):
@app.middleware("http")
async def browser_security(request, call_next):
if request.method not in ("GET", "HEAD", "OPTIONS"):
is_mcp = request.url.path in ("/api/v1/mcp", "/api/v1/mcp/")
if is_mcp:
if not mcp_runtime:
return JSONResponse({"detail": "MCP 未启用"}, status_code=404)
from urllib.parse import urlsplit
from .mcp_api.auth import authenticate
from .mcp_api.server import TOOLS
if request.headers.get("host", "").lower() != urlsplit(settings.public_origin).netloc.lower():
return JSONResponse({"detail": "MCP Host 不被允许"}, status_code=403)
origin = request.headers.get("origin")
if origin and origin.rstrip("/") != settings.public_origin.rstrip("/"):
return JSONResponse({"detail": "MCP Origin 不被允许"}, status_code=403)
scheme, _, secret = request.headers.get("authorization", "").partition(" ")
if scheme.lower() != "bearer" or not secret or len(secret) > 256:
return JSONResponse({"detail": "需要 MCP Bearer 令牌"}, status_code=401,
headers={"WWW-Authenticate": "Bearer"})
try:
async with sessions() as db:
principal = await authenticate(db, secret)
except HTTPException as exc:
return JSONResponse({"detail": exc.detail}, status_code=exc.status_code,
headers={"WWW-Authenticate": "Bearer"})
request.state.mcp_principal = principal
if "research:read" not in principal.scopes:
return JSONResponse({"detail": "缺少读取权限"}, status_code=403)
if request.method == "POST":
body = bytearray()
async for chunk in request.stream():
body.extend(chunk)
if len(body) > 4 * 1024 * 1024:
return JSONResponse({"detail": "MCP 请求过大"}, status_code=413)
# BaseHTTPMiddleware replays cached bytes to the SDK; never log this payload.
request._body = bytes(body)
try:
import json
message = json.loads(body)
except (ValueError, UnicodeDecodeError):
return JSONResponse({"detail": "无效 JSON"}, status_code=400)
if isinstance(message, dict) and message.get("method") == "tools/call":
params = message.get("params")
tool = params.get("name") if isinstance(params, dict) else None
definition = TOOLS.get(tool) if isinstance(tool, str) else None
if definition and definition[2] not in principal.scopes:
return JSONResponse({"detail": "MCP 令牌缺少所需权限"}, status_code=403)
elif request.method not in ("GET", "HEAD", "OPTIONS"):
if request.headers.get("X-WQ-Request") != "1":
return JSONResponse({"detail": "缺少请求校验头"}, status_code=403)
origin = request.headers.get("Origin")
@@ -410,6 +468,9 @@ def create_app(settings=None, wq_client=None, ai_model_factory=None):
await notify_job(runner, "retry_job", result)
return result
if mcp_runtime:
app.mount("/api/v1/mcp", mcp_runtime.app)
app.include_router(mcp_token_router)
app.include_router(backtest_router)
app.include_router(api)
app.include_router(catalog_router)
+1
View File
@@ -0,0 +1 @@
"""Authenticated MCP transport; research behavior lives in research_access."""
+55
View File
@@ -0,0 +1,55 @@
"""Personal access tokens are isolated from browser and upstream credentials."""
import secrets
from dataclasses import dataclass
from datetime import timedelta, timezone
from uuid import uuid4
from fastapi import HTTPException
from sqlalchemy import select
from ..models import Account, Admin, MCPToken, now
from ..security import token_hash
SCOPES = frozenset({"research:read", "research:refresh", "backtests:execute", "backtests:control"})
@dataclass(frozen=True)
class Principal:
token_id: str
admin_id: int
account_id: int
wq_user_id: str
scopes: frozenset[str]
async def create_token(db, name, scopes=None, days=90):
"""Issue a token for the bound account; caller commits and reveals it once."""
scopes = set(scopes if scopes is not None else ["research:read"])
if not name.strip() or len(name) > 100 or not 1 <= days <= 365:
raise ValueError("名称须为 1–100 字,有效期须为 1–365 天")
if not scopes <= SCOPES or "research:read" not in scopes:
raise ValueError("权限无效;所有令牌必须包含 research:read")
account, admin = await db.get(Account, 1), await db.get(Admin, 1)
if not account or not account.wq_user_id or not admin:
raise ValueError("请先初始化系统并确认 WorldQuant 账户身份")
secret = "wqmcp_" + secrets.token_urlsafe(32)
row = MCPToken(
id=str(uuid4()), token_hash=token_hash(secret), name=name.strip(), admin_id=admin.id,
account_id=account.id, wq_user_id=account.wq_user_id, scopes=sorted(scopes),
expires_at=now() + timedelta(days=days),
)
db.add(row)
await db.flush()
return row, secret
async def authenticate(db, secret):
"""Validate every request, including current account binding; return no secrets."""
row = await db.scalar(select(MCPToken).where(MCPToken.token_hash == token_hash(secret)))
if not row or row.revoked_at or row.expires_at.replace(tzinfo=row.expires_at.tzinfo or timezone.utc) <= now():
raise HTTPException(401, "MCP 令牌无效或已过期")
account, admin = await db.get(Account, row.account_id), await db.get(Admin, row.admin_id)
if not account or not admin or account.id != 1 or account.wq_user_id != row.wq_user_id:
raise HTTPException(401, "MCP 令牌账户绑定已失效")
return Principal(row.id, row.admin_id, row.account_id, row.wq_user_id, frozenset(row.scopes))
+127
View File
@@ -0,0 +1,127 @@
"""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)
+82
View File
@@ -0,0 +1,82 @@
"""Cookie-authenticated PAT administration; MCP bearer tokens grant no access here."""
from datetime import timezone
from fastapi import APIRouter, Depends, HTTPException, Query, Request
from pydantic import BaseModel, ConfigDict, Field
from sqlalchemy import func, select
from ..models import Account, MCPToken, now
from ..security import require_auth
from .auth import create_token
router = APIRouter(prefix="/api/v1/mcp-tokens", tags=["mcp-tokens"], dependencies=[Depends(require_auth)])
class TokenInput(BaseModel):
model_config = ConfigDict(extra="forbid")
name: str = Field(min_length=1, max_length=100)
days: int = Field(default=90, ge=1, le=365, strict=True)
scopes: list[str] = Field(default_factory=lambda: ["research:read"], max_length=4)
def token_output(row, account):
"""Return public metadata only, including whether the current binding is usable."""
def timestamp(value):
return value.replace(tzinfo=value.tzinfo or timezone.utc) if value else None
expires = timestamp(row.expires_at)
status = (
"revoked" if row.revoked_at else
"expired" if expires <= now() else
"invalid_binding" if not account or account.wq_user_id != row.wq_user_id else
"active"
)
return {
"id": row.id, "name": row.name, "scopes": row.scopes,
"created_at": timestamp(row.created_at), "expires_at": expires,
"revoked_at": timestamp(row.revoked_at), "status": status,
}
@router.get("")
async def list_tokens(request: Request, limit: int = Query(25, ge=1, le=100), offset: int = Query(0, ge=0)):
async with request.app.state.sessions() as db:
account = await db.get(Account, 1)
owned = (MCPToken.admin_id == 1, MCPToken.account_id == 1)
total = await db.scalar(select(func.count()).select_from(MCPToken).where(*owned))
rows = await db.scalars(select(MCPToken).where(*owned).order_by(
MCPToken.created_at.desc(), MCPToken.id.desc()).offset(offset).limit(limit))
return {
"items": [token_output(row, account) for row in rows], "total": total,
"limit": limit, "offset": offset, "has_more": offset + limit < total,
"enabled": request.app.state.settings.mcp_enabled,
"endpoint": request.app.state.settings.public_origin.rstrip("/") + "/api/v1/mcp/",
"can_create": bool(account and account.wq_user_id),
}
@router.post("", status_code=201)
async def issue_token(body: TokenInput, request: Request):
# The existing single-admin browser session is the authority, never request-supplied IDs.
async with request.app.state.sessions.begin() as db:
try:
row, secret = await create_token(db, body.name, body.scopes, body.days)
except ValueError as exc:
raise HTTPException(422, str(exc)) from exc
result = token_output(row, await db.get(Account, 1))
# Do not expose the secret until the transaction successfully commits.
return {**result, "token": secret}
@router.post("/{token_id}/revoke")
async def revoke_token(token_id: str, request: Request):
async with request.app.state.sessions.begin() as db:
row = await db.scalar(select(MCPToken).where(
MCPToken.id == token_id, MCPToken.admin_id == 1, MCPToken.account_id == 1,
).with_for_update())
if not row:
raise HTTPException(404, "MCP Key 不存在")
row.revoked_at = row.revoked_at or now()
result = token_output(row, await db.get(Account, 1))
return result
+46
View File
@@ -494,3 +494,49 @@ class ResearchStepRun(Base):
created_at: Mapped[datetime] = mapped_column(DateTime(timezone=True), default=now)
updated_at: Mapped[datetime] = mapped_column(DateTime(timezone=True), default=now)
__table_args__ = (UniqueConstraint("run_id", "node_id", "round"),)
class MCPToken(Base):
"""Revocable personal tokens; only the one-way digest is persisted."""
__tablename__ = "mcp_tokens"
id: Mapped[str] = mapped_column(String(36), primary_key=True)
token_hash: Mapped[str] = mapped_column(String(64), unique=True)
name: Mapped[str] = mapped_column(String(100))
admin_id: Mapped[int] = mapped_column(ForeignKey("admins.id"))
account_id: Mapped[int] = mapped_column(ForeignKey("accounts.id"))
wq_user_id: Mapped[str] = mapped_column(String(100))
scopes: Mapped[list] = mapped_column(JSON)
created_at: Mapped[datetime] = mapped_column(DateTime(timezone=True), default=now)
expires_at: Mapped[datetime] = mapped_column(DateTime(timezone=True), index=True)
revoked_at: Mapped[datetime | None] = mapped_column(DateTime(timezone=True))
class ResearchRequest(Base):
"""Successful atomic operations survive retries and credential rotation."""
__tablename__ = "research_requests"
id: Mapped[str] = mapped_column(String(36), primary_key=True)
account_id: Mapped[int] = mapped_column(ForeignKey("accounts.id"))
operation: Mapped[str] = mapped_column(String(50))
idempotency_key: Mapped[str] = mapped_column(String(100))
digest: Mapped[str] = mapped_column(String(64))
business_id: Mapped[str] = mapped_column(String(36))
response: Mapped[dict] = mapped_column(JSON)
created_at: Mapped[datetime] = mapped_column(DateTime(timezone=True), default=now)
__table_args__ = (UniqueConstraint("account_id", "operation", "idempotency_key"),)
class MCPAudit(Base):
"""Minimal call evidence, never raw arguments or authentication material."""
__tablename__ = "mcp_audits"
id: Mapped[str] = mapped_column(String(36), primary_key=True)
token_id: Mapped[str] = mapped_column(ForeignKey("mcp_tokens.id"), index=True)
tool: Mapped[str] = mapped_column(String(100))
request_id: Mapped[str] = mapped_column(String(100))
input_digest: Mapped[str] = mapped_column(String(64))
business_id: Mapped[str | None] = mapped_column(String(100))
result_code: Mapped[str] = mapped_column(String(60))
elapsed_ms: Mapped[int] = mapped_column(Integer)
created_at: Mapped[datetime] = mapped_column(DateTime(timezone=True), default=now, index=True)
+1
View File
@@ -0,0 +1 @@
"""Direct research interface shared by trusted application adapters."""
+170
View File
@@ -0,0 +1,170 @@
"""Bounded direct research inputs; unknown properties are rejected at the interface."""
from datetime import date, datetime
from typing import Annotated, Literal
from pydantic import Field, model_validator
from ..backtests.contracts import Candidate, SimulationSettings
from ..catalog.contracts import CatalogFilters, Scope
from ..schemas import Contract
Identifier = Annotated[str, Field(min_length=1, max_length=100)]
RunId = Annotated[str, Field(min_length=1, max_length=36)]
class Empty(Contract):
pass
class Page(Contract):
limit: int = Field(default=25, ge=1, le=100)
offset: int = Field(default=0, ge=0)
class CompleteSettings(SimulationSettings):
@model_validator(mode="before")
@classmethod
def complete(cls, value):
if isinstance(value, dict) and set(cls.model_fields) - value.keys():
raise ValueError("必须提供每项完整设置;先读取 get_research_capabilities")
return value
model_config = {"json_schema_extra": {"required": list(SimulationSettings.model_fields)}}
class DirectCandidate(Candidate):
settings: CompleteSettings
class Provenance(Contract):
reference: str | None = Field(default=None, max_length=200)
batch_id: str | None = Field(default=None, max_length=200)
hypothesis: str | None = Field(default=None, max_length=2000)
parent_run_id: RunId | None = None
class Submit(Contract):
name: str = Field(min_length=1, max_length=200)
candidates: list[DirectCandidate] = Field(min_length=1, max_length=100)
idempotency_key: Identifier
duplicate_policy: Literal["reject", "rerun"] = "reject"
source: Provenance = Field(default_factory=Provenance)
@model_validator(mode="after")
def unique_ids(self):
if len({c.client_item_id for c in self.candidates}) != len(self.candidates):
raise ValueError("client_item_id 必须唯一")
return self
class Control(Contract):
run_id: RunId
action: Literal["pause", "resume", "stop", "recover"]
expected_version: int = Field(ge=1)
idempotency_key: Identifier
class CatalogSearch(Contract):
filters: CatalogFilters
dataset_id: str | None = Field(default=None, min_length=1, max_length=200)
class Scopes(Contract):
kind: Literal["scopes"]
class SettingOptions(Page):
kind: Literal["settings"]
class Operators(Page):
kind: Literal["operators"]
q: str = Field(default="", max_length=300)
category: str | None = None
class Availability(Contract):
kind: Literal["field_availability"]
field_id: Identifier
scope: Scope
class Metadata(Contract):
query: Annotated[Scopes | SettingOptions | Operators | Availability, Field(discriminator="kind")]
class CatalogRefresh(Contract):
kind: Literal["catalog"]
scope: Scope
dataset_id: str | None = Field(default=None, min_length=1, max_length=200)
class OperatorsRefresh(Contract):
kind: Literal["operators"]
class SettingsRefresh(Contract):
kind: Literal["settings"]
class PnlRefresh(Contract):
kind: Literal["pnl"]
alpha_ids: list[Identifier] = Field(min_length=1, max_length=100)
class Refresh(Contract):
query: Annotated[
CatalogRefresh | OperatorsRefresh | SettingsRefresh | Availability | PnlRefresh,
Field(discriminator="kind"),
]
class JobReference(Page):
job_id: RunId
class History(Page):
source: str | None = Field(default=None, max_length=100)
reference: str | None = Field(default=None, max_length=200)
status: str | None = Field(default=None, max_length=30)
created_from: datetime | None = None
created_to: datetime | None = None
scope: Scope | None = None
q: str = Field(default="", max_length=300)
candidates: list[DirectCandidate] | None = Field(default=None, min_length=1, max_length=100)
@model_validator(mode="after")
def dates(self):
for value in (self.created_from, self.created_to):
if value and not value.tzinfo:
raise ValueError("时间须包含时区")
if self.created_from and self.created_to and self.created_from > self.created_to:
raise ValueError("起始时间不能晚于结束时间")
return self
class RunReference(Contract):
run_id: RunId
after: int | None = Field(default=None, ge=0)
event_limit: int = Field(default=25, ge=1, le=100)
class Results(Page):
run_id: RunId
item_ids: list[RunId] | None = Field(default=None, min_length=1, max_length=100)
class Artifact(Page):
item_id: RunId
kind: Literal["snapshot", "pnl"]
date_from: date | None = None
date_to: date | None = None
@model_validator(mode="after")
def dates(self):
if self.kind == "snapshot" and (self.date_from or self.date_to):
raise ValueError("日期筛选仅用于 PnL")
if self.date_from and self.date_to and self.date_from > self.date_to:
raise ValueError("起始日期不能晚于结束日期")
return self
+134
View File
@@ -0,0 +1,134 @@
"""Historical evidence reads, independent of transport and current Alpha refreshes."""
from collections import Counter
from sqlalchemy import func, select
from ..alphas import number, sanitize
from ..backtests.contracts import fingerprint
from ..models import BacktestItem, BacktestResult, BacktestRun, Pnl
from ..research.serialization import encode_snapshot
def page(items, total, limit, offset):
return {"items": items, "total": total, "limit": limit, "offset": offset,
"has_more": offset + len(items) < total}
def checks_summary(snapshot):
"""Preserve unknown check values; missing checks can never mean passed."""
checks = []
for section in ("is", "os"):
metrics = snapshot.get(section)
if isinstance(metrics, dict) and "checks" in metrics:
raw = metrics["checks"]
checks.extend({"section": section, "raw": c} for c in (raw if isinstance(raw, list) else [raw]))
if "checks" in snapshot:
raw = snapshot["checks"]
checks.extend({"section": "root", "raw": c} for c in (raw if isinstance(raw, list) else [raw]))
counts = Counter({key: 0 for key in ("PASS", "FAIL", "PENDING", "WARNING", "UNKNOWN")})
non_pass = []
for check in checks:
raw = check["raw"]
value = raw.get("result", raw.get("status")) if isinstance(raw, dict) else None
state = value if isinstance(value, str) and value in counts else "UNKNOWN"
counts[state] += 1
if state != "PASS":
non_pass.append({**check, "status": state})
return {"status": "unknown" if not checks else "reported", "counts": dict(counts),
"total": len(checks), "non_pass": non_pass}
def item_summary(item, result):
snapshot = sanitize(result.snapshot) if result else {}
metrics = {}
for section in ("is", "os"):
raw = snapshot.get(section)
raw = raw if isinstance(raw, dict) else {}
metrics[section] = {key: number(raw.get(key)) for key in
("sharpe", "fitness", "returns", "turnover", "margin", "drawdown")}
return encode_snapshot({
**{k: getattr(item, k) for k in (
"id", "run_id", "client_item_id", "expression", "settings", "attempt_id",
"platform_status", "collection_status", "persistence_status", "simulation_id", "alpha_id",
)},
"error": sanitize(item.error), "metrics": metrics,
"missing_metrics_reason": "来源未提供或非有限数字;null 不等于零",
"checks": checks_summary(snapshot),
"result": {"observed_at": result.observed_at, "complete": result.complete} if result else None,
"artifact_reference": {"item_id": item.id},
})
class EvidenceQueries:
def __init__(self, db):
self.db = db
async def history(self, args):
query = select(BacktestItem, BacktestResult, BacktestRun).join(
BacktestRun, BacktestRun.id == BacktestItem.run_id
).outerjoin(BacktestResult, BacktestResult.item_id == BacktestItem.id)
for key in ("source", "reference"):
value = getattr(args, key)
if value is not None:
query = query.where(BacktestRun.source["kind" if key == "source" else key].as_string() == value)
if args.status:
query = query.where(BacktestRun.status == args.status)
if args.created_from:
query = query.where(BacktestRun.created_at >= args.created_from)
if args.created_to:
query = query.where(BacktestRun.created_at <= args.created_to)
if args.scope:
for source, target in (("instrument_type", "instrumentType"), ("region", "region"), ("universe", "universe")):
query = query.where(BacktestItem.settings[target].as_string() == getattr(args.scope, source))
query = query.where(BacktestItem.settings["delay"].as_integer() == args.scope.delay)
if args.q:
escaped = args.q.replace("\\", "\\\\").replace("%", "\\%").replace("_", "\\_")
query = query.where(BacktestItem.expression.ilike(f"%{escaped}%", escape="\\"))
matches = {}
if args.candidates:
for c in args.candidates:
matches.setdefault(fingerprint(c.platform_input()), []).append(c.client_item_id)
query = query.where(BacktestItem.fingerprint.in_(matches))
total = await self.db.scalar(select(func.count()).select_from(query.subquery()))
rows = (await self.db.execute(query.order_by(BacktestRun.created_at.desc(), BacktestRun.id, BacktestItem.ordinal)
.limit(args.limit).offset(args.offset))).all()
items = [{**item_summary(i, r), "source": run.source, "run_status": run.status,
"created_at": run.created_at, "matched_candidates": matches.get(i.fingerprint, []),
"match_type": "exact_input" if args.candidates else "filter"} for i, r, run in rows]
return encode_snapshot(page(items, total, args.limit, args.offset))
async def results(self, args):
query = select(BacktestItem, BacktestResult).outerjoin(
BacktestResult, BacktestResult.item_id == BacktestItem.id
).where(BacktestItem.run_id == args.run_id)
if args.item_ids:
query = query.where(BacktestItem.id.in_(args.item_ids))
total = await self.db.scalar(select(func.count()).select_from(query.subquery()))
rows = (await self.db.execute(query.order_by(BacktestItem.ordinal).limit(args.limit).offset(args.offset))).all()
return {"backtest_run_id": args.run_id, **page([item_summary(i, r) for i, r in rows], total, args.limit, args.offset)}
async def artifact(self, args):
from .service import ResearchError
item = await self.db.get(BacktestItem, args.item_id)
if not item:
raise ResearchError("NOT_FOUND", "候选不存在")
result = await self.db.get(BacktestResult, item.id)
if args.kind == "snapshot":
# Top-level entries retain complete nested values; no hidden string/list truncation.
entries = [{"key": k, "value": v} for k, v in sanitize(result.snapshot).items()] if result else []
return encode_snapshot({"item_id": item.id, "kind": args.kind,
"status": "available" if result else "not_available",
"observed_at": result.observed_at if result else None,
"complete": result.complete if result else False,
**page(entries[args.offset:args.offset + args.limit], len(entries), args.limit, args.offset)})
pnl = await self.db.get(Pnl, item.alpha_id) if item.alpha_id else None
points = pnl.points if pnl else []
points = [p for p in points if
(not args.date_from or p["date"][:10] >= args.date_from.isoformat()) and
(not args.date_to or p["date"][:10] <= args.date_to.isoformat())]
return encode_snapshot({"item_id": item.id, "alpha_id": item.alpha_id, "kind": args.kind,
"status": "available" if pnl else "not_cached", "fetched_at": pnl.fetched_at if pnl else None,
"units": "供应商原始累计值;未提供货币或规模单位",
**page(points[args.offset:args.offset + args.limit], len(points), args.limit, args.offset)})
+222
View File
@@ -0,0 +1,222 @@
"""Direct research operations; caller owns authorization, transaction and wake-up.
The account row serializes mutations with existing HTTP starts. Request records,
previews, runs and control events commit together; failed validation consumes no key.
"""
from collections import Counter
from uuid import uuid4
from sqlalchemy import func, select
from ..backtests.contracts import ControlInput, DraftInput, PreviewInput, Source, StartInput, fingerprint
from ..backtests.service import Backtests
from ..business import Business
from ..catalog.contracts import CatalogJobInput
from ..catalog.platform import platform_options, validate_platform_scope
from ..catalog.research_metadata import ResearchMetadata, availability_key
from ..catalog.service import Catalog
from ..models import Account, Alpha, BacktestItem, Job, JobItem, ResearchRequest, SimulationAttempt, now
from ..research.serialization import encode_snapshot
from ..research.workspace_contracts import FieldAvailabilityInput
from ..schemas import JobInput
from .contracts import DirectCandidate, History
from .queries import EvidenceQueries, page
class ResearchError(Exception):
def __init__(self, code, message, *, retryable=False, retry_after=None, affected_items=None):
super().__init__(message)
self.data = {"code": code, "message": message, "retryable": retryable,
"retry_after": retry_after, "affected_items": affected_items or []}
class ResearchAccess:
def __init__(self, db, principal, client, public_origin):
self.db, self.principal, self.client = db, principal, client
self.public_origin = public_origin.rstrip("/")
self.backtests = Backtests(db)
self.business = Business(db)
self.evidence = EvidenceQueries(db)
self.wake = None
def run_url(self, run_id):
return f"{self.public_origin}/#backtests?run_id={run_id}"
async def capabilities(self, args):
return {**await self.backtests.capabilities(), "max_candidates": 100,
"settings_schema": DirectCandidate.model_json_schema(),
"confirmation": "调用者须已获本批执行授权;直接提交后返回稳定运行 ID",
"duplicate_policies": ["reject", "rerun"], "permissions": sorted(self.principal.scopes),
"metadata_only": True, "actual_platform_allowance": None}
async def catalog(self, args):
data = await Catalog(self.db).search(args.filters, args.dataset_id)
return {**data, "scope": args.filters.model_dump(include={"region", "universe", "delay", "instrument_type"}),
"status": "available" if data["collection_version"] else "not_cached",
"has_more": data["offset"] + len(data["items"]) < data["total"]}
async def metadata(self, args):
q = args.query
metadata = ResearchMetadata(self.db)
if q.kind == "scopes":
return {"source": "worldquant_platform", **await platform_options(self.client)}
if q.kind == "operators":
data = await metadata.operators(q.q, q.category, limit=q.limit, offset=q.offset)
return {**data, "status": "available" if data["fetched_at"] else "not_cached",
"has_more": q.offset + len(data["items"]) < data["total"]}
if q.kind == "settings":
data = await metadata.get("settings")
items = data["content"].get("items", [])
return {"status": "available" if data["fetched_at"] else "not_cached",
"fetched_at": data["fetched_at"], **page(items[q.offset:q.offset+q.limit], len(items), q.limit, q.offset)}
data = await metadata.get(availability_key(q.field_id, q.scope))
return {**data, "status": data["content"].get("status", "unknown")}
async def refresh(self, args):
q = args.query
metadata = ResearchMetadata(self.db, self.client)
if q.kind == "catalog":
await validate_platform_scope(self.client, q.scope)
job = await Catalog(self.db).create_job(CatalogJobInput(scope=q.scope, dataset_id=q.dataset_id))
self.wake = "jobs"
return {"job_id": job.id, "status": job.status}
if q.kind == "pnl":
ids = sorted(set(q.alpha_ids))
existing = set(await self.db.scalars(select(Alpha.id).where(Alpha.id.in_(ids))))
if existing != set(ids):
raise ResearchError("NOT_FOUND", "部分 Alpha 尚未同步", affected_items=sorted(set(ids)-existing))
job = await self.business.create_sync_job(JobInput(kind="pnl_refresh", alpha_ids=ids))
self.wake = "jobs"
return {"job_id": job["id"], "status": job["status"]}
if q.kind == "operators":
data = await metadata.refresh_operators()
elif q.kind == "settings":
data = await metadata.refresh_settings()
else:
data = await metadata.refresh_availability(FieldAvailabilityInput(field_id=q.field_id, scope=q.scope))
# Refresh acknowledgment is bounded; complete content is available through paged reads.
return {"status": "completed", "key": data["key"], "fetched_at": data["fetched_at"],
"read_with": "get_research_metadata"}
async def refresh_job(self, args):
job = await self.db.get(Job, args.job_id)
if not job or job.kind not in {"catalog_sync", "field_sync", "pnl_refresh"}:
raise ResearchError("NOT_FOUND", "研究刷新任务不存在")
result = await self.business.get_job_status(args.job_id)
query = select(JobItem).where(JobItem.job_id == job.id, JobItem.error.is_not(None))
total = await self.db.scalar(select(func.count()).select_from(query.subquery()))
errors = list(await self.db.scalars(query.order_by(JobItem.alpha_id).limit(args.limit).offset(args.offset)))
result.pop("errors", None)
return {**result, "job_id": job.id, "artifact_reference": job.payload,
"errors": page([{"alpha_id": e.alpha_id, "error": e.error} for e in errors], total, args.limit, args.offset)}
async def history(self, args):
return await self.evidence.history(args)
async def previous(self, operation, args):
# PostgreSQL row lock is shared with HTTP start and catalog/job creation.
account = await self.db.scalar(select(Account).where(Account.id == self.principal.account_id).with_for_update())
if not account or account.wq_user_id != self.principal.wq_user_id:
raise ResearchError("ACCOUNT_MISMATCH", "平台账户绑定已变化")
digest = fingerprint(args.model_dump(mode="json", exclude={"idempotency_key"}))
row = await self.db.scalar(select(ResearchRequest).where(
ResearchRequest.account_id == account.id, ResearchRequest.operation == operation,
ResearchRequest.idempotency_key == args.idempotency_key))
if row and row.digest != digest:
raise ResearchError("IDEMPOTENCY_CONFLICT", "幂等键已用于不同内容")
return row, digest
async def remember(self, operation, args, digest, result):
result["_meta"] = {"schema_version": 1, "observed_at": now().isoformat(), "source": "system"}
self.db.add(ResearchRequest(id=str(uuid4()), account_id=self.principal.account_id,
operation=operation, idempotency_key=args.idempotency_key, digest=digest,
business_id=result["backtest_run_id"], response=encode_snapshot(result)))
await self.db.flush()
self.wake = "backtests"
return result
async def validate_settings(self, candidates):
snapshot = await ResearchMetadata(self.db).get("settings")
options = snapshot["content"].get("items", [])
if not snapshot["fetched_at"] or not options:
return {"settings_validation": "unknown", "reason": "设置快照未缓存;未验证平台组合", "field_validation": "unknown"}
invalid = []
for c in candidates:
s = c.settings
matches = [r for r in options if all(r.get(k) == v for k, v in {
"instrument_type": s.instrumentType, "region": s.region, "universe": s.universe, "delay": s.delay}.items())]
if not matches or all(r.get("neutralizations") and s.neutralization not in r["neutralizations"] for r in matches):
invalid.append(c.client_item_id)
if invalid:
raise ResearchError("UNSUPPORTED_SETTINGS", "已缓存平台设置不支持这些组合;可显式刷新后重试", affected_items=invalid)
return {"settings_validation": "cached", "fetched_at": snapshot["fetched_at"], "field_validation": "unknown"}
async def submit(self, args):
previous, digest = await self.previous("submit_backtests", args)
if previous:
return previous.response
seen, within = {}, []
for c in args.candidates:
h = fingerprint(c.platform_input())
if h in seen:
within.append({"client_item_id": c.client_item_id, "duplicate_of": seen[h]})
seen[h] = c.client_item_id
history = await self.evidence.history(History(candidates=args.candidates, limit=100))
if args.duplicate_policy == "reject" and (within or history["total"]):
raise ResearchError("DUPLICATE_INPUT", "发现完整输入重复;未创建运行。重跑须明确 duplicate_policy=rerun",
affected_items={"within_batch": within, "history": history, "read_with": "search_backtests"})
validation = await self.validate_settings(args.candidates)
if args.source.parent_run_id:
await self.backtests.run(args.source.parent_run_id)
source = Source(kind="mcp", **args.source.model_dump())
provenance = {"mcp_token_id": self.principal.token_id, "admin_id": self.principal.admin_id}
# preserve_source prevents the Chatbox-specific generating context rewriting MCP provenance.
backtests = Backtests(self.db, provenance)
preview = await backtests.preview(PreviewInput(inline=DraftInput(
name=args.name, source=source, candidates=args.candidates)), preserve_source=True)
result = await backtests.start(StartInput(preview_id=preview["preview_id"],
idempotency_key="mcp-" + str(uuid4())))
result = {**result, "input_digest": digest, "batch_count": preview["batch_count"],
"duplicates": {"within_batch": within, "historical_matches": history["total"]},
"validation": validation, "web_url": self.run_url(result["backtest_run_id"])}
return await self.remember("submit_backtests", args, digest, result)
async def run(self, args):
result = await self.backtests.run(args.run_id)
attempts = list(await self.db.scalars(select(SimulationAttempt).where(SimulationAttempt.run_id == args.run_id)))
result["submission_counts"] = {
"candidates": result["total"], "attempts": len(attempts),
"post_requests": sum(a.submit_count for a in attempts),
"confirmed_accepted_candidates": sum(len(a.payload) for a in attempts if a.progress_url),
"unknown_acceptance_candidates": sum(len(a.payload) for a in attempts if a.error_code == "submission_unknown" or (a.state == "submitting" and not a.progress_url)),
"actual_platform_consumption": None,
}
if args.after is not None:
result["events"] = await self.backtests.events(args.run_id, args.after, args.event_limit)
return {**result, "web_url": self.run_url(args.run_id)}
async def results(self, args):
await self.backtests.run(args.run_id)
if args.item_ids:
found = set(await self.db.scalars(select(BacktestItem.id).where(
BacktestItem.run_id == args.run_id, BacktestItem.id.in_(args.item_ids))))
if found != set(args.item_ids):
raise ResearchError("NOT_FOUND", "部分候选不属于此运行")
return await self.evidence.results(args)
async def artifact(self, args):
return await self.evidence.artifact(args)
async def control(self, args):
previous, digest = await self.previous("control_backtest", args)
if previous:
return previous.response
before = await self.backtests.run(args.run_id)
states = list(await self.db.scalars(select(SimulationAttempt.state).where(SimulationAttempt.run_id == args.run_id)))
result = await self.backtests.control(args.run_id, ControlInput(action=args.action, version=args.expected_version))
result["impact"] = {"remote_cancelled": False, "attempts_before": dict(Counter(states)),
"indefinite_account_block_cleared": args.action == "resume" and bool(before["scheduler"]["blocked_reason"])
and before["scheduler"]["blocked_until"] is None,
"note": "暂停/停止仅阻止后续提交;已提交模拟继续采集。recover 不重新提交。"}
return await self.remember("control_backtest", args, digest, result)