WebSocket 实时通信与 AI Agent 中断恢复机制深度解析
WebSocket 实时通信与 AI Agent 中断恢复机制深度解析
上一篇我们聊了 LangGraph 的状态图设计和 ReAct Planner 的工程化实现。这篇聚焦一个更棘手的问题：当 AI Agent 需要等待用户输入时，如何在中断后精确恢复到上次的状态？ 本文将以 WebSocket 通信层为切入点，深入剖析 LangGraph 的 checkpoint + interrupt 机制在真实场景中的落地实践。
一、问题本质：AI Agent 不是请求-响应模型
传统的 Web 服务是请求-响应模型：客户端发一个请求，服务端返回一个响应，结束。但 AI Agent 的对话流程是这样的：
用户: "明天8人会议室"
AI: "请问几点开始？持续多久？" ← Agent 主动暂停，等待用户
用户: "下午2点，1小时" ← 可能过了 5 分钟才回复
AI: "找到3间会议室，请选择..." ← Agent 又暂停
用户: "选1" ← 可能过了 10 分钟
AI: "预定成功！"
这里有三个核心挑战：
- 长连接保持：对话可能跨越数分钟，HTTP 请求-响应无法胜任
- 状态持久化：用户回复间隔不确定，Agent 的状态不能一直驻留在内存
- 精确恢复：用户说"选1"时，Agent 必须知道当前处于"等待选择"阶段，且候选列表是什么
二、为什么选 WebSocket 而不是 SSE 或轮询？
| 方案 | 双向通信 | 实时性 | 复杂度 | 适用场景 |
|---|---|---|---|---|
| HTTP 轮询 | 否 | 差 | 低 | 简单通知 |
| SSE（Server-Sent Events） | 否（单向） | 好 | 中 | 服务端推送 |
| WebSocket | 是 | 好 | 中 | 双向实时对话 |
会议室预定是典型的双向实时通信场景：客户端随时发消息，服务端随时推送中断请求。WebSocket 是最自然的选择。
三、通信协议设计
3.1 消息格式
我们定义了一套简洁的 JSON 协议：
客户端 → 服务端：
// 初始化连接
{"type": "init", "session_id": "abc-123"}
// 发送用户消息
{"type": "message", "content": "明天下午2点8人会议室"}
// 心跳
{"type": "ping"}
服务端 → 客户端：
// 欢迎消息
{"type": "welcome", "message": "欢迎使用会议室预定系统！", "session_id": "abc-123"}
// 中断请求（需要用户输入）
{"type": "interrupt", "kind": "clarify", "message": "请问几点开始？", "phase": "COLLECTING_SLOTS"}
{"type": "interrupt", "kind": "select_room", "message": "请选择会议室", "data": {"candidates": [...]}, "phase": "AWAITING_SELECTION"}
{"type": "interrupt", "kind": "confirm", "message": "确认预定？", "phase": "CONFIRMING"}
// 完成
{"type": "complete", "message": "预定成功！", "data": {"booking_id": "BK-001", ...}}
// 心跳响应
{"type": "pong", "timestamp": "2026-06-08T14:30:00"}
3.2 中断类型（Interrupt Kind）
kind 字段是前端渲染的关键：
| kind | 含义 | 前端行为 |
|---|---|---|
clarify |
需要补充信息 | 显示文本输入框 |
select_room |
需要选择会议室 | 渲染候选房间卡片列表 |
confirm |
需要确认 | 显示确认/取消按钮 |
这种设计让前端可以根据 kind 动态切换 UI 组件，而不是千篇一律的聊天框。
四、ConnectionManager：WebSocket 连接的生命周期管理
class ConnectionManager:
"""WebSocket 连接管理器"""
def __init__(self):
self.active_connections: dict[str, WebSocket] = {}
async def connect(self, websocket: WebSocket, session_id: str):
await websocket.accept()
self.active_connections[session_id] = websocket
def disconnect(self, session_id: str):
if session_id in self.active_connections:
del self.active_connections[session_id]
async def send_message(self, session_id: str, message: dict):
if session_id in self.active_connections:
await self.active_connections[session_id].send_json(message)
ConnectionManager 维护了一个 session_id → WebSocket 的映射表。这个设计的隐含假设是：一个 session 同一时刻只有一个活跃连接。这在会议室预定场景中是合理的——用户不会在两个浏览器窗口同时预定同一个会议室。
五、WebSocket Handler：消息循环与异常隔离
5.1 完整的消息处理循环
async def handle_booking_websocket(websocket: WebSocket):
session_id = None
try:
# 1. 接受连接
await websocket.accept()
# 2. 等待初始化消息（握手协议）
init_data = await websocket.receive_json()
if init_data.get('type') != 'init':
await websocket.close(code=1008, reason="Invalid initialization")
return
session_id = init_data.get('session_id', str(uuid.uuid4()))
manager.active_connections[session_id] = websocket
# 3. 创建上下文 & 开始追踪
ctx = BookingContext(session_id=session_id)
session_tracker.start_session(session_id)
# 4. 发送欢迎消息
await manager.send_message(session_id, {
'type': 'welcome',
'message': '欢迎使用会议室预定系统！',
'session_id': session_id
})
# 5. 消息循环
while True:
data = await websocket.receive_json()
if data.get('type') == 'message':
user_message = data.get('content', '')
session_tracker.record_turn(session_id)
# 核心：处理消息（调用 LangGraph）
response = await process_message(ctx, user_message)
await manager.send_message(session_id, response)
# 完成或出错则退出循环
if response.get('type') in ('complete', 'error'):
break
elif data.get('type') == 'ping':
await manager.send_message(session_id, {'type': 'pong'})
except WebSocketDisconnect:
session_tracker.end_session(session_id, status='cancelled')
except Exception as e:
session_tracker.end_session(session_id, status='failed')
finally:
manager.disconnect(session_id)
5.2 握手协议的设计考量
连接建立后的第一条消息必须是 {"type": "init"}，这是一个显式的应用层握手。为什么不在 WebSocket 连接建立时直接传 session_id？
因为 WebSocket 的 URL 传参（如 ws://host/ws?session_id=xxx）有几个问题：
- 敏感信息暴露在 URL 中
- 部分代理/负载均衡器会截断 query string
- 无法在连接建立后动态更新 session_id
应用层握手更灵活，也更安全。
5.3 三层异常隔离
except WebSocketDisconnect:
# 第一层：用户主动断开 → 标记为 cancelled
session_tracker.end_session(session_id, status='cancelled')
except Exception as e:
# 第二层：未知异常 → 标记为 failed + 记录错误
session_tracker.record_error(session_id, error_type='websocket_error', ...)
session_tracker.end_session(session_id, status='failed')
finally:
# 第三层：无论如何都要清理连接
manager.disconnect(session_id)
核心原则：连接可以断，但监控数据不能丢。 即使用户中途关闭浏览器，session_tracker 也能记录这个会话是"cancelled"状态，为后续的成功率统计提供准确数据。
六、核心：process_message 与 LangGraph 的交互
这是整个系统最复杂也最关键的部分。
6.1 图实例的单例缓存
_cached_graph = None
def get_graph():
global _cached_graph
if _cached_graph is None:
_cached_graph = build_graph()
return _cached_graph
LangGraph 图的构建（build_graph()）是一个相对重的操作：注册节点、添加边、编译、初始化 checkpointer。全局缓存避免了每次消息处理都重建图的开销。
注意：图实例是线程安全的，因为状态隔离是通过 thread_id（即 session_id）实现的，不同会话的状态互不干扰。
6.2 从 Checkpoint 恢复状态
async def process_message(ctx: BookingContext, user_message: str) -> dict:
graph = get_graph()
config = {"configurable": {"thread_id": ctx.session_id}}
# 检查是否有之前的状态（从 checkpoint 恢复）
saved_state = graph.get_state(config=config)
if saved_state.values and saved_state.values.get('ctx'):
# 恢复场景：使用 checkpoint 中保存的 ctx
saved_ctx = saved_state.values['ctx']
# 关键：已完成的预订需要重置，允许新的预订
if saved_ctx.done:
saved_ctx.done = False
saved_ctx.all_infos['slots'] = {...} # 重置槽位
saved_ctx.all_infos['candidates'] = []
saved_ctx.set_phase(Phase.COLLECTING_SLOTS)
saved_ctx.actions_executed = []
initial_state = {"ctx": saved_ctx, "user_message": user_message}
else:
# 首次场景：使用新的 ctx
initial_state = {"ctx": ctx, "user_message": user_message}
这段代码解决了两个核心问题：
问题1：用户回复"选1"时，Agent 怎么知道候选列表？
答案在 saved_ctx 中。当图在 await_user 后中断时，checkpoint 会保存完整的 BookingContext，包括 candidates（候选房间列表）、all_infos（所有槽位）等。用户回复时，get_state() 从 checkpoint 读取上次保存的状态，候选列表自然就在里面。
问题2：用户完成一次预定后想再预定一次怎么办？
if saved_ctx.done: 分支处理了这个场景。检测到上次预订已完成，就重置所有状态（槽位清空、候选清空、阶段回退到 COLLECTING_SLOTS），让用户开始新的预订流程。
6.3 图执行与状态收集
# 执行图（stream 模式，可以观察每个节点的更新）
for event in graph.stream(initial_state, config=config):
logger.info(f"节点更新: {list(event.keys())}")
# 获取最终状态
final_state = graph.get_state(config=config)
ctx = final_state.values.get('ctx', ctx)
为什么用 stream 而不是 invoke？
invoke是一次性执行到结束，返回最终状态stream是逐节点 yield，可以观察中间过程
在调试阶段，stream 模式可以清楚地看到每个节点的执行顺序和状态变化。在生产环境中，它也提供了更好的日志可追溯性。
6.4 中断后的响应构建
# 检查是否需要中断（等待用户输入）
if ctx.pending_user_kind:
response = {
'type': 'interrupt',
'kind': ctx.pending_user_kind,
'message': ctx.pending_prompt,
'phase': ctx.get_phase().value
}
if ctx.pending_user_kind == 'select_room':
response['data'] = {
'candidates': ctx.all_infos.get('candidates', [])
}
return response
# 检查是否完成
if ctx.done:
booking_result = ctx.all_infos.get('booking_result', {})
return {
'type': 'complete',
'message': '预定成功！',
'data': {
'booking_id': booking_result.get('booking_id'),
'room_id': booking_result.get('room_id'),
'status': booking_result.get('status')
}
}
响应构建的逻辑是先检查中断，再检查完成。这个顺序很重要——因为 await_user 节点执行后 ctx.done 仍然是 False，如果先检查 done 就会跳过中断响应。
七、LangGraph Checkpoint 机制的内部原理
7.1 Checkpoint 到底保存了什么？
当图在 interrupt_after=["await_user"] 处中断时，LangGraph 会自动将当前的 GraphState 序列化并保存到 checkpointer 中。在我们的实现中，GraphState 包含：
class GraphState(TypedDict):
ctx: BookingContext # 完整的预订上下文
user_message: str # 最后一条用户消息
所以 checkpoint 中保存的是整个 BookingContext 的快照，包括：
- 所有已收集的槽位
- 候选房间列表
- 选中的房间 ID
- 对话历史
- 动作执行记录
- 当前阶段（Phase）
- pending 状态（等待类型和提示文本）
7.2 thread_id 是隔离的关键
config = {"configurable": {"thread_id": ctx.session_id}}
thread_id 是 checkpoint 的隔离键。不同的 session_id 对应不同的 thread，它们的 checkpoint 互不影响。这意味着：
- 用户 A 的预订进度不会影响用户 B
- 同一个用户开两个浏览器窗口（不同的 session_id），各自独立
- 图实例是全局共享的，但状态是按 thread_id 隔离的
7.3 MemorySaver vs PyMySQLSaver
早期开发阶段可以使用 MemorySaver（内存存储），它速度快但进程重启后所有状态丢失。生产环境必须替换为 PyMySQLSaver（MySQL 持久化）。
遵循"数据库操作归口 app/db 模块"的职责划分，checkpointer 的创建封装在 app/db/checkpoint.py 中：
# app/db/checkpoint.py
import pymysql
from langgraph.checkpoint.mysql.pymysql import PyMySQLSaver
from app.config.app_config import app_config
def create_checkpointer() -> PyMySQLSaver:
db_config = app_config.db_meeting
conn = pymysql.connect(
host=db_config.host, port=db_config.port,
user=db_config.user, password=db_config.password,
database=db_config.database,
charset="utf8mb4", autocommit=True,
)
checkpointer = PyMySQLSaver(conn)
checkpointer.setup() # 幂等建表
return checkpointer
graph.py 只需一行导入：
from app.db.checkpoint import create_checkpointer
checkpointer = create_checkpointer()
注意：
langgraph-checkpoint-mysql包中的MySQLSaver类已被移除，正确的同步实现类是langgraph.checkpoint.mysql.pymysql.PyMySQLSaver。
checkpointer.setup() 会在 MySQL 中自动创建以下 4 张表（幂等，重复调用不报错）：
① checkpoint_migrations — 迁移版本记录
| 字段 | 类型 | 说明 |
|---|---|---|
v |
INTEGER PRIMARY KEY | 迁移版本号 |
② checkpoints — 图状态快照（核心表）
| 字段 | 类型 | 说明 |
|---|---|---|
thread_id |
VARCHAR(150) | 会话线程 ID（即 session_id） |
checkpoint_ns |
VARCHAR(2000) | 检查点命名空间（子图嵌套用） |
checkpoint_ns_hash |
BINARY(16) | ns 的 MD5 哈希（主键用） |
checkpoint_id |
VARCHAR(150) | 检查点唯一 ID |
parent_checkpoint_id |
VARCHAR(150) | 父检查点 ID（形成链表） |
type |
VARCHAR(150) | 序列化类型 |
checkpoint |
JSON | 完整的 GraphState 序列化快照 |
metadata |
JSON | 元数据（来源节点、步骤等） |
主键：
(thread_id, checkpoint_ns_hash, checkpoint_id)
③ checkpoint_blobs — 通道数据（大字段）
| 字段 | 类型 | 说明 |
|---|---|---|
thread_id |
VARCHAR(150) | 会话线程 ID |
checkpoint_ns |
VARCHAR(2000) | 命名空间 |
checkpoint_ns_hash |
BINARY(16) | ns 哈希 |
channel |
VARCHAR(150) | 通道名称（如 ctx、user_message） |
version |
VARCHAR(150) | 数据版本号 |
type |
VARCHAR(150) | 序列化类型 |
blob |
LONGBLOB | 序列化后的通道值 |
主键：
(thread_id, checkpoint_ns_hash, channel, version)
④ checkpoint_writes — 待写入记录（WAL 机制）
| 字段 | 类型 | 说明 |
|---|---|---|
thread_id |
VARCHAR(150) | 会话线程 ID |
checkpoint_ns |
VARCHAR(2000) | 命名空间 |
checkpoint_ns_hash |
BINARY(16) | ns 哈希 |
checkpoint_id |
VARCHAR(150) | 关联的检查点 ID |
task_id |
VARCHAR(150) | 任务（节点）ID |
task_path |
VARCHAR(2000) | 任务路径 |
idx |
INTEGER | 写入序号 |
channel |
VARCHAR(150) | 通道名称 |
type |
VARCHAR(150) | 序列化类型 |
blob |
LONGBLOB | 序列化数据 |
主键：
(thread_id, checkpoint_ns_hash, checkpoint_id, task_id, idx)
PyMySQLSaver 将 checkpoint 持久化到 MySQL，实现了：
- 进程重启后状态恢复（用户断线重连不丢失进度）
- 多实例部署时的状态共享（水平扩展无障碍）
- 审计和回溯能力（可查询任意会话的完整状态历史）
八、状态重置：一个容易踩的坑
8.1 问题场景
用户完成一次预定后，说"我还想再预定一个会议室"。此时 checkpoint 中保存的状态是 done=True, phase=BOOKED。如果不重置，图会直接走 finish 分支结束。
8.2 重置策略
if saved_ctx.done:
saved_ctx.done = False
saved_ctx.pending_user_kind = None
saved_ctx.pending_prompt = None
saved_ctx.all_infos['slots'] = {
'date': None, 'start_time': None,
'duration_minutes': None, 'capacity': None,
'building': None, 'organizer_id': None
}
saved_ctx.all_infos['candidates'] = []
saved_ctx.all_infos['selected_room_id'] = None
saved_ctx.all_infos['booking_result'] = None
saved_ctx.all_infos['last_action'] = {}
saved_ctx.set_phase(Phase.COLLECTING_SLOTS)
saved_ctx.actions_executed = []
saved_ctx.last_tool_error = None
重置必须彻底：槽位清空、候选清空、阶段回退、动作记录清空、错误清空。任何一个字段遗漏都可能导致下一轮预订出现诡异的行为。
8.3 一个真实的 Bug
早期版本中，重置时忘记清空 actions_executed。结果用户第二次预订时，ReAct Planner 看到的动作历史包含了第一次预订的记录，LLM 误以为"已经预订过了"，直接返回 finish。这就是状态泄漏的典型表现。
九、Pending 状态的清除时机
pending_user_kind 和 pending_prompt 是两个需要精心管理的中断信号：
# 时机1：用户回复时清除（表示"用户已经响应了"）
if ctx.pending_user_kind:
ctx.pending_user_kind = None
ctx.pending_prompt = None
# 时机2：预订完成时清除（表示"不需要再等用户了"）
if ctx.done:
ctx.pending_user_kind = None
ctx.pending_prompt = None
时机1 在 process_message 的入口处，表示"用户的回复已经收到了，可以清除等待状态"。
时机2 在 process_message 的出口处，表示"预订已完成，不需要再设置中断"。
如果遗漏时机2，即使预订成功，响应中仍然会包含 interrupt 类型的消息，前端会错误地显示输入框。
十、超时会话的自动清理
WebSocket 连接可能因为网络问题突然断开，而没有触发正常的关闭流程。我们实现了一个后台清理任务：
async def cleanup_loop():
"""定期清理超时会话"""
while True:
await asyncio.sleep(60) # 每分钟检查一次
session_tracker.cleanup_timed_out_sessions(timeout_seconds=600)
# 在 FastAPI startup 中启动
asyncio.create_task(cleanup_loop())
SessionTracker 中的清理逻辑：
def cleanup_timed_out_sessions(self, timeout_seconds: int = 600):
current_time = time.time()
timed_out = []
for session_id, session in self._active_sessions.items():
if current_time - session.start_time > timeout_seconds:
timed_out.append(session_id)
for session_id in timed_out:
self.end_session(session_id, status='timeout')
10 分钟超时是一个经验值：太短会导致用户思考时间不够，太长会占用不必要的资源。
十一、完整时序图
Browser WebSocket Handler LangGraph
│ │ │
│──WS Connect───────────────>│ │
│<─Accept────────────────────│ │
│──init{session_id}─────────>│ │
│<─welcome───────────────────│ │
│ │ │
│──message"明天8人"──────────>│ │
│ │──stream(state)────────────>│
│ │ [ingest → planner → │
│ │ await_user → INTERRUPT] │
│ │<─checkpoint saved──────────│
│<─interrupt{clarify}────────│ │
│ │ │
│ ... 5 分钟 ... │ │
│ │ │
│──message"下午2点1小时"─────>│ │
│ │──get_state(thread_id)─────>│
│ │<─restored ctx──────────────│
│ │──stream(restored state)────>│
│ │ [ingest → planner → │
│ │ search → await_user → │
│ │ INTERRUPT] │
│ │<─checkpoint saved──────────│
│<─interrupt{select_room}────│ │
│ │ │
│──message"选1"─────────────>│ │
│ │──get_state → restore ──────>│
│ │──stream → [ingest → │
│ │ planner → book → END] │
│<─complete{booking_id}──────│ │
│ │──end_session(success)──────│
十二、工程经验总结
1. Checkpoint 是多轮对话的基石
没有 checkpoint，每次用户回复都需要从头开始对话。Checkpoint 让 Agent 能够"记住"之前的所有状态，实现真正的多轮交互。
2. 状态重置要彻底
完成一次对话后重新开始，必须重置所有状态字段。任何遗漏都会导致状态泄漏到下一轮对话。建议写一个 reset() 方法，而不是手动逐字段清空。
3. 中断信号是应用层概念
LangGraph 提供了 interrupt_after 机制，但"中断后该返回什么数据给前端"是应用层的责任。pending_user_kind 和 pending_prompt 是我们自定义的中断信号协议。
4. 超时清理不可或缺
WebSocket 连接的断开不一定能被服务端感知到（比如拔网线）。后台定时清理任务确保了监控数据的准确性和资源的及时释放。
5. 图实例复用，状态按 thread 隔离
全局缓存图实例提升性能，通过 thread_id 实现状态隔离。这是一个"无状态服务 + 有状态存储"的经典架构模式。
上一篇：当 ReAct Agent 遇上状态机：LangGraph 在会议室预定系统中的工程化实践
下一篇：AI Agent 的可观测性体系：从 Prometheus 监控到离线评估系统 —— 如何用 Prometheus 指标体系追踪 Agent 的每一次决策，以及如何构建离线评估系统来量化 Agent 的能力。
浙公网安备 33010602011771号