app.py

import asyncio
import multiprocessing
import os
import sys
from contextlib import asynccontextmanager
from typing import Optional
from fastapi import FastAPI, Request
from fastapi.responses import JSONResponse
from starlette.middleware.base import BaseHTTPMiddleware
from starlette.types import ASGIApp
from typing import Callable, Awaitable
from playwright.async_api import Browser as AsyncBrowser, Playwright as AsyncPlaywright
from config import *
from models import WebSocketConnectionManager
from logger import initialize_logging, restore_streams
from browser import _initialize_page_logic, _close_page_logic, load_excluded_models, _handle_initial_model_state_and_storage
from asyncio import Queue, Lock
from . import auth_utils
playwright_manager: Optional[AsyncPlaywright] = None
browser_instance: Optional[AsyncBrowser] = None
page_instance = None
is_playwright_ready = False
is_browser_connected = False
is_page_ready = False
is_initializing = False
global_model_list_raw_json = None
parsed_model_list = []
model_list_fetch_event = None
current_ai_studio_model_id = None
model_switching_lock = None
excluded_model_ids = set()
request_queue = None
processing_lock = None
worker_task = None
page_params_cache = {}
params_cache_lock = None
log_ws_manager = None
STREAM_QUEUE = None
STREAM_PROCESS = None

def _setup_logging():
    import server
    log_level_env = os.environ.get('SERVER_LOG_LEVEL', 'INFO')
    redirect_print_env = os.environ.get('SERVER_REDIRECT_PRINT', 'false')
    server.log_ws_manager = WebSocketConnectionManager()
    return initialize_logging(target_logger=server.logger, ws_manager=server.log_ws_manager, level_name=log_level_env, redirect_output=redirect_print_env)

def _initialize_globals():
    import server
    server.request_queue = Queue()
    server.processing_lock = Lock()
    server.model_switching_lock = Lock()
    server.params_cache_lock = Lock()
    auth_utils.initialize_keys()
    server.logger.info('API keys and global locks initialized.')

def _initialize_proxy_settings():
    import server
    STREAM_PORT = os.environ.get('STREAM_PORT')
    if STREAM_PORT == '0':
        PROXY_SERVER_ENV = os.environ.get('HTTPS_PROXY') or os.environ.get('HTTP_PROXY')
    else:
        PROXY_SERVER_ENV = f'http://127.0.0.1:{STREAM_PORT or 3120}/'
    if PROXY_SERVER_ENV:
        server.PLAYWRIGHT_PROXY_SETTINGS = {'server': PROXY_SERVER_ENV}
        if NO_PROXY_ENV:
            server.PLAYWRIGHT_PROXY_SETTINGS['bypass'] = NO_PROXY_ENV.replace(',', ';')
        server.logger.info(f'Playwright proxy settings configured: {server.PLAYWRIGHT_PROXY_SETTINGS}')
    else:
        server.logger.info('No proxy configured for Playwright.')

async def _wait_for_port(port: int, timeout: float = 10.0, interval: float = 0.3):
    import socket
    deadline = asyncio.get_event_loop().time() + timeout
    while asyncio.get_event_loop().time() < deadline:
        try:
            with socket.create_connection(('127.0.0.1', port), timeout=1):
                return True
        except (ConnectionRefusedError, OSError, socket.timeout):
            await asyncio.sleep(interval)
    return False

async def _start_stream_proxy():
    import proxy
    import server
    STREAM_PORT = os.environ.get('STREAM_PORT')
    if STREAM_PORT != '0':
        port = int(STREAM_PORT or 3120)
        STREAM_PROXY_SERVER_ENV = os.environ.get('UNIFIED_PROXY_CONFIG') or os.environ.get('HTTPS_PROXY') or os.environ.get('HTTP_PROXY')
        server.logger.info(f'Starting STREAM proxy on port {port} with upstream proxy: {STREAM_PROXY_SERVER_ENV}')
        for attempt in range(3):
            current_port = port + attempt
            server.STREAM_QUEUE = multiprocessing.Queue()
            server.STREAM_PROCESS = multiprocessing.Process(target=proxy.start, args=(server.STREAM_QUEUE, current_port, STREAM_PROXY_SERVER_ENV))
            server.STREAM_PROCESS.start()
            server.logger.info(f'STREAM proxy process started on port {current_port}. Waiting for port readiness...')
            if await _wait_for_port(current_port, timeout=30.0):
                server.STREAM_PORT_ACTUAL = current_port
                server.logger.info(f'STREAM proxy port {current_port} is ready.')
                if current_port != port:
                    server.logger.warning(f'STREAM proxy using fallback port {current_port} (requested {port}).')
                return
            else:
                server.logger.warning(f'STREAM proxy port {current_port} not ready, killing process...')
                server.STREAM_PROCESS.terminate()
                server.STREAM_PROCESS.join(timeout=3)
        server.logger.error(f'STREAM proxy failed to start after 3 attempts.')

async def _initialize_browser_and_page():
    import server
    from playwright.async_api import async_playwright
    server.logger.info('Starting Playwright...')
    server.playwright_manager = await async_playwright().start()
    server.is_playwright_ready = True
    server.logger.info('Playwright started.')
    ws_endpoint = os.environ.get('CAMOUFOX_WS_ENDPOINT')
    launch_mode = os.environ.get('LAUNCH_MODE', 'unknown')
    if not ws_endpoint and launch_mode != 'direct_debug_no_browser':
        raise ValueError('CAMOUFOX_WS_ENDPOINT environment variable is missing.')
    if ws_endpoint:
        server.logger.info(f'Connecting to browser at: {ws_endpoint}')
        server.browser_instance = await server.playwright_manager.firefox.connect(ws_endpoint, timeout=30000)
        server.is_browser_connected = True
        server.logger.info(f'Connected to browser: {server.browser_instance.version}')
        server.page_instance, server.is_page_ready = await _initialize_page_logic(server.browser_instance)
        if server.is_page_ready:
            await _handle_initial_model_state_and_storage(server.page_instance)
            server.logger.info('Page initialized successfully.')
        else:
            server.logger.error('Page initialization failed.')
    if not server.model_list_fetch_event.is_set():
        server.model_list_fetch_event.set()

async def _shutdown_resources():
    import server
    logger = server.logger
    logger.info('Shutting down resources...')
    if server.STREAM_PROCESS:
        server.STREAM_PROCESS.terminate()
        logger.info('STREAM proxy terminated.')
    if server.worker_task and (not server.worker_task.done()):
        server.worker_task.cancel()
        try:
            await asyncio.wait_for(server.worker_task, timeout=5.0)
        except (asyncio.TimeoutError, asyncio.CancelledError):
            pass
        logger.info('Worker task stopped.')
    if server.page_instance:
        await _close_page_logic()
    if server.browser_instance and server.browser_instance.is_connected():
        await server.browser_instance.close()
        logger.info('Browser connection closed.')
    if server.playwright_manager:
        await server.playwright_manager.stop()
        logger.info('Playwright stopped.')

@asynccontextmanager
async def lifespan(app: FastAPI):
    """FastAPI application life cycle management"""
    import server
    from api import queue_worker
    original_streams = (sys.stdout, sys.stderr)
    initial_stdout, initial_stderr = _setup_logging()
    logger = server.logger
    _initialize_globals()
    _initialize_proxy_settings()
    load_excluded_models(EXCLUDED_MODELS_FILENAME)
    server.is_initializing = True
    logger.info('Starting AI Studio Proxy Server...')
    try:
        await _start_stream_proxy()
        await _initialize_browser_and_page()
        launch_mode = os.environ.get('LAUNCH_MODE', 'unknown')
        if server.is_page_ready or launch_mode == 'direct_debug_no_browser':
            server.worker_task = asyncio.create_task(queue_worker())
            logger.info('Request processing worker started.')
        else:
            raise RuntimeError('Failed to initialize browser/page, worker not started.')
        logger.info('Server startup complete.')
        server.is_initializing = False
        yield
    except Exception as e:
        if 'Target page, context or browser has been closed' in str(e):
            logger.warning(f'Application startup failed (browser closed): {e}')
        elif 'NS_ERROR_PROXY' in str(e) or 'PROXY_CONNECTION_REFUSED' in str(e):
            logger.warning(f'Application startup failed (proxy error): {e}')
        else:
            logger.critical(f'Application startup failed: {e}', exc_info=True)
        await _shutdown_resources()
        raise RuntimeError(f'Application startup failed: {e}') from e
    finally:
        logger.info('Shutting down server...')
        await _shutdown_resources()
        restore_streams(initial_stdout, initial_stderr)
        restore_streams(*original_streams)
        logger.info('Server shutdown complete.')

class APIKeyAuthMiddleware(BaseHTTPMiddleware):

    def __init__(self, app: ASGIApp):
        super().__init__(app)
        self.excluded_paths = ['/v1/models', '/health', '/docs', '/openapi.json', '/redoc', '/favicon.ico']

    async def dispatch(self, request: Request, call_next: Callable[[Request], Awaitable]):
        if not auth_utils.API_KEYS:
            return await call_next(request)
        if not request.url.path.startswith('/v1/'):
            return await call_next(request)
        for excluded_path in self.excluded_paths:
            if request.url.path == excluded_path or request.url.path.startswith(excluded_path + '/'):
                return await call_next(request)
        api_key = None
        auth_header = request.headers.get('Authorization')
        if auth_header and auth_header.startswith('Bearer '):
            api_key = auth_header[7:]
        if not api_key:
            api_key = request.headers.get('X-API-Key')
        if not api_key or not auth_utils.verify_api_key(api_key):
            return JSONResponse(status_code=401, content={'error': {'message': "Invalid or missing API key. Please provide a valid API key using 'Authorization: Bearer <your_key>' or 'X-API-Key: <your_key>' header.", 'type': 'invalid_request_error', 'param': None, 'code': 'invalid_api_key'}})
        return await call_next(request)

def create_app() -> FastAPI:
    app = FastAPI(
        title='AIStudio2API',
        description='AI Studio 代理服务器 - 将 Google AI Studio 转换为 OpenAI 兼容 API',
        version='1.0.0',
        lifespan=lifespan,
        docs_url='/docs',
        redoc_url='/redoc',
    )
    
    from fastapi.middleware.cors import CORSMiddleware
    app.add_middleware(
        CORSMiddleware,
        allow_origins=["*"],
        allow_credentials=True,
        allow_methods=["*"],
        allow_headers=["*"],
    )
    
    app.add_middleware(APIKeyAuthMiddleware)
    from .routes import get_api_info, health_check, list_models, chat_completions, cancel_request, get_queue_status, websocket_log_endpoint, get_api_keys, add_api_key, test_api_key, delete_api_key, generate_speech, generate_image, generate_video, generate_nano_content
    from fastapi.responses import FileResponse
    
    app.add_api_route('/api/info', get_api_info, methods=['GET'], tags=['System'], summary='获取API配置信息')
    app.add_api_route('/health', health_check, methods=['GET'], tags=['System'], summary='健康检查')
    
    app.add_api_route('/v1/models', list_models, methods=['GET'], tags=['Chat'], summary='获取模型列表', description='返回所有可用的AI模型')
    app.add_api_route('/v1/chat/completions', chat_completions, methods=['POST'], tags=['Chat'], summary='聊天对话', description='OpenAI兼容的聊天接口，支持流式和非流式响应')
    app.add_api_route('/v1/cancel/{req_id}', cancel_request, methods=['POST'], tags=['Chat'], summary='取消请求')
    app.add_api_route('/v1/queue', get_queue_status, methods=['GET'], tags=['Chat'], summary='获取队列状态')
    
    app.add_api_route('/api/keys', get_api_keys, methods=['GET'], tags=['Keys'], summary='获取API密钥列表')
    app.add_api_route('/api/keys', add_api_key, methods=['POST'], tags=['Keys'], summary='添加API密钥')
    app.add_api_route('/api/keys/test', test_api_key, methods=['POST'], tags=['Keys'], summary='测试API密钥')
    app.add_api_route('/api/keys', delete_api_key, methods=['DELETE'], tags=['Keys'], summary='删除API密钥')
    
    app.add_api_route('/generate-speech', generate_speech, methods=['POST'], tags=['Media'], summary='TTS语音合成', description='Gemini 2.5 TTS语音生成')
    app.add_api_route('/v1beta/models/{model}:generateContent', generate_speech, methods=['POST'], tags=['Media'], summary='TTS (v1beta格式)')
    app.add_api_route('/generate-image', generate_image, methods=['POST'], tags=['Media'], summary='Imagen图片生成', description='Imagen 3图片生成')
    app.add_api_route('/v1beta/models/{model}:predict', generate_image, methods=['POST'], tags=['Media'], summary='Imagen (v1beta格式)')
    app.add_api_route('/generate-video', generate_video, methods=['POST'], tags=['Media'], summary='Veo视频生成', description='Veo 2视频生成')
    app.add_api_route('/v1beta/models/{model}:predictLongRunning', generate_video, methods=['POST'], tags=['Media'], summary='Veo (v1beta格式)')
    app.add_api_route('/nano/generate', generate_nano_content, methods=['POST'], tags=['Media'], summary='Nano图片生成', description='Gemini 2.5 Flash原生图片生成')
    app.add_api_route('/v1beta/models/gemini-2.5-flash-image:generateContent', generate_nano_content, methods=['POST'], tags=['Media'], summary='Nano (v1beta格式)')
    
    app.add_api_websocket_route('/ws/logs', websocket_log_endpoint)
    return app