aiagent/backend/app/main.py

"""
天工智能体平台 - FastAPI 主应用
"""
import asyncio
import logging
import os
from logging.handlers import RotatingFileHandler
from pathlib import Path
from fastapi import FastAPI, Request, Response
from fastapi.middleware.cors import CORSMiddleware
from fastapi.exceptions import RequestValidationError
from sqlalchemy.exc import SQLAlchemyError
from app.core.config import settings
from app.core.error_handler import (
    validation_exception_handler,
    api_exception_handler,
    sqlalchemy_exception_handler,
    general_exception_handler
)
from app.core.exceptions import BaseAPIException
from app.core.database import init_db
from app.core.rate_limiter import RateLimiterMiddleware
from app.core.behavior_middleware import BehaviorCollectionMiddleware
from app.core.security_headers import SecurityHeadersMiddleware
from app.core.metrics import setup_metrics

_SECURITY_WARNED = False


def _check_security_config() -> None:
    """启动时检查敏感配置，输出安全警告。"""
    global _SECURITY_WARNED
    if _SECURITY_WARNED:
        return
    _SECURITY_WARNED = True

    warnings: list[str] = []

    # JWT 密钥检查
    if settings.JWT_SECRET_KEY == "dev-jwt-secret-key-change-in-production" or \
       "change-in-production" in settings.JWT_SECRET_KEY.lower():
        warnings.append("JWT_SECRET_KEY 使用默认值，生产环境必须更换为随机字符串")

    # 应用密钥
    if settings.SECRET_KEY == "dev-secret-key-change-in-production" or \
       "change-in-production" in settings.SECRET_KEY.lower():
        warnings.append("SECRET_KEY 使用默认值，生产环境必须更换")

    # API Key 检查
    api_keys_set = bool(
        settings.OPENAI_API_KEY
        or settings.DEEPSEEK_API_KEY
        or settings.ANTHROPIC_API_KEY
    )
    if not api_keys_set:
        warnings.append("未配置任何 AI API Key (OPENAI / DEEPSEEK / ANTHROPIC)，LLM 调用将失败")

    # 数据库密码检查
    db_url = settings.DATABASE_URL or ""
    if "change" in db_url.lower() or "CHANGE_ME" in db_url:
        warnings.append("DATABASE_URL 疑似使用默认密码，请检查数据库凭据")

    # HSTS 检查（仅在生产提醒）
    if not settings.HSTS_ENABLED and not settings.DEBUG:
        warnings.append("HSTS 未启用，生产环境建议启用 HSTS 以强制 HTTPS")

    # 生产环境检查
    if settings.ENVIRONMENT == "prod":
        if settings.DEBUG:
            warnings.append("生产环境 DEBUG 必须为 False")
        if not settings.HSTS_ENABLED:
            warnings.append("生产环境建议启用 HSTS (HSTS_ENABLED=True)")
        if not api_keys_set:
            warnings.append("生产环境必须配置至少一个 AI API Key")
    elif settings.ENVIRONMENT == "staging":
        if settings.DEBUG:
            warnings.append("预发布环境建议关闭 DEBUG")

    for w in warnings:
        logger.warning("[安全] %s", w)

    if not warnings:
        logger.info("[安全] 配置检查通过，未发现明显安全风险")

# 配置日志：同时输出到控制台和文件（带轮转）
_log_dir = Path(settings.LOG_DIR)
_log_dir.mkdir(parents=True, exist_ok=True)

_stream_handler = logging.StreamHandler()
_stream_handler.setFormatter(logging.Formatter('%(asctime)s - %(name)s - %(levelname)s - %(message)s'))

_file_handler = RotatingFileHandler(
    _log_dir / "app.log",
    maxBytes=settings.LOG_MAX_BYTES,
    backupCount=settings.LOG_BACKUP_COUNT,
    encoding="utf-8"
)
_file_handler.setFormatter(logging.Formatter('%(asctime)s - %(name)s - %(levelname)s - %(message)s'))

logging.basicConfig(
    level=getattr(logging, settings.LOG_LEVEL.upper(), logging.INFO),
    handlers=[_stream_handler, _file_handler]
)
logger = logging.getLogger(__name__)

app = FastAPI(
    title=settings.APP_NAME,
    version=settings.APP_VERSION,
    description="""
## 天工智能体平台 API

一个支持可视化工作流设计和自主智能Agent配置的AI平台。

### 主要功能

* **用户认证** - 用户注册、登录、JWT认证
* **工作流管理** - 工作流的创建、读取、更新、删除、执行
* **工作流版本管理** - 版本保存、版本列表、版本回滚
* **执行管理** - 工作流执行、执行记录查询、执行状态监控
* **执行日志** - 详细的执行日志记录和查询
* **数据源管理** - 多种数据源的连接和管理
* **WebSocket实时推送** - 执行状态实时更新

### 认证方式

大部分API需要JWT认证。请先通过 `/api/v1/auth/login` 获取token，然后在请求头中添加：

```
Authorization: Bearer <your_token>
```

### API版本

当前API版本：v1

### 文档

* **Swagger UI**: `/docs` - 交互式API文档
* **ReDoc**: `/redoc` - 可读性更好的API文档
    """,
    docs_url="/docs",
    redoc_url="/redoc",
    openapi_tags=[
        {
            "name": "auth",
            "description": "用户认证相关API，包括注册、登录、获取用户信息等。"
        },
        {
            "name": "workflows",
            "description": "工作流管理API，包括工作流的CRUD操作、执行、版本管理等。"
        },
        {
            "name": "executions",
            "description": "执行管理API，包括执行记录的创建、查询、状态获取等。"
        },
        {
            "name": "execution-logs",
            "description": "执行日志API，包括日志查询、日志统计等。"
        },
        {
            "name": "data-sources",
            "description": "数据源管理API，包括数据源的CRUD操作、连接测试、数据查询等。"
        },
        {
            "name": "websocket",
            "description": "WebSocket API，用于实时推送执行状态。"
        },
        {
            "name": "goals",
            "description": "目标管理API — Main Agent 数字员工的目标 CRUD、启停、任务树查看。"
        },
        {
            "name": "tasks",
            "description": "任务管理API — 目标分解后的子任务 CRUD、依赖检查、审批。"
        }
    ]
)

# CORS 配置
cors_origins = [origin.strip() for origin in settings.CORS_ORIGINS.split(",")]
# 添加调试日志
logger.info(f"CORS允许的源: {cors_origins}")
app.add_middleware(
    CORSMiddleware,
    allow_origins=cors_origins,
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
    expose_headers=["*"],
)

# API 限流中间件（在 CORS 之后、日志之前）
app.add_middleware(RateLimiterMiddleware)

# 安全响应头中间件（HSTS + X-Frame-Options + X-Content-Type-Options 等）
app.add_middleware(SecurityHeadersMiddleware)

# 用户行为自动采集中间件
app.add_middleware(BehaviorCollectionMiddleware)

# 注册全局异常处理器
app.add_exception_handler(RequestValidationError, validation_exception_handler)
app.add_exception_handler(BaseAPIException, api_exception_handler)
app.add_exception_handler(SQLAlchemyError, sqlalchemy_exception_handler)
app.add_exception_handler(Exception, general_exception_handler)

# Prometheus 指标收集 + /metrics 端点
setup_metrics(app)

# 请求日志中间件
@app.middleware("http")
async def log_requests(request: Request, call_next):
    """记录请求日志 + 写操作审计"""
    import time
    start_time = time.time()

    # 记录请求
    logger.info(f"{request.method} {request.url.path} - 客户端: {request.client.host if request.client else 'unknown'}")

    try:
        response = await call_next(request)
        process_time = time.time() - start_time

        # 记录响应
        logger.info(
            f"{request.method} {request.url.path} - "
            f"状态码: {response.status_code} - "
            f"耗时: {process_time:.3f}s"
        )

        # 对写操作自动记录审计日志（fire-and-forget）
        if request.method.upper() in ("POST", "PUT", "PATCH", "DELETE"):
            _schedule_audit_log(request, response.status_code, process_time)

        return response
    except Exception as e:
        process_time = time.time() - start_time
        logger.error(
            f"{request.method} {request.url.path} - "
            f"异常: {str(e)} - "
            f"耗时: {process_time:.3f}s"
        )
        raise


def _schedule_audit_log(request: Request, status_code: int, duration: float) -> None:
    """fire-and-forget 写入审计日志，不影响请求响应"""
    import asyncio
    asyncio.ensure_future(_write_audit_log(request, status_code, duration))


async def _write_audit_log(request: Request, status_code: int, duration: float) -> None:
    try:
        from app.core.security import decode_access_token
        from app.core.database import SessionLocal
        from app.models.user import User
        from app.models.audit_log import AuditLog

        # 提取用户（从 Authorization header 解码 JWT）
        username = "anonymous"
        user_id = ""
        auth_header = request.headers.get("Authorization", "")
        if auth_header.startswith("Bearer "):
            token = auth_header[7:]
            try:
                payload = decode_access_token(token)
                if payload:
                    uid = payload.get("sub")
                    if uid:
                        db = SessionLocal()
                        try:
                            user = db.query(User).filter(User.id == uid).first()
                            if user:
                                username = user.username
                                user_id = user.id
                        finally:
                            db.close()
            except Exception:
                pass

        # 推断资源类型
        path = request.url.path
        method = request.method.upper()
        resource_type = _infer_resource_type(path)
        action = _infer_action(method, path)

        # 跳过非 API 路径
        if action == "OTHER":
            return

        # 写入审计日志
        db = SessionLocal()
        try:
            audit = AuditLog(
                user_id=user_id,
                username=username,
                action=action,
                resource_type=resource_type,
                ip_address=request.client.host if request.client else None,
                status="success" if 200 <= status_code < 400 else "failure",
            )
            db.add(audit)
            db.commit()
        except Exception:
            db.rollback()
        finally:
            db.close()
    except Exception:
        pass  # 审计失败不影响业务


def _infer_resource_type(path: str) -> str:
    mapping = {
        "agents": "agent",
        "workflows": "workflow",
        "executions": "execution",
        "auth": "auth",
        "users": "user",
        "permissions": "permission",
        "alert-rules": "alert_rule",
        "knowledge": "knowledge_base",
        "data-sources": "data_source",
        "model-configs": "model_config",
        "tools": "tool",
        "uploads": "upload",
        "plugins": "plugin",
        "tasks": "task",
        "goals": "goal",
        "agent-schedules": "schedule",
        "notifications": "notification",
        "templates": "template",
        "agent-chat": "agent_chat",
        "agent-monitoring": "monitoring",
        "system-logs": "system_log",
        "audit-logs": "audit_log",
    }
    parts = path.strip("/").split("/")
    for i, part in enumerate(parts):
        if part in mapping:
            return mapping[part]
    return "other"


def _infer_action(method: str, path: str) -> str:
    path_lower = path.lower()
    if "login" in path_lower:
        return "LOGIN"
    if method == "POST":
        return "CREATE"
    if method in ("PUT", "PATCH"):
        return "UPDATE"
    if method == "DELETE":
        return "DELETE"
    return "OTHER"

@app.get("/")
async def root():
    """根路径"""
    return {
        "message": "欢迎使用天工智能体平台 API",
        "version": settings.APP_VERSION,
        "docs": "/docs"
    }

@app.get("/health")
async def health_check():
    """健康检查（含内置工具是否已注册，便于排查「可动手 Agent」）。"""
    from app.services.tool_registry import tool_registry

    names = tool_registry.builtin_tool_names()
    count = tool_registry.builtin_tool_count()
    file_agent_core = {"file_write", "file_read", "system_info"}
    subset_ok = file_agent_core.issubset(set(names))
    expected_ok = count >= 10
    tools_ready = expected_ok and subset_ok

    return {
        "status": "healthy",
        "checks": {
            "builtin_tools_ready": tools_ready,
            "builtin_tools_count_ok": expected_ok,
            "file_agent_core_ready": subset_ok,
        },
        "builtin_tools": {
            "count": count,
            "names": names,
            "expected_min_count": 10,
        },
        "notes": {
            "celery": "工作流/Agent 执行通常在 Celery Worker 中跑；Worker 日志中也应出现「内置工具就绪」且 count 应与 API 一致。若仅 API 有工具而 Worker 无，会出现 LLM 无法真正调用 file_write。",
        },
    }

# 应用启动时初始化数据库和工具
@app.on_event("startup")
async def startup_event():
    """应用启动事件"""
    # 启用 JSON 结构化日志（ELK 采集）
    try:
        from app.core.logging_config import setup_json_logging
        setup_json_logging()
    except Exception as e:
        logger.warning(f"JSON 日志启用失败: {e}")

    try:
        logger.info("正在初始化数据库...")
        init_db()
        logger.info("数据库初始化完成")
    except Exception as e:
        logger.error(f"数据库初始化失败: {e}")
        # 不抛出异常，允许应用继续启动
    
    # 注册内置工具（与 Celery Worker 共用 app.core.tools_bootstrap）
    try:
        from app.core.tools_bootstrap import ensure_builtin_tools_registered

        ensure_builtin_tools_registered()
        from app.services.tool_registry import tool_registry

        logger.info("内置工具注册完成（count=%s）", tool_registry.builtin_tool_count())
    except Exception as e:
        logger.error(f"内置工具注册失败: {e}")
        # 不抛出异常，允许应用继续启动

    # 加载自定义工具（从数据库同步到注册表）
    try:
        from app.core.database import SessionLocal
        db = SessionLocal()
        try:
            tool_registry.load_tools_from_db(db)
            logger.info("自定义工具加载完成（count=%s）", len(tool_registry._tool_schemas) - tool_registry.builtin_tool_count())
        finally:
            db.close()
    except Exception as e:
        logger.error(f"自定义工具加载失败: {e}")

    # 启动飞书长连接（在主事件循环中运行）
    try:
        from app.services.feishu_ws_handler import start_ws_client
        asyncio.ensure_future(start_ws_client())
    except Exception as e:
        logger.error(f"飞书长连接启动失败: {e}")

    # 启动橙子飞书长连接
    try:
        from app.services.orange_ws_handler import start_ws_client as start_orange_ws
        asyncio.ensure_future(start_orange_ws())
    except Exception as e:
        logger.error(f"橙子长连接启动失败: {e}")

    # 启动苏瑶飞书长连接
    try:
        from app.services.suyao_ws_handler import start_ws_client as start_suyao_ws
        asyncio.ensure_future(start_suyao_ws())
    except Exception as e:
        logger.error(f"苏瑶长连接启动失败: {e}")

    # 启动甜甜飞书长连接（苏瑶3号）
    try:
        from app.services.tiantian_ws_handler import start_ws_client as start_tiantian_ws
        asyncio.ensure_future(start_tiantian_ws())
    except Exception as e:
        logger.error(f"甜甜长连接启动失败: {e}")

    # 启动灵犀飞书长连接（学习助手）
    try:
        from app.services.lingxi_ws_handler import start_ws_client as start_lingxi_ws
        asyncio.ensure_future(start_lingxi_ws())
    except Exception as e:
        logger.error(f"灵犀长连接启动失败: {e}")

    # 启动人参果飞书长连接（AI学习助手 — KG+RAG理想版）
    try:
        from app.services.renshenguo_ws_handler import start_ws_client as start_renshenguo_ws
        asyncio.ensure_future(start_renshenguo_ws())
    except Exception as e:
        logger.error(f"人参果长连接启动失败: {e}")

    # 启动人参果1号飞书长连接（AI学习助手 — 行为约束版，禁止主动消息）
    try:
        from app.services.renshenguo2_ws_handler import start_ws_client as start_renshenguo2_ws
        asyncio.ensure_future(start_renshenguo2_ws())
    except Exception as e:
        logger.error(f"人参果1号长连接启动失败: {e}")

    # 模板市场预置注入（若市场为空，注入 8 个预置模板）
    try:
        from app.core.database import SessionLocal
        from app.services.marketplace_bootstrap import bootstrap_preset_templates

        db = SessionLocal()
        try:
            count = bootstrap_preset_templates(db)
            if count > 0:
                logger.info("模板市场预置注入完成，新增 %s 个模板", count)
        finally:
            db.close()
    except Exception as e:
        logger.error(f"模板市场预置注入失败: {e}")

    # 安全配置检查
    _check_security_config()

    # 启动内置定时任务调度器（每分钟检查到期任务）
    try:
        from app.services.agent_schedule_service import run_scheduler_loop
        asyncio.ensure_future(run_scheduler_loop())
        logger.info("定时任务调度器已启动（内置模式）")
    except Exception as e:
        logger.error(f"定时任务调度器启动失败: {e}")

# 注册路由
from app.api import auth, workspaces, uploads, workflows, executions, websocket, execution_logs, data_sources, agents, platform_templates, model_configs, webhooks, template_market, batch_operations, collaboration, permissions, monitoring, alert_rules, node_test, node_templates, tools, agent_chat, agent_branches, agent_monitoring, knowledge_base, knowledge_dashboard, agent_schedules, notifications, feishu_bind, approval, orchestration_templates, plugins, agent_market, goals, tasks, system_logs, audit_logs, feedback, agent_swarm, push, voice, fcm, scene_contracts, teams, agent_memory

app.include_router(auth.router)
app.include_router(workspaces.router)
app.include_router(uploads.router)
app.include_router(workflows.router)
app.include_router(executions.router)
app.include_router(websocket.router)
app.include_router(execution_logs.router)
app.include_router(data_sources.router)
app.include_router(agents.router)
app.include_router(platform_templates.router)
app.include_router(model_configs.router)
app.include_router(webhooks.router)
app.include_router(template_market.router)
app.include_router(batch_operations.router)
app.include_router(collaboration.router)
app.include_router(permissions.router)
app.include_router(monitoring.router)
app.include_router(alert_rules.router)
app.include_router(node_test.router)
app.include_router(node_templates.router)
app.include_router(tools.router)
app.include_router(agent_branches.router)
app.include_router(agent_chat.router)
app.include_router(agent_monitoring.router)
app.include_router(knowledge_base.router)
app.include_router(agent_schedules.router)
app.include_router(notifications.router)
app.include_router(feishu_bind.router)
app.include_router(approval.router)
app.include_router(plugins.router)
app.include_router(orchestration_templates.router)
app.include_router(agent_market.router)
app.include_router(goals.router)
app.include_router(tasks.router)
app.include_router(system_logs.router)
app.include_router(audit_logs.router)
app.include_router(feedback.router)
app.include_router(agent_swarm.router)
app.include_router(push.router)
app.include_router(voice.router)
app.include_router(fcm.router)
app.include_router(knowledge_dashboard.router)
app.include_router(scene_contracts.router)
app.include_router(teams.router)
app.include_router(agent_memory.router)

if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=8000)