409 lines
13 KiB
Python
409 lines
13 KiB
Python
"""
|
||
Dev Agent MCP Server（HTTP 模式）
|
||
将 Dev Agent 暴露为 MCP 服务器，CodeBuddy 通过 HTTP POST MCP 调用
|
||
|
||
直接用 FastAPI 处理 MCP JSON-RPC 消息，不依赖 mcp SDK 的 HTTP transport，
|
||
避免各种版本兼容问题。
|
||
"""
|
||
import os
|
||
import sys
|
||
import asyncio
|
||
import json
|
||
import uuid
|
||
from dotenv import load_dotenv
|
||
|
||
# 确保项目根目录在 path 里
|
||
PROJECT_ROOT = os.path.abspath(os.path.join(os.path.dirname(__file__), ".."))
|
||
if PROJECT_ROOT not in sys.path:
|
||
sys.path.insert(0, PROJECT_ROOT)
|
||
|
||
from mcp.server import Server
|
||
from mcp.types import Tool, TextContent
|
||
from fastapi import FastAPI, Request, Response
|
||
import uvicorn
|
||
|
||
from .task_manager import TaskManager
|
||
from .agent_runner import AgentRunner
|
||
|
||
|
||
# 加载 .env
|
||
load_dotenv(os.path.join(PROJECT_ROOT, "my_agent", ".env"))
|
||
|
||
# 配置
|
||
HOST = os.getenv("MCP_SERVER_HOST", "0.0.0.0")
|
||
PORT = int(os.getenv("MCP_SERVER_PORT", "8001"))
|
||
DATA_DIR = os.getenv("MCP_DATA_DIR", os.path.join(PROJECT_ROOT, "data"))
|
||
MCP_PATH = "/mcp" # MCP 端点路径
|
||
|
||
# 初始化组件
|
||
task_manager = TaskManager(store_dir=os.path.join(DATA_DIR, "tasks"))
|
||
agent_runner = AgentRunner()
|
||
|
||
mcp_server = Server("dev-agent-mcp-server")
|
||
|
||
|
||
# --- MCP 工具定义 ---
|
||
|
||
@mcp_server.list_tools()
|
||
async def list_tools():
|
||
return [
|
||
Tool(
|
||
name="submit_task",
|
||
description=(
|
||
"提交一个开发任务给 Dev Agent 执行。任务将异步执行，"
|
||
"提交后返回 task_id，用 get_task_status 查询进度。"
|
||
),
|
||
inputSchema={
|
||
"type": "object",
|
||
"properties": {
|
||
"description": {
|
||
"type": "string",
|
||
"description": "任务的详细描述，要做什么开发工作",
|
||
},
|
||
"project_path": {
|
||
"type": "string",
|
||
"description": "项目的本地路径，agent 将在此目录下工作",
|
||
},
|
||
"requirements": {
|
||
"type": "string",
|
||
"description": "（可选）额外的要求或约束条件",
|
||
},
|
||
},
|
||
"required": ["description", "project_path"],
|
||
},
|
||
),
|
||
Tool(
|
||
name="get_task_status",
|
||
description="查询任务的当前状态（pending/running/completed/failed/cancelled）",
|
||
inputSchema={
|
||
"type": "object",
|
||
"properties": {
|
||
"task_id": {
|
||
"type": "string",
|
||
"description": "任务 ID",
|
||
},
|
||
},
|
||
"required": ["task_id"],
|
||
},
|
||
),
|
||
Tool(
|
||
name="get_task_result",
|
||
description="获取任务的执行结果（完成后调用）",
|
||
inputSchema={
|
||
"type": "object",
|
||
"properties": {
|
||
"task_id": {
|
||
"type": "string",
|
||
"description": "任务 ID",
|
||
},
|
||
},
|
||
"required": ["task_id"],
|
||
},
|
||
),
|
||
Tool(
|
||
name="get_task_log",
|
||
description="获取任务的执行日志",
|
||
inputSchema={
|
||
"type": "object",
|
||
"properties": {
|
||
"task_id": {
|
||
"type": "string",
|
||
"description": "任务 ID",
|
||
},
|
||
},
|
||
"required": ["task_id"],
|
||
},
|
||
),
|
||
Tool(
|
||
name="cancel_task",
|
||
description="取消一个正在执行或等待中的任务",
|
||
inputSchema={
|
||
"type": "object",
|
||
"properties": {
|
||
"task_id": {
|
||
"type": "string",
|
||
"description": "任务 ID",
|
||
},
|
||
},
|
||
"required": ["task_id"],
|
||
},
|
||
),
|
||
Tool(
|
||
name="list_tasks",
|
||
description="列出所有任务，可按状态过滤",
|
||
inputSchema={
|
||
"type": "object",
|
||
"properties": {
|
||
"status": {
|
||
"type": "string",
|
||
"description": "（可选）按状态过滤：pending/running/completed/failed/cancelled",
|
||
},
|
||
"limit": {
|
||
"type": "integer",
|
||
"description": "（可选）返回数量限制，默认 20",
|
||
"default": 20,
|
||
},
|
||
},
|
||
},
|
||
),
|
||
]
|
||
|
||
|
||
# --- MCP 工具实现 ---
|
||
|
||
@mcp_server.call_tool()
|
||
async def call_tool(name: str, arguments: dict):
|
||
if name == "submit_task":
|
||
return await _submit_task(arguments)
|
||
elif name == "get_task_status":
|
||
return _get_task_status(arguments)
|
||
elif name == "get_task_result":
|
||
return _get_task_result(arguments)
|
||
elif name == "get_task_log":
|
||
return _get_task_log(arguments)
|
||
elif name == "cancel_task":
|
||
return await _cancel_task(arguments)
|
||
elif name == "list_tasks":
|
||
return _list_tasks(arguments)
|
||
else:
|
||
return [TextContent(type="text", text=f"错误：未知工具 {name}")]
|
||
|
||
|
||
async def _submit_task(args: dict):
|
||
description = args.get("description", "")
|
||
project_path = args.get("project_path", "")
|
||
requirements = args.get("requirements", "")
|
||
|
||
if not description:
|
||
return [TextContent(type="text", text="错误：description 不能为空")]
|
||
if not project_path:
|
||
return [TextContent(type="text", text="错误：project_path 不能为空")]
|
||
if not os.path.isdir(project_path):
|
||
return [TextContent(type="text", text=f"错误：项目路径不存在 {project_path}")]
|
||
|
||
task = await task_manager.submit_task(
|
||
description=description,
|
||
project_path=project_path,
|
||
requirements=requirements,
|
||
)
|
||
|
||
return [TextContent(
|
||
type="text",
|
||
text=(
|
||
f"任务已提交\n"
|
||
f"任务ID: {task['id']}\n"
|
||
f"状态: {task['status']}\n"
|
||
f"描述: {description[:100]}\n"
|
||
f"项目: {project_path}\n"
|
||
f"\n"
|
||
f"请使用 get_task_status 查询进度。"
|
||
),
|
||
)]
|
||
|
||
|
||
def _get_task_status(args: dict):
|
||
task_id = args.get("task_id", "")
|
||
task = task_manager.get_task(task_id)
|
||
if not task:
|
||
return [TextContent(type="text", text=f"错误：任务不存在 {task_id}")]
|
||
|
||
return [TextContent(
|
||
type="text",
|
||
text=(
|
||
f"任务状态\n"
|
||
f"任务ID: {task['id']}\n"
|
||
f"状态: {task['status']}\n"
|
||
f"描述: {task['description'][:100]}\n"
|
||
f"创建时间: {_format_time(task.get('created_at'))}\n"
|
||
f"更新时间: {_format_time(task.get('updated_at'))}\n"
|
||
),
|
||
)]
|
||
|
||
|
||
def _get_task_result(args: dict):
|
||
task_id = args.get("task_id", "")
|
||
task = task_manager.get_task(task_id)
|
||
if not task:
|
||
return [TextContent(type="text", text=f"错误：任务不存在 {task_id}")]
|
||
|
||
result = task.get("result")
|
||
status = task["status"]
|
||
|
||
if status in ("pending", "running"):
|
||
return [TextContent(
|
||
type="text",
|
||
text=(
|
||
f"任务尚未完成（状态：{status}），"
|
||
f"请稍后再试或使用 get_task_status 查询进度。"
|
||
),
|
||
)]
|
||
|
||
if not result:
|
||
return [TextContent(type="text", text=f"任务结果为空，状态：{status}")]
|
||
|
||
if isinstance(result, dict):
|
||
summary = result.get("summary", str(result))
|
||
tool_count = result.get("tool_calls_count", 0)
|
||
full = result.get("full_response", "")
|
||
|
||
return [TextContent(
|
||
type="text",
|
||
text=(
|
||
f"任务结果（{status}）\n"
|
||
f"{'='*40}\n"
|
||
f"{summary}\n"
|
||
f"{'='*40}\n"
|
||
f"工具调用次数: {tool_count}\n"
|
||
f"\n"
|
||
f"--- 完整回复 ---\n"
|
||
f"{full[:5000]}"
|
||
f"\n{'...' if len(full) > 5000 else ''}"
|
||
),
|
||
)]
|
||
|
||
return [TextContent(type="text", text=str(result))]
|
||
|
||
|
||
def _get_task_log(args: dict):
|
||
task_id = args.get("task_id", "")
|
||
task = task_manager.get_task(task_id)
|
||
if not task:
|
||
return [TextContent(type="text", text=f"错误：任务不存在 {task_id}")]
|
||
|
||
logs = task.get("logs", [])
|
||
if not logs:
|
||
return [TextContent(type="text", text="暂无日志")]
|
||
|
||
lines = []
|
||
for log in logs[-50:]:
|
||
lines.append(f"[{log['time']}] {log['message']}")
|
||
|
||
return [TextContent(type="text", text="\n".join(lines))]
|
||
|
||
|
||
async def _cancel_task(args: dict):
|
||
task_id = args.get("task_id", "")
|
||
success = await task_manager.cancel_task(task_id)
|
||
if success:
|
||
return [TextContent(type="text", text=f"任务 {task_id} 已取消")]
|
||
else:
|
||
return [TextContent(type="text", text=f"取消失败：任务不存在或已结束")]
|
||
|
||
|
||
def _list_tasks(args: dict):
|
||
status = args.get("status")
|
||
limit = int(args.get("limit", 20))
|
||
tasks = task_manager.list_tasks(status=status)
|
||
tasks = tasks[:limit]
|
||
|
||
if not tasks:
|
||
return [TextContent(type="text", text="没有找到任务")]
|
||
|
||
lines = [f"任务列表（共 {len(tasks)} 个）:"]
|
||
for t in tasks:
|
||
lines.append(
|
||
f" [{t['status']}] {t['id']} - {t['description'][:50]} "
|
||
f"({_format_time(t.get('created_at'))})"
|
||
)
|
||
|
||
return [TextContent(type="text", text="\n".join(lines))]
|
||
|
||
|
||
def _format_time(ts: float = None) -> str:
|
||
import time
|
||
if not ts:
|
||
return "-"
|
||
return time.strftime("%Y-%m-%d %H:%M:%S", time.localtime(ts))
|
||
|
||
|
||
# --- HTTP Server ---
|
||
|
||
def create_fastapi_app() -> FastAPI:
|
||
"""创建 FastAPI 应用，处理 MCP JSON-RPC 请求"""
|
||
fastapi_app = FastAPI(title="Dev Agent MCP Server")
|
||
|
||
# 存储 session：session_id -> (read_stream, write_stream, session_task)
|
||
sessions = {}
|
||
|
||
async def _get_or_create_session(session_id: str):
|
||
"""获取或创建 MCP session（简单的内存会话管理）"""
|
||
if session_id not in sessions:
|
||
from anyio.streams.memory import MemoryObjectReceiveStream, MemoryObjectSendStream
|
||
from mcp.server.session import ServerSession
|
||
|
||
read_stream_writer, read_stream = MemoryObjectSendStream(100), MemoryObjectReceiveStream(100)
|
||
write_stream, write_stream_reader = MemoryObjectSendStream(100), MemoryObjectReceiveStream(100)
|
||
|
||
session = ServerSession(read_stream, write_stream)
|
||
task = asyncio.create_task(
|
||
mcp_server.run(
|
||
read_stream, write_stream,
|
||
mcp_server.create_initialization_options(),
|
||
)
|
||
)
|
||
sessions[session_id] = (read_stream_writer, write_stream_reader, task, session)
|
||
|
||
return sessions[session_id]
|
||
|
||
@fastapi_app.post(MCP_PATH)
|
||
async def handle_mcp(request: Request):
|
||
"""处理 MCP JSON-RPC 请求"""
|
||
body = await request.json()
|
||
|
||
# 简单处理：单条请求（非批量）
|
||
# 从 header 获取 session_id，没有就创建新的
|
||
session_id = request.headers.get("mcp-session-id") or str(uuid.uuid4())
|
||
read_stream_writer, write_stream_reader, task, session = await _get_or_create_session(session_id)
|
||
|
||
# 把请求写入 read_stream
|
||
await read_stream_writer.send(body)
|
||
|
||
# 等待响应（简单地从 write_stream 读一条）
|
||
response = await write_stream_reader.receive()
|
||
|
||
# 返回响应
|
||
return Response(
|
||
content=json.dumps(response),
|
||
media_type="application/json",
|
||
headers={"mcp-session-id": session_id},
|
||
)
|
||
|
||
@fastapi_app.get("/health")
|
||
async def health():
|
||
return {"status": "ok", "server": "dev-agent-mcp-server"}
|
||
|
||
return fastapi_app
|
||
|
||
|
||
async def _task_executor(task: dict) -> dict:
|
||
"""任务执行器，交给 TaskManager 调用"""
|
||
return await agent_runner.run_task(task)
|
||
|
||
|
||
async def main():
|
||
"""启动 MCP Server"""
|
||
print("=" * 50)
|
||
print("Dev Agent MCP Server 启动中...")
|
||
print(f" 监听地址: {HOST}:{PORT}")
|
||
print(f" 数据目录: {DATA_DIR}")
|
||
print(f" MCP 端点: http://{HOST}:{PORT}{MCP_PATH}")
|
||
print(f" 健康检查: http://{HOST}:{PORT}/health")
|
||
print("=" * 50)
|
||
|
||
# 启动任务管理器
|
||
await task_manager.start(executor=_task_executor)
|
||
|
||
# 启动 HTTP server
|
||
fastapi_app = create_fastapi_app()
|
||
config = uvicorn.Config(fastapi_app, host=HOST, port=PORT, log_level="info")
|
||
server = uvicorn.Server(config)
|
||
|
||
try:
|
||
await server.serve()
|
||
finally:
|
||
await task_manager.stop()
|
||
|
||
|
||
if __name__ == "__main__":
|
||
asyncio.run(main())
|