codex-app-server/cloud/mcp_server.py

"""云端任务管理的 MCP 入口：把 REST API 翻成 MCP tools。

方向说明（别跟任务内 MCP 搞混）：项目内 agent 用的 MCP 是
`[mcp_servers.*]`（给 agent 接外部工具）；本文件是反方向——
让外部 MCP 客户端（编辑器、agent）用 tools/list / tools/call
管云端任务，底层一律走已有的 REST API，不另开一套任务逻辑。

协议：stdio 上每行一个 JSON-RPC 2.0 请求，支持 initialize、
ping、tools/list、tools/call；notifications/* 不回包（按 MCP
惯例）。tools/call 的结果一律是 content=[{type: text,
text: <REST 响应 JSON>}]，REST 报错时 isError=true，文本里
含 HTTP 状态与服务端 error，不伪造成功。

鉴权：只认调用方自己的 X-API-Key，从环境变量
CODEX_CLOUD_API_KEY 读（缺了就启动失败，绝不读文件、
绝不打印 key）；BASE 从 CODEX_CLOUD_BASE_URL 读，缺省
https://api.diante.me。MCP 服务本身不存任何凭证、不建
用户、不绕过服务端权限——服务端按 key 归属做 viewer/
contributor 等权限校验，403 原样透传。

工具（不包管理端）：
  health、list_projects、list_tasks、get_task、create_task、
  cancel_task、get_task_cost、get_task_trace、get_task_files、
  get_task_tree、
  get_task_result_full、retry_task、merge_task。
create_task 只传 project_id + prompt + 可选 executor/model/
effort/approval/priority，其余字段走服务端缺省；cancel_task
只做取消，不做真 kill 之外的任何事（真 kill 由 worker 的
kill_requests 机制处理，不在这里另起一套）；merge_task 与
retry_task 是写操作，调用方 key 必须有相应项目权限，服务端
按状态机拒绝非法转换（例如未完成任务 merge 会 409），MCP
层不绕过；get_task_files 的 path 参数是 worktree 内相对
路径，服务端 _safe_worktree_path 防逃逸，MCP 层不另做路径
展开，绝不传绝对路径。
"""

import json
import os
import sys
import urllib.error
import urllib.parse
import urllib.request

__all__ = ["MCPServer", "main"]


def _tool(name, description, properties, required):
    return {
        "name": name,
        "description": description,
        "inputSchema": {
            "type": "object",
            "properties": properties,
            "required": required,
        },
    }


TOOLS = [
    _tool("health", "云端控制面健康检查（不含任务数据）。",
          {}, []),
    _tool("list_projects", "列出当前 key 可见的项目。",
          {}, []),
    _tool("list_tasks", "列出某项目的任务（需 project_id）。",
          {"project_id": {"type": "string"}}, ["project_id"]),
    _tool("get_task", "取单个任务详情（状态、结果、成本字段等）。",
          {"task_id": {"type": "string"}}, ["task_id"]),
    _tool("create_task", "在某项目建一个任务（需 project_id 与 prompt）。",
          {"project_id": {"type": "string"},
           "prompt": {"type": "string"},
           "executor": {"type": "string"},
           "model": {"type": "string"},
           "effort": {"type": "string"},
           "approval": {"type": "string"},
           "priority": {"type": "integer"}},
          ["project_id", "prompt"]),
    _tool("cancel_task", "取消一个任务（标 interrupted，容器真杀由 worker 处理）。",
          {"task_id": {"type": "string"}}, ["task_id"]),
    _tool("get_task_cost", "取单个任务的成本明细。",
          {"task_id": {"type": "string"}}, ["task_id"]),
    _tool("get_task_trace", "取单个任务的链路条目（摘要，完整结果在任务详情）。",
          {"task_id": {"type": "string"}}, ["task_id"]),
    _tool("get_task_files", "列/读任务 worktree 里的文件；不传 path 列顶层，传相对路径读文件或列目录。",
          {"task_id": {"type": "string"},
           "path": {"type": "string"}}, ["task_id"]),
    _tool("get_task_tree", "取任务 worktree 的结构化目录树（只含名、类型、大小，"
          "不读文件内容）；不传 path 列 worktree 根，"
          "传相对路径列该子目录；max_depth 1..5 缺省 3，max_entries 1..500 缺省 200，截断时 truncated=True。",
          {"task_id": {"type": "string"},
           "path": {"type": "string"},
           "max_depth": {"type": "integer"},
           "max_entries": {"type": "integer"}}, ["task_id"]),
    _tool("get_task_result_full", "取单个任务的完整结果文本（result_text 只是摘要）。",
          {"task_id": {"type": "string"}}, ["task_id"]),
    _tool("retry_task", "重试一个失败/中断的任务（服务端状态机校验）。",
          {"task_id": {"type": "string"}}, ["task_id"]),
    _tool("merge_task", "合并一个已完成任务的改动（服务端状态机校验，未完成会拒绝）。",
          {"task_id": {"type": "string"}}, ["task_id"]),
    _tool("pause_task", "暂停一个进行中任务（服务端状态机校验）。",
          {"task_id": {"type": "string"}}, ["task_id"]),
    _tool("goal_resume_task", "断点续跑一个 goal 任务（服务端状态机校验）。",
          {"task_id": {"type": "string"}}, ["task_id"]),
    _tool("create_followup", "给进行中任务发后续消息；"
          "mode=queued 排队进队尾不打断本轮，direct_send 立即送达且本轮结束后优先开下一轮。",
          {"task_id": {"type": "string"},
           "body": {"type": "string"},
           "mode": {"type": "string"}}, ["task_id", "body"]),
    _tool("get_task_agents", "取单个任务的 agent 分工条目。",
          {"task_id": {"type": "string"}}, ["task_id"]),
    _tool("list_workers", "列出 worker 聚合状态（普通用户只见在线/离线数，管理员见完整清单）。",
          {}, []),
    _tool("get_metrics", "取控制面可观测性聚合：任务计数、成功率、平均耗时、成本合计（只统计调用者是成员的项目）。",
          {}, []),
    _tool("list_project_members", "列出某项目的成员。",
          {"project_id": {"type": "string"}}, ["project_id"]),
    _tool("get_models", "取模型一览（价目与默认模型，服务端口径）。",
          {}, []),
    _tool("list_project_memories", "列出某项目的项目级记忆（此前任务沉淀的经验/约定，"
          "新任务会自动注入；top_k 可选，默认 20）。",
          {"project_id": {"type": "string"},
           "top_k": {"type": "integer"}}, ["project_id"]),
    _tool("create_project_memory", "给某项目记一条项目级记忆（kind 与 content 必填）。",
          {"project_id": {"type": "string"},
           "kind": {"type": "string"},
           "content": {"type": "string"}}, ["project_id", "kind", "content"]),
    _tool("supersede_project_memory", "废弃某项目的一条记忆，可同时给替代内容（不传则只废弃）。",
          {"project_id": {"type": "string"},
           "memory_id": {"type": "string"},
           "replacement_content": {"type": "string"}}, ["project_id", "memory_id"]),
]


class MCPServer:
    """无状态的 REST→MCP 翻译层；transport 可注入，单测用假 transport。"""

    def __init__(self, base_url, api_key, transport=None):
        if not api_key:
            raise ValueError("CODEX_CLOUD_API_KEY 未设置")
        self.base_url = base_url.rstrip("/")
        self.api_key = api_key
        self.transport = transport or self._http

    def _http(self, method, path, body=None):
        url = self.base_url + path
        data = json.dumps(body).encode("utf-8") if body is not None else None
        req = urllib.request.Request(url, data=data, method=method)
        req.add_header("X-API-Key", self.api_key)
        if data is not None:
            req.add_header("Content-Type", "application/json")
        try:
            with urllib.request.urlopen(req, timeout=60) as resp:
                raw = resp.read().decode("utf-8") or "{}"
                return resp.status, json.loads(raw)
        except urllib.error.HTTPError as exc:
            raw = exc.read().decode("utf-8") or "{}"
            try:
                payload = json.loads(raw)
            except ValueError:
                payload = {"raw": raw[:500]}
            return exc.code, payload
        except Exception as exc:  # noqa: BLE001 - 网络层一律转错误结果
            return 0, {"error": str(exc)[:500]}

    def call_tool(self, name, args):
        """返回 (is_error, payload_dict)。未知工具、缺参一律 is_error。"""
        args = args or {}
        if name == "health":
            return self._wrap(*self.transport("GET", "/health"))
        if name == "list_projects":
            return self._wrap(*self.transport("GET", "/projects"))
        if name == "list_tasks":
            if not args.get("project_id"):
                return True, {"error": "project_id required"}
            q = urllib.parse.urlencode({"project_id": args["project_id"]})
            return self._wrap(*self.transport("GET", "/tasks?" + q))
        if name == "get_task":
            if not args.get("task_id"):
                return True, {"error": "task_id required"}
            return self._wrap(*self.transport(
                "GET", "/tasks/" + args["task_id"]))
        if name == "create_task":
            if not args.get("project_id") or not args.get("prompt"):
                return True, {"error": "project_id and prompt required"}
            body = {"project_id": args["project_id"],
                    "prompt": args["prompt"]}
            for key in ("executor", "model", "effort", "approval",
                        "priority"):
                if args.get(key) is not None:
                    body[key] = args[key]
            return self._wrap(*self.transport("POST", "/tasks", body))
        if name == "cancel_task":
            if not args.get("task_id"):
                return True, {"error": "task_id required"}
            return self._wrap(*self.transport(
                "POST", "/tasks/" + args["task_id"] + "/cancel", {}))
        if name == "get_task_cost":
            if not args.get("task_id"):
                return True, {"error": "task_id required"}
            return self._wrap(*self.transport(
                "GET", "/tasks/" + args["task_id"] + "/cost"))
        if name == "get_task_trace":
            if not args.get("task_id"):
                return True, {"error": "task_id required"}
            return self._wrap(*self.transport(
                "GET", "/tasks/" + args["task_id"] + "/trace"))
        if name == "get_task_files":
            if not args.get("task_id"):
                return True, {"error": "task_id required"}
            path = "/tasks/" + args["task_id"] + "/files"
            if args.get("path"):
                path += "?" + urllib.parse.urlencode(
                    {"path": args["path"]})
            return self._wrap(*self.transport("GET", path))
        if name == "get_task_tree":
            if not args.get("task_id"):
                return True, {"error": "task_id required"}
            query = {}
            if args.get("path"):
                query["path"] = args["path"]
            if args.get("max_depth") is not None:
                query["max_depth"] = args["max_depth"]
            if args.get("max_entries") is not None:
                query["max_entries"] = args["max_entries"]
            path = "/tasks/" + args["task_id"] + "/tree"
            if query:
                path += "?" + urllib.parse.urlencode(query)
            return self._wrap(*self.transport("GET", path))
        if name == "get_task_result_full":
            if not args.get("task_id"):
                return True, {"error": "task_id required"}
            return self._wrap(*self.transport(
                "GET", "/tasks/" + args["task_id"] + "/result-full"))
        if name == "retry_task":
            if not args.get("task_id"):
                return True, {"error": "task_id required"}
            return self._wrap(*self.transport(
                "POST", "/tasks/" + args["task_id"] + "/retry", {}))
        if name == "merge_task":
            if not args.get("task_id"):
                return True, {"error": "task_id required"}
            return self._wrap(*self.transport(
                "POST", "/tasks/" + args["task_id"] + "/merge", {}))
        if name == "pause_task":
            if not args.get("task_id"):
                return True, {"error": "task_id required"}
            return self._wrap(*self.transport(
                "POST", "/tasks/" + args["task_id"] + "/pause", {}))
        if name == "goal_resume_task":
            if not args.get("task_id"):
                return True, {"error": "task_id required"}
            return self._wrap(*self.transport(
                "POST", "/tasks/" + args["task_id"] + "/goal-resume", {}))
        if name == "create_followup":
            if not args.get("task_id") or not args.get("body"):
                return True, {"error": "task_id and body required"}
            body = {"body": args["body"]}
            if args.get("mode") is not None:
                body["mode"] = args["mode"]
            return self._wrap(*self.transport(
                "POST", "/tasks/" + args["task_id"] + "/followups", body))
        if name == "get_task_agents":
            if not args.get("task_id"):
                return True, {"error": "task_id required"}
            return self._wrap(*self.transport(
                "GET", "/tasks/" + args["task_id"] + "/agents"))
        if name == "list_workers":
            return self._wrap(*self.transport("GET", "/workers"))
        if name == "get_metrics":
            return self._wrap(*self.transport("GET", "/metrics"))
        if name == "list_project_members":
            if not args.get("project_id"):
                return True, {"error": "project_id required"}
            return self._wrap(*self.transport(
                "GET", "/projects/" + args["project_id"] + "/members"))
        if name == "get_models":
            return self._wrap(*self.transport("GET", "/models"))
        if name == "list_project_memories":
            if not args.get("project_id"):
                return True, {"error": "project_id required"}
            path = "/projects/" + args["project_id"] + "/identity-memories"
            if args.get("top_k") is not None:
                path += "?" + urllib.parse.urlencode(
                    {"top_k": args["top_k"]})
            return self._wrap(*self.transport("GET", path))
        if name == "create_project_memory":
            if not args.get("project_id") or not args.get("kind") \
                    or not args.get("content"):
                return True, {"error": "project_id, kind and content required"}
            return self._wrap(*self.transport(
                "POST", "/projects/" + args["project_id"]
                + "/identity-memories",
                {"kind": args["kind"], "content": args["content"]}))
        if name == "supersede_project_memory":
            if not args.get("project_id") or not args.get("memory_id"):
                return True, {"error": "project_id and memory_id required"}
            body = {}
            if args.get("replacement_content") is not None:
                body["replacement_content"] = args["replacement_content"]
            return self._wrap(*self.transport(
                "POST", "/projects/" + args["project_id"]
                + "/identity-memories/" + args["memory_id"] + "/supersede",
                body))
        return True, {"error": "unknown tool: %s" % name}

    @staticmethod
    def _wrap(status, payload):
        if status == 0 or status >= 400:
            return True, {"http_status": status, "error": payload}
        return False, payload

    def handle(self, request):
        """处理一个 JSON-RPC 请求，返回响应 dict 或 None（notification）。"""
        if not isinstance(request, dict):
            return {"jsonrpc": "2.0", "id": None,
                    "error": {"code": -32600, "message": "invalid request"}}
        method = request.get("method", "")
        req_id = request.get("id")
        if method.startswith("notifications/"):
            return None
        if method == "initialize":
            return {"jsonrpc": "2.0", "id": req_id, "result": {
                "protocolVersion": request.get("params", {}).get(
                    "protocolVersion", "2024-11-05"),
                "capabilities": {"tools": {}},
                "serverInfo": {"name": "codex-cloud-mcp",
                               "version": "0.1.0"}}}
        if method == "ping":
            return {"jsonrpc": "2.0", "id": req_id, "result": {}}
        if method == "tools/list":
            return {"jsonrpc": "2.0", "id": req_id,
                    "result": {"tools": TOOLS}}
        if method == "tools/call":
            params = request.get("params", {}) or {}
            is_error, payload = self.call_tool(
                params.get("name"), params.get("arguments") or {})
            return {"jsonrpc": "2.0", "id": req_id, "result": {
                "content": [{"type": "text",
                             "text": json.dumps(payload,
                                                ensure_ascii=False)}],
                "isError": is_error}}
        return {"jsonrpc": "2.0", "id": req_id,
                "error": {"code": -32601,
                          "message": "method not found: %s" % method}}

    def serve(self, instream, outstream):
        """stdio 主循环：每行一个请求，EOF 结束；坏行回 -32700 不炸循环。"""
        for line in instream:
            line = line.strip()
            if not line:
                continue
            try:
                request = json.loads(line)
            except ValueError:
                outstream.write(json.dumps(
                    {"jsonrpc": "2.0", "id": None,
                     "error": {"code": -32700,
                               "message": "parse error"}}) + "\n")
                outstream.flush()
                continue
            response = self.handle(request)
            if response is not None:
                outstream.write(json.dumps(response,
                                           ensure_ascii=False) + "\n")
                outstream.flush()


def main(argv=None):
    argv = list(sys.argv[1:] if argv is None else argv)
    base_url = os.environ.get("CODEX_CLOUD_BASE_URL",
                              "https://api.diante.me")
    api_key = os.environ.get("CODEX_CLOUD_API_KEY", "")
    try:
        server = MCPServer(base_url, api_key)
    except ValueError as exc:
        sys.stderr.write("mcp_server: %s\n" % exc)
        return 2
    server.serve(sys.stdin, sys.stdout)
    return 0


if __name__ == "__main__":
    raise SystemExit(main())