paper-doc/backend/app/mcp/server.py

"""Build and run the paper-doc MCP server.

One build, three transports
---------------------------
``stdio`` is what a client spawns: Claude Code, Codex and the DeepSeek Harness
all launch a command and speak JSON-RPC over its stdin/stdout. Nothing may be
printed to stdout by anything in this process — a stray ``print`` corrupts the
protocol stream, so diagnostics go to stderr (the SDK's own logging already
does, see ``configure_logging``).

``streamable-http`` is what a client *on another machine* connects to: one
long-lived process serving ``POST /mcp``, optionally behind a bearer token.
``sse`` is the older HTTP shape, kept for clients that have not moved yet.

Nothing is registered here
--------------------------
The tools are registered by :func:`app.mcp.tools.register_all`, so this module
stays about transports and configuration. It is also the only module in the
package that knows a transport exists at all, which keeps the tool layer
testable by calling its functions directly.

The instructions below are sent once, at ``initialize``. They are the model's
map of the tool surface, so they describe the *workflow* and leave the tool
descriptions to describe the tools.
"""

from __future__ import annotations

import argparse
import os
import sys
from typing import Any

from mcp.server.mcpserver import MCPServer
from mcp.server.transport_security import TransportSecuritySettings

from app.core.config import get_settings
from app.mcp.tools import register_all

#: Version reported at ``initialize``. Bumped when the tool surface changes in
#: a way a client could notice, since it is the only version a client sees.
SERVER_VERSION = "1.0.0"

INSTRUCTIONS = """\
paper-doc 论文写作服务：管理论文（paper）、段落（paragraph）、句子（sentence）、
模板（template）和字段库（field）。工具名前缀对应这五类。

典型写作流程：
1. paper_list / paper_get 找到论文；没有就 paper_create（建议带 template_id）。
2. paper_outline 看段落骨架：每段有 position 和标题（如 “1. Introduction”）。
3. 写入正文：整篇用 paper_write_text（Markdown 标题分段），单段用 paragraph_write，
   多段一起用 paper_write。段落可以用 heading 指定，服务端会解析成 position。
4. 微调用 sentence_add / sentence_update / sentence_delete，不必重写整段。
5. 验证用 paper_document（text 形式）或 paragraph_get。

关键规则：
- paper_sentence 是一行一句，写入时句子以数组给出最稳妥；给 text 时用 split 指定切句方式。
- 段落由模板的 position 决定；换模板不丢正文，模板没定义的位置会显示为“未设定”。
- 删除论文会连带句子与引用且不可恢复，paper_delete 需要 confirm=true。
- 模板正被论文使用时不能删；字段正被模板使用时不能删——两者都会说明原因和解决办法。
"""


def build_server() -> MCPServer:
    """Assemble the server with every paper-doc tool registered."""
    server = MCPServer(
        name="paper-doc",
        title="paper-doc 论文写作",
        version=SERVER_VERSION,
        instructions=INSTRUCTIONS,
        # The library's INFO stream includes one line per refused call, which
        # is useful in a log and noise on a terminal shared with the protocol.
        log_level="WARNING",
    )
    register_all(server)
    return server


# --- HTTP transports ---------------------------------------------------------


class BearerTokenMiddleware:
    """Require ``Authorization: Bearer <token>`` on every MCP request.

    A token, not an OAuth flow: this server is reached over a private network by
    a client holding a shared secret, and the SDK's OAuth support would need a
    full authorization server to say the same thing. When no token is
    configured the check is skipped entirely — which is the right default for
    ``stdio`` and for a loopback bind, and the reason
    :func:`_serve_http` warns when a non-loopback host is used without one.
    """

    def __init__(self, app: Any, token: str) -> None:
        self.app = app
        self.token = token

    async def __call__(self, scope: Any, receive: Any, send: Any) -> None:
        if scope["type"] == "http":
            headers = {
                key.lower(): value for key, value in (scope.get("headers") or [])
            }
            presented = headers.get(b"authorization", b"").decode("latin-1")
            if presented != f"Bearer {self.token}":
                await send(
                    {
                        "type": "http.response.start",
                        "status": 401,
                        "headers": [
                            (b"content-type", b"application/json"),
                            (b"www-authenticate", b"Bearer"),
                        ],
                    }
                )
                await send(
                    {
                        "type": "http.response.body",
                        "body": b'{"error":"unauthorized"}',
                    }
                )
                return
        await self.app(scope, receive, send)


def _transport_security(host: str) -> TransportSecuritySettings | None:
    """DNS-rebinding protection, off for a non-loopback bind.

    The SDK turns it on by default for ``127.0.0.1``, where it costs nothing:
    the client is local and sends the right ``Host``. A server bound to a LAN
    address is reached by whatever name the client knows it by, so the
    allow-list would reject every request with a confusing 421 — the token is
    the gate there instead.
    """
    if host in ("127.0.0.1", "localhost", "::1"):
        return None
    return TransportSecuritySettings(enable_dns_rebinding_protection=False)


def _serve_http(
    server: MCPServer,
    *,
    transport: str,
    host: str,
    port: int,
    path: str | None,
    token: str | None,
) -> None:
    """Build a Starlette app for an HTTP transport and serve it with uvicorn."""
    import uvicorn

    security = _transport_security(host)

    if transport == "sse":
        app = server.sse_app(host=host, transport_security=security)
    else:
        app = server.streamable_http_app(
            streamable_http_path=path or "/mcp",
            host=host,
            transport_security=security,
        )

    if token:
        app = BearerTokenMiddleware(app, token)
    elif host not in ("127.0.0.1", "localhost", "::1"):
        print(
            f"[paper-doc-mcp] 警告：绑定在 {host} 且没有设置 MCP_HTTP_TOKEN，"
            "网络内任何人都能改论文",
            file=sys.stderr,
        )

    uvicorn.run(app, host=host, port=port, log_level="info")


# --- command line ------------------------------------------------------------


def build_parser() -> argparse.ArgumentParser:
    """The CLI, shared by the script entry point and the module entry point."""
    parser = argparse.ArgumentParser(
        prog="paper-doc-mcp",
        description="paper-doc 的 MCP 服务：论文、段落、句子、模板、字段的增删改查。",
    )
    parser.add_argument(
        "--transport",
        default=os.getenv("MCP_TRANSPORT", "stdio"),
        choices=["stdio", "http", "streamable-http", "sse"],
        help="stdio（默认，给 Claude Code / Codex / Harness 本地拉起）或 http",
    )
    parser.add_argument(
        "--host",
        default=os.getenv("MCP_HOST", "127.0.0.1"),
        help="HTTP 监听地址，默认 127.0.0.1；对外提供服务时用 0.0.0.0",
    )
    parser.add_argument(
        "--port",
        type=int,
        default=int(os.getenv("MCP_PORT", "8931")),
        help="HTTP 端口，默认 8931",
    )
    parser.add_argument(
        "--path",
        default=os.getenv("MCP_PATH"),
        help="streamable-http 的路径，默认 /mcp",
    )
    parser.add_argument(
        "--token",
        default=os.getenv("MCP_HTTP_TOKEN"),
        help="HTTP 传输的 Bearer 口令；对外暴露时务必设置",
    )
    parser.add_argument(
        "--check",
        action="store_true",
        help="只做自检：连数据库、列出工具名，然后退出（不启动服务）",
    )
    return parser


def main(argv: list[str] | None = None) -> int:
    """Entry point of ``scripts/mcp_server.py`` and ``python -m app.mcp``."""
    args = build_parser().parse_args(argv)
    server = build_server()

    if args.check:
        from app.mcp.selfcheck import run_selfcheck

        return run_selfcheck(server)

    settings = get_settings()
    print(
        f"[paper-doc-mcp] transport={args.transport} db={settings.safe_database_url}",
        file=sys.stderr,
    )

    if args.transport == "stdio":
        server.run(transport="stdio")
    else:
        _serve_http(
            server,
            transport=args.transport,
            host=args.host,
            port=args.port,
            path=args.path,
            token=args.token,
        )
    return 0


__all__ = ["BearerTokenMiddleware", "INSTRUCTIONS", "SERVER_VERSION", "build_server", "main"]