浅念
3d6064dc14
- test_lead_push_cc: 两处"双钉"注释与模块头改为如实的跨文件互补分工 (常量同源钉本文件,字节级措辞钉 test_lead_push_service.py / test_email_rendering.py) - chat_rest: _server_ip/_merge_client_info docstring 与调用点注释恢复 上游半角标点原文(逐字同步,便于将来 diff 式跟上游) - chat.py WS 自报注释补 REST 侧交叉引用,闭合悬空不对称说明 全量回归 459 passed(不变) Co-Authored-By: Claude Code <noreply@anthropic.com>
139 lines
6.0 KiB
Python
139 lines
6.0 KiB
Python
import json
|
||
import logging
|
||
|
||
from fastapi import APIRouter, WebSocket, WebSocketDisconnect
|
||
|
||
from app.engine.models import ConversationSource, ConversationStatus
|
||
from app.engine.routes.runtime import get_runtime
|
||
from app.engine.services.conn_limiter import DEFAULT_WS_CONN_LIMIT
|
||
from app.engine.services.conn_limiter import limiter
|
||
from app.engine.services.conversation import ConversationManager
|
||
|
||
logger = logging.getLogger(__name__)
|
||
router = APIRouter()
|
||
|
||
|
||
@router.websocket("/ws/chat")
|
||
async def ws_chat(websocket: WebSocket):
|
||
# Task 8: 每 (租户, IP) 并发连接上限——必须在 accept 前拒绝(close 代
|
||
# accept,客户端握手即断)。IP 提取与下方 client_info 同一来源(XFF 优先,
|
||
# scope client 兜底),握手期 headers/scope 已就绪、不依赖首条 init 消息;
|
||
# 客户端自报的 client_info.ip 仍按原逻辑只进会话记录,不参与限流。
|
||
runtime = get_runtime(websocket)
|
||
client_ip = websocket.client.host if websocket.client else ""
|
||
forwarded = websocket.headers.get("x-forwarded-for", "").split(",")[0].strip()
|
||
if forwarded:
|
||
client_ip = forwarded
|
||
|
||
limit = runtime.config.data["engine"].get(
|
||
"ws_conn_limit", DEFAULT_WS_CONN_LIMIT)
|
||
if not limiter.try_acquire(runtime.config.slug, client_ip, limit=limit):
|
||
await websocket.close(code=1013, reason="too many connections")
|
||
return
|
||
|
||
session_id = None
|
||
|
||
try:
|
||
await websocket.accept()
|
||
# Task 5b: 租户运行时经 state/默认租户兜底解析(过渡期无中间件);
|
||
# manager 由装配线从 runtime 造(db 会话 + runtime 引用)
|
||
async with runtime.session_factory() as db:
|
||
manager = ConversationManager(db, runtime=runtime)
|
||
|
||
# Collect client info(client_ip/forwarded 已在握手期提取,复用)
|
||
user_agent = websocket.headers.get("user-agent", "")
|
||
client_info = {"ip": client_ip, "user_agent": user_agent}
|
||
|
||
# First message: init or resume
|
||
init_data = await websocket.receive_text()
|
||
init_msg = json.loads(init_data)
|
||
|
||
# Task 10: 客户端自报 site(空域名站点如 es 靠嵌入代码 data-site
|
||
# 接入)作会话站点标——只加记录,不参与业务逻辑;写入点在下方
|
||
# 新建会话处(优先级对齐 REST,见那里注释)
|
||
site_slug = str(init_msg.get("site") or "").strip()
|
||
|
||
if "session_id" in init_msg:
|
||
conv = await manager.get_by_session(init_msg["session_id"])
|
||
if conv and conv.status == ConversationStatus.active:
|
||
session_id = conv.session_id
|
||
await websocket.send_text(json.dumps({
|
||
"type": "resumed",
|
||
"session_id": session_id,
|
||
}))
|
||
else:
|
||
session_id = None
|
||
|
||
# page_token: reuse active conversation from same page session
|
||
if not session_id and "page_token" in init_msg and init_msg["page_token"]:
|
||
conv = await manager.get_by_page_token(init_msg["page_token"])
|
||
if conv:
|
||
session_id = conv.session_id
|
||
await websocket.send_text(json.dumps({
|
||
"type": "resumed",
|
||
"session_id": session_id,
|
||
}))
|
||
|
||
if not session_id:
|
||
source = ConversationSource(
|
||
init_msg.get("source", "standalone")
|
||
)
|
||
merged_info = {**client_info, **(init_msg.get("client_info") or {})}
|
||
# 服务端 XFF 优先于客户端自报 ip（客户端自报可伪造，可伪造面最小）；
|
||
# 无 XFF 时保留客户端自报，user_agent 等其他字段不受影响。
|
||
# （REST 侧已同步上游用直连 host，WS 保持自报——见
|
||
# tests/test_chat_rest_ip.py）
|
||
if forwarded:
|
||
merged_info["ip"] = forwarded
|
||
# Task 10 复审: 专用 site 字段同样赢过 init_msg.client_info
|
||
# 自报的 site(与 REST ?site= 覆盖 req.client_info 的优先级
|
||
# 对齐——专用通道是嵌入代码显式声明,可信度高于任意自报 dict)
|
||
if site_slug:
|
||
merged_info["site"] = site_slug
|
||
conv = await manager.create_conversation(
|
||
source=source,
|
||
client_info=merged_info,
|
||
page_token=init_msg.get("page_token"),
|
||
)
|
||
session_id = conv.session_id
|
||
await websocket.send_text(json.dumps({
|
||
"type": "init",
|
||
"session_id": session_id,
|
||
}))
|
||
|
||
# Chat loop
|
||
while True:
|
||
data = await websocket.receive_text()
|
||
msg = json.loads(data)
|
||
content = msg.get("content", "").strip()
|
||
if not content:
|
||
continue
|
||
|
||
ended = False
|
||
|
||
async for event in manager.stream_process_message(session_id, content):
|
||
if event["type"] == "error":
|
||
await websocket.send_text(json.dumps({
|
||
"type": "error",
|
||
"message": event["message"],
|
||
}))
|
||
break
|
||
|
||
await websocket.send_text(json.dumps(event, ensure_ascii=False))
|
||
|
||
if event["type"] == "evaluated" and event.get("ended"):
|
||
ended = True
|
||
break
|
||
|
||
if ended:
|
||
break
|
||
|
||
except WebSocketDisconnect:
|
||
logger.info("Client disconnected: %s", session_id)
|
||
except Exception:
|
||
logger.exception("WebSocket error for session %s", session_id)
|
||
finally:
|
||
# Task 8: 正常关闭与异常断开都归还名额——acquire 成功后所有路径
|
||
# 必经此处,不泄漏(accept 失败/循环中抛错同样覆盖)
|
||
limiter.release(runtime.config.slug, client_ip)
|