app/main.py

"""FastAPI 入口：组装上游 provider + 挂载 OpenAI / Anthropic 兼容路由。"""
from __future__ import annotations

import asyncio
import logging
from contextlib import asynccontextmanager
from pathlib import Path

import httpx
from fastapi import FastAPI

from app.account import AccountPool, FailReason, set_cooldown_policy
from app.adapters.anthropic_messages import router as anthropic_router
from app.adapters.openai_chat import router as openai_chat_router
from app.adapters.openai_images import router as openai_images_router
from app.adapters.openai_models import router as openai_models_router
from app.adapters.openai_responses import router as openai_responses_router
from app.admin import router as admin_router
from app.config import get_settings
from app.http_log import RequestResponseLogMiddleware
from app.logging_setup import setup_logging
from app.upstream import get_provider

logger = logging.getLogger(__name__)

async def _watch_pool(app: FastAPI, interval: float) -> None:
    """后台任务：周期扫描 account/*.json，把注册机/手动新写的账号热加载进池。

    2026-08-08：注册机（独立进程）写盘新账号时网关不知情，之前要手动 /admin/reload
    或重启才感知。本任务定期 reload：新号进池 + 建 provider；已删号（余额<=0 删除）
    同步移除。与 /admin/reload 等价，但自动。
    """
    while True:
        await asyncio.sleep(interval)
        try:
            st = app.state
            pool: AccountPool = st.pool
            providers: dict[str, object] = st.providers
            before = {a.name for a in pool.all()}
            pool.reload()
            after = {a.name for a in pool.all()}
            new = after - before
            removed = before - after
            changed = False
            for name in sorted(new):
                acc = next(a for a in pool.all() if a.name == name)
                providers[name] = get_provider(acc, st.settings, st.http_client)
                changed = True
            for name in sorted(removed):
                providers.pop(name, None)
                changed = True
            if changed:
                logger.info("pool hot-reload: +%d new, -%d removed (total=%d)",
                            len(new), len(removed), len(after))
        except Exception:  # noqa: BLE001
            logger.warning("pool hot-reload failed", exc_info=True)


@asynccontextmanager
async def lifespan(app: FastAPI):
    settings = get_settings()
    setup_logging(settings)
    # 同步可恢复失效的处理策略与冷却时长
    seconds_map: dict[FailReason, float] = {}
    if settings.cooldown_seconds_quota is not None:
        seconds_map[FailReason.QUOTA_EXHAUSTED] = settings.cooldown_seconds_quota
    if settings.cooldown_seconds_cf is not None:
        seconds_map[FailReason.CF_CHALLENGE] = settings.cooldown_seconds_cf
    set_cooldown_policy(
        settings.quota_exhausted_action,
        seconds=settings.cooldown_seconds,
        seconds_map=seconds_map,
    )
    # proxy 未配置时 httpx 直连（proxy=None）
    http_client = httpx.AsyncClient(
        timeout=settings.request_timeout,
        proxy=settings.effective_proxy(),
    )
    # 加载账号池；为每个账号建独立的 UpstreamProvider（共享 http_client）
    pool = AccountPool.load(Path(settings.account_dir))
    providers: dict[str, object] = {}
    for acc in pool.all():
        providers[acc.name] = get_provider(acc, settings, http_client)
    app.state.settings = settings
    app.state.http_client = http_client
    app.state.pool = pool
    app.state.providers = providers
    # 面板生成的 API Key 库（多 key，文件存储；config 的 gateway_api_key 兼容并存）
    from app.api_keys import load_keys
    app.state.api_keys = load_keys(Path(settings.log_dir))
    # token 用量存储（SQLite，logs/usage.db；面板 /admin/usage 聚合查询）
    try:
        from app import usage_store
        usage_store.init(Path(settings.log_dir) / "usage.db")
    except Exception:  # noqa: BLE001
        logger.warning("usage_store init failed (usage 统计将不可用)", exc_info=True)
    n_all = len(pool.all())
    n_ok = pool.available_count() if hasattr(pool, "available_count") else n_all
    logger.info(
        "startup accounts=%d available=%d proxy=%s log_file=%s timeout=%.0fs",
        n_all,
        n_ok,
        settings.effective_proxy() or "(direct)",
        f"{settings.log_dir}/{settings.log_filename}" if settings.log_enabled else "(disabled)",
        settings.request_timeout,
    )
    # 账号池热加载后台任务（注册机写盘的新号自动进池，不用重启/手动 reload）
    watch_task: asyncio.Task | None = None
    watch_interval = float(getattr(settings, "pool_watch_interval", 0.0) or 0.0)
    if watch_interval > 0:
        watch_task = asyncio.create_task(_watch_pool(app, watch_interval))
        logger.info("pool hot-reload enabled (every %.0fs)", watch_interval)
    try:
        yield
    finally:
        logger.info("shutdown")
        if watch_task is not None:
            watch_task.cancel()
        await http_client.aclose()

app = FastAPI(title="anuma2api", version="0.1.0", lifespan=lifespan)
app.add_middleware(RequestResponseLogMiddleware)

app.include_router(openai_models_router)
app.include_router(openai_chat_router)
app.include_router(openai_responses_router)
app.include_router(openai_images_router)
app.include_router(anthropic_router)
app.include_router(admin_router)

@app.get("/healthz")
async def healthz() -> dict[str, str]:
    return {"status": "ok"}