worldquant-alpha-system/backend/app/backtests/service.py

"""Transactional research interface. Callers own authorization and commit boundaries."""

from collections import Counter, defaultdict
from uuid import uuid4

from fastapi import HTTPException
from fastapi.encoders import jsonable_encoder
from sqlalchemy import func, select, update

from ..models import (
    Account,
    BacktestConfig,
    BacktestDraft,
    BacktestEvent,
    BacktestItem,
    BacktestPreview,
    BacktestResult,
    BacktestRun,
    SimulationAttempt,
    now,
)
from .contracts import Candidate, DraftInput, PreviewInput, Source, fingerprint, group_key


def uid():
    return str(uuid4())


async def event(db, run, kind, payload):
    """Append a run-local cursor under the run row lock, in the result's transaction."""
    run.event_seq += 1
    run.updated_at = now()
    db.add(BacktestEvent(run_id=run.id, seq=run.event_seq, kind=kind, payload=payload))


async def locked_run(db, run_id):
    run = await db.scalar(select(BacktestRun).where(BacktestRun.id == run_id).with_for_update())
    if not run:
        raise HTTPException(404, "回测运行不存在")
    return run


async def refresh_status(db, run):
    await db.flush()
    states = list(await db.scalars(select(SimulationAttempt.state).where(SimulationAttempt.run_id == run.id)))
    if any(s in ("needs_review", "collection_failed") for s in states):
        run.status = "needs_review"
    elif all(s in ("completed", "failed", "skipped") for s in states):
        run.status = (
            "stopped"
            if run.control == "stopped"
            else "completed_with_errors"
            if "failed" in states
            else "completed"
        )
    elif run.control == "paused":
        run.status = "paused"
    elif run.control == "stopped":
        run.status = "stopping"
    elif any(s in ("submitting", "submitted", "collecting") for s in states):
        run.status = "running"
    else:
        run.status = "queued"


class Backtests:
    def __init__(self, db, ai_context=None):
        self.db = db
        self.ai_context = ai_context or {}

    async def config(self):
        row = await self.db.get(BacktestConfig, 1)
        return jsonable_encoder(
            {
                k: getattr(row, k)
                for k in ("concurrency", "batch_size", "version", "blocked_reason", "blocked_until")
            }
        )

    async def configure(self, body):
        result = await self.db.execute(
            update(BacktestConfig)
            .where(BacktestConfig.id == 1, BacktestConfig.version == body.version)
            .values(
                concurrency=body.concurrency,
                batch_size=body.batch_size,
                version=BacktestConfig.version + 1,
            )
        )
        if result.rowcount != 1:
            raise HTTPException(409, "调度配置已变化，请刷新后重试")
        return await self.config()

    async def capabilities(self):
        return {
            "alpha_types": ["REGULAR", "SUPER"],
            "languages": ["FASTEXPR"],
            "instrument_types": ["EQUITY"],
            "settings_schema": Candidate.model_json_schema(),
            "scheduler": await self.config(),
            "max_candidates": 10000,
            "remote_cancel": False,
            "automatic_history_reuse": False,
            "confirmation": "每个固定运行确认一次；启动后返回 ID，不循环等待",
            "mapping": "完整输入匹配；证据不足待核对，不按 children 顺序匹配",
            "limits": "并发和批大小为本地配置，并非平台可用额度；账户外提交不在预算内",
        }

    async def bind_preparations(self, body):
        from ..preparations.service import Preparations
        from ..research.expressions import analyze
        if not body.preparation_refs and not body.input_ids:
            return
        if any(c.alpha_type == "SUPER" for c in body.candidates):
            raise HTTPException(422, "SUPER 组件快照不能使用字段数据准备集合")
        await Preparations(self.db).bind(body)
        snapshots = [await Preparations(self.db).snapshot(i) for i in body.input_ids]
        for candidate in body.candidates:
            scope = dict(instrument_type=candidate.settings.instrumentType, region=candidate.settings.region,
                         universe=candidate.settings.universe, delay=candidate.settings.delay)
            if any(s["scope"] != scope for s in snapshots):
                raise HTTPException(422, "数据准备集合与回测范围不一致")
            fields = {}
            for snapshot in snapshots:
                for field, kind in snapshot["field_types"].items():
                    if field in fields and fields[field] != kind:
                        raise HTTPException(422, "输入字段类型冲突")
                    fields[field] = kind
            validation = analyze(candidate.expression, fields)
            if validation["syntax"] or validation["types"]:
                raise HTTPException(422, "；".join(validation["syntax"] + validation["types"]))
        body.source.input_snapshot_ids = body.input_ids
        body.source.input_snapshot_id = body.input_ids[0] if len(body.input_ids) == 1 else None

    async def save_draft(self, body, draft_id=None):
        await self.bind_preparations(body)
        data = body.model_dump(mode="json", exclude={"version", "preparation_refs", "input_ids"})
        if draft_id:
            changed = await self.db.execute(
                update(BacktestDraft)
                .where(BacktestDraft.id == draft_id, BacktestDraft.version == body.version)
                .values(
                    **data,
                    version=BacktestDraft.version + 1,
                    updated_at=now(),
                )
            )
            if changed.rowcount != 1:
                raise HTTPException(409, "草稿已变化或不存在；保留当前编辑并重新载入")
        else:
            draft_id = uid()
            self.db.add(BacktestDraft(id=draft_id, **data))
        await self.db.flush()
        return await self.draft(draft_id)

    async def drafts(self, limit=25, offset=0, q="", sort="updated_at", direction="desc"):
        query = select(BacktestDraft)
        if q:
            query = query.where(BacktestDraft.name.contains(q, autoescape=True))
        column = {"name": BacktestDraft.name, "updated_at": BacktestDraft.updated_at}[sort]
        order = column.asc() if direction == "asc" else column.desc()
        rows = (
            await self.db.scalars(
                query
                .order_by(order, BacktestDraft.id)
                .limit(limit)
                .offset(offset)
            )
        ).all()
        return {
            "items": [
                jsonable_encoder(
                    {
                        "id": r.id,
                        "version": r.version,
                        "name": r.name,
                        "total": len(r.candidates),
                        "updated_at": r.updated_at,
                    }
                )
                for r in rows
            ],
            "total": await self.db.scalar(select(func.count()).select_from(query.subquery())),
            "limit": limit,
            "offset": offset,
        }

    async def draft(self, draft_id):
        row = await self.db.get(BacktestDraft, draft_id)
        if not row:
            raise HTTPException(404, "候选草稿不存在")
        return jsonable_encoder(
            {k: getattr(row, k) for k in ("id", "version", "name", "source", "candidates", "updated_at")}
        )

    async def preview(self, body, *, preserve_source=False):
        """Fix inputs; new chatbox candidates inherit trusted generating-run provenance.

        Existing draft references and server-side subsets/reruns retain their
        producer. ai_context separately identifies whoever starts the execution.
        """
        if body.inline:
            await self.bind_preparations(body.inline)
            data = body.inline.model_dump(mode="json")
            if self.ai_context and not preserve_source:
                data["source"] = {
                    **data["source"],
                    "kind": "chatbox",
                    "reference": self.ai_context["conversation_id"],
                    "research_id": self.ai_context["ai_run_id"],
                    "parent_run_id": None,
                }
        else:
            draft = await self.db.scalar(
                select(BacktestDraft).where(BacktestDraft.id == body.draft_id).with_for_update()
            )
            if not draft or draft.version != body.draft_version:
                raise HTTPException(409, "候选草稿已变化，请重新准备预览")
            candidates = draft.candidates
            if body.selection is not None:
                selection = set(body.selection)
                candidates = [c for c in candidates if c["client_item_id"] in selection]
                if len(candidates) != len(selection):
                    raise HTTPException(422, "选择包含不属于当前草稿的候选")
            data = {"name": draft.name, "source": draft.source, "candidates": candidates}
        from ..superalpha.service import validate_source
        await validate_source(self.db, data["source"], data["candidates"])
        candidates = DraftInput.model_validate(data).model_dump(mode="json")["candidates"]
        config = await self.db.get(BacktestConfig, 1)
        groups = defaultdict(list)
        hashes = []
        for i, c in enumerate(candidates):
            groups[group_key(c)].append(i)
            hashes.append(fingerprint(Candidate.model_validate(c).platform_input()))
        # Query hashes in bounded chunks, including SQLite's bind-parameter limit.
        existing = set()
        for index in range(0, len(hashes), 400):
            existing.update(
                await self.db.scalars(
                    select(BacktestItem.fingerprint)
                    .where(BacktestItem.fingerprint.in_(hashes[index : index + 400]))
                    .distinct()
                )
            )
        seen, duplicates = set(), []
        for c, h in zip(candidates, hashes):
            if h in seen or h in existing:
                duplicates.append(
                    {
                        "client_item_id": c["client_item_id"],
                        "historical": h in existing,
                        "within_preview": h in seen,
                    }
                )
            seen.add(h)
        batches = []
        for indices in groups.values():
            local_batches = []
            for index in indices:
                if candidates[index]["alpha_type"] == "SUPER":
                    batches.append([index])
                    continue
                batch = next(
                    (
                        b
                        for b in local_batches
                        if len(b) < config.batch_size and all(hashes[i] != hashes[index] for i in b)
                    ),
                    None,
                )
                if batch is None:
                    batch = []
                    local_batches.append(batch)
                batch.append(index)
            batches.extend(local_batches)
        row = BacktestPreview(
            id=uid(),
            name=data["name"],
            source=data["source"],
            candidates=candidates,
            batches=batches,
            batch_size=config.batch_size,
            digest=fingerprint({"candidates": candidates, "source": data["source"]}),
            duplicates=duplicates,
            ai_context=self.ai_context,
        )
        self.db.add(row)
        await self.db.flush()
        return await self.get_preview(row.id)

    async def get_preview(self, preview_id, limit=25, offset=0):
        row = await self.db.get(BacktestPreview, preview_id)
        if not row:
            raise HTTPException(404, "回测预览不存在")
        return jsonable_encoder(
            {
                "preview_id": row.id,
                "version": row.version,
                "name": row.name,
                "source": row.source,
                "digest": row.digest,
                "total": len(row.candidates),
                "batch_count": len(row.batches),
                "batch_size": row.batch_size,
                "duplicate_count": len(row.duplicates),
                "duplicates": row.duplicates[offset : offset + limit],
                "items": row.candidates[offset : offset + limit],
                "limit": limit,
                "offset": offset,
                "has_more": offset + limit < len(row.candidates),
                "created_at": row.created_at,
            }
        )

    async def start(self, body):
        # One account row serializes all starts; unique keys remain the final DB invariant.
        account = await self.db.scalar(select(Account).where(Account.id == 1).with_for_update())
        previous = await self.db.scalar(
            select(BacktestRun).where(BacktestRun.idempotency_key == body.idempotency_key)
        )
        if previous:
            if previous.preview_id != body.preview_id or body.version != 1:
                raise HTTPException(409, "幂等键已用于另一份预览")
            return await self.run(previous.id)
        preview = await self.db.get(BacktestPreview, body.preview_id)
        if not preview or preview.version != body.version:
            raise HTTPException(409, "预览不存在或版本不匹配")
        previous = await self.db.scalar(select(BacktestRun).where(BacktestRun.preview_id == preview.id))
        if previous:
            return await self.run(previous.id)
        if not account or not account.wq_user_id or account.connection_status != "connected":
            raise HTTPException(409, "请先连接并确认 WorldQuant 账户身份")
        run = BacktestRun(
            id=uid(),
            preview_id=preview.id,
            idempotency_key=body.idempotency_key,
            name=preview.name,
            source=preview.source,
            total=len(preview.candidates),
            batch_size=preview.batch_size,
            ai_context=self.ai_context or preview.ai_context,
            event_seq=0,
        )
        self.db.add(run)
        await self.db.flush()
        for n, indices in enumerate(preview.batches):
            candidates = [Candidate.model_validate(preview.candidates[i]) for i in indices]
            attempt = SimulationAttempt(
                id=uid(), run_id=run.id, ordinal=n, payload=[c.platform_input() for c in candidates]
            )
            self.db.add(attempt)
            await self.db.flush()
            for i, c in zip(indices, candidates):
                self.db.add(
                    BacktestItem(
                        id=uid(),
                        run_id=run.id,
                        attempt_id=attempt.id,
                        ordinal=i,
                        client_item_id=c.client_item_id,
                        expression=c.expression,
                        alpha_type=c.alpha_type, selection=c.selection, combo=c.combo,
                        settings=c.settings.model_dump(),
                        fingerprint=fingerprint(c.platform_input()),
                    )
                )
        await event(self.db, run, "created", {"total": run.total, "batch_count": len(preview.batches)})
        await self.db.flush()
        return await self.run(run.id)

    async def runs(self, limit=25, offset=0, source=None, reference=None, research_id=None, q="", sort="created_at", direction="desc", alpha_type=None):
        query = select(BacktestRun)
        if alpha_type:
            query = query.where(BacktestRun.id.in_(select(BacktestItem.run_id).where(BacktestItem.alpha_type == alpha_type)))
        if q:
            query = query.where(BacktestRun.name.contains(q, autoescape=True))
        column = {"name": BacktestRun.name, "created_at": BacktestRun.created_at}[sort]
        order = column.asc() if direction == "asc" else column.desc()
        for key, value in (("kind", source), ("reference", reference), ("research_id", research_id)):
            if value:
                query = query.where(BacktestRun.source[key].as_string() == value)
        total = await self.db.scalar(select(func.count()).select_from(query.subquery()))
        rows = (
            await self.db.scalars(
                query.order_by(order, BacktestRun.id).limit(limit).offset(offset)
            )
        ).all()
        return {
            "items": [await self.run(r.id) for r in rows],
            "total": total,
            "limit": limit,
            "offset": offset,
        }

    async def sources(self):
        kinds = await self.db.scalars(select(BacktestRun.source["kind"].as_string()).distinct())
        return sorted({kind for kind in kinds if kind} | {"chatbox", "manual"})

    async def run(self, run_id):
        row = await self.db.get(BacktestRun, run_id)
        if not row:
            raise HTTPException(404, "回测运行不存在")
        groups = (
            await self.db.execute(
                select(
                    BacktestItem.platform_status,
                    BacktestItem.collection_status,
                    BacktestItem.persistence_status,
                    func.count(),
                )
                .where(BacktestItem.run_id == run_id)
                .group_by(
                    BacktestItem.platform_status,
                    BacktestItem.collection_status,
                    BacktestItem.persistence_status,
                )
            )
        ).all()
        counts = {"platform": Counter(), "collection": Counter(), "persistence": Counter()}
        for p, c, s, n in groups:
            for key, value in (("platform", p), ("collection", c), ("persistence", s)):
                counts[key][value] += n
        return jsonable_encoder(
            {
                "backtest_run_id": row.id,
                **{
                    k: getattr(row, k)
                    for k in (
                        "preview_id",
                        "name",
                        "source",
                        "ai_context",
                        "control",
                        "status",
                        "version",
                        "total",
                        "batch_size",
                        "created_at",
                        "updated_at",
                    )
                },
                "counts": counts,
                "cursor": row.event_seq,
                "scheduler": await self.config(),
            }
        )

    async def results(self, run_id, limit=25, offset=0):
        run = await self.run(run_id)
        rows = (
            await self.db.execute(
                select(BacktestItem, BacktestResult)
                .outerjoin(BacktestResult, BacktestResult.item_id == BacktestItem.id)
                .where(BacktestItem.run_id == run_id)
                .order_by(BacktestItem.ordinal)
                .limit(limit)
                .offset(offset)
            )
        ).all()
        return jsonable_encoder(
            {
                "backtest_run_id": run_id,
                "total": run["total"],
                "limit": limit,
                "offset": offset,
                "items": [
                    {
                        **{
                            k: getattr(i, k)
                            for k in (
                                "id",
                                "client_item_id",
                                "expression", "alpha_type", "selection", "combo",
                                "settings",
                                "attempt_id",
                                "platform_status",
                                "collection_status",
                                "persistence_status",
                                "simulation_id",
                                "alpha_id",
                                "error",
                            )
                        },
                        "result": {
                            "snapshot": r.snapshot,
                            "observed_at": r.observed_at,
                            "complete": r.complete,
                        }
                        if r
                        else None,
                    }
                    for i, r in rows
                ],
            }
        )

    async def events(self, run_id, after=0, limit=100):
        await self.run(run_id)
        rows = (
            await self.db.scalars(
                select(BacktestEvent)
                .where(BacktestEvent.run_id == run_id, BacktestEvent.seq > after)
                .order_by(BacktestEvent.seq)
                .limit(limit + 1)
            )
        ).all()
        return jsonable_encoder(
            {
                "items": [
                    {"seq": r.seq, "kind": r.kind, "payload": r.payload, "created_at": r.created_at}
                    for r in rows[:limit]
                ],
                "next_cursor": rows[min(len(rows), limit) - 1].seq if rows else after,
                "has_more": len(rows) > limit,
            }
        )

    async def attempts(self, run_id):
        await self.run(run_id)
        rows = (
            await self.db.scalars(
                select(SimulationAttempt)
                .where(SimulationAttempt.run_id == run_id)
                .order_by(SimulationAttempt.ordinal)
            )
        ).all()
        return jsonable_encoder(
            [
                {
                    k: getattr(a, k)
                    for k in (
                        "id",
                        "state",
                        "ordinal",
                        "progress_url",
                        "remote_complete",
                        "children",
                        "error",
                        "error_code",
                        "poll_count",
                        "submit_count",
                        "next_poll_at",
                    )
                }
                for a in rows
            ]
        )

    async def control(self, run_id, body):
        run = await locked_run(self.db, run_id)
        if run.version != body.version:
            raise HTTPException(409, "运行控制已变化，请重新确认")
        attempts = (
            await self.db.scalars(select(SimulationAttempt).where(SimulationAttempt.run_id == run_id))
        ).all()
        if body.action == "recover":
            for a in attempts:
                if a.state in ("needs_review", "collection_failed") and a.progress_url:
                    if len(a.children) != len(a.payload):
                        # Re-enumerate missing children while retaining collected receipts/results.
                        a.children = []
                    a.state, a.poll_count, a.next_poll_at, a.error, a.error_code = (
                        "submitted",
                        0,
                        None,
                        None,
                        None,
                    )
            # Recovery never clears uncertain submissions or creates a new POST.
        elif body.action == "resume":
            if run.control == "stopped":
                raise HTTPException(409, "已停止的剩余项不能恢复，请生成重跑预览")
            run.control = "active"
            config = await self.db.scalar(
                select(BacktestConfig).where(BacktestConfig.id == 1).with_for_update()
            )
            # An explicit resume may clear an indefinite quota block, never a Retry-After deadline.
            if config.blocked_until is None:
                config.blocked_reason = None
        elif body.action == "pause":
            if run.control == "stopped":
                raise HTTPException(409, "该运行已经停止")
            run.control = "paused"
        else:
            run.control = "stopped"
            for a in attempts:
                if a.state == "queued":
                    a.state = "skipped"
                    await self.db.execute(
                        update(BacktestItem)
                        .where(BacktestItem.attempt_id == a.id)
                        .values(
                            platform_status="skipped",
                            collection_status="not_required",
                            persistence_status="not_required",
                        )
                    )
        run.version += 1
        await refresh_status(self.db, run)
        await event(self.db, run, "control", {"action": body.action, "control": run.control})
        await self.db.flush()
        return await self.run(run_id)

    async def rerun(self, run_id, body):
        run = await locked_run(self.db, run_id)
        rows = (
            await self.db.scalars(
                select(BacktestItem).where(BacktestItem.run_id == run_id).order_by(BacktestItem.ordinal)
            )
        ).all()
        selected = [r for r in rows if r.id in set(body.item_ids)]
        if len(selected) != len(set(body.item_ids)):
            raise HTTPException(422, "重跑项不属于指定运行")
        if any(r.platform_status not in ("completed", "failed", "skipped") for r in selected):
            raise HTTPException(409, "仍在执行或结果未知的项须先核对，不能直接重跑")
        return await self.preview(
            PreviewInput(
                inline=DraftInput(
                    name=f"{run.name[:190]} · 重跑",
                    source=Source.model_validate({**run.source, "parent_run_id": run.id}),
                    candidates=[
                        Candidate(
                            client_item_id=r.client_item_id, expression=r.expression, settings=r.settings,
                            alpha_type=r.alpha_type, selection=r.selection, combo=r.combo
                        )
                        for r in selected
                    ],
                )
            ),
            preserve_source=True,
        )

    async def attach_reference(self, attempt_id, body):
        """Record a human-supplied original simulation; collection still verifies its input."""
        a = await self.db.get(SimulationAttempt, attempt_id)
        if not a:
            raise HTTPException(404, "执行尝试不存在")
        run = await locked_run(self.db, a.run_id)
        if run.version != body.version or a.state != "needs_review" or a.progress_url:
            raise HTTPException(409, "执行状态已变化或已有平台引用，请重新读取")
        duplicate = await self.db.scalar(
            select(SimulationAttempt.id).where(SimulationAttempt.progress_url == body.progress_url)
        )
        if duplicate:
            raise HTTPException(409, "此模拟引用已经关联其他执行尝试")
        a.progress_url, a.state, a.error, a.error_code = body.progress_url, "submitted", None, None
        a.next_poll_at, a.poll_count = None, 0
        run.version += 1
        await self.db.execute(
            update(BacktestItem)
            .where(BacktestItem.attempt_id == a.id)
            .values(platform_status="submitted", error=None)
        )
        await refresh_status(self.db, run)
        await event(
            self.db, run, "reference_attached", {"attempt_id": a.id, "progress_url": body.progress_url}
        )
        return await self.run(run.id)

    async def subset(self, preview_id, body):
        parent = await self.db.get(BacktestPreview, preview_id)
        if not parent:
            raise HTTPException(404, "预览不存在")
        excluded = set(body.exclude_ids)
        if not excluded.issubset({c["client_item_id"] for c in parent.candidates}):
            raise HTTPException(422, "排除集合包含未知候选")
        candidates = [c for c in parent.candidates if c["client_item_id"] not in excluded]
        if not candidates:
            raise HTTPException(422, "至少保留一条候选")
        return await self.preview(
            PreviewInput(inline=DraftInput(name=parent.name, source=parent.source, candidates=candidates)),
            preserve_source=True,
        )