main.py

"""Onyx to OpenAI API proxy server."""
import json
import logging
import time
import traceback
import uuid
import hmac
import html
from contextlib import asynccontextmanager
from typing import Optional

import httpx
from fastapi import FastAPI, Header, HTTPException, Request
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import HTMLResponse, JSONResponse, StreamingResponse

import config
import onyx
from status_manager import status_manager
from auth_manager import auth_manager

logging.basicConfig(
    level=getattr(logging, config.LOG_LEVEL),
    format="%(asctime)s %(name)s %(levelname)s %(message)s",
)
logger = logging.getLogger("onyxtoopenaicodex")

http_client: Optional[httpx.AsyncClient] = None


@asynccontextmanager
async def lifespan(app):
    global http_client
    # trust_env=True 允许读取本地的 HTTP_PROXY / HTTPS_PROXY 翻墙软件环境变量
    http_client = httpx.AsyncClient(timeout=float(config.REQUEST_TIMEOUT), trust_env=True)
    logger.info("Server starting on port %s", config.PORT)
    logger.info("Available models: %s", len(config.MODEL_MAP))
    yield
    if http_client:
        await http_client.aclose()


app = FastAPI(title="Onyx2OpenAI", version="1.0.0", lifespan=lifespan)
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=False,
    allow_methods=["*"],
    allow_headers=["*"],
)


@app.exception_handler(Exception)
async def global_exception_handler(request: Request, exc: Exception):
    logger.error("Global Error: %s\n%s", str(exc), traceback.format_exc())
    return JSONResponse(
        status_code=500,
        content={"error": {"message": "Internal server error", "type": "server_error"}},
    )


def verify_auth(authorization: Optional[str] = None):
    if not config.API_KEY:
        return
    if not authorization or not authorization.startswith("Bearer "):
        raise HTTPException(401, "Missing Authorization header")
    token = authorization.split(" ", 1)[1]
    if not hmac.compare_digest(token, config.API_KEY):
        raise HTTPException(401, "Invalid token")


@app.get("/favicon.ico", include_in_schema=False)
async def favicon():
    return HTMLResponse(content="")


@app.get("/")
async def root():
    # 重定向到仪表盘或返回简单消息
    return HTMLResponse("<script>window.location.href='/dashboard'</script>")


@app.get("/dashboard", response_class=HTMLResponse)
async def dashboard(key: Optional[str] = None):
    # 如果设置了 API_KEY，则需要通过 ?key=xxx 验证才能访问控制面板
    if config.API_KEY and not hmac.compare_digest(key or "", config.API_KEY):
        return HTMLResponse(
            "<h2 style='text-align:center;margin-top:20%;color:#ef4444;font-family:sans-serif;'>"
            "🔒 需要密码才能访问控制面板<br><small style='color:#888;'>请在 URL 后加上 ?key=你的密码</small></h2>",
            status_code=403
        )
    stats = await status_manager.get_stats()
    
    # 账号状态 HTML 生成
    accounts = auth_manager.accounts
    status_dict = auth_manager.status
    legacy = auth_manager.legacy_cookies
    
    cookie_status_html = ""
    expired_count = 0
    total_accounts = len(accounts) + len(legacy)
    
    for email, pwd in accounts:
        acc_status = status_dict.get(email, {})
        has_tried = "fastapiusersauth" in acc_status
        is_valid = acc_status.get("valid", False)
        reason = acc_status.get("reason", "")
        
        if not has_tried:
            status_text = "⚪ 等待首次请求/未分配"
            status_class = ""
        elif not is_valid:
            expired_count += 1
            status_text = f"🔴 失效 ({reason})" if reason else "🔴 失效/登录中"
            status_class = "expired"
        else:
            status_text = "🟢 正常可用"
            status_class = "active"
            
        cookie_val = acc_status.get("fastapiusersauth", "")
        display_val = f"{cookie_val[:10]}...{cookie_val[-10:]}" if cookie_val else "未分配 Token"
        
        cookie_status_html += f'''
        <div class="cookie-card {status_class}">
            <span class="cookie-index">{html.escape(email)}</span>
            <span class="cookie-status">{html.escape(status_text)}</span>
            <div class="cookie-val">{html.escape(display_val)}</div>
        </div>
        '''
        
    for i, c in enumerate(legacy):
        cookie_status_html += f'''
        <div class="cookie-card active">
            <span class="cookie-index">Legacy Cookie #{i}</span>
            <span class="cookie-status">🟢 未知/旧配置</span>
            <div class="cookie-val">{html.escape(c[:10] + "..." + c[-10:])}</div>
        </div>
        '''

    html_content = f"""
    <!DOCTYPE html>
    <html>
    <head>
        <title>Onyx2OpenAI 控制面板</title>
        <meta charset="utf-8">
        <meta name="viewport" content="width=device-width, initial-scale=1">
        <style>
            :root {{
                --bg: #0f172a;
                --card: #1e293b;
                --text: #f8fafc;
                --primary: #38bdf8;
                --success: #22c55e;
                --danger: #ef4444;
            }}
            body {{
                font-family: 'Inter', system-ui, sans-serif;
                background: var(--bg);
                color: var(--text);
                margin: 0;
                padding: 2rem;
                display: flex;
                flex-direction: column;
                align-items: center;
            }}
            .container {{
                max-width: 900px;
                width: 100%;
            }}
            h1 {{
                color: var(--primary);
                font-size: 2.5rem;
                margin-bottom: 2rem;
                text-align: center;
                background: linear-gradient(to right, #38bdf8, #818cf8);
                -webkit-background-clip: text;
                -webkit-text-fill-color: transparent;
            }}
            .stats-grid {{
                display: grid;
                grid-template-columns: repeat(auto-fit, minmax(200px, 1fr));
                gap: 1.5rem;
                margin-bottom: 3rem;
            }}
            .stat-card {{
                background: var(--card);
                padding: 1.5rem;
                border-radius: 1rem;
                text-align: center;
                border: 1px solid rgba(255,255,255,0.1);
                box-shadow: 0 10px 15px -3px rgba(0,0,0,0.1);
            }}
            .stat-val {{
                font-size: 2rem;
                font-weight: bold;
                color: var(--primary);
                margin-bottom: 0.5rem;
            }}
            .stat-label {{
                font-size: 0.875rem;
                opacity: 0.7;
            }}
            .cookie-list {{
                display: grid;
                gap: 1rem;
            }}
            .cookie-card {{
                background: var(--card);
                padding: 1.25rem;
                border-radius: 0.75rem;
                display: flex;
                align-items: center;
                justify-content: space-between;
                border-left: 4px solid var(--success);
            }}
            .cookie-card.expired {{
                border-left-color: var(--danger);
                opacity: 0.8;
            }}
            .cookie-index {{ font-weight: bold; }}
            .cookie-val {{ font-family: monospace; opacity: 0.5; font-size: 0.75rem; }}
            .footer {{
                margin-top: 4rem;
                text-align: center;
                opacity: 0.5;
                font-size: 0.875rem;
            }}
            .btn-reset {{
                background: var(--card);
                border: 1px solid var(--primary);
                color: var(--primary);
                padding: 0.5rem 1rem;
                border-radius: 0.5rem;
                cursor: pointer;
                text-decoration: none;
                transition: all 0.2s;
            }}
            .btn-reset:hover {{
                background: var(--primary);
                color: var(--bg);
            }}
        </style>
    </head>
    <body>
        <div class="container">
            <h1>Onyx2OpenAI Status</h1>
            
            <div class="stats-grid">
                <div class="stat-card">
                    <div class="stat-val">{stats['total_requests']}</div>
                    <div class="stat-label">总请求数</div>
                </div>
                <div class="stat-card">
                    <div class="stat-val" style="color: var(--success)">{stats['success_counts']}</div>
                    <div class="stat-label">成功</div>
                </div>
                <div class="stat-card">
                    <div class="stat-val" style="color: var(--danger)">{stats['failure_counts']}</div>
                    <div class="stat-label">API失败</div>
                </div>
                <div class="stat-card">
                    <div class="stat-val">{expired_count} / {total_accounts}</div>
                    <div class="stat-label">失效/耗尽账号</div>
                </div>
            </div>

            <h2 style="margin-bottom: 1.5rem">账号与凭证节点健康状况</h2>
            <div class="cookie-list">
                {cookie_status_html}
            </div>

            <div class="footer">
                <form action="/dashboard/reset{('?key=' + html.escape(key)) if key else ''}" method="post" style="display: inline;">
                    <button type="submit" class="btn-reset">清除所有保存的账户凭证并重试</button>
                </form>
                <div style="margin-top: 1rem">Server v1.1 | 运行正常</div>
            </div>
        </div>
    </body>
    </html>
    """
    return html_content


@app.post("/dashboard/reset")
async def reset_cookies(key: Optional[str] = None):
    # 重置操作也需要密码验证
    if config.API_KEY and not hmac.compare_digest(key or "", config.API_KEY):
        raise HTTPException(403, "Unauthorized")
    import os
    if os.path.exists(auth_manager.status_file):
        try:
            os.remove(auth_manager.status_file)
        except Exception:
            pass
    async with auth_manager.lock:
        auth_manager.status.clear()
    await status_manager.reset_expired_cookies()
    return HTMLResponse("<script>alert('已重置所有账号鉴权缓存，下次请求将全部重新登录！'); window.location.href='/dashboard'</script>")


@app.get("/health")
async def health():
    return {"status": "ok", "version": "1.0.0", "models": len(config.MODEL_MAP)}


@app.get("/v1/models")
async def list_models(authorization: Optional[str] = Header(None)):
    verify_auth(authorization)
    data = [
        {
            "id": model_name,
            "object": "model",
            "created": 1700000000,
            "owned_by": "onyx",
        }
        for model_name in config.MODEL_MAP
    ]
    return {"object": "list", "data": data}


@app.post("/v1/chat/completions")
async def chat_completions(request: Request, authorization: Optional[str] = Header(None)):
    verify_auth(authorization)

    try:
        body = await request.json()
    except Exception as e:
        raise HTTPException(400, f"Invalid JSON: {e}")

    messages = body.get("messages", [])
    model_name = body.get("model", "claude-opus-4.6")
    stream = body.get("stream", True)
    include_reasoning = body.get("include_reasoning", True)

    logger.info("Request: model=%s, messages=%s, stream=%s", model_name, len(messages), stream)

    if stream:
        # 流式模式：统计在生成器内部完成，因为此处只是创建 StreamingResponse 对象
        return await _stream_response(messages, model_name, include_reasoning)
    else:
        try:
            response = await _non_stream_response(messages, model_name, include_reasoning)
            await status_manager.record_request(success=True)
            return response
        except Exception as e:
            await status_manager.record_request(success=False)
            raise e


async def _stream_response(messages, model_name, include_reasoning):
    response_id = f"chatcmpl-{uuid.uuid4()}"

    async def generate():
        had_error = False
        try:
            async for item_type, content in onyx.stream_chat(http_client, messages, model_name):
                if item_type == "thinking" and include_reasoning:
                    chunk = {
                        "id": response_id,
                        "object": "chat.completion.chunk",
                        "created": int(time.time()),
                        "model": model_name,
                        "choices": [{
                            "index": 0,
                            "delta": {"reasoning_content": content},
                            "finish_reason": None,
                        }],
                    }
                    yield f"data: {json.dumps(chunk, ensure_ascii=False)}\n\n"
                elif item_type == "text":
                    chunk = {
                        "id": response_id,
                        "object": "chat.completion.chunk",
                        "created": int(time.time()),
                        "model": model_name,
                        "choices": [{
                            "index": 0,
                            "delta": {"content": content},
                            "finish_reason": None,
                        }],
                    }
                    yield f"data: {json.dumps(chunk, ensure_ascii=False)}\n\n"
        except Exception as e:
            had_error = True
            logger.error("Stream error: %s", e, exc_info=True)
            err_msg = "上游服务暂时不可用，请稍后重试"
            error_event = {"error": {"message": err_msg, "type": "upstream_error"}}
            yield f"data: {json.dumps(error_event, ensure_ascii=False)}\n\n"
        finally:
            await status_manager.record_request(success=not had_error)

        if not had_error:
            end_chunk = {
                "id": response_id,
                "object": "chat.completion.chunk",
                "created": int(time.time()),
                "model": model_name,
                "choices": [{"index": 0, "delta": {}, "finish_reason": "stop"}],
            }
            yield f"data: {json.dumps(end_chunk)}\n\n"
        yield "data: [DONE]\n\n"

    return StreamingResponse(
        generate(),
        media_type="text/event-stream",
        headers={"Cache-Control": "no-cache", "Connection": "keep-alive"},
    )


async def _non_stream_response(messages, model_name, include_reasoning):
    response_id = f"chatcmpl-{uuid.uuid4()}"
    text_content, thinking_content = await onyx.full_chat(http_client, messages, model_name)

    message = {"role": "assistant", "content": text_content}
    if include_reasoning and thinking_content:
        message["reasoning_content"] = thinking_content

    return JSONResponse({
        "id": response_id,
        "object": "chat.completion",
        "created": int(time.time()),
        "model": model_name,
        "choices": [{
            "index": 0,
            "message": message,
            "finish_reason": "stop",
        }],
        "usage": {"prompt_tokens": 0, "completion_tokens": 0, "total_tokens": 0},
    })


if __name__ == "__main__":
    import uvicorn
    print("=" * 50)
    print("Onyx2OpenAI Server v1.0")
    print("=" * 50)
    print(f"Address: http://127.0.0.1:{config.PORT}")
    print("Models:  /v1/models")
    print("Chat:    /v1/chat/completions")
    print("=" * 50)
    uvicorn.run(app, host="0.0.0.0", port=config.PORT)