Create main.py

main.py

@@ -0,0 +1,799 @@
+"""
+Gemini API 兼容中间件服务器
+透明代理 - 将 HTTP 请求通过 WebSocket 转发给 WSClient 处理
+
+版本: 1.0.0
+协议: Gemini Compatible WebSocket Proxy Protocol
+"""
+
+import asyncio
+import json
+import time
+import uuid
+import logging
+from dataclasses import dataclass, field
+from enum import Enum
+from typing import Optional, Dict, List, Any, AsyncGenerator
+from contextlib import asynccontextmanager
+from datetime import datetime
+
+from fastapi import FastAPI, WebSocket, WebSocketDisconnect, HTTPException, Request, Path
+from fastapi.responses import StreamingResponse, JSONResponse
+from pydantic import BaseModel
+import uvicorn
+
+
+# ============================================================
+# 日志配置
+# ============================================================
+
+class ColoredFormatter(logging.Formatter):
+    """彩色日志格式化器"""
+    
+    COLORS = {
+        'DEBUG': '\033[36m',     # 青色
+        'INFO': '\033[32m',      # 绿色
+        'WARNING': '\033[33m',   # 黄色
+        'ERROR': '\033[31m',     # 红色
+        'CRITICAL': '\033[35m',  # 紫色
+    }
+    RESET = '\033[0m'
+    
+    def format(self, record):
+        color = self.COLORS.get(record.levelname, self.RESET)
+        record.levelname = f"{color}{record.levelname}{self.RESET}"
+        record.msg = f"{color}{record.msg}{self.RESET}"
+        return super().format(record)
+
+
+def setup_logging():
+    """配置日志"""
+    handler = logging.StreamHandler()
+    handler.setFormatter(ColoredFormatter(
+        fmt='%(asctime)s | %(levelname)-17s | %(message)s',
+        datefmt='%Y-%m-%d %H:%M:%S'
+    ))
+    
+    logger = logging.getLogger("middleware")
+    logger.setLevel(logging.DEBUG)
+    logger.addHandler(handler)
+    return logger
+
+
+log = setup_logging()
+
+
+# ============================================================
+# 配置
+# ============================================================
+
+class Config:
+    HOST = "0.0.0.0"
+    PORT = 8000
+    API_KEY = "sk-123456"              # 修改为你的密钥
+    REQUEST_TIMEOUT = 120              # 请求超时（秒）
+    HEARTBEAT_INTERVAL = 30            # 心跳间隔（秒）
+    REGISTER_TIMEOUT = 10              # 注册超时（秒）
+    LOG_BODY_MAX_LENGTH = 500          # 日志中 body 最大显示长度
+
+
+# ============================================================
+# 协议定义
+# ============================================================
+
+class MessageType(str, Enum):
+    # 连接管理
+    REGISTER = "register"
+    REGISTER_ACK = "register_ack"
+    # 请求
+    REQUEST = "request"
+    # 响应
+    RESPONSE = "response"
+    CHUNK = "chunk"
+    END = "end"
+    ERROR = "error"
+    # 控制
+    ABORT = "abort"
+    PING = "ping"
+    PONG = "pong"
+
+
+# ============================================================
+# 请求上下文
+# ============================================================
+
+class RequestStatus(str, Enum):
+    PENDING = "pending"
+    STREAMING = "streaming"
+    COMPLETED = "completed"
+    ERROR = "error"
+    ABORTED = "aborted"
+
+
+@dataclass
+class RequestContext:
+    """请求上下文，追踪每个请求的状态"""
+    id: str
+    created_at: float
+    is_stream: bool
+    model: str
+    status: RequestStatus = RequestStatus.PENDING
+    # 非流式请求：用 Future 等待完整响应
+    response_future: asyncio.Future = field(default_factory=lambda: asyncio.get_running_loop().create_future())
+    # 流式请求：用 Queue 传递 chunks
+    chunk_queue: asyncio.Queue = field(default_factory=asyncio.Queue)
+    # 统计
+    chunk_count: int = 0
+    
+    def elapsed_ms(self) -> int:
+        """返回请求耗时（毫秒）"""
+        return int((time.time() - self.created_at) * 1000)
+
+
+# ============================================================
+# 请求管理器
+# ============================================================
+
+class RequestManager:
+    """管理所有待处理的请求"""
+    
+    def __init__(self):
+        self.pending_requests: Dict[str, RequestContext] = {}
+        self._lock = asyncio.Lock()
+    
+    async def create_request(self, is_stream: bool, model: str) -> RequestContext:
+        """创建新的请求上下文"""
+        request_id = str(uuid.uuid4())
+        ctx = RequestContext(
+            id=request_id,
+            created_at=time.time(),
+            is_stream=is_stream,
+            model=model
+        )
+        async with self._lock:
+            self.pending_requests[request_id] = ctx
+        
+        log.debug(f"[ReqMgr] 创建请求 | id={request_id[:8]}... | stream={is_stream} | model={model}")
+        return ctx
+    
+    def get_request(self, request_id: str) -> Optional[RequestContext]:
+        return self.pending_requests.get(request_id)
+    
+    async def wait_for_response(self, request_id: str, timeout: float) -> Dict:
+        """等待非流式请求的完整响应"""
+        ctx = self.pending_requests.get(request_id)
+        if not ctx:
+            raise ValueError(f"Request {request_id} not found")
+        
+        try:
+            log.debug(f"[ReqMgr] 等待响应 | id={request_id[:8]}... | timeout={timeout}s")
+            result = await asyncio.wait_for(ctx.response_future, timeout=timeout)
+            log.info(f"[ReqMgr] 收到响应 | id={request_id[:8]}... | elapsed={ctx.elapsed_ms()}ms")
+            return result
+        except asyncio.TimeoutError:
+            ctx.status = RequestStatus.ERROR
+            log.error(f"[ReqMgr] 请求超时 | id={request_id[:8]}... | elapsed={ctx.elapsed_ms()}ms")
+            raise TimeoutError(f"Request {request_id} timed out after {timeout}s")
+        finally:
+            await self._cleanup_request(request_id)
+    
+    async def wait_for_stream(self, request_id: str, timeout: float) -> AsyncGenerator[Dict, None]:
+        """等待流式请求的数据块"""
+        ctx = self.pending_requests.get(request_id)
+        if not ctx:
+            raise ValueError(f"Request {request_id} not found")
+        
+        ctx.status = RequestStatus.STREAMING
+        log.debug(f"[ReqMgr] 开始流式接收 | id={request_id[:8]}...")
+        
+        try:
+            while True:
+                try:
+                    chunk = await asyncio.wait_for(ctx.chunk_queue.get(), timeout=timeout)
+                    if chunk is None:  # 流结束信号
+                        log.info(f"[ReqMgr] 流结束 | id={request_id[:8]}... | chunks={ctx.chunk_count} | elapsed={ctx.elapsed_ms()}ms")
+                        break
+                    ctx.chunk_count += 1
+                    yield chunk
+                except asyncio.TimeoutError:
+                    ctx.status = RequestStatus.ERROR
+                    log.error(f"[ReqMgr] 流超时 | id={request_id[:8]}... | chunks={ctx.chunk_count}")
+                    raise TimeoutError(f"Stream {request_id} timed out")
+        finally:
+            await self._cleanup_request(request_id)
+    
+    def resolve_request(self, request_id: str, response: Dict):
+        """解决非流式请求"""
+        ctx = self.pending_requests.get(request_id)
+        if ctx and not ctx.response_future.done():
+            ctx.status = RequestStatus.COMPLETED
+            ctx.response_future.set_result(response)
+            log.debug(f"[ReqMgr] 请求已解决 | id={request_id[:8]}...")
+    
+    def push_chunk(self, request_id: str, chunk: Dict):
+        """推送流式数据块"""
+        ctx = self.pending_requests.get(request_id)
+        if ctx and ctx.is_stream:
+            ctx.chunk_queue.put_nowait(chunk)
+    
+    def end_stream(self, request_id: str, final_body: Optional[Dict] = None):
+        """结束流式响应"""
+        ctx = self.pending_requests.get(request_id)
+        if ctx:
+            ctx.status = RequestStatus.COMPLETED
+            if final_body:
+                ctx.chunk_queue.put_nowait(final_body)
+            ctx.chunk_queue.put_nowait(None)  # 结束信号
+    
+    def fail_request(self, request_id: str, error: Dict):
+        """标记请求失败"""
+        ctx = self.pending_requests.get(request_id)
+        if ctx:
+            ctx.status = RequestStatus.ERROR
+            log.error(f"[ReqMgr] 请求失败 | id={request_id[:8]}... | error={error}")
+            if not ctx.response_future.done():
+                ctx.response_future.set_exception(Exception(json.dumps(error)))
+            if ctx.is_stream:
+                ctx.chunk_queue.put_nowait(None)
+    
+    def abort_request(self, request_id: str):
+        """中止请求"""
+        ctx = self.pending_requests.get(request_id)
+        if ctx:
+            ctx.status = RequestStatus.ABORTED
+            log.warning(f"[ReqMgr] 请求中止 | id={request_id[:8]}...")
+            if not ctx.response_future.done():
+                ctx.response_future.set_exception(Exception("Request aborted"))
+            if ctx.is_stream:
+                ctx.chunk_queue.put_nowait(None)
+    
+    async def fail_all_requests(self, error_message: str):
+        """使所有待处理请求失败"""
+        async with self._lock:
+            count = len(self.pending_requests)
+            if count > 0:
+                log.warning(f"[ReqMgr] 批量失败 | count={count} | reason={error_message}")
+            for request_id in list(self.pending_requests.keys()):
+                self.fail_request(request_id, {
+                    "error": {
+                        "code": 503,
+                        "message": error_message,
+                        "status": "UNAVAILABLE"
+                    }
+                })
+    
+    async def _cleanup_request(self, request_id: str):
+        """清理已完成的请求"""
+        async with self._lock:
+            self.pending_requests.pop(request_id, None)
+    
+    @property
+    def pending_count(self) -> int:
+        return len(self.pending_requests)
+
+
+# ============================================================
+# WebSocket 管理器
+# ============================================================
+
+class WebSocketManager:
+    """管理 WSClient 连接"""
+    
+    def __init__(self):
+        self.active_connection: Optional[WebSocket] = None
+        self.client_id: Optional[str] = None
+        self.client_version: Optional[str] = None
+        self.models: List[str] = []
+        self.max_concurrent: int = 1
+        self.connected_at: Optional[float] = None
+        self._lock = asyncio.Lock()
+        self._heartbeat_task: Optional[asyncio.Task] = None
+    
+    async def register(self, websocket: WebSocket, payload: Dict):
+        """注册 WSClient 连接"""
+        async with self._lock:
+            # 关闭旧连接
+            if self.active_connection and self.active_connection != websocket:
+                log.warning("[WSMgr] 关闭旧连接...")
+                try:
+                    await self.active_connection.close()
+                except:
+                    pass
+            
+            self.active_connection = websocket
+            self.client_id = payload.get("clientId", "unknown")
+            self.client_version = payload.get("clientVersion", "unknown")
+            self.models = payload.get("models", [])
+            self.max_concurrent = payload.get("maxConcurrent", 1)
+            self.connected_at = time.time()
+            
+            # 发送注册确认
+            ack_message = {
+                "id": str(uuid.uuid4()),
+                "type": MessageType.REGISTER_ACK.value,
+                "timestamp": int(time.time() * 1000),
+                "payload": {
+                    "success": True,
+                    "serverId": "gemini-middleware-001",
+                    "config": {
+                        "heartbeatInterval": Config.HEARTBEAT_INTERVAL * 1000,
+                        "requestTimeout": Config.REQUEST_TIMEOUT * 1000
+                    }
+                }
+            }
+            await websocket.send_json(ack_message)
+            
+            log.info(f"[WSMgr] ✅ WSClient 已注册")
+            log.info(f"[WSMgr]    ├─ clientId: {self.client_id}")
+            log.info(f"[WSMgr]    ├─ version: {self.client_version}")
+            log.info(f"[WSMgr]    ├─ models: {self.models}")
+            log.info(f"[WSMgr]    └─ maxConcurrent: {self.max_concurrent}")
+            
+            # 启动心跳
+            if self._heartbeat_task:
+                self._heartbeat_task.cancel()
+            self._heartbeat_task = asyncio.create_task(self._heartbeat_loop())
+    
+    async def unregister(self, websocket: WebSocket):
+        """注销连接"""
+        async with self._lock:
+            if self.active_connection == websocket:
+                uptime = int(time.time() - self.connected_at) if self.connected_at else 0
+                log.warning(f"[WSMgr] ❌ WSClient 已断开 | uptime={uptime}s")
+                
+                self.active_connection = None
+                self.client_id = None
+                self.models = []
+                self.connected_at = None
+                
+                if self._heartbeat_task:
+                    self._heartbeat_task.cancel()
+                    self._heartbeat_task = None
+    
+    def is_available(self) -> bool:
+        return self.active_connection is not None
+    
+    async def send_message(self, message: Dict):
+        """发送消息到 WSClient"""
+        if self.active_connection:
+            await self.active_connection.send_json(message)
+            log.debug(f"[WSMgr] 发送消息 | type={message.get('type')} | id={message.get('id', '')[:8]}...")
+    
+    async def send_request(self, request_id: str, is_stream: bool, body: Dict):
+        """发送请求到 WSClient"""
+        message = {
+            "id": request_id,
+            "type": MessageType.REQUEST.value,
+            "timestamp": int(time.time() * 1000),
+            "stream": is_stream,
+            "body": body
+        }
+        await self.send_message(message)
+        
+        body_preview = self._truncate_body(body)
+        log.info(f"[WSMgr] 📤 发送请求 | id={request_id[:8]}... | stream={is_stream}")
+        log.debug(f"[WSMgr]    └─ body: {body_preview}")
+    
+    async def send_abort(self, request_id: str, reason: str = "client_disconnected"):
+        """发送取消请求"""
+        message = {
+            "id": request_id,
+            "type": MessageType.ABORT.value,
+            "timestamp": int(time.time() * 1000),
+            "reason": reason
+        }
+        await self.send_message(message)
+        log.warning(f"[WSMgr] 🚫 发送取消 | id={request_id[:8]}... | reason={reason}")
+    
+    async def _heartbeat_loop(self):
+        """心跳循环"""
+        while True:
+            try:
+                await asyncio.sleep(Config.HEARTBEAT_INTERVAL)
+                if self.active_connection:
+                    ping_id = str(uuid.uuid4())
+                    await self.send_message({
+                        "id": ping_id,
+                        "type": MessageType.PING.value,
+                        "timestamp": int(time.time() * 1000)
+                    })
+                    log.debug(f"[WSMgr] 💓 发送心跳 | id={ping_id[:8]}...")
+            except asyncio.CancelledError:
+                break
+            except Exception as e:
+                log.error(f"[WSMgr] 心跳错误: {e}")
+    
+    def _truncate_body(self, body: Dict) -> str:
+        """截断 body 用于日志显示"""
+        s = json.dumps(body, ensure_ascii=False)
+        if len(s) > Config.LOG_BODY_MAX_LENGTH:
+            return s[:Config.LOG_BODY_MAX_LENGTH] + "..."
+        return s
+    
+    def get_status(self) -> Dict:
+        """获取连接状态"""
+        return {
+            "connected": self.is_available(),
+            "clientId": self.client_id,
+            "clientVersion": self.client_version,
+            "models": self.models,
+            "maxConcurrent": self.max_concurrent,
+            "uptime": int(time.time() - self.connected_at) if self.connected_at else 0
+        }
+
+
+# ============================================================
+# 全局实例
+# ============================================================
+
+request_manager = RequestManager()
+ws_manager = WebSocketManager()
+
+
+# ============================================================
+# FastAPI 应用
+# ============================================================
+
+@asynccontextmanager
+async def lifespan(app: FastAPI):
+    log.info("=" * 60)
+    log.info("  Gemini API 兼容中间件")
+    log.info("=" * 60)
+    log.info(f"  HTTP 端点: http://{Config.HOST}:{Config.PORT}")
+    log.info(f"  WebSocket: ws://{Config.HOST}:{Config.PORT}/ws")
+    log.info(f"  API 文档:  http://{Config.HOST}:{Config.PORT}/docs")
+    log.info("=" * 60)
+    yield
+    log.info("[Server] 服务关闭")
+
+
+app = FastAPI(
+    title="Gemini API Compatible Middleware",
+    description="透明代理 - 将 Gemini API 请求通过 WebSocket 转发给 WSClient",
+    version="1.0.0",
+    lifespan=lifespan
+)
+
+
+# ============================================================
+# API Key 校验 - 支持多种传递方式
+# ============================================================
+
+def extract_api_key(request: Request) -> Optional[str]:
+    """
+    从多个位置提取 API Key，按优先级：
+    1. Header: x-goog-api-key
+    2. Header: Authorization: Bearer <key>
+    3. Query: ?key=<key>
+    """
+    # 方式1: x-goog-api-key header
+    api_key = request.headers.get("x-goog-api-key")
+    if api_key:
+        log.debug(f"[Auth] 从 x-goog-api-key header 获取 key")
+        return api_key
+    
+    # 方式2: Authorization header (Bearer token)
+    auth_header = request.headers.get("authorization", "")
+    if auth_header.lower().startswith("bearer "):
+        api_key = auth_header[7:].strip()
+        log.debug(f"[Auth] 从 Authorization header 获取 key")
+        return api_key
+    
+    # 方式3: Query parameter
+    api_key = request.query_params.get("key")
+    if api_key:
+        log.debug(f"[Auth] 从 query parameter 获取 key")
+        return api_key
+    
+    return None
+
+
+def verify_api_key(request: Request):
+    """从请求中校验 API Key"""
+    api_key = extract_api_key(request)
+    
+    # 调试日志
+    log.debug(f"[Auth] 提取到的 Key: {api_key[:10] + '...' if api_key else 'None'}")
+    
+    if api_key != Config.API_KEY:
+        log.warning(f"[Auth] ⛔ 认证失败 | key={api_key[:10] + '...' if api_key else 'None'}")
+        raise HTTPException(
+            status_code=401,
+            detail={
+                "error": {
+                    "code": 401,
+                    "message": "Invalid API key. Provide via 'x-goog-api-key' header, 'Authorization: Bearer <key>' header, or '?key=<key>' query parameter.",
+                    "status": "UNAUTHENTICATED"
+                }
+            }
+        )
+    log.debug("[Auth] ✓ API Key 验证通过")
+
+
+# ============================================================
+# HTTP 路由 - Gemini API 兼容
+# ============================================================
+
+@app.get("/")
+async def root():
+    """根路径 - 服务信息"""
+    return {
+        "service": "Gemini API Compatible Middleware",
+        "version": "1.0.0",
+        "status": "running",
+        "wsClientConnected": ws_manager.is_available(),
+        "endpoints": {
+            "models": "/v1beta/models",
+            "generateContent": "/v1beta/models/{model}:generateContent",
+            "streamGenerateContent": "/v1beta/models/{model}:streamGenerateContent",
+            "health": "/health",
+            "websocket": "/ws"
+        }
+    }
+
+
+@app.get("/health")
+async def health_check():
+    """健康检查"""
+    status = ws_manager.get_status()
+    return {
+        "status": "ok" if status["connected"] else "degraded",
+        "timestamp": datetime.now().isoformat(),
+        "wsClient": status,
+        "pendingRequests": request_manager.pending_count
+    }
+
+
+@app.get("/v1beta/models")
+async def list_models(request: Request):
+    """列出可用模型"""
+    verify_api_key(request)
+    
+    log.info("[API] GET /v1beta/models")
+    
+    if not ws_manager.is_available():
+        log.warning("[API] WSClient 未连接，返回空模型列表")
+        return {"models": []}
+    
+    models = [
+        {
+            "name": f"models/{model}",
+            "displayName": model,
+            "supportedGenerationMethods": ["generateContent", "streamGenerateContent"]
+        }
+        for model in ws_manager.models
+    ]
+    
+    log.info(f"[API] 返回 {len(models)} 个模型")
+    return {"models": models}
+
+
+@app.get("/v1beta/models/{model}")
+async def get_model(
+    request: Request,
+    model: str = Path(..., description="模型名称")
+):
+    """获取模型信息"""
+    verify_api_key(request)
+    
+    log.info(f"[API] GET /v1beta/models/{model}")
+    
+    # 移除可能的 "models/" 前缀
+    model_name = model.replace("models/", "")
+    
+    if not ws_manager.is_available():
+        raise HTTPException(status_code=503, detail={
+            "error": {"code": 503, "message": "Service unavailable", "status": "UNAVAILABLE"}
+        })
+    
+    if model_name not in ws_manager.models:
+        raise HTTPException(status_code=404, detail={
+            "error": {"code": 404, "message": f"Model '{model}' not found", "status": "NOT_FOUND"}
+        })
+    
+    return {
+        "name": f"models/{model_name}",
+        "displayName": model_name,
+        "supportedGenerationMethods": ["generateContent", "streamGenerateContent"]
+    }
+
+
+@app.post("/v1beta/models/{model}:generateContent")
+async def generate_content(
+    request: Request,
+    model: str = Path(..., description="模型名称")
+):
+    """生成内容 - 非流式"""
+    verify_api_key(request)
+    
+    model_name = model.replace("models/", "")
+    log.info(f"[API] POST /v1beta/models/{model_name}:generateContent")
+    
+    # 检查 WSClient 是否可用
+    if not ws_manager.is_available():
+        log.error("[API] WSClient 未连接")
+        raise HTTPException(status_code=503, detail={
+            "error": {"code": 503, "message": "WSClient not connected", "status": "UNAVAILABLE"}
+        })
+    
+    # 读取请求体（透传）
+    body = await request.json()
+    log.debug(f"[API] 请求体: {ws_manager._truncate_body(body)}")
+    
+    # 创建请求上下文
+    ctx = await request_manager.create_request(is_stream=False, model=model_name)
+    
+    try:
+        # 发送请求到 WSClient
+        await ws_manager.send_request(ctx.id, is_stream=False, body=body)
+        
+        # 等待响应
+        response_body = await request_manager.wait_for_response(ctx.id, Config.REQUEST_TIMEOUT)
+        
+        log.info(f"[API] ✅ 请求完成 | id={ctx.id[:8]}... | elapsed={ctx.elapsed_ms()}ms")
+        return JSONResponse(content=response_body)
+        
+    except TimeoutError:
+        log.error(f"[API] ⏱️ 请求超时 | id={ctx.id[:8]}...")
+        raise HTTPException(status_code=504, detail={
+            "error": {"code": 504, "message": "Request timeout", "status": "DEADLINE_EXCEEDED"}
+        })
+    except Exception as e:
+        log.error(f"[API] ❌ 请求失败 | id={ctx.id[:8]}... | error={e}")
+        # 尝试解析错误 JSON
+        try:
+            error_detail = json.loads(str(e))
+        except:
+            error_detail = {"error": {"code": 500, "message": str(e), "status": "INTERNAL"}}
+        raise HTTPException(status_code=500, detail=error_detail)
+
+
+@app.post("/v1beta/models/{model}:streamGenerateContent")
+async def stream_generate_content(
+    request: Request,
+    model: str = Path(..., description="模型名称")
+):
+    """生成内容 - 流式"""
+    verify_api_key(request)
+    
+    model_name = model.replace("models/", "")
+    log.info(f"[API] POST /v1beta/models/{model_name}:streamGenerateContent")
+    
+    # 检查 WSClient 是否可用
+    if not ws_manager.is_available():
+        log.error("[API] WSClient 未连接")
+        raise HTTPException(status_code=503, detail={
+            "error": {"code": 503, "message": "WSClient not connected", "status": "UNAVAILABLE"}
+        })
+    
+    # 读取请求体（透传）
+    body = await request.json()
+    log.debug(f"[API] 请求体: {ws_manager._truncate_body(body)}")
+    
+    # 创建请求上下文
+    ctx = await request_manager.create_request(is_stream=True, model=model_name)
+    
+    # 发送请求到 WSClient
+    await ws_manager.send_request(ctx.id, is_stream=True, body=body)
+    
+    # 返回流式响应
+    return StreamingResponse(
+        stream_generator(ctx.id, request),
+        media_type="text/event-stream",
+        headers={
+            "Cache-Control": "no-cache",
+            "Connection": "keep-alive",
+            "X-Accel-Buffering": "no"
+        }
+    )
+
+
+async def stream_generator(request_id: str, http_request: Request) -> AsyncGenerator[str, None]:
+    """生成 SSE 流式响应"""
+    ctx = request_manager.get_request(request_id)
+    
+    try:
+        async for chunk_body in request_manager.wait_for_stream(request_id, Config.REQUEST_TIMEOUT):
+            # 检查客户端是否断开
+            if await http_request.is_disconnected():
+                log.warning(f"[Stream] 客户端断开 | id={request_id[:8]}...")
+                await ws_manager.send_abort(request_id, "client_disconnected")
+                request_manager.abort_request(request_id)
+                break
+            
+            # 输出 SSE 格式（透传 body）
+            yield f"data: {json.dumps(chunk_body)}\n\n"
+        
+        log.info(f"[Stream] ✅ 流完成 | id={request_id[:8]}...")
+        
+    except TimeoutError:
+        log.error(f"[Stream] ⏱️ 流超时 | id={request_id[:8]}...")
+        error_body = {"error": {"code": 504, "message": "Stream timeout", "status": "DEADLINE_EXCEEDED"}}
+        yield f"data: {json.dumps(error_body)}\n\n"
+    except Exception as e:
+        log.error(f"[Stream] ❌ 流错误 | id={request_id[:8]}... | error={e}")
+        error_body = {"error": {"code": 500, "message": str(e), "status": "INTERNAL"}}
+        yield f"data: {json.dumps(error_body)}\n\n"
+
+
+# ============================================================
+# WebSocket 端点
+# ============================================================
+
+@app.websocket("/ws")
+async def websocket_endpoint(websocket: WebSocket):
+    """WSClient WebSocket 连接端点"""
+    await websocket.accept()
+    client_ip = websocket.client.host if websocket.client else "unknown"
+    log.info(f"[WS] 🔌 新连接 | ip={client_ip}")
+    log.info(f"[WS]    └─ 等待 REGISTER 消息 (timeout={Config.REGISTER_TIMEOUT}s)...")
+    
+    try:
+        while True:
+            data = await websocket.receive_json()
+            await handle_ws_message(websocket, data)
+    
+    except WebSocketDisconnect:
+        log.warning(f"[WS] 连接断开 | ip={client_ip}")
+    except Exception as e:
+        log.error(f"[WS] 连接错误 | ip={client_ip} | error={e}")
+    finally:
+        await ws_manager.unregister(websocket)
+        await request_manager.fail_all_requests("WSClient disconnected")
+
+
+async def handle_ws_message(websocket: WebSocket, data: Dict):
+    """处理 WSClient 消息"""
+    msg_type = data.get("type", "")
+    msg_id = data.get("id", "")
+    timestamp = data.get("timestamp", 0)
+    
+    # 简短日志
+    log.debug(f"[WS] 📩 收到消息 | type={msg_type} | id={msg_id[:8] if msg_id else 'N/A'}...")
+    
+    # 根据消息类型处理
+    if msg_type == MessageType.REGISTER.value or msg_type == "register":
+        payload = data.get("payload", {})
+        await ws_manager.register(websocket, payload)
+    
+    elif msg_type == MessageType.RESPONSE.value or msg_type == "response":
+        body = data.get("body", {})
+        log.info(f"[WS] 📥 收到响应 | id={msg_id[:8]}...")
+        log.debug(f"[WS]    └─ body: {ws_manager._truncate_body(body)}")
+        request_manager.resolve_request(msg_id, body)
+    
+    elif msg_type == MessageType.CHUNK.value or msg_type == "chunk":
+        body = data.get("body", {})
+        index = data.get("index", 0)
+        log.debug(f"[WS] 📦 收到数据块 | id={msg_id[:8]}... | index={index}")
+        request_manager.push_chunk(msg_id, body)
+    
+    elif msg_type == MessageType.END.value or msg_type == "end":
+        body = data.get("body")
+        log.info(f"[WS] 🏁 流结束 | id={msg_id[:8]}...")
+        request_manager.end_stream(msg_id, body)
+    
+    elif msg_type == MessageType.ERROR.value or msg_type == "error":
+        body = data.get("body", {})
+        log.error(f"[WS] ⚠️ 收到错误 | id={msg_id[:8]}... | body={body}")
+        request_manager.fail_request(msg_id, body)
+    
+    elif msg_type == MessageType.PONG.value or msg_type == "pong":
+        log.debug(f"[WS] 💓 收到心跳响应 | id={msg_id[:8]}...")
+    
+    else:
+        log.warning(f"[WS] ❓ 未知消息类型 | type={msg_type}")
+
+
+# ============================================================
+# 启动
+# ============================================================
+
+if __name__ == "__main__":
+    uvicorn.run(
+        "main:app",
+        host=Config.HOST,
+        port=Config.PORT,
+        reload=True,
+        log_level="warning"  # 使用我们自己的日志
+    )