wehub-resource-sync
151 lines
4.9 KiB
Python
151 lines
4.9 KiB
Python
"""WebSocket 端点 — 实时推送事件到前端。
|
|
|
|
取代旧架构的 5 秒 HTTP 轮询，改为：
|
|
- 客户端 WebSocket 连接
|
|
- 服务端订阅 Redis Pub/Sub 频道
|
|
- 实时推送事件（状态变更、Agent 思考流、心跳等）
|
|
"""
|
|
|
|
import asyncio
|
|
import json
|
|
import logging
|
|
|
|
import redis.asyncio as aioredis
|
|
from fastapi import APIRouter, WebSocket, WebSocketDisconnect
|
|
|
|
from ..config import get_settings
|
|
from ..services.event_bus import get_event_bus
|
|
|
|
log = logging.getLogger("edict.ws")
|
|
router = APIRouter()
|
|
|
|
# 活跃连接管理
|
|
_connections: set[WebSocket] = set()
|
|
|
|
|
|
@router.websocket("/ws")
|
|
async def websocket_endpoint(ws: WebSocket):
|
|
"""主 WebSocket 端点 — 推送所有事件。"""
|
|
await ws.accept()
|
|
_connections.add(ws)
|
|
log.info(f"WebSocket connected. Total: {len(_connections)}")
|
|
|
|
# 创建独立的 Redis Pub/Sub 连接
|
|
settings = get_settings()
|
|
pubsub_redis = aioredis.from_url(settings.redis_url, decode_responses=True)
|
|
pubsub = pubsub_redis.pubsub()
|
|
|
|
# 订阅所有 edict 频道
|
|
await pubsub.psubscribe("edict:pubsub:*")
|
|
|
|
try:
|
|
# 并发：监听 Redis Pub/Sub + 客户端消息
|
|
await asyncio.gather(
|
|
_relay_events(pubsub, ws),
|
|
_handle_client_messages(ws),
|
|
)
|
|
except WebSocketDisconnect:
|
|
log.info("WebSocket disconnected")
|
|
except Exception as e:
|
|
log.error(f"WebSocket error: {e}")
|
|
finally:
|
|
_connections.discard(ws)
|
|
await pubsub.punsubscribe("edict:pubsub:*")
|
|
await pubsub_redis.aclose()
|
|
log.info(f"WebSocket cleaned up. Remaining: {len(_connections)}")
|
|
|
|
|
|
async def _relay_events(pubsub, ws: WebSocket):
|
|
"""从 Redis Pub/Sub 接收事件，推送到 WebSocket。"""
|
|
async for message in pubsub.listen():
|
|
if message["type"] == "pmessage":
|
|
channel = message["channel"]
|
|
data = message["data"]
|
|
|
|
# 提取 topic 名
|
|
topic = channel.replace("edict:pubsub:", "") if channel.startswith("edict:pubsub:") else channel
|
|
|
|
try:
|
|
event_data = json.loads(data) if isinstance(data, str) else data
|
|
await ws.send_json({
|
|
"type": "event",
|
|
"topic": topic,
|
|
"data": event_data,
|
|
})
|
|
except Exception as e:
|
|
log.warning(f"Failed to relay event: {e}")
|
|
break
|
|
|
|
|
|
async def _handle_client_messages(ws: WebSocket):
|
|
"""处理客户端发送的消息（心跳、订阅过滤等）。"""
|
|
while True:
|
|
try:
|
|
data = await ws.receive_json()
|
|
msg_type = data.get("type", "")
|
|
|
|
if msg_type == "ping":
|
|
await ws.send_json({"type": "pong"})
|
|
elif msg_type == "subscribe":
|
|
# 前端可请求只订阅特定 topic（未来扩展）
|
|
topics = data.get("topics", [])
|
|
log.debug(f"Client subscribe request: {topics}")
|
|
await ws.send_json({"type": "subscribed", "topics": topics})
|
|
else:
|
|
log.debug(f"Unknown client message: {msg_type}")
|
|
|
|
except WebSocketDisconnect:
|
|
raise
|
|
except Exception:
|
|
break
|
|
|
|
|
|
@router.websocket("/ws/task/{task_id}")
|
|
async def task_websocket(ws: WebSocket, task_id: str):
|
|
"""单任务 WebSocket — 只推送与特定任务相关的事件。"""
|
|
await ws.accept()
|
|
_connections.add(ws)
|
|
|
|
settings = get_settings()
|
|
pubsub_redis = aioredis.from_url(settings.redis_url, decode_responses=True)
|
|
pubsub = pubsub_redis.pubsub()
|
|
await pubsub.psubscribe("edict:pubsub:*")
|
|
|
|
try:
|
|
async for message in pubsub.listen():
|
|
if message["type"] == "pmessage":
|
|
data = message["data"]
|
|
try:
|
|
event_data = json.loads(data) if isinstance(data, str) else data
|
|
payload = event_data.get("payload", {})
|
|
if isinstance(payload, str):
|
|
payload = json.loads(payload)
|
|
|
|
# 只转发与此任务相关的事件
|
|
if payload.get("task_id") == task_id:
|
|
topic = message["channel"].replace("edict:pubsub:", "")
|
|
await ws.send_json({
|
|
"type": "event",
|
|
"topic": topic,
|
|
"data": event_data,
|
|
})
|
|
except Exception:
|
|
continue
|
|
except WebSocketDisconnect:
|
|
pass
|
|
finally:
|
|
_connections.discard(ws)
|
|
await pubsub.punsubscribe("edict:pubsub:*")
|
|
await pubsub_redis.aclose()
|
|
|
|
|
|
async def broadcast(event: dict):
|
|
"""向所有连接的 WebSocket 客户端广播事件（服务端内部调用用）。"""
|
|
dead = set()
|
|
for ws in _connections:
|
|
try:
|
|
await ws.send_json(event)
|
|
except Exception:
|
|
dead.add(ws)
|
|
_connections -= dead
|