# Fix: Server-side assistant message persistence

**Date:** 2026-04-17
**Status:** Implemented

---

## Problem

Workspace 聊天消息在后端进程重启后丢失。用户完整看到 assistant 回复数十秒后，因 CI 触发 K8s 滚动更新，pod 被替换，消息消失。

### Root cause

messages_v2 的写入依赖前端 fire-and-forget POST：

```
Agent 生成完毕
  → SSE complete 事件
    → 前端 appendMessagesToAPI(assistantMsg)   ← 唯一持久化路径
      → POST /api/workspace/{id}/messages
```

当后端 pod 在 complete 事件之后、POST 到达之前重启（或网络抖动），POST 会收到 502/connection-reset，被 `catch(e) { console.warn(...) }` 静默吞掉。

### 证据

对 workspace `ws-1776411106708-e0c19bfac` 的 DB 查询：

| 数据 | 结果 |
|------|------|
| `token_usage` 记录 | 存在（说明 agent 正常完成） |
| `messages_v2` 行数 | 1（只有 user 消息，0 条 assistant 消息） |
| `messages_v2_meta.msg_count` | 1 |

backend 完成了生成并记录了 token，但 assistant 消息从未写入 DB。

### 影响场景

| 场景 | 是否丢消息 |
|------|-----------|
| K8s 滚动更新 (Recreate strategy) | 是 — POST 502 |
| 用户关闭浏览器 mid-generation | 是 — POST 永远不会发出 |
| 网络抖动 | 是 — POST 失败无重试 |
| 正常使用 | 否 |

---

## Design

### 核心思路

将 assistant 消息的持久化从前端 POST 移到后端服务端，在 agent 生成过程中以 **tool 粒度** 增量 flush 到 SQLite。

```
                    Before                              After
              ┌──────────────┐                   ┌──────────────┐
              │   Frontend   │                   │   Frontend   │
              │  POST assist │ ← 唯一路径        │  POST user   │ ← 仅 user/system
              └──────┬───────┘                   └──────┬───────┘
                     │                                  │
              ┌──────▼───────┐                   ┌──────▼───────┐
              │   Backend    │                   │   Backend    │
              │  (no write)  │                   │ flush on:    │
              └──────────────┘                   │  • tool done │
                                                 │  • prompt end│
                                                 │  • cancel    │
                                                 └──────┬───────┘
                                                        │
                                                 ┌──────▼───────┐
                                                 │   SQLite     │
                                                 │ messages_v2  │
                                                 └──────────────┘
```

### Flush 时机

| 时机 | 触发点 | 说明 |
|------|--------|------|
| Tool 完成 | `ToolCallProgress` (state=completed) | 每完成一个 tool call 就 flush，保证即使中途断开也有部分内容 |
| Prompt 结束 | `ACPSession.prompt()` return | 最终 flush，包含完整内容 |
| Cancel | `ACPSession.cancel_prompt()` | flush 已累积的内容后再取消 agent |

### 消息 ID 协调

Backend 生成 `assistant-{timestamp}` 作为消息 ID，通过新增的 `message_id` SSE 事件发送给前端：

```
Backend                              Frontend
  │                                    │
  │ ─── SSE: message_id {id} ────────► │  更新 assistantMessageId
  │ ─── SSE: text {...} ────────────► │  追加到同一消息
  │ ─── SSE: tool {...} ────────────► │
  │     └─ flush to DB (upsert)        │
  │ ─── SSE: complete ──────────────► │  不再 POST assistant
  │     └─ final flush to DB           │
```

### Upsert 语义

新增 `_db_upsert_message_v2()` 使用 `INSERT ... ON CONFLICT DO UPDATE`，替代原来的 `INSERT OR IGNORE`（append-only）：

```sql
INSERT INTO messages_v2
  (id, workspace_id, role, content, parts, issue_info, created_at)
VALUES (?, ?, ?, ?, ?, NULL, ?)
ON CONFLICT(workspace_id, id) DO UPDATE SET
  content=excluded.content, parts=excluded.parts
```

同一消息多次 flush 安全 — 内容单调增长，后写覆盖前写。

### SQLite 并发安全

- WAL mode 已启用（多读单写）
- 所有写操作通过 `_db_lock = threading.Lock()` 序列化
- `flush_assistant_message` 是同步函数，通过 `asyncio.to_thread()` 在线程池中执行，不阻塞事件循环
- Tool 粒度的 flush 频率较低（通常 3-10 次/prompt），不会造成锁争用

---

## Implementation

### 新增 DB 函数

**`backend/investibook/db/queries.py`** — `_db_upsert_message_v2()`

单条消息 upsert，同时更新 `messages_v2_meta` 的 `updated_at` 和 `msg_count`。

### WebACPClient 累积器

**`backend/investibook/core/acp.py`** — 新增字段和方法：

```python
# 累积器字段
self._msg_id: str | None = None         # 当前 assistant 消息 ID
self._msg_text_parts: list[str] = []     # 文本片段
self._msg_parts: Dict[str, dict] = {}    # parts (text/reasoning/tool)
self._msg_parts_order: list[str] = []    # 保持插入顺序

# 方法
def start_assistant_message(msg_id: str)   # 开始累积
def flush_assistant_message()              # 同步 upsert 到 DB
def reset_assistant_message()              # 清空累积器
```

`session_update()` 中各事件类型的累积逻辑：

| SSE 事件 | 累积行为 | 触发 flush |
|----------|---------|-----------|
| `AgentMessageChunk` (text) | 追加到 `_msg_text_parts`，维护 text part | 否 |
| `AgentThoughtChunk` (thought) | 创建/追加 reasoning part | 否 |
| `ToolCallStart` | 创建 tool part entry | 否 |
| `ToolCallProgress` (completed) | 更新 tool part 状态 | **是** |

### SSE message_id 事件

**`backend/investibook/api/chat.py`** — 在 `generate()` 和 `run_chat_and_wait()` 中：

```python
assistant_msg_id = f"assistant-{int(time.time() * 1000)}"
session.client.start_assistant_message(assistant_msg_id)
yield _sse({"type": "message_id", "id": assistant_msg_id})
```

### 前端变更

**`frontend/components/ChatPanel/index.tsx`**:

1. `appendMessagesToAPI` 增加 3 次重试 + 指数退避（1s/2s/4s），4xx 不重试
2. 处理 `message_id` SSE 事件，用 backend 分配的 ID 替换前端临时 ID
3. 移除 5 处 assistant 消息的 `appendMessagesToAPI` 调用（complete、error、timeout、onerror、handleStop）
4. 保留 user 消息和 system 消息的 POST（backend 不负责这些）

**`frontend/components/ChatPanel/types.ts`**:

```typescript
| { type: 'message_id'; id: string }
```

---

## Files changed

| File | Change |
|------|--------|
| `backend/investibook/db/queries.py` | 新增 `_db_upsert_message_v2()` |
| `backend/investibook/core/acp.py` | 新增累积器字段 + 3 个方法；`session_update()` 累积逻辑；tool completed 触发 flush；`prompt()` / `cancel_prompt()` 最终 flush |
| `backend/investibook/api/chat.py` | 生成 assistant message ID，发送 `message_id` SSE 事件，调用 `start_assistant_message()` |
| `frontend/components/ChatPanel/index.tsx` | 重试逻辑、`message_id` handler、移除 assistant POST |
| `frontend/components/ChatPanel/types.ts` | 新增 `message_id` 事件类型 |

---

## What's still on the frontend

| 消息类型 | 持久化方 | 说明 |
|----------|---------|------|
| User message | 前端 POST | backend 只收到 prompt 纯文本，没有完整 Message 对象 |
| System message | 前端 POST | issue-ready 等系统消息由前端构造 |
| Welcome message | 前端 POST | clearChat 时写入 |
| **Assistant message** | **后端 flush** | 不再依赖前端 |

---

## Verification

1. 正常对话：assistant 消息应在 DB 中可查到，`messages_v2_meta.msg_count` 正确
2. 关闭浏览器 mid-generation：重新打开后消息应保留到最后一个完成的 tool call
3. 模拟 pod 重启：assistant 消息不丢失
4. Cancel prompt：已生成的部分内容应被保存
5. `run_chat_and_wait`（飞书/自动化路径）：同样受益于服务端 flush