浅念
d041077bb6
- main.py: 删全局 engine/async_session;lifespan 装配 TenantRegistry (seed_profiles → 仓库根 tenants/,路径锚定),预热默认租户 jiancai 挂 app.state.default_runtime,逐 active 租户(库文件已存在)跑 sweep_stale_pending,退出 dispose_all。app.state.registry 预置即沿用 (测试注入缝) - routes/runtime.py(过渡,Task 7 删): get_runtime——state.tenant 优先, 缺省落 default_runtime;chat.py/chat_rest.py 全部改走它,消除 main↔routes 循环依赖 - ConversationManager 构造线: 新增 runtime 可选参(持 session_factory/ compiled_wall/config 引用;prompts 仍模块级,Task 6;业务接线 Task 9), fire_pending_pushes 随任务传 session_factory(后台写回派发租户的库) - lead_push_service: 去 app.main import;写库入口显式收 session_factory, 模块级 async_session=None 仅作测试注入兜底 - conftest: client 夹具 = 每测试独立 Registry(tmp data_dir)+ 真实 lifespan 在 TestClient 循环内预热;删 _dispose_main_engine; async_db 保持 NullPool(跨 asyncio.run 循环安全,刻意不换 Registry 引擎) - 测试仅动夹具/桩签名(断言零改动): wiring 桩补 session_factory 形参; history-404 用例改经真实 WS 路由造 widget 会话(引擎常驻连接绑定 TestClient 循环,测试线程不得直开 session_factory) Co-Authored-By: Claude Code <noreply@anthropic.com>
110 lines
4.3 KiB
Python
110 lines
4.3 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.conversation import ConversationManager
|
|
|
|
logger = logging.getLogger(__name__)
|
|
router = APIRouter()
|
|
|
|
|
|
@router.websocket("/ws/chat")
|
|
async def ws_chat(websocket: WebSocket):
|
|
await websocket.accept()
|
|
session_id = None
|
|
|
|
try:
|
|
# Task 5b: 租户运行时经 state/默认租户兜底解析(过渡期无中间件);
|
|
# manager 由装配线从 runtime 造(db 会话 + runtime 引用)
|
|
runtime = get_runtime(websocket)
|
|
async with runtime.session_factory() as db:
|
|
manager = ConversationManager(db, runtime=runtime)
|
|
|
|
# Collect client info
|
|
client_ip = websocket.client.host if websocket.client else ""
|
|
forwarded = websocket.headers.get("x-forwarded-for", "").split(",")[0].strip()
|
|
if forwarded:
|
|
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)
|
|
|
|
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 等其他字段不受影响
|
|
if forwarded:
|
|
merged_info["ip"] = forwarded
|
|
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)
|