TraceStudio-dist/scripts/manage.py

#!/usr/bin/env python3
"""
TraceStudio 项目管家
统一管理开发、调试、发布的命令行工具。
"""
import argparse
import os
import sys
import shutil
import subprocess
import time
from pathlib import Path

# --- 路径配置 ---
# 自动判断是在源码根目录，还是在 dist 目录下
CURRENT_SCRIPT_DIR = Path(__file__).resolve().parent
if (CURRENT_SCRIPT_DIR.parent / "server" / "main.py").exists():
    ROOT = CURRENT_SCRIPT_DIR.parent
else:
    ROOT = Path.cwd()

DIST_DIR = ROOT / "dist"
ENV_FILE = ROOT / ".env"
CLOUD_DIR = ROOT / "cloud"


# --- 辅助函数 ---


def print_step(msg):
    print(f"\n[TraceStudio] ➤ {msg}")


def ensure_env():
    """确保 .env 文件存在，用于 Docker Compose 读取变量"""
    if ENV_FILE.exists():
        return

    print_step("初始化环境配置 (.env)...")
    if not CLOUD_DIR.exists():
        CLOUD_DIR.mkdir(parents=True, exist_ok=True)

    # Windows 下路径处理
    local_cloud = str(CLOUD_DIR)

    content = (
        f"# 自动生成的本地开发配置\n"
        f"LOCAL_CLOUD_ROOT={local_cloud}\n"
        f"PYTHONUNBUFFERED=1\n"
        f"AUTO_START=1\n"
    )
    with open(ENV_FILE, "w", encoding="utf-8") as f:
        f.write(content)
    print(f"已创建 .env，LOCAL_CLOUD_ROOT 指向: {local_cloud}")


def run_concurrent(commands, cwd=ROOT):
    """
    并发运行多个进程，并处理退出信号
    commands: list of {'cmd': list, 'name': str, 'cwd': Path, 'silent': bool}
    """
    processes = []
    try:
        for item in commands:
            cmd = item["cmd"]
            name = item["name"]
            work_dir = item.get("cwd", cwd)
            # 新增 silent 标记：如果为 True，则丢弃该进程的输出
            is_silent = item.get("silent", False)

            print(f">>> 启动 {name} ... {'(日志已隐藏)' if is_silent else ''}")
            
            # Windows 平台特殊处理
            is_shell = True if sys.platform == "win32" else False
            if isinstance(cmd, list) and len(cmd) > 0 and cmd[0] == "npm" and sys.platform == "win32":
                is_shell = True
                cmd = " ".join(cmd)

            # 配置输出流
            # 如果 silent=True，指向 DEVNULL；否则为 None (继承父进程控制台输出)
            stdout_dest = subprocess.DEVNULL if is_silent else None
            stderr_dest = subprocess.DEVNULL if is_silent else None

            p = subprocess.Popen(
                cmd, 
                cwd=str(work_dir), 
                shell=is_shell,
                stdout=stdout_dest,
                stderr=stderr_dest
            )
            processes.append((name, p))
            time.sleep(1) 

        print_step("所有服务已启动 (按 Ctrl+C 停止)")
        print("提示：Docker 运行在后台，请使用 Docker Desktop 或 'docker logs' 查看日志。")

        # 监控进程
        while True:
            time.sleep(1)
            for name, p in processes:
                if p.poll() is not None:
                    # 如果进程意外退出
                    print(f"\n!!! {name} 意外退出 (Code: {p.returncode})")
                    # 如果是静默模式退出的，通常用户不知道发生了什么，给个提示
                    if item.get("silent", False):
                        print(f"提示：{name} 是静默启动的，请手动运行命令检查报错原因。")
                    raise KeyboardInterrupt

    except KeyboardInterrupt:
        print("\n\n正在停止所有服务...")
        for name, p in processes:
            if p.poll() is None:
                print(f"停止 {name}...")
                p.terminate()
                try:
                    p.wait(timeout=3)
                except subprocess.TimeoutExpired:
                    p.kill()
        print("服务已停止。")


# --- 命令实现 ---


def cmd_setup(args):
    ensure_env()
    print("环境配置检查完毕。")


def cmd_dev(args):
    """混合开发模式: Docker (Server/Web) + Local Agent"""
    ensure_env()
    print_step("启动混合开发模式 (Docker Backend + Local Agent)")

    cmds = [
        {
            "name": "Docker Compose",
            "cmd": ["docker", "compose", "-f", "docker-compose.dev.yml", "up", "--build"],
            "cwd": ROOT,
            "silent": True, # <--- 关键修改：隐藏 Docker 日志
        },
        {
            "name": "Local Agent",
            "cmd": [sys.executable, "agent/main.py"],
            "cwd": ROOT,
            "silent": False, # Agent 日志保留，方便调试 Exe 调用
        },
    ]
    run_concurrent(cmds)


def cmd_local(args):
    """全本地模式: Local Server + Agent + Web"""
    ensure_env()
    print_step("启动全本地调试模式 (Local Host)")

    cmds = [
        {
            "name": "Backend Server",
            "cmd": [sys.executable, "-m", "uvicorn", "server.main:app", "--reload", "--port", "8000"],
            "cwd": ROOT / "server",
        },
        {"name": "Local Agent", "cmd": [sys.executable, "agent/main.py"], "cwd": ROOT},
        {"name": "Web Frontend", "cmd": ["npm", "run", "dev"], "cwd": ROOT / "web"},
    ]
    run_concurrent(cmds)


def cmd_build(args):
    """构建发布包"""
    print_step(f"清理并构建发布包到: {DIST_DIR}")
    if DIST_DIR.exists():
        # 遍历 DIST_DIR 下的所有文件和文件夹
        for item in DIST_DIR.iterdir():
            # 【关键】核心保护逻辑：如果名字是 .git，直接跳过
            if item.name == ".git":
                continue 
            try:
                if item.is_dir():
                    #如果是文件夹，递归删除
                    shutil.rmtree(item)
                else:
                    # 如果是文件，直接删除
                    item.unlink()
            except Exception as e:
                print(f"!!! 删除失败: {item} - {e}")
                # Windows 下有时候会因为文件被占用删不掉，视情况决定是否要 raise 阻断流程
    else:
        # 如果目录压根不存在，才需要 mkdir
        DIST_DIR.mkdir(parents=True, exist_ok=True)

    ignore = shutil.ignore_patterns(
        ".git", ".gitignore", ".dockerignore", "__pycache__", "node_modules",
        ".venv", "dist", ".vscode", ".devcontainer", "*.pyc", "docker-compose.dev.yml", "tests"
    )

    # 1. 复制核心代码
    for item in ["server", "web", "agent", "scripts"]:
        src = ROOT / item
        dst = DIST_DIR / item
        if src.is_dir():
            shutil.copytree(src, dst, ignore=ignore)

    shutil.copytree(ROOT / "cloud/custom_nodes", DIST_DIR / "cloud/custom_nodes", ignore=ignore)

    # 2. 复制生产配置
    prod_compose = ROOT / "docker-compose.yml"
    if prod_compose.exists():
        shutil.copy2(prod_compose, DIST_DIR / "docker-compose.yml")

    # 3. 复制管理脚本
    if not (DIST_DIR / "scripts").exists():
        (DIST_DIR / "scripts").mkdir()
    shutil.copy2(Path(__file__), DIST_DIR / "scripts" / "manage.py")

    print(f"构建完成！位置: {DIST_DIR}")


def cmd_prod(args):
    """【发布模式】生产环境运行"""
    ensure_env()
    compose_file = ROOT / 'docker-compose.yml'
    
    if not compose_file.exists():
        print(f"错误: 找不到 {compose_file}，请确保这是发布包目录。")
        return

    print_step("启动生产环境 (Docker Prod + Agent)")
    cmds = [
        {
            'name': 'Docker Compose (Prod)',
            'cmd': ['docker', 'compose', 'up', '--build'],
            'cwd': ROOT,
            'silent': False # 生产环境下建议保留日志以便排查
        },
        {
            'name': 'Local Agent',
            'cmd': [sys.executable, 'agent/main.py'],
            'cwd': ROOT
        }
    ]
    run_concurrent(cmds)


# --- 主入口 ---


def main():
    parser = argparse.ArgumentParser(description="TraceStudio 项目管理工具")
    subparsers = parser.add_subparsers(dest="command", help="可用命令")

    subparsers.add_parser("setup", help="生成 .env 配置文件")
    subparsers.add_parser("dev", help="启动 Docker开发环境(静默) + 本地Agent")
    subparsers.add_parser("local", help="启动全本地环境 (Server+Web+Agent)")
    subparsers.add_parser("build", help="构建发布包 (dist)")
    subparsers.add_parser("prod", help="启动生产发布包")

    args = parser.parse_args()

    mapping = {
        "setup": lambda x: ensure_env(),
        "dev": cmd_dev,
        "local": cmd_local,
        "build": cmd_build,
        "prod": cmd_prod
    }

    func = mapping.get(args.command)
    if func:
        func(args)
    else:
        parser.print_help()


if __name__ == "__main__":
    main()