worldquant-alpha-system/backend/tests/test_mcp_templates.py

"""External-model templates share the real asset/expansion path, using synthetic research."""

import asyncio
import json
from copy import deepcopy
from unittest.mock import Mock

import httpx
import httpx2
import pytest
from mcp import ClientSession
from mcp.client.streamable_http import streamable_http_client
from sqlalchemy import func, select

from app.mcp_api.auth import SCOPES, authenticate
from app.mcp_api.server import MCPResearchServer
from app.models import (
    Account,
    BacktestItem,
    BacktestResult,
    BacktestRun,
    Job,
    MCPAudit,
    ResearchAsset,
    ResearchRequest,
    ResearchRevision,
)
from tests.test_backtests import candidate, execute
from tests.test_mcp import ENDPOINT, credentials, invoke, submission
from tests.test_mcp import mcp_app as mcp_app_fixture
from tests.test_research_workspace import catalog as catalog_fixture
from tests.test_research_workspace import research_input as research_input_fixture

mcp_test_app = mcp_app_fixture
catalog = catalog_fixture
research_input = research_input_fixture
TOOL = "create_research_template"


@pytest.fixture
async def app(mcp_test_app):
    # The catalog fixture connects this same synthetic account; keep its email consistent.
    async with mcp_test_app.state.sessions.begin() as db:
        (await db.get(Account, 1)).email = "test@example.com"
    return mcp_test_app


def template_request(item_id):
    return {
        "template": {
            "name": "研究后的字段排序模板",
            "description": "比较相同类型字段及偏移参数，检验研究假设的稳定性。",
            "expression": "rank({field}) + {offset}",
            "variables": {
                "field": {"kind": "field", "field_type": "MATRIX", "values": ["TEST_FIN_001", "TEST_FIN_002"]},
                "offset": {"kind": "integer", "values": [0, 1, 5]},
            },
        },
        "hypothesis": "来源结果值得进一步检验，字段替换和参数组合仍需独立回测。",
        "source_item_ids": [item_id],
        "reference": "synthetic-research-round-1",
        "idempotency_key": "template-1",
    }


@pytest.fixture
async def completed_source(app):
    principal, _ = await credentials(app)
    run = await invoke(app, principal, "submit_backtests", submission(items=[
        candidate() | {"expression": "rank(TEST_FIN_001) + 0"},
    ]))
    await execute(app, app.state.runner.backtests, run["backtest_run_id"])
    result = await invoke(app, principal, "get_backtest_results", {"run_id": run["backtest_run_id"]})
    item_id = result["items"][0]["id"]
    async with app.state.sessions.begin() as db:
        row = await db.get(BacktestResult, item_id)
        row.snapshot = {**row.snapshot, "is": {"sharpe": 1.9, "fitness": 0.9, "checks": [
            {"name": "LOW_FITNESS", "result": "FAIL"}, {"name": "FUTURE_CHECK", "result": "NEW_STATUS"},
        ]}}
    result = await invoke(app, principal, "get_backtest_results", {"run_id": run["backtest_run_id"]})
    return result["items"][0]


async def test_sdk_template_creation_frozen_evidence_and_web_expansion(app, logged_in, completed_source, research_input, monkeypatch):
    principal, secret = await credentials(app, {"research:read", "research:write"})
    body = template_request(completed_source["id"])
    async with app.state.sessions() as db:
        job_count = await db.scalar(select(func.count()).select_from(Job))
    factory = Mock(side_effect=AssertionError("Template persistence must not call a model"))
    monkeypatch.setattr(app.state.ai, "model_factory", factory)
    app.state.runner.wake.clear()
    app.state.runner.backtests.wake.clear()
    async with httpx2.AsyncClient(transport=httpx2.ASGITransport(app=app),
                                 headers={"Authorization": f"Bearer {secret}"}) as http:
        async with streamable_http_client("http://testserver/api/v1/mcp/", http_client=http) as streams:
            async with ClientSession(streams[0], streams[1]) as client:
                await client.initialize()
                listed = {tool.name: tool for tool in (await client.list_tools()).tools}
                tool = listed[TOOL]
                assert "submit_backtests" not in listed
                assert not tool.annotations.read_only_hint and not tool.annotations.destructive_hint
                assert tool.annotations.idempotent_hint and not tool.annotations.open_world_hint
                assert {"template", "hypothesis", "idempotency_key"} <= set(tool.input_schema["required"])
                assert tool.input_schema["additionalProperties"] is False
                caps = await client.call_tool("get_research_capabilities", {})
                assert caps.structured_content["templates"]["create_with"] == TOOL
                result = await client.call_tool(TOOL, body)
                assert not result.is_error, result
                saved = result.structured_content
                assert json.loads(result.content[0].text) == saved
    assert saved["id"] == saved["template_id"] and saved["version"] == 1
    assert saved["combination_count"] == "6" and saved["web_url"] == "http://testserver/#templates"
    assert saved["validation"]["expanded_candidates"] == "not_validated"
    assert saved["provenance"]["source"] == {"kind": "mcp", "reference": body["reference"]}
    assert saved["provenance"]["hypothesis"] == body["hypothesis"]
    assert saved["provenance"]["source_items"] == [completed_source]
    factory.assert_not_called()
    assert not app.state.runner.wake.is_set() and not app.state.runner.backtests.wake.is_set()
    async with app.state.sessions.begin() as db:
        audit = await db.scalar(select(MCPAudit).where(MCPAudit.tool == TOOL))
        assert audit.business_id == saved["id"] and audit.result_code == "OK"
        assert secret not in json.dumps(audit.__dict__, default=str)
        row = await db.get(BacktestResult, completed_source["id"])
        row.snapshot = {**row.snapshot, "is": {"sharpe": 999}}
    # The browser reads the same library; later source changes do not rewrite the saved evidence.
    listing = (await logged_in.get("/api/v1/research/assets?kind=template")).json()
    assert listing["items"][0]["id"] == saved["id"]
    stored = (await logged_in.get(f'/api/v1/research/assets/{saved["id"]}?version=1')).json()
    assert stored["provenance"]["source_items"][0]["metrics"]["is"]["sharpe"] == 1.9
    assert stored["provenance"]["source_items"][0]["checks"]["counts"]["FAIL"] == 1
    assert stored["provenance"]["source_items"][0]["checks"]["counts"]["UNKNOWN"] == 1
    expanded = await logged_in.post("/api/v1/research/experiments", json={
        "asset_id": saved["id"], "version": saved["version"], "input_ids": [research_input["id"]],
        "hypothesis": body["hypothesis"], "settings": completed_source["settings"], "limit": 100,
    })
    assert expanded.status_code == 201, expanded.text
    experiment = expanded.json()
    assert {c["expression"] for c in experiment["candidates"]} == {
        f"rank({field}) + {offset}" for field in ["TEST_FIN_001", "TEST_FIN_002"] for offset in [0, 1, 5]
    }
    assert all("validation" not in c for c in experiment["candidates"])
    assert experiment["evidence"]["template"]["provenance"] == stored["provenance"]
    async with app.state.sessions() as db:
        assert await db.scalar(select(func.count()).select_from(BacktestRun)) == 1
        assert await db.scalar(select(func.count()).select_from(Job)) == job_count


async def test_template_idempotency_concurrency_rotation_and_name_conflict(app, completed_source):
    principal, _ = await credentials(app, {"research:read", "research:write"})
    body = template_request(completed_source["id"])
    first, repeated = await asyncio.gather(invoke(app, principal, TOOL, body), invoke(app, principal, TOOL, body))
    assert first == repeated
    rotated, _ = await credentials(app, {"research:read", "research:write"})
    restarted = MCPResearchServer(app.state.sessions, app.state.runner, app.state.settings)
    replay = await restarted.invoke(rotated, TOOL, body)
    assert not replay.is_error and replay.structured_content == first
    conflict = await app.state.mcp.invoke(principal, TOOL, body | {"hypothesis": "另一假设"})
    assert conflict.structured_content["error"]["code"] == "IDEMPOTENCY_CONFLICT"
    same_name = await app.state.mcp.invoke(principal, TOOL, body | {"idempotency_key": "new-key"})
    assert same_name.structured_content["error"]["code"] == "TEMPLATE_NAME_CONFLICT"
    async with app.state.sessions() as db:
        assert await db.scalar(select(func.count()).select_from(ResearchAsset)) == 1
        assert await db.scalar(select(func.count()).select_from(ResearchRevision)) == 1
        requests = list(await db.scalars(select(ResearchRequest).where(ResearchRequest.operation == TOOL)))
        assert len(requests) == 1 and requests[0].business_id == first["id"]


async def test_template_invalid_inputs_and_missing_sources_are_atomic(app, completed_source):
    principal, _ = await credentials(app)
    valid = template_request(completed_source["id"])
    variants = [
        {"source_item_ids": [completed_source["id"]] * 21},
        {"source_item_ids": [completed_source["id"]] * 2}, {"hypothesis": " "},
        {"force": True}, {"idempotency_key": ""},
        {"template": valid["template"] | {"expression": "rank({missing})"}},
        {"template": valid["template"] | {"category": "fragment"}},
        {"template": valid["template"] | {"name": " "}},
        {"template": valid["template"] | {"expression": " ", "variables": {}}},
    ]
    missing_type = deepcopy(valid["template"])
    missing_type["variables"]["field"].pop("field_type")
    variants.append({"template": missing_type})
    for changed in variants:
        failed = await app.state.mcp.invoke(principal, TOOL, valid | changed)
        assert failed.is_error and failed.structured_content["error"]["code"] == "INVALID_INPUT"
    missing = await app.state.mcp.invoke(principal, TOOL, valid | {"source_item_ids": [completed_source["id"], "missing"]})
    assert missing.structured_content["error"]["code"] == "NOT_FOUND"
    assert missing.structured_content["error"]["affected_items"] == ["missing"]
    async with app.state.sessions() as db:
        assert await db.scalar(select(func.count()).select_from(ResearchAsset)) == 0
        assert await db.scalar(select(func.count()).select_from(ResearchRequest).where(ResearchRequest.operation == TOOL)) == 0
    await invoke(app, principal, TOOL, valid)


@pytest.mark.parametrize("state", ["pending", "collection_failed", "unsaved", "incomplete_snapshot", "missing_snapshot"])
async def test_template_rejects_incomplete_source_without_consuming_key(app, completed_source, state):
    principal, _ = await credentials(app)
    body = template_request(completed_source["id"])
    async with app.state.sessions.begin() as db:
        item = await db.get(BacktestItem, completed_source["id"])
        result = await db.get(BacktestResult, item.id)
        if state == "pending":
            item.platform_status = "pending"
        elif state == "collection_failed":
            item.collection_status = "failed"
        elif state == "unsaved":
            item.persistence_status = "pending"
        elif state == "incomplete_snapshot":
            result.complete = False
        else:
            await db.delete(result)
    failed = await app.state.mcp.invoke(principal, TOOL, body)
    assert failed.is_error and failed.structured_content["error"]["code"] == "SOURCE_NOT_READY"
    assert failed.structured_content["error"]["affected_items"] == body["source_item_ids"]
    async with app.state.sessions() as db:
        assert await db.scalar(select(func.count()).select_from(ResearchAsset)) == 0
        assert await db.scalar(select(func.count()).select_from(ResearchRequest).where(ResearchRequest.operation == TOOL)) == 0


@pytest.mark.parametrize("scopes", [{"research:read"}, SCOPES - {"research:write"}])
async def test_template_write_permission_not_granted_to_old_keys(app, scopes):
    _, secret = await credentials(app, scopes)
    async with httpx.AsyncClient(transport=httpx.ASGITransport(app=app), base_url="http://testserver",
        headers={"Authorization": f"Bearer {secret}", "Accept": "application/json, text/event-stream"}) as http:
        listed = await http.post(ENDPOINT, json={"jsonrpc": "2.0", "id": 1, "method": "tools/list"})
        assert TOOL not in {tool["name"] for tool in listed.json()["result"]["tools"]}
        denied = await http.post(ENDPOINT, json={"jsonrpc": "2.0", "id": 2, "method": "tools/call",
            "params": {"name": TOOL, "arguments": template_request("missing")}})
        assert denied.status_code == 403


async def test_browser_can_issue_template_only_and_all_permissions(app, logged_in):
    for scopes in [{"research:read", "research:write"}, SCOPES]:
        response = await logged_in.post("/api/v1/mcp-tokens", json={"name": "template writer", "scopes": sorted(scopes)})
        assert response.status_code == 201, response.text
        async with app.state.sessions() as db:
            principal = await authenticate(db, response.json()["token"])
            assert principal.scopes == scopes


async def test_mcp_template_version_is_idempotent_and_preserves_history(app, completed_source):
    principal, _ = await credentials(app, {"research:read", "research:write"})
    body = template_request(completed_source["id"])
    created = await invoke(app, principal, TOOL, body)
    content = deepcopy(body["template"])
    content["variables"]["field"] = {"kind": "field", "field_type": "MATRIX", "description": "数据准备中的矩阵字段"}
    update = {**body, "template": content, "template_id": created["id"],
              "expected_version": 1, "idempotency_key": "version-2"}
    tool = "create_research_template_version"
    first = await invoke(app, principal, tool, update)
    replay = await invoke(app, principal, tool, update)
    assert first == replay and first["version"] == 2
    assert first["combination_count"] is None
    assert first["provenance"]["parent_template"] == {"id": created["id"], "version": 1}
    read = await invoke(app, principal, "get_research_template", {"template_id": created["id"], "version": 1})
    assert read["content"]["variables"]["field"]["values"] == ["TEST_FIN_001", "TEST_FIN_002"]
    listed = await invoke(app, principal, "search_research_templates", {"q": content["name"]})
    assert listed["items"][0]["version"] == 2
    stale = await app.state.mcp.invoke(principal, tool, update | {"idempotency_key": "stale-version"})
    assert stale.is_error and stale.structured_content["error"]["code"] == "CONFLICT"
    async with app.state.sessions() as db:
        assert await db.scalar(select(func.count()).select_from(ResearchRevision)) == 2
    denied, _ = await credentials(app, {"research:read"})
    from fastapi import HTTPException
    with pytest.raises(HTTPException) as forbidden:
        await app.state.mcp.invoke(denied, tool, update)
    assert forbidden.value.status_code == 403


async def test_template_without_result_sources_can_be_created_and_versioned(app):
    principal, _ = await credentials(app, {"research:read", "research:write"})
    body = template_request("unused")
    body.pop("source_item_ids")
    created = await invoke(app, principal, TOOL, body)
    assert created["provenance"]["source_items"] == []
    assert created["validation"]["source_evidence"] == "not_provided"
    revised = await invoke(app, principal, "create_research_template_version", {
        **body, "template_id": created["id"], "expected_version": 1, "idempotency_key": "no-source-v2",
        "source_item_ids": [],
    })
    assert revised["version"] == 2 and revised["provenance"]["source_items"] == []
    async with app.state.sessions() as db:
        assert await db.scalar(select(func.count()).select_from(BacktestRun)) == 0


async def template_collection(app, principal, research_input, count=130):
    body = template_request("unused")
    body.pop("source_item_ids")
    body["template"]["variables"]["field"]["values"] = ["TEST_FIN_001"]
    body["template"]["variables"]["offset"]["values"] = list(range(count))
    saved = await invoke(app, principal, TOOL, body)
    args = {
        "template_id": saved["id"], "version": saved["version"],
        "preparation_refs": [{"id": research_input["preparation_id"], "version": research_input["preparation_version"]}],
        "settings": candidate()["settings"], "limit": count, "idempotency_key": "expand-collection",
    }
    return saved, args


async def test_template_collection_mcp_paging_execution_and_replay(app, logged_in, research_input):
    from sqlalchemy import delete

    from app.models import BacktestPreview, CatalogResource, ResearchExperiment

    principal, _ = await credentials(app)
    saved, args = await template_collection(app, principal, research_input)
    async with app.state.sessions.begin() as db:
        await db.execute(delete(CatalogResource))
    app.state.runner.backtests.wake.clear()
    first, retry = await asyncio.gather(*[invoke(app, principal, "expand_research_template", args) for _ in range(2)])
    assert first == retry and first["total"] == 130 and len(first["items"]) == 25
    assert first["has_more"] and not first["starts_backtests"]
    assert first["template"]["id"] == saved["id"] and first["template"]["version"] == 1
    assert not app.state.runner.backtests.wake.is_set()
    collected = []
    for offset in (0, 100):
        result = await invoke(app, principal, "get_template_candidates", {
            "experiment_id": first["experiment_id"], "limit": 100, "offset": offset,
        })
        collected += result["items"]
    assert len(collected) == 130 and not result["has_more"]
    assert all("validation" not in item for item in collected)
    async with app.state.sessions() as db:
        assert await db.scalar(select(func.count()).select_from(ResearchExperiment)) == 1
        assert await db.scalar(select(func.count()).select_from(BacktestRun)) == 0
    args = {"experiment_id": first["experiment_id"], "candidate_ids": [c["client_item_id"] for c in collected],
            "idempotency_key": "execute-collection"}
    run, replay = await asyncio.gather(*[invoke(app, principal, "start_template_backtest", args) for _ in range(2)])
    assert run == replay and run["total"] == 130
    assert run["source"]["kind"] == "template" and run["source"]["research_id"] == first["experiment_id"]
    assert run["source"]["input_snapshot_ids"]
    assert app.state.runner.backtests.wake.is_set()
    rotated, _ = await credentials(app)
    assert await invoke(app, rotated, "start_template_backtest", args) == run
    conflict = await app.state.mcp.invoke(principal, "start_template_backtest", args | {"candidate_ids": ["c1"]})
    assert conflict.is_error and conflict.structured_content["error"]["code"] == "IDEMPOTENCY_CONFLICT"
    async with app.state.sessions() as db:
        assert await db.scalar(select(func.count()).select_from(BacktestRun)) == 1
        assert await db.scalar(select(func.count()).select_from(BacktestPreview)) == 1
        row = await db.get(BacktestRun, run["backtest_run_id"])
        assert row.ai_context["mcp_token_id"] == principal.token_id
        audit = await db.scalar(select(MCPAudit).where(MCPAudit.tool == "expand_research_template"))
        assert audit.business_id == first["experiment_id"]
    record = (await logged_in.get(f"/api/v1/research/experiments/{first['experiment_id']}")).json()
    assert record["backtest_run_ids"] == [run["backtest_run_id"]]


async def test_template_mcp_failures_are_atomic_and_do_not_consume_keys(app, research_input):
    from app.models import BacktestPreview, ResearchExperiment

    principal, _ = await credentials(app)
    saved, args = await template_collection(app, principal, research_input, 2)
    for changed in [args | {"settings": args["settings"] | {"region": "EUR"}},
                    args | {"preparation_refs": [{**args["preparation_refs"][0], "version": 999}]}]:
        result = await app.state.mcp.invoke(principal, "expand_research_template", changed)
        assert result.is_error
    async with app.state.sessions() as db:
        assert await db.scalar(select(func.count()).select_from(ResearchExperiment)) == 0
        assert not await db.scalar(select(ResearchRequest).where(ResearchRequest.operation == "expand_research_template"))
    collection = await invoke(app, principal, "expand_research_template", args)
    conflict = await app.state.mcp.invoke(principal, "expand_research_template", args | {"seed": 1})
    assert conflict.structured_content["error"]["code"] == "IDEMPOTENCY_CONFLICT"
    request = {"experiment_id": collection["experiment_id"], "candidate_ids": ["c1"], "idempotency_key": "execute"}
    for ids in [[], ["unknown"], ["c1", "c1"]]:
        result = await app.state.mcp.invoke(principal, "start_template_backtest", request | {"candidate_ids": ids})
        assert result.is_error
    result = await app.state.mcp.invoke(principal, "start_template_backtest", request | {"expression": "rank(other)"})
    assert result.is_error and result.structured_content["error"]["code"] == "INVALID_INPUT"
    async with app.state.sessions() as db:
        assert await db.scalar(select(func.count()).select_from(BacktestRun)) == 0
        assert await db.scalar(select(func.count()).select_from(BacktestPreview)) == 0
        assert not await db.scalar(select(ResearchRequest).where(ResearchRequest.operation == "start_template_backtest"))
    await invoke(app, principal, "start_template_backtest", request)


async def test_template_tool_discovery_and_execute_permissions(app, research_input):
    from fastapi import HTTPException

    principal, secret = await credentials(app, {"research:read", "research:write"})
    async with httpx.AsyncClient(transport=httpx.ASGITransport(app=app), base_url="http://testserver",
                               headers={"Authorization": f"Bearer {secret}", "Accept": "application/json, text/event-stream"}) as http:
        listed = (await http.post(ENDPOINT, json={"jsonrpc": "2.0", "id": 1, "method": "tools/list"})).json()
        tools = {t["name"]: t for t in listed["result"]["tools"]}
        assert "expand_research_template" in tools and "get_template_candidates" in tools
        assert "start_template_backtest" not in tools
        assert "source_item_ids" not in tools[TOOL]["inputSchema"]["required"]
        assert tools["expand_research_template"]["annotations"]["idempotentHint"]
    _, args = await template_collection(app, principal, research_input, 2)
    collection = await invoke(app, principal, "expand_research_template", args)
    with pytest.raises(HTTPException) as denied:
        await app.state.mcp.invoke(principal, "start_template_backtest", {
            "experiment_id": collection["experiment_id"], "candidate_ids": ["c1"], "idempotency_key": "denied",
        })
    assert denied.value.status_code == 403
    reader, _ = await credentials(app, {"research:read"})
    await invoke(app, reader, "get_template_candidates", {"experiment_id": collection["experiment_id"]})
    with pytest.raises(HTTPException):
        await app.state.mcp.invoke(reader, "expand_research_template", args)