agent1.4版本 精简代码
This commit is contained in:
parent
81acd16783
commit
d3e0861a40
@ -110,6 +110,10 @@ def main():
|
||||
api_server = create_api_server()
|
||||
fastapi_app = api_server.get_fast_api_app()
|
||||
|
||||
# 挂载 A2A 网关任务接收端点(POST /tasks/{request_id})
|
||||
from task_receiver import create_task_router
|
||||
fastapi_app.include_router(create_task_router(dev_app))
|
||||
|
||||
# 向 A2A 网关注册并启动心跳(注册失败不阻塞服务启动)
|
||||
gateway_ok = gateway_client.register_agent(dev_app.name, f"http://127.0.0.1:{PORT}")
|
||||
if gateway_ok:
|
||||
|
||||
@ -111,8 +111,8 @@ def main():
|
||||
fastapi_app = api_server.get_fast_api_app()
|
||||
|
||||
# 挂载 A2A 网关任务接收端点(POST /tasks/{request_id})
|
||||
from agents.my_agent.task_receiver import router as task_router
|
||||
fastapi_app.include_router(task_router)
|
||||
from task_receiver import create_task_router
|
||||
fastapi_app.include_router(create_task_router(dev_app))
|
||||
|
||||
# 向 A2A 网关注册并启动心跳(注册失败不阻塞服务启动)
|
||||
gateway_ok = gateway_client.register_agent(dev_app.name, f"http://127.0.0.1:{PORT}")
|
||||
|
||||
@ -131,7 +131,7 @@ async def chat(session_id: str | None = None):
|
||||
print("再见!")
|
||||
break
|
||||
|
||||
print("花花: ", end="", flush=True)
|
||||
print("agent: ", end="", flush=True)
|
||||
|
||||
async def _agent_task():
|
||||
"""运行 agent 并流式输出,返回是否完成"""
|
||||
|
||||
@ -1,120 +0,0 @@
|
||||
"""A2A 网关任务接收端点:接收网关主动推送的任务,后台执行 agent,完成后回传结果。
|
||||
|
||||
契约(网关 relay.py dispatch_command 推送):
|
||||
POST {endpoint}/tasks/{request_id}
|
||||
body: {"auth": GATEWAY_AUTH, "request_id": str, "payload": {...}}
|
||||
成功响应 202(立即确认),执行完成后由后台线程回传网关 /api/agent/result。
|
||||
"""
|
||||
import asyncio
|
||||
import logging
|
||||
import threading
|
||||
|
||||
from fastapi import APIRouter, HTTPException, Request
|
||||
from fastapi.responses import JSONResponse
|
||||
|
||||
import gateway_client
|
||||
from agents.my_agent.app import dev_app
|
||||
from google.adk.runners import InMemoryRunner
|
||||
from google.genai import types as genai_types
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
router = APIRouter(tags=["tasks"])
|
||||
|
||||
# 与 api_server.py 一致的导入路径(独立运行时兜底)
|
||||
import os
|
||||
import sys
|
||||
|
||||
PROJECT_ROOT = os.path.dirname(os.path.abspath(__file__))
|
||||
if PROJECT_ROOT not in sys.path:
|
||||
sys.path.insert(0, PROJECT_ROOT)
|
||||
|
||||
|
||||
def _payload_to_prompt(payload: dict) -> str:
|
||||
"""将网关任务 payload 转换为 agent 的用户指令。"""
|
||||
if not payload:
|
||||
return "请执行任务并汇报结果。"
|
||||
if "prompt" in payload and payload["prompt"]:
|
||||
return str(payload["prompt"])
|
||||
if "cmd" in payload and payload["cmd"]:
|
||||
return f"请执行以下命令并汇报执行结果:\n{payload['cmd']}"
|
||||
# 兜底:序列化整个 payload
|
||||
return "请根据以下任务载荷执行并汇报结果:\n" + str(payload)
|
||||
|
||||
|
||||
async def _run_agent_once(prompt: str, request_id: str) -> tuple[str, str]:
|
||||
"""用 ADK InMemoryRunner 运行一次 agent,返回 (接受任务后的首条回复, 最终总结)。
|
||||
|
||||
dev_app 是 App 容器(根 agent 非裸 LlmAgent),run_async 不会自动创建
|
||||
session,需先用 runner.session_service 显式创建。
|
||||
"""
|
||||
runner = InMemoryRunner(app=dev_app)
|
||||
session_id = f"task-{request_id}"
|
||||
await runner.session_service.create_session(
|
||||
app_name=runner.app_name,
|
||||
user_id="gateway",
|
||||
session_id=session_id,
|
||||
)
|
||||
texts: list[str] = []
|
||||
final_text = ""
|
||||
async for event in runner.run_async(
|
||||
user_id="gateway",
|
||||
session_id=session_id,
|
||||
new_message=genai_types.Content(role="user", parts=[genai_types.Part(text=prompt)]),
|
||||
):
|
||||
if event.is_final_response():
|
||||
if event.content and event.content.parts:
|
||||
for part in event.content.parts:
|
||||
text = getattr(part, "text", None)
|
||||
if text:
|
||||
final_text += text
|
||||
break
|
||||
if event.content and event.content.parts:
|
||||
for part in event.content.parts:
|
||||
text = getattr(part, "text", None)
|
||||
if text:
|
||||
texts.append(text)
|
||||
# 首条回复 = 接受任务后的第一条回应;最终总结 = final response
|
||||
reply = (texts[0] if texts else final_text).strip()
|
||||
summary = final_text.strip() or "\n".join(t for t in texts if t).strip() or "(无输出)"
|
||||
return reply, summary
|
||||
|
||||
|
||||
def _execute_and_report(request_id: str, payload: dict) -> None:
|
||||
"""后台线程:执行 agent,成功后回传 success,异常回传 failed。"""
|
||||
try:
|
||||
prompt = _payload_to_prompt(payload)
|
||||
reply, summary = asyncio.run(_run_agent_once(prompt, request_id))
|
||||
gateway_client.report_result(
|
||||
request_id,
|
||||
agent_id=dev_app.name,
|
||||
status="success",
|
||||
progress=100,
|
||||
result={"reply": reply, "output": summary},
|
||||
)
|
||||
except Exception as e:
|
||||
logger.exception("agent task failed request=%s", request_id)
|
||||
gateway_client.report_result(
|
||||
request_id,
|
||||
agent_id=dev_app.name,
|
||||
status="failed",
|
||||
progress=100,
|
||||
error_info=str(e),
|
||||
)
|
||||
|
||||
|
||||
@router.post("/tasks/{request_id}")
|
||||
async def receive_task(request_id: str, request: Request):
|
||||
"""接收网关推送的任务,立即 202 确认,后台执行。"""
|
||||
body = await request.json()
|
||||
if body.get("auth") != gateway_client.GATEWAY_AUTH:
|
||||
raise HTTPException(status_code=401, detail="invalid auth")
|
||||
payload = body.get("payload") or {}
|
||||
threading.Thread(
|
||||
target=_execute_and_report,
|
||||
args=(request_id, payload),
|
||||
name=f"task-{request_id[:8]}",
|
||||
daemon=True,
|
||||
).start()
|
||||
logger.info("task received request=%s payload=%s", request_id, payload)
|
||||
return JSONResponse(status_code=202, content={"ok": True, "request_id": request_id, "status": "accepted"})
|
||||
@ -110,6 +110,10 @@ def main():
|
||||
api_server = create_api_server()
|
||||
fastapi_app = api_server.get_fast_api_app()
|
||||
|
||||
# 挂载 A2A 网关任务接收端点(POST /tasks/{request_id})
|
||||
from task_receiver import create_task_router
|
||||
fastapi_app.include_router(create_task_router(dev_app))
|
||||
|
||||
# 向 A2A 网关注册并启动心跳(注册失败不阻塞服务启动)
|
||||
gateway_ok = gateway_client.register_agent(dev_app.name, f"http://127.0.0.1:{PORT}")
|
||||
if gateway_ok:
|
||||
|
||||
@ -28,6 +28,82 @@ HEARTBEAT_INTERVAL = int(os.getenv("HEARTBEAT_INTERVAL", "10"))
|
||||
|
||||
_client = httpx.Client(timeout=5.0)
|
||||
|
||||
# 取消/停止指令下发:网关复用 CLI 的 /api/cli/events 通道(event=task_stop)
|
||||
# Agent 通过同一个 cli_session_id 订阅,收到匹配自身正在执行 request_id 的 task_stop 时置停止标志。
|
||||
_stop_registry: dict[str, threading.Event] = {}
|
||||
|
||||
|
||||
def mark_stop_requested(request_id: str) -> None:
|
||||
"""标记某任务需要停止(收到 task_stop 后调用)。"""
|
||||
ev = _stop_registry.get(request_id)
|
||||
if ev is None:
|
||||
ev = threading.Event()
|
||||
_stop_registry[request_id] = ev
|
||||
ev.set()
|
||||
|
||||
|
||||
def clear_stop_requested(request_id: str) -> None:
|
||||
"""任务开始执行前清除停止标志。"""
|
||||
ev = _stop_registry.get(request_id)
|
||||
if ev is not None:
|
||||
ev.clear()
|
||||
|
||||
|
||||
def is_stop_requested(request_id: str) -> bool:
|
||||
"""判断某任务是否已被要求停止(供执行循环轮询检查)。"""
|
||||
ev = _stop_registry.get(request_id)
|
||||
return ev is not None and ev.is_set()
|
||||
|
||||
|
||||
def wait_stop(request_id: str, timeout: float = 0.5) -> bool:
|
||||
"""等待停止标志,返回 True 表示已收到停止请求。执行循环可用它做可中断 sleep。"""
|
||||
ev = _stop_registry.get(request_id)
|
||||
if ev is None:
|
||||
ev = threading.Event()
|
||||
_stop_registry[request_id] = ev
|
||||
return ev.wait(timeout)
|
||||
|
||||
|
||||
def _sse_subscribe_poll(cli_session_id: str) -> None:
|
||||
"""后台线程:订阅网关 /api/cli/events 通道,识别 task_stop 指令并标记停止。"""
|
||||
url = f"{GATEWAY_URL}/api/cli/events?cli_session_id={cli_session_id}&auth={GATEWAY_AUTH}"
|
||||
while True:
|
||||
try:
|
||||
with _client.stream("GET", url, timeout=None) as resp:
|
||||
if resp.status_code != 200:
|
||||
logger.warning("sse subscribe failed status=%s", resp.status_code)
|
||||
time.sleep(HEARTBEAT_INTERVAL)
|
||||
continue
|
||||
for line in resp.iter_lines():
|
||||
if not line or not line.startswith("data:"):
|
||||
continue
|
||||
try:
|
||||
import json
|
||||
evt = json.loads(line[len("data:"):].strip())
|
||||
except Exception:
|
||||
continue
|
||||
if evt.get("event") == "task_stop":
|
||||
rid = evt.get("request_id")
|
||||
if rid:
|
||||
mark_stop_requested(rid)
|
||||
logger.info("stop requested received request=%s", rid)
|
||||
except Exception as e:
|
||||
logger.warning("sse subscribe loop error err=%s", e)
|
||||
time.sleep(HEARTBEAT_INTERVAL)
|
||||
|
||||
|
||||
def start_stop_listener(cli_session_id: str) -> threading.Thread:
|
||||
"""启动 SSE 停止指令订阅线程(daemon)。"""
|
||||
t = threading.Thread(
|
||||
target=_sse_subscribe_poll,
|
||||
args=(cli_session_id,),
|
||||
name=f"gateway-sse-{cli_session_id[:8]}",
|
||||
daemon=True,
|
||||
)
|
||||
t.start()
|
||||
logger.info("stop listener started session=%s", cli_session_id)
|
||||
return t
|
||||
|
||||
|
||||
def _parse_tags() -> list[str]:
|
||||
raw = os.getenv("AGENT_TAGS", "code")
|
||||
@ -62,7 +138,7 @@ def _heartbeat_once(agent_id: str, current_load: int) -> bool:
|
||||
try:
|
||||
resp = _client.post(
|
||||
f"{GATEWAY_URL}/api/agent/heartbeat",
|
||||
json={"agent_id": agent_id, "current_load": current_load},
|
||||
json={"auth": GATEWAY_AUTH, "agent_id": agent_id, "current_load": current_load},
|
||||
)
|
||||
if resp.status_code == 200:
|
||||
return True
|
||||
@ -93,7 +169,7 @@ def start_heartbeat(agent_id: str) -> threading.Thread:
|
||||
def unregister_agent(agent_id: str) -> bool:
|
||||
"""向网关注销本 Agent。"""
|
||||
try:
|
||||
resp = _client.post(f"{GATEWAY_URL}/api/agent/unregister", json={"agent_id": agent_id})
|
||||
resp = _client.post(f"{GATEWAY_URL}/api/agent/unregister", json={"auth": GATEWAY_AUTH, "agent_id": agent_id})
|
||||
if resp.status_code in (200, 404):
|
||||
logger.info("unregistered from gateway agent=%s", agent_id)
|
||||
return True
|
||||
@ -117,6 +193,7 @@ def report_result(request_id: str, agent_id: str, status: str = "success",
|
||||
error_info: 错误信息(失败时必填)
|
||||
"""
|
||||
body = {
|
||||
"auth": GATEWAY_AUTH,
|
||||
"request_id": request_id,
|
||||
"agent_id": agent_id,
|
||||
"status": status,
|
||||
|
||||
@ -1 +0,0 @@
|
||||
# mcp_server package
|
||||
@ -1,153 +0,0 @@
|
||||
"""
|
||||
封装 ADK Agent 调用
|
||||
将 Dev Agent 的执行包装为可被任务管理器调用的异步函数
|
||||
"""
|
||||
import os
|
||||
import sys
|
||||
|
||||
# 确保项目根目录在 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 google.adk.runners import Runner
|
||||
from google.adk.sessions import InMemorySessionService
|
||||
from google.genai import types
|
||||
|
||||
|
||||
class AgentRunner:
|
||||
def __init__(self):
|
||||
# 延迟导入 agent,避免循环导入
|
||||
from agents.my_agent.agent import root_agent
|
||||
self.agent = root_agent
|
||||
self.session_service = InMemorySessionService()
|
||||
self._runner: Runner = None
|
||||
|
||||
def _get_runner(self) -> Runner:
|
||||
if self._runner is None:
|
||||
self._runner = Runner(
|
||||
agent=self.agent,
|
||||
app_name="dev_agent_server",
|
||||
session_service=self.session_service,
|
||||
auto_create_session=True,
|
||||
)
|
||||
return self._runner
|
||||
|
||||
async def run_task(self, task: dict) -> dict:
|
||||
"""
|
||||
执行一个开发任务,返回结构化结果
|
||||
|
||||
Args:
|
||||
task: 任务字典,包含 description, project_path, requirements 等
|
||||
|
||||
Returns:
|
||||
结构化的任务结果
|
||||
"""
|
||||
task_id = task["id"]
|
||||
description = task["description"]
|
||||
project_path = task.get("project_path", "")
|
||||
requirements = task.get("requirements", "")
|
||||
|
||||
# 构建给 agent 的提示词
|
||||
prompt = self._build_prompt(description, project_path, requirements)
|
||||
|
||||
print(f"[AgentRunner] 执行任务 {task_id}: {description[:60]}...")
|
||||
|
||||
runner = self._get_runner()
|
||||
session_id = f"task_{task_id}"
|
||||
user_id = "task_manager"
|
||||
|
||||
all_text = [] # 收集所有 model 消息中的文本
|
||||
tool_calls = [] # 收集所有 function_call
|
||||
final_text = [] # 最终响应的文本
|
||||
event_count = 0
|
||||
error_msg = None
|
||||
|
||||
try:
|
||||
async for event in runner.run_async(
|
||||
user_id=user_id,
|
||||
session_id=session_id,
|
||||
new_message=types.Content(
|
||||
role="user",
|
||||
parts=[types.Part(text=prompt)],
|
||||
),
|
||||
):
|
||||
event_count += 1
|
||||
|
||||
# 跳过没有 content 的事件
|
||||
if not event.content:
|
||||
continue
|
||||
|
||||
# 只收集 model 角色的消息(不是 user / function_response 等)
|
||||
content = event.content
|
||||
role = getattr(content, "role", "")
|
||||
if role != "model":
|
||||
continue
|
||||
|
||||
# 遍历 parts 收集文本和 function_call
|
||||
parts = getattr(content, "parts", [])
|
||||
for part in parts:
|
||||
# 文本
|
||||
if hasattr(part, "text") and part.text:
|
||||
all_text.append(part.text)
|
||||
# 最终回复(没有 function_call 的 model 消息)
|
||||
if event.is_final_response():
|
||||
final_text.append(part.text)
|
||||
|
||||
# function_call
|
||||
if hasattr(part, "function_call") and part.function_call:
|
||||
fc = part.function_call
|
||||
tool_calls.append({
|
||||
"name": fc.name,
|
||||
"args": dict(fc.args) if hasattr(fc, "args") else {},
|
||||
})
|
||||
|
||||
result_text = "\n".join(final_text) if final_text else "\n".join(all_text)
|
||||
status = "success" if result_text.strip() else "empty"
|
||||
|
||||
except Exception as e:
|
||||
error_msg = f"{type(e).__name__}: {e}"
|
||||
result_text = f"任务执行出错:{error_msg}"
|
||||
status = "error"
|
||||
|
||||
return {
|
||||
"summary": self._extract_summary(result_text),
|
||||
"full_response": result_text,
|
||||
"tool_calls_count": len(tool_calls),
|
||||
"tool_calls_sample": tool_calls[:10],
|
||||
"event_count": event_count,
|
||||
"status": status,
|
||||
"error": error_msg,
|
||||
}
|
||||
|
||||
def _build_prompt(self, description: str, project_path: str, requirements: str) -> str:
|
||||
"""构建给 agent 的任务指令"""
|
||||
parts = [
|
||||
"你需要完成以下开发任务:",
|
||||
"",
|
||||
f"**任务描述:**{description}",
|
||||
]
|
||||
if project_path:
|
||||
parts.append(f"**项目路径:**{project_path}")
|
||||
if requirements:
|
||||
parts.append(f"**额外要求:**{requirements}")
|
||||
parts.extend([
|
||||
"",
|
||||
"请按照你的工作流程执行:",
|
||||
"1. 浏览项目结构,理解上下文",
|
||||
"2. 编写或修改代码",
|
||||
"3. 运行编译/构建验证",
|
||||
"4. 完成后,按照你规定的报告格式输出结果",
|
||||
"",
|
||||
"请开始执行。",
|
||||
])
|
||||
return "\n".join(parts)
|
||||
|
||||
def _extract_summary(self, text: str) -> str:
|
||||
"""从 agent 回复中提取摘要"""
|
||||
# 简单处理:取前 1000 字符作为摘要
|
||||
if not text:
|
||||
return "(无响应)"
|
||||
if len(text) <= 1000:
|
||||
return text
|
||||
return text[:1000] + "...(已截断)"
|
||||
@ -1,408 +0,0 @@
|
||||
"""
|
||||
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())
|
||||
@ -1,25 +0,0 @@
|
||||
"""
|
||||
Dev Agent MCP Server 启动脚本
|
||||
供 CodeBuddy stdio MCP 调用,使用绝对路径确保可以从任何目录启动
|
||||
"""
|
||||
import os
|
||||
import sys
|
||||
|
||||
# 项目根目录(脚本所在目录的上一级)
|
||||
PROJECT_ROOT = os.path.abspath(os.path.join(os.path.dirname(__file__), ".."))
|
||||
|
||||
# 确保项目根目录在 Python path 里
|
||||
if PROJECT_ROOT not in sys.path:
|
||||
sys.path.insert(0, PROJECT_ROOT)
|
||||
|
||||
# 切换到项目目录,确保相对路径正确
|
||||
os.chdir(PROJECT_ROOT)
|
||||
|
||||
# 强制 UTF-8
|
||||
os.environ["PYTHONUTF8"] = "1"
|
||||
|
||||
# 导入并运行
|
||||
from mcp_server.stdio_server import main
|
||||
import asyncio
|
||||
|
||||
asyncio.run(main())
|
||||
@ -1,337 +0,0 @@
|
||||
"""
|
||||
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())
|
||||
@ -1,164 +0,0 @@
|
||||
"""
|
||||
异步任务管理器
|
||||
负责任务的提交、调度、状态管理
|
||||
"""
|
||||
import asyncio
|
||||
import uuid
|
||||
import time
|
||||
import os
|
||||
from typing import Dict, List, Optional, Callable, Awaitable
|
||||
from .task_store import TaskStore
|
||||
|
||||
|
||||
# 任务状态
|
||||
STATUS_PENDING = "pending"
|
||||
STATUS_RUNNING = "running"
|
||||
STATUS_COMPLETED = "completed"
|
||||
STATUS_FAILED = "failed"
|
||||
STATUS_CANCELLED = "cancelled"
|
||||
|
||||
|
||||
class TaskManager:
|
||||
def __init__(self, store_dir: str = "./data/tasks", max_concurrent: int = 3):
|
||||
self.store = TaskStore(store_dir)
|
||||
self.max_concurrent = max_concurrent
|
||||
self._tasks: Dict[str, dict] = {}
|
||||
self._running = 0
|
||||
self._semaphore = asyncio.Semaphore(max_concurrent)
|
||||
self._worker_task: Optional[asyncio.Task] = None
|
||||
self._queue: asyncio.Queue = asyncio.Queue()
|
||||
self._executor: Optional[Callable[[dict], Awaitable[dict]]] = None
|
||||
|
||||
async def start(self, executor: Callable[[dict], Awaitable[dict]]):
|
||||
"""启动任务管理器,executor 是实际执行任务的异步函数"""
|
||||
self._executor = executor
|
||||
self._worker_task = asyncio.create_task(self._worker_loop())
|
||||
print(f"[TaskManager] 已启动,最大并发: {self.max_concurrent}")
|
||||
|
||||
async def stop(self):
|
||||
"""停止任务管理器"""
|
||||
if self._worker_task:
|
||||
self._worker_task.cancel()
|
||||
try:
|
||||
await self._worker_task
|
||||
except asyncio.CancelledError:
|
||||
pass
|
||||
print("[TaskManager] 已停止")
|
||||
|
||||
async def submit_task(self, description: str, project_path: str,
|
||||
requirements: str = "", extra: dict = None) -> dict:
|
||||
"""提交一个新任务"""
|
||||
task_id = str(uuid.uuid4())[:8]
|
||||
now = time.time()
|
||||
task = {
|
||||
"id": task_id,
|
||||
"description": description,
|
||||
"project_path": project_path,
|
||||
"requirements": requirements,
|
||||
"status": STATUS_PENDING,
|
||||
"created_at": now,
|
||||
"updated_at": now,
|
||||
"result": None,
|
||||
"logs": [],
|
||||
"extra": extra or {},
|
||||
}
|
||||
self._tasks[task_id] = task
|
||||
self.store.save(task)
|
||||
await self._queue.put(task_id)
|
||||
print(f"[TaskManager] 任务已提交: {task_id} - {description[:50]}")
|
||||
return task
|
||||
|
||||
def get_task(self, task_id: str) -> Optional[dict]:
|
||||
"""获取任务详情"""
|
||||
# 优先从内存取,没有再从文件读
|
||||
if task_id in self._tasks:
|
||||
return self._tasks[task_id]
|
||||
return self.store.load(task_id)
|
||||
|
||||
def get_task_status(self, task_id: str) -> Optional[str]:
|
||||
task = self.get_task(task_id)
|
||||
return task["status"] if task else None
|
||||
|
||||
def list_tasks(self, status: str = None) -> List[dict]:
|
||||
"""列出所有任务,可按状态过滤"""
|
||||
tasks = list(self._tasks.values())
|
||||
# 加上磁盘上的任务
|
||||
disk_tasks = self.store.list_all()
|
||||
disk_ids = {t["id"] for t in tasks}
|
||||
for t in disk_tasks:
|
||||
if t["id"] not in disk_ids:
|
||||
tasks.append(t)
|
||||
if status:
|
||||
tasks = [t for t in tasks if t["status"] == status]
|
||||
tasks.sort(key=lambda t: t.get("created_at", 0), reverse=True)
|
||||
return tasks
|
||||
|
||||
def append_log(self, task_id: str, message: str):
|
||||
"""追加任务日志"""
|
||||
task = self._tasks.get(task_id)
|
||||
if not task:
|
||||
return
|
||||
if "logs" not in task:
|
||||
task["logs"] = []
|
||||
task["logs"].append({
|
||||
"time": time.strftime("%Y-%m-%d %H:%M:%S"),
|
||||
"message": message,
|
||||
})
|
||||
if len(task["logs"]) > 500:
|
||||
task["logs"] = task["logs"][-500:]
|
||||
# 异步持久化(这里直接同步写,简单起见)
|
||||
self.store.save(task)
|
||||
|
||||
async def cancel_task(self, task_id: str) -> bool:
|
||||
"""取消任务"""
|
||||
task = self._tasks.get(task_id)
|
||||
if not task:
|
||||
return False
|
||||
if task["status"] in (STATUS_COMPLETED, STATUS_FAILED, STATUS_CANCELLED):
|
||||
return False
|
||||
task["status"] = STATUS_CANCELLED
|
||||
task["updated_at"] = time.time()
|
||||
self.store.save(task)
|
||||
print(f"[TaskManager] 任务已取消: {task_id}")
|
||||
return True
|
||||
|
||||
async def _worker_loop(self):
|
||||
"""后台 worker,从队列取任务执行"""
|
||||
while True:
|
||||
try:
|
||||
task_id = await self._queue.get()
|
||||
async with self._semaphore:
|
||||
await self._execute_task(task_id)
|
||||
except asyncio.CancelledError:
|
||||
break
|
||||
except Exception as e:
|
||||
print(f"[TaskManager] Worker 异常: {e}")
|
||||
await asyncio.sleep(1)
|
||||
|
||||
async def _execute_task(self, task_id: str):
|
||||
"""执行单个任务"""
|
||||
task = self._tasks.get(task_id)
|
||||
if not task or task["status"] == STATUS_CANCELLED:
|
||||
return
|
||||
|
||||
task["status"] = STATUS_RUNNING
|
||||
task["updated_at"] = time.time()
|
||||
self.store.save(task)
|
||||
print(f"[TaskManager] 开始执行: {task_id}")
|
||||
|
||||
try:
|
||||
result = await self._executor(task)
|
||||
task["result"] = result
|
||||
# 检查是否已被取消
|
||||
if task["status"] == STATUS_CANCELLED:
|
||||
return
|
||||
task["status"] = STATUS_COMPLETED
|
||||
print(f"[TaskManager] 任务完成: {task_id}")
|
||||
except Exception as e:
|
||||
task["status"] = STATUS_FAILED
|
||||
task["result"] = {"error": str(e)}
|
||||
self.append_log(task_id, f"执行失败: {e}")
|
||||
print(f"[TaskManager] 任务失败: {task_id} - {e}")
|
||||
finally:
|
||||
task["updated_at"] = time.time()
|
||||
self.store.save(task)
|
||||
@ -1,56 +0,0 @@
|
||||
"""
|
||||
任务持久化存储
|
||||
使用 JSON 文件存储任务数据
|
||||
"""
|
||||
import os
|
||||
import json
|
||||
import time
|
||||
from typing import Dict, List, Optional
|
||||
|
||||
|
||||
class TaskStore:
|
||||
def __init__(self, store_dir: str):
|
||||
self.store_dir = os.path.abspath(store_dir)
|
||||
os.makedirs(self.store_dir, exist_ok=True)
|
||||
|
||||
def _task_path(self, task_id: str) -> str:
|
||||
return os.path.join(self.store_dir, f"{task_id}.json")
|
||||
|
||||
def save(self, task: dict) -> None:
|
||||
task["updated_at"] = time.time()
|
||||
path = self._task_path(task["id"])
|
||||
with open(path, "w", encoding="utf-8") as f:
|
||||
json.dump(task, f, ensure_ascii=False, indent=2)
|
||||
|
||||
def load(self, task_id: str) -> Optional[dict]:
|
||||
path = self._task_path(task_id)
|
||||
if not os.path.exists(path):
|
||||
return None
|
||||
with open(path, "r", encoding="utf-8") as f:
|
||||
return json.load(f)
|
||||
|
||||
def list_all(self) -> List[dict]:
|
||||
tasks = []
|
||||
for filename in os.listdir(self.store_dir):
|
||||
if filename.endswith(".json"):
|
||||
task_id = filename[:-5]
|
||||
task = self.load(task_id)
|
||||
if task:
|
||||
tasks.append(task)
|
||||
tasks.sort(key=lambda t: t.get("created_at", 0), reverse=True)
|
||||
return tasks
|
||||
|
||||
def append_log(self, task_id: str, log_line: str) -> None:
|
||||
task = self.load(task_id)
|
||||
if not task:
|
||||
return
|
||||
if "logs" not in task:
|
||||
task["logs"] = []
|
||||
task["logs"].append({
|
||||
"time": time.strftime("%Y-%m-%d %H:%M:%S"),
|
||||
"message": log_line,
|
||||
})
|
||||
# 日志最多保留 500 条
|
||||
if len(task["logs"]) > 500:
|
||||
task["logs"] = task["logs"][-500:]
|
||||
self.save(task)
|
||||
@ -1 +0,0 @@
|
||||
# mcp_tools package
|
||||
@ -1 +0,0 @@
|
||||
# command_executor package
|
||||
@ -1,127 +0,0 @@
|
||||
"""
|
||||
终端命令执行 MCP Server
|
||||
通过 MCP 协议暴露 run_command 命令,供 Dev Agent 使用
|
||||
"""
|
||||
import asyncio
|
||||
import subprocess
|
||||
import shlex
|
||||
import os
|
||||
import sys
|
||||
|
||||
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
|
||||
|
||||
from mcp.server import Server
|
||||
from mcp.server.stdio import stdio_server
|
||||
from mcp.types import Tool, TextContent
|
||||
|
||||
|
||||
def log(msg):
|
||||
"""写日志到 stderr,不污染 MCP stdio 协议通道"""
|
||||
print(msg, file=sys.stderr, flush=True)
|
||||
|
||||
|
||||
app = Server("command-executor")
|
||||
|
||||
|
||||
@app.list_tools()
|
||||
async def list_tools():
|
||||
return [
|
||||
Tool(
|
||||
name="run_command",
|
||||
description="在终端中执行一条命令,返回输出结果。适用于编译、构建、运行测试等场景。",
|
||||
inputSchema={
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"command": {
|
||||
"type": "string",
|
||||
"description": "要执行的命令,如 'npm run build'、'npm test' 等",
|
||||
},
|
||||
"cwd": {
|
||||
"type": "string",
|
||||
"description": "命令执行的工作目录,默认使用当前目录",
|
||||
},
|
||||
"timeout": {
|
||||
"type": "integer",
|
||||
"description": "超时时间(秒),默认 300 秒",
|
||||
"default": 300,
|
||||
},
|
||||
},
|
||||
"required": ["command"],
|
||||
},
|
||||
),
|
||||
]
|
||||
|
||||
|
||||
@app.call_tool()
|
||||
async def call_tool(name: str, arguments: dict):
|
||||
if name != "run_command":
|
||||
raise ValueError(f"Unknown tool: {name}")
|
||||
|
||||
cmd = arguments.get("command", "")
|
||||
cwd = arguments.get("cwd") or os.getcwd()
|
||||
timeout = int(arguments.get("timeout", 300))
|
||||
|
||||
if not cmd:
|
||||
return [TextContent(type="text", text="错误:命令不能为空")]
|
||||
|
||||
log(f" [command] {cmd}")
|
||||
log(f" [cwd] {cwd}")
|
||||
log(f" [timeout] {timeout}s")
|
||||
log(f" [PATH] {os.environ.get('PATH', 'N/A')[:200]}")
|
||||
log(f" [where npm] {__import__('shutil').which('npm.cmd')}")
|
||||
|
||||
try:
|
||||
# 使用异步 subprocess,避免阻塞 asyncio 事件循环
|
||||
# 在 Windows 上,npm 是 .cmd 文件,需要通过 shell 执行
|
||||
proc = await asyncio.create_subprocess_shell(
|
||||
cmd,
|
||||
cwd=cwd,
|
||||
stdout=asyncio.subprocess.PIPE,
|
||||
stderr=asyncio.subprocess.PIPE,
|
||||
)
|
||||
stdout_bytes, stderr_bytes = await asyncio.wait_for(
|
||||
proc.communicate(),
|
||||
timeout=timeout,
|
||||
)
|
||||
|
||||
stdout = stdout_bytes.decode("utf-8", errors="replace")
|
||||
stderr = stderr_bytes.decode("utf-8", errors="replace")
|
||||
|
||||
output_parts = []
|
||||
if stdout:
|
||||
output_parts.append(f"[stdout]\n{stdout}")
|
||||
if stderr:
|
||||
output_parts.append(f"[stderr]\n{stderr}")
|
||||
|
||||
output = "\n".join(output_parts) if output_parts else "(无输出)"
|
||||
|
||||
max_len = 10000
|
||||
if len(output) > max_len:
|
||||
output = output[:max_len] + f"\n\n...(输出已截断,共 {len(output)} 字符)"
|
||||
|
||||
status = "成功" if proc.returncode == 0 else f"失败 (退出码 {proc.returncode})"
|
||||
log(f" [result] {status}")
|
||||
return [TextContent(type="text", text=f"命令执行{status}\n{output}")]
|
||||
|
||||
except asyncio.TimeoutError:
|
||||
# 超时后杀进程
|
||||
proc.kill()
|
||||
await proc.wait()
|
||||
log(f" [error] 超时")
|
||||
return [TextContent(type="text", text=f"命令执行超时({timeout}秒): {cmd}")]
|
||||
except Exception as e:
|
||||
log(f" [error] {e}")
|
||||
return [TextContent(type="text", text=f"命令执行出错: {e}")]
|
||||
|
||||
|
||||
async def main():
|
||||
log("MCP 命令执行服务器启动中...")
|
||||
async with stdio_server() as (read_stream, write_stream):
|
||||
await app.run(
|
||||
read_stream, write_stream,
|
||||
app.create_initialization_options()
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
asyncio.run(main())
|
||||
Loading…
Reference in New Issue
Block a user