引言
CountBot 的 Web UI 通过 WebSocket 实现实时双向通信，支持 LLM 流式响应推送、工具调用状态通知和任务取消等功能。本文将分析其 WebSocket 架构设计和流式传输优化策略。
消息协议设计
客户端消息
class ClientMessage(BaseModel):
type: str = Field(..., description="消息类型")
session_id: str = Field(..., alias="sessionId")
content: str | None = Field(None)
服务端消息类型
class MessageChunk(ServerMessage):
"""流式文本块"""
type: str = "message_chunk"
content: str
class ToolCall(ServerMessage):
"""工具调用通知"""
type: str = "tool_call"
tool: str
arguments: dict[str, Any]
message_id: int | None
class ToolResult(ServerMessage):
"""工具执行结果"""
type: str = "tool_result"
tool: str
result: str
message_id: int | None
class MessageComplete(ServerMessage):
"""消息完成通知"""
type: str = "message_complete"
message_id: str | None
这种类型化的消息协议让前端可以精确地处理不同类型的事件。
连接管理
取消令牌机制
_session_cancel_tokens: dict[str, CancellationToken] = {}
def get_cancel_token(session_id: str) -> CancellationToken:
"""获取或创建会话的取消令牌"""
if session_id in _session_cancel_tokens:
old_token = _session_cancel_tokens[session_id]
if old_token.is_cancelled:
del _session_cancel_tokens[session_id]
if session_id not in _session_cancel_tokens:
_session_cancel_tokens[session_id] = CancellationToken()
return _session_cancel_tokens[session_id]
def cancel_session(session_id: str) -> bool:
"""取消会话的处理"""
if session_id in _session_cancel_tokens:
_session_cancel_tokens[session_id].cancel()
return True
return False
当用户点击"停止生成"按钮时，前端发送取消消息，后端通过取消令牌中断 AgentLoop 的处理。
流式响应处理器
基础流式处理器
class StreamingResponseHandler:
def __init__(self, session_id, chunk_size=50, delay_ms=0, on_progress=None):
self.session_id = session_id
self.chunk_size = chunk_size
self.delay_ms = delay_ms
self.total_sent = 0
self.chunk_count = 0
async def stream_text(self, text: str) -> int:
"""分块推送文本"""
for i in range(0, len(text), self.chunk_size):
chunk = text[i:i + self.chunk_size]
count = await send_message_chunk(self.session_id, chunk)
if self.delay_ms > 0:
await asyncio.sleep(self.delay_ms / 1000.0)
return sent_count
async def stream_iterator(self, iterator: AsyncIterator[str]) -> int:
"""流式推送异步迭代器内容"""
async for chunk in iterator:
if chunk:
await send_message_chunk(self.session_id, chunk)
缓冲流式处理器
class BufferedStreamingHandler:
def __init__(self, session_id, buffer_size=100, flush_interval_ms=100):
self.buffer = ""
self.buffer_size = buffer_size
self.flush_interval_ms = flush_interval_ms
async def write(self, text: str):
self.buffer += text
# 缓冲区满或超时则刷新
if len(self.buffer) >= self.buffer_size or time_since_flush >= self.flush_interval_ms:
await self.flush()
async def flush(self):
if self.buffer:
await send_message_chunk(self.session_id, self.buffer)
self.buffer = ""
缓冲处理器通过合并小块来减少 WebSocket 帧数，优化网络开销。两种策略的选择：
- 基础模式：适合已经分好块的内容（如 LLM 流式输出）
- 缓冲模式：适合频繁的小块写入（如逐字符输出）
便捷函数
async def stream_response(session_id, iterator, use_buffer=True, buffer_size=100):
"""流式推送响应（自动选择处理器）"""
if use_buffer:
handler = BufferedStreamingHandler(session_id, buffer_size)
else:
handler = StreamingResponseHandler(session_id)
await handler.stream_iterator(iterator)
return handler.get_stats()
前端 WebSocket 集成
前端通过 WebSocket 接收不同类型的消息并分别处理：
// 前端消息处理（概念代码）
ws.onmessage = (event) => {
const msg = JSON.parse(event.data)
switch (msg.type) {
case 'message_chunk':
appendToChat(msg.content) // 追加文本
break
case 'tool_call':
showToolCallIndicator(msg.tool, msg.arguments) // 显示工具调用
break
case 'tool_result':
updateToolResult(msg.tool, msg.result) // 更新工具结果
break
case 'message_complete':
finalizeMessage(msg.message_id) // 完成消息
break
}
}
性能考量
背压控制
当客户端处理速度跟不上服务端推送速度时，WebSocket 的 TCP 层会自动提供背压。CountBot 通过 delay_ms 参数提供了应用层的速率控制。
连接复用
同一个会话的多次交互复用同一个 WebSocket 连接，避免频繁的连接建立和断开开销。
统计监控
每个流式处理器都提供 get_stats() 方法，返回发送统计信息，便于性能监控和调优。
总结
CountBot 的 WebSocket 实现展示了如何在 AI 应用中构建高效的实时通信层。通过类型化的消息协议、灵活的流式处理器和优雅的取消机制，为用户提供了流畅的交互体验。
