TermControlAgent/app/schemas/websocket.py

"""
WebSocket相关的Pydantic模型
"""
from pydantic import BaseModel, Field, validator
from typing import Optional, Any, Dict, List
from enum import Enum

class WebSocketClientState(str, Enum):
    """WebSocket客户端状态"""
    DISCONNECTED = "disconnected"
    CONNECTING = "connecting"
    CONNECTED = "connected"
    RECONNECTING = "reconnecting"
    ERROR = "error"

class ChannelState(str, Enum):
    """Channel状态"""
    IDLE = "idle"
    CONNECTING = "connecting"
    CONNECTED = "connected"
    DISCONNECTED = "disconnected"
    ERROR = "error"

# ==================== 请求模型 ====================

class CreateWebSocketClientRequest(BaseModel):
    """创建WebSocket客户端请求"""
    name: str = Field(..., min_length=1, max_length=50, description="客户端名称")
    url: str = Field(..., description="WebSocket服务器URL")
    heartbeat_interval: int = Field(default=120, ge=5, le=300, description="心跳间隔(秒)")
    
    @validator('url')
    def validate_url(cls, v):
        if not v.startswith(('ws://', 'wss://')):
            raise ValueError('URL必须以ws://或wss://开头')
        return v

class SendMessageRequest(BaseModel):
    """发送消息请求"""
    client_name: str = Field(..., min_length=1, max_length=50, description="客户端名称")
    message_type: str = Field(..., min_length=1, max_length=100, description="消息类型")
    data: Any = Field(..., description="消息数据")
    channel_name: str = Field(default="default", min_length=1, max_length=50, description="Channel名称")

class BroadcastMessageRequest(BaseModel):
    """广播消息请求"""
    message_type: str = Field(..., min_length=1, max_length=100, description="消息类型")
    data: Any = Field(..., description="消息数据")
    channel_name: str = Field(default="default", min_length=1, max_length=50, description="Channel名称")

class SubscribeChannelRequest(BaseModel):
    """订阅Channel请求"""
    client_name: str = Field(..., min_length=1, max_length=50, description="客户端名称")
    channel_name: str = Field(..., min_length=1, max_length=50, description="Channel名称")

class UnsubscribeChannelRequest(BaseModel):
    """取消订阅Channel请求"""
    client_name: str = Field(..., min_length=1, max_length=50, description="客户端名称")
    channel_name: str = Field(..., min_length=1, max_length=50, description="Channel名称")

# ==================== 响应模型 ====================

class WebSocketClientInfo(BaseModel):
    """WebSocket客户端信息"""
    name: str
    url: str
    state: WebSocketClientState
    is_connected: bool
    reconnect_attempts: int
    handler_count: int
    created_at: str
    subscribed_channels: List[str] = Field(default_factory=list, description="已订阅的Channel列表")

class ChannelInfo(BaseModel):
    """Channel信息"""
    name: str
    state: ChannelState
    queue_size: int
    max_size: int
    is_full: bool
    subscriber_count: int
    filter_count: int
    connection_count: int
    created_at: str
    last_message_at: Optional[str]
    subscribers: List[str] = Field(default_factory=list, description="订阅者列表")

class AdapterInfo(BaseModel):
    """适配器信息"""
    client_name: str
    channel_name: str
    client_connected: bool
    channel_connected: bool
    send_task_running: bool
    created_at: str

class MessageResponse(BaseModel):
    """消息响应"""
    client: str
    type: str
    channel: str
    timestamp: str
    message_id: str

class BroadcastResponse(BaseModel):
    """广播响应"""
    type: str
    channel: str
    results: Dict[str, bool]
    success_count: int
    total_count: int
    timestamp: str

class WebSocketStats(BaseModel):
    """WebSocket统计信息"""
    created_at: str
    client_count: int
    channel_count: int
    adapter_count: int
    total_messages_sent: int
    total_messages_received: int
    clients: Dict[str, Dict[str, Any]]
    channels: Dict[str, Dict[str, Any]]
    adapters: Dict[str, Dict[str, Any]]

# ==================== 通用响应模型 ====================

class SuccessResponse(BaseModel):
    """成功响应"""
    success: bool = True
    message: str
    data: Optional[Any] = None
    timestamp: str

class ErrorResponse(BaseModel):
    """错误响应"""
    success: bool = False
    message: str
    error_code: Optional[str] = None
    details: Optional[Any] = None
    timestamp: str

# ==================== 配置模型 ====================

class WebSocketConfig(BaseModel):
    """WebSocket配置"""
    default_channels: List[str] = Field(default_factory=lambda: ["default", "system", "events"], description="默认Channel列表")
    max_channel_size: int = Field(default=1000, ge=1, le=10000, description="默认Channel最大大小")
    heartbeat_interval: int = Field(default=120, ge=5, le=300, description="心跳间隔(秒)")
    reconnect_attempts: int = Field(default=5, ge=1, le=20, description="重连尝试次数")
    reconnect_delay: float = Field(default=1.0, ge=0.1, le=10.0, description="重连延迟(秒)")