docs/plan/2026-05-30-sse-to-websocket-migration.md
Rama main · réplica de lectura del repositorio original.
# SSE → WebSocket 多任务流式输出改造方案
> Note: The `bash` tool has been removed. `shell` is now the only command execution tool. See `2026-06-03-shell-unification-design.md`.
> 解决桌面客户端切换任务时流式输出断连、无法并行监听多任务的问题。
## 1. 问题分析
### 当前架构
```
桌面客户端 后端
┌─────────────┐ ┌──────────────────┐
│ TaskView │ EventSource │ SseEmitter │
│ (单实例复用) │◄──────────────│ Registry │
│ │ 1:1 绑定 session │ (1 emitter/session) │
│ useChat │ │ │
│ (局部状态) │ │ AgentLoop │
│ messages[] │ │ (线程池执行) │
└─────────────┘ └──────────────────┘
```
**三个核心问题：**
| # | 问题 | 根因 | 影响 |
|---|------|------|------|
| 1 | 切走即断流 | `loadSession()` 调用 `cleanup()` 关闭 EventSource | 无法后台监听任务进度 |
| 2 | 切回不重连 | `restoreSession()` 只拉历史消息，不重建 SSE | 回来后看不到实时输出 |
| 3 | 消息随组件销毁 | `messages` 是 composable 局部 ref，不在 Pinia store | 切任务 = 清空消息数组 |
### 不受影响的部分
- 后端 agent 执行不依赖客户端连接，断开后继续运行
- 消息最终持久化到 `message` 表，不会丢失最终结果
- 侧边栏 phase 指示器有 30s 轮询兜底（但不够实时）
---
## 2. 目标架构
```
桌面客户端 后端
┌─────────────┐ ┌──────────────────┐
│ TaskView │ │ StreamingWs │
│ (渲染层) │ │ Handler │
│ ▲ │ WebSocket │ ▲ │
│ │ │◄══════════════│ │ │
│ Pinia Store │ 1:N 多路复用 │ AgentLoop │
│ Map<sid, │ │ (per-session) │
│ msgs[]> │ │ │
└─────────────┘ └──────────────────┘
切任务 = 切换 activeSessionId → store 渲染对应 key 的消息
连接始终在线，不因切任务断开
```
**核心变化：**
1. **一条 WebSocket 连接**：用户登录后建立，生命周期跟随会话而非单个 task
2. **消息路由**：服务端按 sessionId 标记每条事件，客户端按 sessionId 分发到 Pinia store
3. **切任务零开销**：`loadSession()` 不再需要 `cleanup()`，只切换 `activeSessionId`
---
## 3. 协议设计
### 3.1 WebSocket 连接
```
URL: ws(s)://{host}/ws/stream?token={jwt}
生命周期: 登录 → 建立连接 → 持续在线 → 登出/关闭标签页 → 断开
心跳: 客户端每 30s 发 {"type":"ping"}，服务端回 {"type":"pong"}
断线重连: 客户端指数退避重连（1s → 2s → 4s → ... → 30s cap）
```
### 3.2 消息格式（服务端 → 客户端）
所有事件统一封装为：
```json
{
"type": "content_delta | tool_call_start | tool_call_result | activity | todo_updated | session_status | session_list_update | context_window | compaction_start | compaction_end | message_end | error | pong",
"sessionId": 123,
"data": { ... }
}
```
`data` 字段复用现有 SSE 事件的 payload，不改格式。示例：
```json
// content_delta
{ "type": "content_delta", "sessionId": 42, "data": { "delta": "Hello" } }
// tool_call_start
{ "type": "tool_call_start", "sessionId": 42, "data": { "tool_call_id": "tc_1", "tool_name": "bash", "arguments": "{\"command\":\"ls\"}" } }
// session_status
{ "type": "session_status", "sessionId": 42, "data": { "phase": "RUNNING" } }
// session_list_update（其他 session 的状态变化）
{ "type": "session_list_update", "sessionId": 99, "data": { "phase": "COMPLETED" } }
// message_end
{ "type": "message_end", "sessionId": 42, "data": { "prompt_tokens": 100, "completion_tokens": 50, "total_tokens": 150 } }
```
### 3.3 客户端 → 服务端消息
```json
// 订阅：告诉服务端"我要收这个 session 的实时事件"
// 进入任务页面时发送，不触发任何执行
{ "type": "subscribe", "sessionId": 42 }
// 取消订阅：离开任务页面时发送（可选，服务端在连接断开时自动清理）
{ "type": "unsubscribe", "sessionId": 42 }
// 发送消息并触发 agent 执行（替代 POST /messages + GET /stream 两步流程）
// 服务端收到后：保存消息 → 缓存内容 → 提交 AgentLoop → 通过 WS 流式返回事件
{ "type": "send_message", "sessionId": 42, "data": { "content": "请帮我...", "eventId": "uuid" } }
// 心跳
{ "type": "ping" }
// 取消正在执行的 agent
{ "type": "cancel", "sessionId": 42 }
```
**subscribe 与 send_message 的关系：**
- `subscribe` 是**被动监听**：只声明"我要收事件"，不触发执行。用于：打开已在运行的 session 时接收其实时输出。
- `send_message` 是**主动触发**：保存消息 + 启动 AgentLoop。隐含 subscribe 语义——不需要先 subscribe 再 send_message。
- 时序保障：`send_message` 的事件回调绑定在创建它的那次 AgentLoop 上，不会因为 subscribe 时序问题丢失事件。
**重连后的追赶机制：**
客户端重连后对 RUNNING 状态的 session 重新 subscribe，服务端发送 `session_snapshot` 事件补发当前状态：
```json
// 服务端在 subscribe 时检查 session 是否正在运行，如果是则发送一次快照
{
"type": "session_snapshot",
"sessionId": 42,
"data": {
"phase": "RUNNING",
"currentToolCall": { "tool_call_id": "tc_3", "tool_name": "bash", "status": "running" },
"partialContent": "已完成的部分文本输出...",
"todos": [{ "id": "t1", "content": "分析代码", "status": "completed" }]
}
}
```
断线期间的 `content_delta`、`tool_call_start` 等实时事件**不缓冲不补发**——客户端通过 `fetchMessages()` REST 接口拉取已持久化的完整消息历史，`session_snapshot` 只补充当前正在执行的那一轮的实时状态。
---
## 4. 后端改造
### 4.1 新增文件
| 文件 | 说明 |
|------|------|
| `session/ws/WsEvent.java` | WS 事件 DTO，包含 type、sessionId、data 字段 |
| `session/ws/StreamingWsHandler.java` | WebSocket 主处理器，替代 SSE 流式推送 |
| `session/ws/StreamingWsRegistry.java` | 连接注册表，管理 userId → WebSocket 映射 |
| `session/ws/WsStreamingEventListener.java` | 事件监听器，实现 AgentEventListener 接口，将回调转为 WS 消息 |
### 4.2 改造文件
#### `WebSocketConfig.java`
新增 `/ws/stream` 端点：
```java
@Override
public void registerWebSocketHandlers(WebSocketHandlerRegistry registry) {
registry.addHandler(localToolHandler, "/ws/local-tool")
.setAllowedOrigins("*");
registry.addHandler(streamingWsHandler, "/ws/stream") // 新增
.setAllowedOrigins("*");
}
```
#### `SecurityConfig.java`
`/ws/stream` 已被现有 `/ws/**` 规则覆盖，无需改动。
#### `SessionController.java`
**删除旧端点**：`POST /v1/sessions/{id}/messages`（line 168）和 `GET /v1/sessions/{id}/stream`（line 192）整体删除，逻辑迁入 `StreamingWsHandler`。
`SessionController` 保留其余端点（创建 session、获取消息列表、删除 session 等 REST API）。
当前流程：
```
POST /messages → 保存消息 → 缓存内容 → 返回 {eventId}
GET /stream → 创建 SseEmitter → 提交 AgentLoop 到线程池
```
新流程：**agent 执行触发统一走 WS**，`send_message` 替代 POST /messages + GET /stream 两步：
```
客户端 WS send_message
→ StreamingWsHandler.onTextMessage()
→ 校验 session 存在且有权限
→ 如果 executionMode=LOCAL，校验 LocalToolSessionRegistry 有对应连接
→ 保存用户消息到 DB
→ 缓存内容到 Caffeine
→ 创建 WsStreamingEventListener（绑定 userId + sessionId + registry）
→ 提交 AgentLoop 到线程池
→ AgentLoop 事件回调 → WsStreamingEventListener → WS 推送给客户端
```
旧的 `POST /v1/sessions/{id}/messages` 和 `GET /v1/sessions/{id}/stream` 端点**删除**，不再保留。
**AgentLoop 生命周期（与当前 SSE 路径完全一致）：**
```
send_message 到达
→ agentExecutor.submit(() -> {
sessionService.updatePhase(sessionId, "RUNNING")
wsRegistry.send(userId, session_status(RUNNING))
AgentLoop loop = new AgentLoop(sessionId, content, listener, ...)
loop.execute() // 阻塞直到 LLM 不再调用工具
sessionService.updatePhase(sessionId, "COMPLETED")
wsRegistry.send(userId, session_status(COMPLETED))
wsRegistry.send(userId, message_end(...))
})
```
线程池配置沿用当前 `agentExecutor`（core=8, max=16, queue=200）。同一 session 不会并发执行——`send_message` 处理时检查 session phase，RUNNING 状态拒绝新请求（同当前逻辑）。
#### `SseEmitterRegistry.java` — 删除
SSE 路径整体移除，`StreamingWsRegistry` 完全替代。
#### `AgentEventListener` 实现
新增 `WsStreamingEventListener`，将每个事件回调转为 WS 消息发送：
```java
public class WsStreamingEventListener implements AgentEventListener {
private final StreamingWsRegistry registry;
private final Long sessionId;
private final Long userId;
@Override
public void onContentDelta(String delta) {
registry.send(userId, new WsEvent("content_delta", sessionId, Map.of("delta", delta)));
}
@Override
public void onToolCallStart(String toolCallId, String toolName, String arguments) {
registry.send(userId, new WsEvent("tool_call_start", sessionId,
Map.of("tool_call_id", toolCallId, "tool_name", toolName, "arguments", arguments)));
}
// ... 其他事件同理
}
```
#### `SessionController` 广播改造
当前 `broadcastPhaseChange()` 只推给单个 session 的 SSE emitter。
**广播范围定义：** `session_list_update` 应推送给**该用户的所有已连接客户端**（多设备场景），而非所有在线用户。理由：
- session 列表是按用户隔离的，A 用户看不到 B 用户的 session
- 侧边栏只显示当前用户的 session，所以只有 session 创建者关心状态变化
- 实现：`StreamingWsRegistry.send(userId, event)` 只发给指定用户的连接集合
```java
// 在 AgentLoop 生命周期的关键节点广播：
// session phase 变更时（RUNNING/COMPLETED/FAILED）
registry.send(userId, new WsEvent("session_list_update", sessionId, Map.of("phase", "RUNNING")));
```
如果未来支持**共享 session**（多人协作），再改为按 session 成员列表广播。当前无此需求。
### 4.3 连接管理
```
StreamingWsRegistry
├── ConcurrentHashMap<Long, Set<WebSocketSession>> // userId → 连接集合
├── ConcurrentHashMap<WebSocketSession, Long> // 连接 → userId（反查）
├── ConcurrentHashMap<Long, Set<Long>> // userId → 订阅的 sessionId 集合
│
├── register(session, userId) // 连接建立时
├── unregister(session) // 连接断开时
├── subscribe(userId, sessionId) // 客户端发送 subscribe 时
├── unsubscribe(userId, sessionId) // 客户端发送 unsubscribe 时
├── send(userId, event) // 发送给指定用户的所有连接
└── broadcast(sessionId, event) // 广播给订阅了该 session 的所有用户
```
### 4.4 并发安全
- AgentLoop 线程池与 WS 发送线程分离：AgentLoop 回调 → 事件入队 → WS 异步发送
- WS 发送使用 `WebSocketSession.sendMessage()`（异步，非阻塞 AgentLoop）
- 用户多设备登录：同一 userId 可能有多个 WS 连接，`send()` 遍历发送
---
## 5. 前端改造
### 5.1 新增文件
| 文件 | 说明 |
|------|------|
| `composables/useStreamWS.ts` | WebSocket 连接管理，替代 useSSE.ts |
### 5.2 改造文件
#### `stores/session.ts` — 新增消息缓存
```typescript
// 新增状态
const sessionMessages = ref<Map<string, ChatMessage[]>>(new Map())
const sessionTodos = ref<Map<string, TodoItem[]>>(new Map())
// 新增 actions
function appendDelta(sessionId: string, delta: string) { ... }
function addToolCallStart(sessionId: string, event: ToolCallStartEvent) { ... }
function addToolCallResult(sessionId: string, event: ToolCallResultEvent) { ... }
function setMessages(sessionId: string, messages: ChatMessage[]) { ... }
function updateTodos(sessionId: string, todos: TodoItem[]) { ... }
// 新增 getter
const activeMessages = computed(() =>
sessionMessages.value.get(activeSessionId.value ?? '') ?? []
)
```
#### `composables/useStreamWS.ts` — 全局 WebSocket 管理
```typescript
// 在 TaskView 或 App 级别初始化一次，全局共享
export function useStreamWS() {
const ws = ref<WebSocket | null>(null)
const connected = ref(false)
const reconnectDelay = ref(1000)
function connect() {
const token = authStore.token
const wsUrl = apiBase.replace(/^http/, 'ws').replace(/\/api$/, '') + `/ws/stream?token=${token}`
ws.value = new WebSocket(wsUrl)
ws.value.onopen = () => {
connected.value = true
reconnectDelay.value = 1000
// 重新订阅当前活跃 session
subscribe(sessionStore.activeSessionId)
}
ws.value.onmessage = (event) => {
const msg = JSON.parse(event.data)
routeEvent(msg) // 按 type + sessionId 分发到 store
}
ws.value.onclose = () => {
connected.value = false
setTimeout(connect, reconnectDelay.value)
reconnectDelay.value = Math.min(reconnectDelay.value * 2, 30000)
}
}
function routeEvent(msg: WsEvent) {
const { type, sessionId, data } = msg
switch (type) {
case 'content_delta':
sessionStore.appendDelta(sessionId, data.delta)
break
case 'tool_call_start':
sessionStore.addToolCallStart(sessionId, data)
break
case 'tool_call_result':
sessionStore.addToolCallResult(sessionId, data)
break
case 'session_status':
sessionStore.updateSessionPhase(sessionId, data.phase)
break
case 'session_list_update':
sessionStore.updateSessionPhase(sessionId, data.phase)
break
case 'message_end':
sessionStore.markMessageComplete(sessionId, data)
break
// ... 其他事件
}
}
function subscribe(sessionId: string | null) {
if (sessionId && ws.value?.readyState === WebSocket.OPEN) {
ws.value.send(JSON.stringify({ type: 'subscribe', sessionId }))
}
}
function sendMessage(sessionId: string, content: string, eventId: string) {
ws.value?.send(JSON.stringify({
type: 'send_message', sessionId, data: { content, eventId }
}))
}
return { connected, connect, subscribe, sendMessage, cancel }
}
```
#### `composables/useChat.ts` — 简化
改造后 `useChat` 不再持有 `messages`、`sse` 等局部状态，变为薄封装层：
```typescript
export function useChat() {
const { sendMessage: wsSend, subscribe, cancel } = useStreamWS()
const sessionStore = useSessionStore()
// 消息直接从 store 读
const messages = computed(() => sessionStore.activeMessages)
const todos = computed(() => sessionStore.activeTodos)
async function sendMessage(content: string) {
// 1. 乐观更新：立即在 store 中添加用户消息
sessionStore.addUserMessage(sessionStore.activeSessionId, content)
// 2. 如果是新 session，先创建
if (!sessionStore.activeSessionId) {
await createSession()
}
// 3. 通过 WS 发送
const eventId = crypto.randomUUID()
wsSend(sessionStore.activeSessionId, content, eventId)
}
// loadSession 只需加载历史 + 订阅，不再 cleanup
async function loadSession(sid: string) {
sessionStore.setActiveSession(sid)
subscribe(sid) // 告诉服务端我要收这个 session 的事件
await fetchMessages(sid) // 拉历史消息到 store
}
// 不再需要 cleanup()，WS 连接是全局的
return { messages, todos, sendMessage, loadSession, cancel }
}
```
#### `views/task/TaskView.vue` — 简化
```typescript
// 删除：
// - useSSE.ts 导入和所有相关调用
// - cleanup() 中的 sse.stop()、SSE 相关逻辑
// - messages.value = [] 等手动状态重置（消息从 store 读）
// - 30s 轮询 fetchSessions（session_list_update 已通过 WS 实时推送）
// 改造 watch：
watch(sessionIdParam, (newSid, oldSid) => {
if (newSid && newSid !== oldSid) {
loadSession(newSid) // 只切换订阅 + 拉历史，不断开连接
}
})
// 模板不变，messages 从 store 的 computed 读取
```
#### `electron/main.cjs` — 本地工具 WS 保留不变
`/ws/local-tool` 连接保持现状，它负责工具执行桥接，与流式输出无关。
---
## 6. 改造步骤
一次性完成，不分阶段。
```
后端：
1. 新建 StreamingWsHandler + StreamingWsRegistry + WsStreamingEventListener
2. WebSocketConfig 注册 /ws/stream 端点
3. StreamingWsHandler 实现消息路由：
- subscribe → 注册到 registry 的订阅集合
- send_message → 校验权限 + LOCAL 连接 → 保存消息 → 提交 AgentLoop
- cancel → 取消对应 session 的 AgentLoop
- ping → 回 pong
4. WsStreamingEventListener 接入 AgentLoop，事件通过 registry 推送给 userId
5. session_list_update 广播：phase 变更时推送给 session 创建者的所有连接
6. 删除 SessionController 中的 POST /messages + GET /stream 端点
7. 删除 SseEmitterRegistry.java
前端：
8. 新建 useStreamWS.ts，实现全局 WS 连接管理
9. stores/session.ts 新增 sessionMessages Map 和相关 actions
10. 改造 useChat.ts：messages 从 store computed 读取，sendMessage 走 WS
11. 改造 TaskView.vue：loadSession 只切订阅，删除 SSE/cleanup 逻辑和 30s 轮询
12. 删除 useSSE.ts
验证：
13. 多任务并行执行 + 切换 + 返回，确认流式输出不中断
14. 断线重连：断网 → 恢复 → session_snapshot 补发 → 消息不丢失
15. 50+ 并发 session 的 WS 广播性能
```
---
## 7. 边界场景处理
### 7.1 断线重连
```
客户端断线
→ 自动重连（指数退避 1s → 30s）
→ 重连成功后重新 subscribe 所有正在查看的 session
→ 服务端检查 session 状态：
IDLE/COMPLETED/FAILED → 无需补偿，历史消息通过 REST 拉取
RUNNING → 发送 session_snapshot 补发当前执行状态，之后继续实时推送
```
**断线期间丢失的事件类型：**
| 事件类型 | 是否持久化 | 断线后能否恢复 | 说明 |
|---------|-----------|--------------|------|
| `content_delta` | 否 | 不可恢复 | 实时打字效果，最终完整文本通过 ASSISTANT 消息持久化 |
| `tool_call_start` | 否 | 不可恢复 | 工具开始执行的通知，最终 TOOL 消息持久化了结果 |
| `tool_call_result` | 部分 | 通过 REST 恢复 | result 内容持久化为 TOOL 消息，但 summary/status 字段丢失 |
| `activity` | 否 | 不可恢复 | 中间活动通知，无持久化 |
| `todo_updated` | 否 | 通过 session_snapshot 恢复 | 快照中包含最新 todos 状态 |
| `context_window` | 否 | 不可恢复 | 信息性事件，不影响功能 |
| `compaction_start/end` | 否 | 不可恢复 | 信息性事件，不影响功能 |
| `session_status` | 是 | 通过 REST 恢复 | phase 持久化在 session 表 |
| `message_end` | 是 | 通过 REST 恢复 | 完整消息持久化在 message 表 |
**结论：服务端不需要缓冲事件**，因为：
1. 最终结果（USER/ASSISTANT/TOOL 消息）已持久化，`fetchMessages()` 可完整恢复
2. 实时事件（delta、activity）是"打字效果"，丢失不影响最终内容
3. 中间状态（todos、当前工具调用）通过 `session_snapshot` 在重连时补发
### 7.2 多设备登录
同一用户可能在多台设备登录。`StreamingWsRegistry` 按 userId 管理连接集合，事件广播到该用户的所有连接。
### 7.3 浏览器标签页（非 Electron）
桌面端用 Electron 走原生 WS。如果未来有 Web 版，浏览器端也直接用 `new WebSocket()`，逻辑完全一致。不需要 SockJS/STOMP。
---
## 8. 风险评估
| 风险 | 概率 | 影响 | 缓解措施 |
|------|------|------|---------|
| WS 连接被 Nginx/网关超时断开 | 中 | 流式中断 | 心跳 30s + 客户端自动重连 |
| 大量并发 session 导致 WS 消息风暴 | 低 | 前端渲染卡顿 | 只推送 subscribe 的 session 事件，非全量广播 |
| AgentLoop 事件回调阻塞在 WS 发送 | 低 | agent 执行变慢 | WS 发送异步化，不阻塞回调线程 |
| 浏览器 WS 连接数限制（同域） | 极低 | 连接失败 | 一条连接多路复用，不受 6 连接限制影响 |
---
## 9. 文件变更清单
### 新增
```
backend/src/main/java/cn/etarch/mao/session/ws/
├── StreamingWsHandler.java
├── StreamingWsRegistry.java
└── WsStreamingEventListener.java
desktop/src/composables/
└── useStreamWS.ts
```
### 改造
```
backend/
├── config/WebSocketConfig.java — 注册 /ws/stream 端点
└── session/service/SessionService.java — 广播 phase 变更到 WS
desktop/
├── stores/session.ts — 新增 sessionMessages Map
├── composables/useChat.ts — 简化，消息从 store 读，sendMessage 走 WS
└── views/task/TaskView.vue — 删除 SSE 逻辑，loadSession 只切订阅
```
### 删除
```
backend/
├── session/SseEmitterRegistry.java — SSE 注册表，整体移除
└── session/controller/SessionController.java — 删除 POST /messages + GET /stream 端点
desktop/
└── composables/useSSE.ts — SSE 客户端，整体移除
```
---
## 10. 附：SSE vs WebSocket 对比
| 维度 | 当前 SSE | 改造后 WebSocket |
|------|---------|-----------------|
| 连接数 | 每 session 1 条 | 全局 1 条 |
| 切任务 | 断开旧连接 + 建新连接 | 只切换渲染目标 |
| 后台监听 | 不支持 | 支持（subscribe 机制） |
| 双向通信 | 不支持 | 支持（可发送消息、取消） |
| 浏览器兼容性 | 好 | 好（IE 不支持但已不需考虑） |
| 代理/网关穿透 | 好 | 需确认 Nginx proxy_read_timeout（默认 60s，心跳 30s 可覆盖） |
| 实现复杂度 | 低 | 中 |
| 多任务体验 | 差（断流） | 好（无缝切换） |