Agent/api_server.py

"""
API Server - FastAPI 应用入口

聚合所有模块的 API 路由：
- GET  /api/traces — 查询（trace/api.py）
- POST /api/traces — 执行控制（trace/run_api.py，需配置 Runner）
- WS   /api/traces/{id}/watch — 实时推送（trace/websocket.py）
- GET  /api/experiences — 经验查询（trace/run_api.py，需配置 Runner）
"""

import json
import logging
import os
import fcntl
from ipaddress import ip_address
from pathlib import Path

from dotenv import load_dotenv

load_dotenv()
from fastapi import FastAPI, Request, WebSocket
from fastapi.responses import JSONResponse
from fastapi.middleware.cors import CORSMiddleware
from fastapi.middleware.trustedhost import TrustedHostMiddleware
import uvicorn

from cyber_agent.trace import FileSystemTraceStore
from cyber_agent.trace.api import router as api_router, set_trace_store as set_api_trace_store
from cyber_agent.trace.run_api import router as run_router, experiences_router, set_runner
from cyber_agent.trace.websocket import router as ws_router, set_trace_store as set_ws_trace_store
from cyber_agent.trace.examples_api import router as examples_router
from cyber_agent.trace.logs_websocket import router as logs_router, setup_websocket_logging
from cyber_agent.trace.upload_api import router as upload_router, set_trace_store as set_upload_trace_store


# ===== 日志配置 =====

logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s [%(levelname)s] %(name)s: %(message)s"
)
logger = logging.getLogger(__name__)


# ===== 单用户本地开发边界 =====

LOCAL_DEVELOPMENT_ORIGINS = (
    "http://127.0.0.1:3000",
    "http://localhost:3000",
)
LOCAL_DEVELOPMENT_HOSTS = ("127.0.0.1", "localhost")
_writer_lock_file = None


def _env_flag(name: str, *, default: bool = False) -> bool:
    """解析显式开关；未配置和未识别的值均按关闭处理。"""
    value = os.getenv(name)
    if value is None:
        return default
    return value.strip().lower() in {"1", "true", "yes", "on"}


def _is_local_websocket_origin(websocket: WebSocket) -> bool:
    """浏览器 WebSocket 必须来自本地前端。"""
    return websocket.headers.get("origin") in LOCAL_DEVELOPMENT_ORIGINS


def _is_loopback_host(host: str | None) -> bool:
    """Only numeric IPv4/IPv6 loopback peers are trusted."""
    if not host:
        return False
    try:
        return ip_address(host).is_loopback
    except ValueError:
        return False


def _is_loopback_websocket_peer(websocket: WebSocket) -> bool:
    return _is_loopback_host(websocket.client.host if websocket.client else None)


def _acquire_single_writer_lock():
    """Fail startup when another API process owns this Trace directory."""
    global _writer_lock_file
    if _writer_lock_file is not None:
        return
    lock_path = Path(".trace").resolve() / ".api-writer.lock"
    lock_path.parent.mkdir(parents=True, exist_ok=True)
    handle = lock_path.open("a+")
    try:
        fcntl.flock(handle.fileno(), fcntl.LOCK_EX | fcntl.LOCK_NB)
    except BlockingIOError as exc:
        handle.close()
        raise RuntimeError(
            f"another Agent API writer already owns {lock_path.parent}"
        ) from exc
    _writer_lock_file = handle


def _release_single_writer_lock():
    global _writer_lock_file
    if _writer_lock_file is None:
        return
    fcntl.flock(_writer_lock_file.fileno(), fcntl.LOCK_UN)
    _writer_lock_file.close()
    _writer_lock_file = None

# 设置WebSocket日志推送
setup_websocket_logging(level=logging.INFO)


# ===== FastAPI 应用 =====

app = FastAPI(
    title="Agent API",
    description="Agent 查询 + 执行 API",
    version="1.0.0"
)

# CORS 配置（允许前端跨域访问）
app.add_middleware(
    CORSMiddleware,
    allow_origins=list(LOCAL_DEVELOPMENT_ORIGINS),
    allow_credentials=False,
    allow_methods=["*"],
    allow_headers=["*"],
)

# 拒绝非本地 Host，防止误暴露和 DNS rebinding。
app.add_middleware(
    TrustedHostMiddleware,
    allowed_hosts=list(LOCAL_DEVELOPMENT_HOSTS),
)


@app.middleware("http")
async def reject_non_local_browser_origins(request: Request, call_next):
    """Fail closed unless the TCP peer and optional browser Origin are local."""
    if not _is_loopback_host(request.client.host if request.client else None):
        return JSONResponse(
            status_code=403,
            content={"detail": "Loopback client required"},
        )
    origin = request.headers.get("origin")
    if origin is not None and origin not in LOCAL_DEVELOPMENT_ORIGINS:
        return JSONResponse(
            status_code=403,
            content={"detail": "Local development origin required"},
        )
    return await call_next(request)


# ===== 初始化存储 =====

# 使用文件系统存储（支持跨进程和持久化）
trace_store = FileSystemTraceStore(base_path=".trace")

# 注入到 step_tree 模块
set_api_trace_store(trace_store)
set_ws_trace_store(trace_store)
set_upload_trace_store(trace_store)


# ===== 可选：配置 Runner（启用执行 API）=====

# 如需启用 POST /api/traces（新建/运行/停止/反思），取消以下注释并配置 LLM：

from cyber_agent.core.runner import AgentRunner
from cyber_agent.llm import create_openrouter_llm_call

runner = AgentRunner(
    trace_store=trace_store,
    llm_call=create_openrouter_llm_call(model="anthropic/claude-sonnet-4.5"),
)
set_runner(runner)


# ===== 注册路由 =====

# Examples API（GET /api/examples）
app.include_router(examples_router)

# Trace 上传 API（POST /api/traces/upload）
app.include_router(upload_router)

# Trace 执行 API（POST + GET /running，需配置 Runner）
# 注意：run_router 必须在 api_router 之前注册，否则 GET /running 会被 /{trace_id} 捕获
app.include_router(run_router)

# 经验 API（GET /api/experiences，需配置 Runner）
app.include_router(experiences_router)

# Trace 查询 API（GET）
app.include_router(api_router)

# Trace WebSocket（实时推送）
app.include_router(ws_router)

# Logs WebSocket（日志推送）
app.include_router(logs_router)


@app.on_event("startup")
async def on_startup():
    """服务器启动时执行状态对齐"""
    _acquire_single_writer_lock()
    for trace in await trace_store.list_traces(limit=10_000):
        await trace_store.reconcile_trace_integrity(trace.trace_id)
    from cyber_agent.trace.run_api import reconcile_traces
    await reconcile_traces()


@app.on_event("shutdown")
async def on_shutdown():
    _release_single_writer_lock()

@app.websocket("/ws_ping")
async def ws_ping(websocket: WebSocket):
    if not _is_loopback_websocket_peer(websocket):
        await websocket.close(code=1008, reason="Loopback client required")
        return
    if not _is_local_websocket_origin(websocket):
        await websocket.close(code=1008, reason="Local development origin required")
        return
    await websocket.accept()
    await websocket.send_text("pong")
    await websocket.close()


# ===== 健康检查 =====

@app.get("/health")
async def health_check():
    """健康检查"""
    return {
        "status": "ok",
        "service": "Agent Step Tree API",
        "version": "1.0.0",
        "deployment_mode": "local_development",
        "network_scope": "loopback_only",
        "authentication": "none",
    }


# ===== 启动服务 =====

if __name__ == "__main__":
    logger.info("Starting API server...")
    uvicorn.run(
        "api_server:app",
        host="127.0.0.1",
        port=8000,
        reload=_env_flag("TRACE_API_RELOAD"),
        log_level="info"
    )