TraceStudio-dev/server/main.py

"""
TraceStudio 主启动文件
规范化架构 v2.0
"""
from fastapi import FastAPI
from fastapi.middleware.cors import CORSMiddleware
import uvicorn
import yaml
from pathlib import Path

from server.app.api import endpoints_graph, endpoints_upload
from server.api_spec import mock_schema
from server.app.core.user_manager import init_base_structure
from server.app.core.node_loader import reload_custom_nodes
from server.app.core.cache_manager import CacheManager
import tempfile
import asyncio
# 加载系统配置
def load_config():
    config_path = Path(__file__).parent / "system_config.yaml"
    if config_path.exists():
        with open(config_path, 'r', encoding='utf-8') as f:
            return yaml.safe_load(f)
    return {
        "server": {
            "host": "127.0.0.1",
            "port": 8000,
            "reload": True
        },
        "security": {
            "allowed_ips": ["*"]
        }
    }

config = load_config()

# 创建 FastAPI 应用
app = FastAPI(
    title="TraceStudio Server",
    description="Unreal Insights 数据分析后端 - 规范化架构",
    version="2.0.0"
)

# 配置 CORS
allow_origins = config.get("security", {}).get("allowed_ips", ["*"])
if "*" in allow_origins:
    allow_origins = ["*"]

app.add_middleware(
    CORSMiddleware,
    allow_origins=allow_origins,
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

# 注册路由
app.include_router(endpoints_graph.router, prefix="/api", tags=["Graph"])
app.include_router(endpoints_upload.router, prefix="/api", tags=["Files"])
# register mock schema endpoint for manifest
mock_schema.register(app)
# app.include_router(endpoints_custom_nodes.router, tags=["Custom Nodes"])  # 暂时禁用

@app.on_event("startup")
async def startup_event():
    """应用启动事件"""
    
    # 加载自定义节点
    result = reload_custom_nodes()
    print(f"🔌 自定义节点: {result['loaded']}/{result['total']} 个已加载")
    
    init_base_structure()
    # 在应用启动阶段异步初始化缓存以避免阻塞启动（尤其在 reload/子进程场景）
    try:
        print("🗂️ 调度缓存初始化任务（内存 + 磁盘）...")
        disk_dir = Path(tempfile.gettempdir()) / "tracestudio_cache"

        mem_task = asyncio.create_task(asyncio.to_thread(CacheManager.init_memory_cache, 100, None))
        disk_task = asyncio.create_task(asyncio.to_thread(CacheManager.init_disk_cache, disk_dir, None))

        async def _watch_init(m_t, d_t, d_path):
            try:
                await m_t
                print("🗂️ 内存缓存初始化完成")
            except Exception as e:
                print("⚠️ 内存缓存初始化失败:", e)
            try:
                await d_t
                print(f"🗂️ 磁盘缓存初始化完成，目录: {d_path}")
            except Exception as e:
                print("⚠️ 磁盘缓存初始化失败:", e)

        asyncio.create_task(_watch_init(mem_task, disk_task, disk_dir))
        print(f"🗂️ 缓存初始化任务已调度，磁盘缓存目录: {disk_dir}")
    except Exception as e:
        print("⚠️ 缓存初始化任务调度失败:", e)
    print("🚀 TraceStudio Server v2.0 启动完成")
    print(f"📡 服务地址: http://{config['server']['host']}:{config['server']['port']}")
    print(f"📖 API 文档: http://{config['server']['host']}:{config['server']['port']}/docs")

@app.get("/")
async def root():
    """健康检查"""
    return {
        "status": "running",
        "service": "TraceStudio Server",
        "version": "2.0.0",
        "architecture": "normalized"
    }

if __name__ == "__main__":
    print("🚀 启动 TraceStudio 服务器...")
    # 当直接运行 main.py 时也初始化缓存（保持向后兼容）
    try:
        CacheManager.init_memory_cache(max_size=100, ttl=None)
        CacheManager.init_disk_cache(Path(tempfile.gettempdir()) / "tracestudio_cache", ttl=None)
    except Exception:
        pass
    uvicorn.run(
        "main:app",
        host=config["server"]["host"],
        port=config["server"]["port"],
        reload=config["server"]["reload"]
    )