359 lines
17 KiB
Python
359 lines
17 KiB
Python
"""云端任务管理的 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_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_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_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())
|