338 lines
10 KiB
Python
338 lines
10 KiB
Python
"""
|
||
Dev Agent MCP Server(stdio 模式)
|
||
通过标准输入输出与 MCP 客户端通信,适合 CodeBuddy 本地使用。
|
||
|
||
用法:python -m mcp_server.stdio_server
|
||
"""
|
||
import os
|
||
import sys
|
||
import asyncio
|
||
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 mcp.server.stdio import stdio_server
|
||
|
||
from .task_manager import TaskManager
|
||
from .agent_runner import AgentRunner
|
||
|
||
|
||
# 加载 .env
|
||
load_dotenv(os.path.join(PROJECT_ROOT, "my_agent", ".env"))
|
||
|
||
# 配置
|
||
DATA_DIR = os.getenv("MCP_DATA_DIR", os.path.join(PROJECT_ROOT, "data"))
|
||
|
||
# 初始化组件
|
||
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))
|
||
|
||
|
||
async def _task_executor(task: dict) -> dict:
|
||
"""任务执行器,交给 TaskManager 调用"""
|
||
return await agent_runner.run_task(task)
|
||
|
||
|
||
async def main():
|
||
"""启动 stdio MCP Server"""
|
||
# 日志写 stderr,不污染 stdout(MCP 协议通道)
|
||
print("Dev Agent MCP Server (stdio) 启动中...", file=sys.stderr)
|
||
|
||
# 启动任务管理器
|
||
await task_manager.start(executor=_task_executor)
|
||
|
||
try:
|
||
async with stdio_server() as (read_stream, write_stream):
|
||
await mcp_server.run(
|
||
read_stream, write_stream,
|
||
mcp_server.create_initialization_options(),
|
||
)
|
||
finally:
|
||
await task_manager.stop()
|
||
print("Dev Agent MCP Server 已停止", file=sys.stderr)
|
||
|
||
|
||
if __name__ == "__main__":
|
||
asyncio.run(main())
|