PetAgent/token_optimization_plan.md

# Token 计费系统优化计划

## 一、问题现状

### 1.1 已修复的问题（前置修改 + 本轮优化）

| 编号 | 文件 | 修改内容 | 状态 |
|------|------|----------|------|
| M1 | `voice_chat_router.py` | usage 扣费失败时重置 `token_charged=False` 并发送 `token_update` | 已完成 |
| M2 | `voice_chat_router.py` | `asr_start` 事件新增 token 耗尽检查 + break | 已完成 |
| M3 | `chatStore.js` | `handleVoiceMessage` 和 `handleRealtimeMessage` 新增 `token_update` case | 已完成 |
| O1 | `voice_chat_router.py` | usage 扣费失败 + 余额=0 时 break 退出事件循环 | 已完成 |
| O2 | `chat_router.py` | 文字聊天 WebSocket 扣费失败 + 余额=0 时发 error 到前端 | 已完成 |
| O3 | `chatStore.js` | `sendMessage` 新增 `_tokenExhausted` 检查 + `token_update` case | 已完成 |
| O4 | `voice_chat_router.py` | `asr_start` 时重置 `token_charged=False`，每轮对话触发扣费 | 已完成 |

### 1.2 仍存在的问题

#### 问题 A：实时通话 token 卡住不动

**现象**：后台 500 token → 实时聊天几分钟后降至 300 → 继续聊十几分钟，token 一直停在 300。火山引擎实际消费 7.8 元。

**根因分析**：

1. `receive_events_loop` 中 `_charge_user_tokens` 扣费失败时不会 break，火山引擎继续工作但本地不扣费
2. `token_charged` 在首次扣费成功后设为 True，跳过后续所有 usage 扣费（设计上是防重复，但如果火山引擎一轮对话只有一次 usage event，只会扣一次费）
3. 火山引擎 usage event 通常在会话结束时才发送，不是持续的，所以"卡住不动"是因为只收到一次 usage
4. **扣费失败无前端通知**：`_charge_user_tokens` 返回失败时，旧代码只 logger.warning，不发送任何消息到前端

#### 问题 B：按住说话在 token 耗尽时仍可用

**现象**：后台剩 100 token，使用按住说话还能得到 AI 文字回复，token 不变。

**根因分析**：

1. 修改 M2 已在 `asr_start` 加了 token 检查，**本轮对话已修复**
2. 但 `asr_start` 的 `break` 只退出事件循环，不断开 WebSocket，前端可能还在发音频数据
3. voice_chat 的 `session_started` 时做了 token 检查，但那只检查一次，不检查后续轮次

#### 问题 C：文字聊天 token 不限制

**现象**：文字聊天聊了半天，后台 100 token 始终不变化。

**根因分析**：

1. `chat_router.py` WebSocket 端点（L263-274）扣费逻辑存在，但 `charge_tokens` 失败时**不发送 error 到前端**
2. 前端 `sendMessage` 的 `case 'error'` 处理：收到 `insufficient_tokens` 后调用 `fail()` → 显示 "消息发送失败"，但**消息已经生成并返回了**（AI 回复已经通过 `message_chunk` 流式发出）
3. `fallback_amount=1`，如果 `tokens_used` 没有正确返回，每次只扣 1 token
4. 前端 `sendMessage` 没有 `token_update` case，无法接收 token 更新通知

---

## 二、优化方案

### 优化 1：usage 扣费失败时终止会话（优先级：P0）

**文件**：`backend/api/voice_chat_router.py`
**位置**：`receive_events_loop` 函数，`usage` 事件的 else 分支（L1150-1166）

**改动内容**：

在扣费失败且余额为 0 时，增加 `break` 退出事件循环：

```python
# 当前代码（L1150-1166）：
else:
    session_state['token_charged'] = False
    logger.warning(
        f"Failed to charge tokens for user {user_id} in {charge_mode}: "
        f"{charge_result.get('message')}, token_charged reset"
    )
    try:
        await websocket.send_text(json.dumps({
            "type": "token_update",
            "remaining_tokens": charge_result.get("available_tokens", 0),
            "consumed": 0,
            "mode": charge_mode,
            "warning": charge_result.get("message", "扣费失败"),
            "timestamp": datetime.now().isoformat()
        }))
    except Exception:
        pass

# 优化后：在 try/except 块之后增加：
    remaining = charge_result.get("available_tokens", 0)
    if remaining <= 0:
        logger.warning(
            f"[VOICE_TOKEN_EXHAUSTED] user_id={user_id} mode={charge_mode} "
            f"remaining={remaining}, breaking event loop"
        )
        break
```

**预期效果**：余额为 0 时，事件循环立即退出，`_persist_voice_session_history` 在 finally 中执行，火山引擎的后续事件不再被转发。

**验证标准**：
- [ ] 给测试用户设置 1 个 token，发起实时通话
- [ ] 首次 usage 扣费后，观察后端日志出现 `[VOICE_TOKEN_EXHAUSTED]`
- [ ] 前端收到 `token_update` 消息，显示 warning
- [ ] 事件循环退出后，前端 `asr_start` 检查触发 break（双重保护）

---

### 优化 2：文字聊天扣费失败时通知前端（优先级：P0）

**文件**：`backend/api/chat_router.py`
**位置**：`websocket_chat_endpoint` 函数，stats chunk 扣费逻辑（L270-274）

**改动内容**：

```python
# 当前代码（L270-275）：
if not charge_result.get("success"):
    logger.warning(
        f"Failed to charge tokens for user {user_id} in text_chat: "
        f"{charge_result.get('message')}"
    )
remaining_tokens = charge_result.get("available_tokens")

# 优化后：
if not charge_result.get("success"):
    logger.warning(
        f"Failed to charge tokens for user {user_id} in text_chat: "
        f"{charge_result.get('message')}"
    )
    # 余额耗尽时通知前端
    remaining = charge_result.get("available_tokens", 0)
    if remaining <= 0:
        await websocket.send_text(json.dumps({
            "type": "error",
            "error": "insufficient_tokens",
            "message": "Token不足，请充值",
            "available_tokens": remaining,
            "timestamp": datetime.now().isoformat()
        }))
remaining_tokens = charge_result.get("available_tokens")
```

**预期效果**：文字聊天余额耗尽时，前端收到 `insufficient_tokens` 错误，停止后续发送。

**验证标准**：
- [ ] 设置用户 token 为 1，发送一条文字消息
- [ ] 后端日志显示扣费失败
- [ ] 前端收到 error 消息
- [ ] 再发送消息时，前端在发送前检查 token 状态（需配合前端优化 4）

---

### 优化 3：前端文字聊天发送前检查 token（优先级：P1）

**文件**：`frontend/src/stores/chatStore.js`
**位置**：`sendMessage` 函数（L108 附近）

**改动内容**：

在 `sendMessage` 中增加本地 token 状态检查，并增加 `token_update` 处理：

```javascript
// 在 sendMessage 函数开头增加：
// 当前代码 L126 之前插入：
// 如果之前收到过 token 耗尽通知，阻止发送
// (通过模块级变量追踪)

let _tokenExhausted = false  // 模块级变量

// 在 switch(data.type) 中增加 case：
case 'token_update':
  {
    const remaining = data.remaining_tokens ?? data.data?.remaining_tokens ?? null
    if (remaining !== null && remaining <= 0) {
      _tokenExhausted = true
    }
  }
  break
case 'error':
  if (data.error === 'insufficient_tokens') {
    _tokenExhausted = true
  }
  fail(new Error(data.message || data.error || '聊天失败'))
  break
```

在 `socket.onopen` 发送消息前增加检查：

```javascript
socket.onopen = () => {
  if (_tokenExhausted) {
    fail(new Error('Token已耗尽，请充值'))
    return
  }
  socket.send(JSON.stringify({
    type: 'chat_message',
    content,
    pet_id: petId,
    background_id: bgId,
    conversation_id: conversationId
  }))
}
```

**预期效果**：前端在 token 耗尽后自动阻止发送新消息。

**验证标准**：
- [ ] token 耗尽后，文字聊天按钮显示禁用或发送时提示
- [ ] 不再产生无效的 AI 回复请求

---

### 优化 4：voice_chat asr_start break 后的清理（优先级：P1）

**文件**：`backend/api/voice_chat_router.py`
**位置**：`receive_events_loop` 函数，`asr_start` 事件处理（L1096-1111）

**改动内容**：

`break` 后，事件循环退出，但 WebSocket 仍连接。需确保前端收到错误后主动断开。当前修改 M2 已发送 error 消息，前端 `handleVoiceMessage` 的 `case 'error'` 中 `insufficient_tokens` 分支会调用 `disconnectVoiceWebSocket()`。**这部分已通过 M3 修复。**

但需确认：`handleVoiceMessage` 中 `token_update` case 的 warning 处理不会干扰断开逻辑。当前 M3 的 `token_update` 只是 console.log，不会断开连接，这是正确的。

**无需额外代码修改，仅需验证。**

**验证标准**：
- [ ] 语音聊天 token 耗尽后，前端自动断开 voice WebSocket
- [ ] 不会出现"连接还在但无法使用"的僵死状态

---

### 优化 5：voice_chat 的 token_charged 跨轮次重置（优先级：P2）

**文件**：`backend/api/voice_chat_router.py`
**位置**：`websocket_voice_chat` 函数中 `start_session` 消息处理（L392-402）

**改动内容**：

当前 `start_session` 已重置 `token_charged=False`（L400）。但实时通话的 `start_call`（L625）也已重置。**无需修改。**

但需确认：同一 WebSocket 连接中，用户多次交互（多轮对话）时，`token_charged` 不应在 usage 事件之间被错误地设为 True 后跳过后续扣费。

**分析**：火山引擎端到端模型的 usage event 在每轮对话结束时发送一次。`token_charged=True` 会跳过同一会话内的重复扣费。这意味着：如果一轮对话产生 2 个 usage event，第二个会被跳过。

**建议**：在每次收到 `asr_start` 时重置 `token_charged=False`，确保每轮对话都会触发扣费。

```python
# voice_chat_router.py receive_events_loop 中 asr_start 处理（L1096-1111）增加：
if event.event_type == 'asr_start':
    logger.debug(f"ASR started: {event.data}")
    session_state['token_charged'] = False  # 新增：每轮对话重置扣费标记
    _tc = await _check_user_has_tokens(user_id)
    if not _tc.get("has_tokens"):
        # ...现有逻辑...
```

**预期效果**：每轮对话都会触发 usage 扣费，不会因 `token_charged=True` 跳过。

**验证标准**：
- [ ] 设置用户 100 token，发起语音聊天
- [ ] 进行 3 轮对话，观察后端日志出现 3 次 `[VOICE_CHARGE_OK]`
- [ ] token 余额递减，不是只扣一次

---

### 优化 6：REST API 端点的 token 检查（优先级：P2）

**文件**：`backend/api/chat_router.py`
**位置**：`send_message_endpoint` 函数（L53-111）

**改动内容**：

当前 REST API 端点在发送前检查 token（L69-78），但扣费失败后仍然返回响应（L95-106）。需在扣费失败时返回 402：

```python
# 当前代码（L95-106）：
if not charge_result.get("success"):
    logger.warning(...)
# 当前代码继续返回响应

# 优化后：
if not charge_result.get("success"):
    logger.warning(...)
    # 如果余额为0，返回 402
    if charge_result.get("available_tokens", 0) <= 0:
        return SendMessageResponse(
            response=result["response"],
            conversation_id=result["conversation_id"],
            tokens_used=result["tokens_used"],
            remaining_tokens=0
        )
```

**预期效果**：REST API 扣费失败时前端可感知余额为 0。

---

## 三、实施顺序

```
第 1 阶段（P0，立即执行）：
  优化 1 → 优化 2
  验证：实时通话和文字聊天的 token 耗尽终止

第 2 阶段（P1，本周内）：
  优化 3 → 优化 4
  验证：前端主动阻止 + WebSocket 清理

第 3 阶段（P2，下周）：
  优化 5 → 优化 6
  验证：跨轮次扣费 + REST API 一致性
```

---

## 四、架构层面的长期建议

### 4.1 当前架构的根本问题

当前扣费是**事后结算**模式：火山引擎先消费，usage event 返回后才扣费。这导致：

- 本地 token 余额与实际消费不同步
- 无法在用户耗尽 token 时实时停止火山引擎的消费
- 火山引擎的费用与本地 token 扣费可能不一致

### 4.2 建议的架构改进

```
短期（可立即做）：
  ├── asr_start 时检查 token → 耗尽则 break + 通知前端
  ├── usage 扣费失败时 break → 防止继续消费
  └── 前端 token 耗尽时禁用发送按钮

中期（需要火山引擎 API 支持）：
  ├── 在火山引擎 session 配置中设置 max_tokens
  ├── 或在 receive_events_loop 中监控 usage 累计值
  └── 超过 token 预算时主动关闭 session

长期（架构优化）：
  ├── 预扣费模式：开始对话前先扣费，结束后按实际用量多退少补
  ├── 或 token 预算模式：每次对话设定 token 上限
  └── 与火山引擎的计费对齐：建立本地 token 与火山引擎费用的映射关系
```

### 4.3 建议的 token 扣费流程（新架构）

```
用户发起语音对话
  │
  ├─→ [asr_start] 检查 token 余额
  │     ├─ 余额 > 0 → 继续
  │     └─ 余额 = 0 → 发送 error，break，前端断开
  │
  ├─→ [asr_result] 记录用户消息
  │
  ├─→ [chat_response] AI 生成回复
  │
  ├─→ [usage] 收到实际用量
  │     ├─ 扣费成功 → 发送 token_update，设置 token_charged=True
  │     └─ 扣费失败 → 发送 token_update(warning)，break
  │
  └─→ [下一轮 asr_start] 重置 token_charged=False，重新检查
```

---

## 五、文件修改清单

| 优先级 | 文件 | 修改位置 | 改动描述 | 状态 |
|--------|------|----------|----------|------|
| P0 | `backend/api/voice_chat_router.py` | `receive_events_loop` usage else 分支 | 扣费失败 + 余额=0 时 break | 已完成 |
| P0 | `backend/api/chat_router.py` | `websocket_chat_endpoint` stats chunk | 扣费失败 + 余额=0 时发 error | 已完成 |
| P1 | `frontend/src/stores/chatStore.js` | `sendMessage` 函数 | 增加 token 耗尽检查 + token_update case | 已完成 |
| P1 | `backend/api/voice_chat_router.py` | `receive_events_loop` asr_start | 每轮对话重置 token_charged | 已完成 |
| P2 | `backend/api/chat_router.py` | REST `/send` 端点 | 已有 pre-check 返回 402，无需额外修改 | 已有 |

---

## 六、风险评估

| 风险 | 影响 | 缓解措施 |
|------|------|----------|
| break 过于激进导致会话中断 | 用户体验下降 | break 前发送明确的 error 消息，前端优雅处理 |
| token_charged 重置导致多扣费 | 用户投诉 | 在 ChatLog 中记录每次扣费，可追溯 |
| 前端 _tokenExhausted 状态未同步 | 用户无法发送 | 定期刷新 token 状态，或在页面加载时重新检查 |
| 火山引擎 usage event 延迟 | 扣费不及时 | 在 asr_start 时预检查，usage 到达后确认扣费 |