125 lines
4.2 KiB
Python
125 lines
4.2 KiB
Python
import subprocess
|
||
import time
|
||
import os
|
||
import asyncio
|
||
|
||
# 全局常量
|
||
AGENT_SCRIPT = "compliance-mcp-agent/agent_main_loop.py"
|
||
|
||
async def log_subprocess_output(stream, prefix):
|
||
"""异步读取并打印子进程的输出流。"""
|
||
while True:
|
||
try:
|
||
line = await stream.readline()
|
||
if line:
|
||
print(f"[{prefix}] {line.decode().strip()}")
|
||
else:
|
||
# End of stream
|
||
break
|
||
except Exception as e:
|
||
print(f"Error reading stream for {prefix}: {e}")
|
||
break
|
||
|
||
async def start_servers():
|
||
"""异步启动所有MCP服务器，并捕获它们的日志。"""
|
||
print("="*20 + " Starting MCP Servers " + "="*20)
|
||
|
||
server_processes = {}
|
||
log_tasks = []
|
||
|
||
server_scripts = [
|
||
"servers/APICallerServer.py",
|
||
"servers/SchemaValidatorServer.py",
|
||
"servers/DMSProviderServer.py",
|
||
"servers/TestManagerServer.py",
|
||
]
|
||
|
||
project_root = os.path.dirname(os.path.abspath(__file__))
|
||
|
||
for script_path_rel in server_scripts:
|
||
script_path_abs = os.path.join(project_root, script_path_rel)
|
||
server_name = os.path.splitext(os.path.basename(script_path_abs))[0]
|
||
server_dir = os.path.dirname(script_path_abs)
|
||
|
||
print(f"Starting server: {script_path_rel}...")
|
||
|
||
env = os.environ.copy()
|
||
|
||
process = await asyncio.create_subprocess_exec(
|
||
"uv", "run", "python", os.path.basename(script_path_abs),
|
||
stdout=asyncio.subprocess.PIPE,
|
||
stderr=asyncio.subprocess.PIPE,
|
||
env=env,
|
||
cwd=server_dir
|
||
)
|
||
|
||
server_processes[server_name] = process
|
||
print(f" -> Started {script_path_rel} with PID: {process.pid}")
|
||
|
||
log_tasks.append(asyncio.create_task(log_subprocess_output(process.stdout, server_name)))
|
||
log_tasks.append(asyncio.create_task(log_subprocess_output(process.stderr, f"{server_name}-ERROR")))
|
||
|
||
print(f"\nAll {len(server_scripts)} servers are running in the background.")
|
||
return server_processes, log_tasks
|
||
|
||
async def run_agent(agent_script_path):
|
||
"""异步运行Agent主循环并实时打印其输出。"""
|
||
print("\n" + "="*20 + " Running Agent " + "="*20)
|
||
|
||
process = await asyncio.create_subprocess_exec(
|
||
"uv", "run", "python", agent_script_path,
|
||
stdout=asyncio.subprocess.PIPE,
|
||
stderr=asyncio.subprocess.PIPE
|
||
)
|
||
|
||
agent_stdout_task = asyncio.create_task(log_subprocess_output(process.stdout, "Agent"))
|
||
agent_stderr_task = asyncio.create_task(log_subprocess_output(process.stderr, "Agent-ERROR"))
|
||
|
||
await process.wait()
|
||
|
||
await asyncio.gather(agent_stdout_task, agent_stderr_task)
|
||
|
||
print(f"\nAgent process finished with return code: {process.returncode}")
|
||
|
||
|
||
def cleanup_servers(processes):
|
||
"""停止所有服务器进程。"""
|
||
print("\n" + "="*20 + " Cleaning up Servers " + "="*20)
|
||
for name, process in processes.items():
|
||
if process.returncode is None: # 仅当进程仍在运行时才终止
|
||
print(f"Terminating server {name} (PID: {process.pid})...")
|
||
try:
|
||
process.terminate()
|
||
print(f" -> Terminated signal sent.")
|
||
except ProcessLookupError:
|
||
print(f" -> Process {name} (PID: {process.pid}) already gone.")
|
||
else:
|
||
print(f" -> Server {name} (PID: {process.pid}) already finished with code {process.returncode}.")
|
||
|
||
print("Cleanup complete.")
|
||
|
||
async def main():
|
||
"""
|
||
主入口点，异步运行所有测试。
|
||
"""
|
||
server_processes, log_tasks = await start_servers()
|
||
|
||
print("\nWaiting for servers to initialize...")
|
||
await asyncio.sleep(8)
|
||
|
||
try:
|
||
await run_agent(AGENT_SCRIPT)
|
||
finally:
|
||
cleanup_servers(server_processes)
|
||
|
||
# 取消仍在运行的日志任务
|
||
for task in log_tasks:
|
||
task.cancel()
|
||
await asyncio.gather(*log_tasks, return_exceptions=True)
|
||
print("Log watchers terminated.")
|
||
|
||
if __name__ == "__main__":
|
||
try:
|
||
asyncio.run(main())
|
||
except KeyboardInterrupt:
|
||
print("\nMain process interrupted by user. Cleaning up...") |