xiaoski/PicoBot
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

merge: streaming turn architecture

Browse Source
This commit is contained in:
xiaoxixi 2026-07-17 17:47:41 +08:00
commit 7a5d95e786
48 changed files with 6119 additions and 653 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

16
AGENTS.md
View File

@ -47,6 +47,8 @@ Channel → MessageBus.inbound → Gateway processor → SessionManager → per-
↑ │ ↑ │
└── Channel ← OutboundDispatcher ← per-conversation lane ← MessageBus.outbound └── Channel ← OutboundDispatcher ← per-conversation lane ← MessageBus.outbound
AgentLoop → TurnEvent → Session TurnController → latest TurnSnapshot → DeliveryCoordinator → per-turn TurnSink → Channel
WebSocket/Channel → MessageBus.control → Gateway processor → SessionManager (dialog operations) WebSocket/Channel → MessageBus.control → Gateway processor → SessionManager (dialog operations)
Scheduler → SessionManager scheduled execution → AgentLoop → Scheduler delivery policy → SessionManager/MessageBus Scheduler → SessionManager scheduled execution → AgentLoop → Scheduler delivery policy → SessionManager/MessageBus
``` ```
@ -59,9 +61,10 @@ Scheduler → SessionManager scheduled execution → AgentLoop → Scheduler del
| `client` | TUI rendering, WebSocket client for CLI chat | `App`, `run()` | | `client` | TUI rendering, WebSocket client for CLI chat | `App`, `run()` |
| `channels` | External integrations (Feishu, CLI chat) | `ChannelManager`, `Channel` trait | | `channels` | External integrations (Feishu, CLI chat) | `ChannelManager`, `Channel` trait |
| `bus` | Bounded async queues and ordered outbound delivery lanes | `MessageBus`, `OutboundDispatcher`, `InboundMessage`, `OutboundMessage`, `ControlMessage` | | `bus` | Bounded async queues and ordered outbound delivery lanes | `MessageBus`, `OutboundDispatcher`, `InboundMessage`, `OutboundMessage`, `ControlMessage` |
| `session` | Conversation lifecycle, dialog operations, per-session serialization, persistence coordination | `SessionManager`, `Session` | | `session` | Conversation lifecycle, dialog operations, per-session serialization, Turn state, persistence coordination | `SessionManager`, `Session`, `TurnController` |
| `agent` | LLM call loop, tool execution, context compression | `AgentLoop` | | `agent` | LLM call loop, tool execution, context compression, semantic Turn events | `AgentLoop`, `TurnEvent` |
| `providers` | LLM API clients (OpenAI-compatible, Anthropic) | `LLMProvider` trait, factory `create_provider()` | | `providers` | Native LLM streams normalized into text/reasoning/tool/usage chunks | `LLMProvider`, `ProviderChunk`, `create_provider()` |
| `delivery` | Snapshot projection, latest-wins throttling, terminal retry, per-turn sink lifecycle | `DeliveryCoordinator`, `TurnDeliveryService`, `PresentationPolicy` |
| `tools` | Agent tools (bash, file ops, http, web, get_skill) | `ToolRegistry`, `Tool` trait | | `tools` | Agent tools (bash, file ops, http, web, get_skill) | `ToolRegistry`, `Tool` trait |
| `skills` | Skills loading, management, and prompt building | `SkillsLoader`, `Skill` | | `skills` | Skills loading, management, and prompt building | `SkillsLoader`, `Skill` |
| `storage` | SQLite persistence for sessions and messages | `Storage`, `SessionMeta`, `MessageMeta` | | `storage` | SQLite persistence for sessions and messages | `Storage`, `SessionMeta`, `MessageMeta` |
@ -75,9 +78,11 @@ Scheduler → SessionManager scheduled execution → AgentLoop → Scheduler del
### Functional Boundaries ### Functional Boundaries
- **Channels** only send/receive messages via `MessageBus`; they know nothing about sessions or LLM - **Channels** publish inbound messages through `MessageBus`; outbound writes arrive through `OutboundDispatcher` or a per-turn `TurnSink`. They know nothing about sessions or LLM
- **MessageBus** owns bounded inbound/outbound/control queues; outbound routing and retries belong to `OutboundDispatcher`, not the queue - **MessageBus** owns bounded inbound/outbound/control queues; outbound routing and retries belong to `OutboundDispatcher`, not the queue
- **SessionManager** owns session state, dialog operations, context construction, per-session work queues, and persistence coordination - **SessionManager** owns session state, dialog operations, context construction, per-session work queues, and persistence coordination
- **TurnController** is the only owner of active Turn state; snapshots are complete latest-wins values, not token queues, and `Completed` is published only after atomic persistence succeeds
- **DeliveryCoordinator** projects active Turn snapshots without mutating history; it owns `TurnSink` lifecycle but no platform message IDs, which remain private to each sink
- **WorkManager** owns the single active plan per session, item state transitions, plan versions, and plan-change events; plans are optional and absent from ordinary chat context - **WorkManager** owns the single active plan per session, item state transitions, plan versions, and plan-change events; plans are optional and absent from ordinary chat context
- **Scheduler** supports legacy direct-delivery jobs and managed `task`/`monitor` jobs; managed agents cannot call `send_message`, and `on_alert` suppresses only healthy informational results - **Scheduler** supports legacy direct-delivery jobs and managed `task`/`monitor` jobs; managed agents cannot call `send_message`, and `on_alert` suppresses only healthy informational results
- **AgentLoop** is stateless across turns; it receives prepared history, calls LLM providers, executes tools, and returns one result - **AgentLoop** is stateless across turns; it receives prepared history, calls LLM providers, executes tools, and returns one result
@ -87,12 +92,15 @@ Scheduler → SessionManager scheduled execution → AgentLoop → Scheduler del
- **WebUI/TUI file transfer** streams bytes over authenticated HTTP and sends only short-lived upload IDs/attachment metadata over WebSocket; messages persist local media paths without guaranteeing later availability, and client responses must never expose those paths - **WebUI/TUI file transfer** streams bytes over authenticated HTTP and sends only short-lived upload IDs/attachment metadata over WebSocket; messages persist local media paths without guaranteeing later availability, and client responses must never expose those paths
- **WebUI authentication** protects every management API and `/ws`; only static pairing assets, public health/status, and pairing submission may bypass device auth. Pair-code issuance requires both a real loopback peer and the filesystem-held admin token; never put bearer tokens in URLs or logs - **WebUI authentication** protects every management API and `/ws`; only static pairing assets, public health/status, and pairing submission may bypass device auth. Pair-code issuance requires both a real loopback peer and the filesystem-held admin token; never put bearer tokens in URLs or logs
- **Providers** are pure HTTP clients; no bus/session/channel awareness - **Providers** are pure HTTP clients; no bus/session/channel awareness
- **Provider reasoning state** is private replay data: persist it, replay it only to the matching provider, and never expose it to clients, channels, or logs
- **Tools** are executed by `AgentLoop`; they receive raw arguments and normally return text. Tools that produce model-consumable media use the structured `execute_with_media` side channel; model capability checks and provider content-block serialization stay outside tools - **Tools** are executed by `AgentLoop`; they receive raw arguments and normally return text. Tools that produce model-consumable media use the structured `execute_with_media` side channel; model capability checks and provider content-block serialization stay outside tools
### Concurrency and Lifecycle Invariants ### Concurrency and Lifecycle Invariants
- Messages in one session are processed serially through a bounded queue; different sessions may run concurrently - Messages in one session are processed serially through a bounded queue; different sessions may run concurrently
- Outbound messages are ordered per `(channel, chat_id)`; a slow destination must not block unrelated destinations - Outbound messages are ordered per `(channel, chat_id)`; a slow destination must not block unrelated destinations
- Active Turn delivery and ordinary outbound delivery share the same per-`(channel, chat_id)` write lock; never enqueue token deltas into MessageBus
- Slow Turn consumers may skip intermediate snapshots but must receive an explicit bounded terminal delivery; shutdown must call sink abort so platform cleanup remains possible
- Never hold a Session mutex across model, network, or database I/O unless a documented invariant requires it - Never hold a Session mutex across model, network, or database I/O unless a documented invariant requires it
- Slow work derived from session state must validate `worker_generation`/`state_version` before committing results - Slow work derived from session state must validate `worker_generation`/`state_version` before committing results
- Related durable mutations use Storage transaction APIs; persistence failure must not leave silent memory/database divergence - Related durable mutations use Storage transaction APIs; persistence failure must not leave silent memory/database divergence

23
README.md
View File

@ -9,9 +9,10 @@ PicoBot 是一个用 Rust 编写的个人 AI 助手运行时。它在本地启
## 适合做什么 ## 适合做什么
- 在终端里和本地 AI 助手持续对话。 - 在终端里和本地 AI 助手持续对话。
- 在浏览器中聊天，并查看日志、任务和记忆，修改运行配置与助手档案。 - 在 TUI 或浏览器中实时查看正文、思考过程和工具执行状态，并在完成后收敛到持久化历史。
- 在浏览器中查看日志、任务和记忆，修改运行配置与助手档案。
- 复杂任务可创建 session 级 Todo 计划，把不同子项并行委托给多个子 Agent；聊天页侧栏实时显示进度。 - 复杂任务可创建 session 级 Todo 计划，把不同子项并行委托给多个子 Agent；聊天页侧栏实时显示进度。
- 将同一套 Agent 能力接入飞书/Lark。 - 将同一套 Agent 能力接入飞书/Lark，并可选用单张卡片实时更新回复。
- 让 Agent 使用本地文件、Shell、搜索、HTTP、浏览器、MCP 工具完成任务。 - 让 Agent 使用本地文件、Shell、搜索、HTTP、浏览器、MCP 工具完成任务。
- 把长期偏好、事实和历史摘要存成可检索记忆。 - 把长期偏好、事实和历史摘要存成可检索记忆。
- 用 Cron 定时执行任务，并把结果发回目标渠道。 - 用 Cron 定时执行任务，并把结果发回目标渠道。
@ -143,7 +144,7 @@ picobot pair
WebUI 随二进制嵌入，不需要 Node.js、npm 或单独部署静态文件，提供： WebUI 随二进制嵌入，不需要 Node.js、npm 或单独部署静态文件，提供：
- 在线聊天、会话创建/切换、历史回放、Markdown 消息、可折叠工具调用卡片，以及基于 Gateway 实时命令清单的 `/` 斜杠命令补全。 - 在线聊天、会话创建/切换、历史回放、流式 Markdown、独立思考区、实时工具状态、可折叠历史工具调用卡片，以及基于 Gateway 实时命令清单的 `/` 斜杠命令补全。
- 文件选择、拖放和剪贴板图片上传；消息中的附件可预览或下载。附件按服务端路径引用，原文件移动或删除后历史附件可能不可用。 - 文件选择、拖放和剪贴板图片上传；消息中的附件可预览或下载。附件按服务端路径引用，原文件移动或删除后历史附件可能不可用。
- 可持久化的浅色/深色主题，首次访问时跟随系统偏好。 - 可持久化的浅色/深色主题，首次访问时跟随系统偏好。
- Cron 定时任务、最近运行记录和后台子任务状态。 - Cron 定时任务、最近运行记录和后台子任务状态。
@ -189,7 +190,7 @@ picobot service uninstall
unit 位于 `~/.config/systemd/user/picobot.service`，以执行 `service install` 时的当前目录作为初始工作目录。服务异常退出时由 systemd 自动重启；`stop` 和 `restart` 会通过 SIGTERM 触发 Gateway 的有界优雅关停。 unit 位于 `~/.config/systemd/user/picobot.service`，以执行 `service install` 时的当前目录作为初始工作目录。服务异常退出时由 systemd 自动重启；`stop` 和 `restart` 会通过 SIGTERM 触发 Gateway 的有界优雅关停。
TUI 会把一个随机客户端标识保存到 `~/.picobot/tui_client_id`，因此关闭并重新打开客户端后会恢复同一组 dialog 和最近使用的会话。界面支持历史回放、会话列表与归档筛选、命令补全、Unicode/中文编辑、括号粘贴、多行输入和文件传输。 TUI 会把一个随机客户端标识保存到 `~/.picobot/tui_client_id`，因此关闭并重新打开客户端后会恢复同一组 dialog 和最近使用的会话。界面支持流式正文、独立思考与工具状态、历史回放、会话列表与归档筛选、命令补全、Unicode/中文编辑、括号粘贴、多行输入和文件传输。
常用快捷键： 常用快捷键：
@ -209,11 +210,11 @@ WebUI/TUI 上传文件默认保存到 `~/.picobot/media/cli_chat`，单文件上
## 运行时数据流 ## 运行时数据流
用户消息进入 PicoBot 后，会被转换为统一的 inbound message，经由 MessageBus 交给 SessionManager。SessionManager 选择当前 dialog、组装上下文、调用 AgentLoop；AgentLoop 调用模型和工具，最终响应通过 outbound bus 回到原渠道。 用户消息进入 PicoBot 后，会被转换为统一的 inbound message，经由 MessageBus 交给 SessionManager。SessionManager 选择当前 dialog、组装上下文并创建活动 Turn；AgentLoop 消费 Provider 原生流、执行工具并发出结构化事件，TurnController 将它们归约成可丢中间帧的完整快照。DeliveryCoordinator 把快照投影给 TUI、WebUI 或 Channel，最终消息在 SQLite 原子提交成功后才进入 `Completed`。
详细时序和失败语义见 [架构文档：消息与控制数据流](docs/ARCHITECTURE.md#4-消息与控制数据流)。 详细时序和失败语义见 [架构文档：消息与控制数据流](docs/ARCHITECTURE.md#4-消息与控制数据流)。
同一 session 的普通消息由专属有界队列串行处理，不同 session 可以并发；出站消息按 `(channel, chat_id)` 分 lane 保序，慢渠道不会阻塞其他目标。Gateway 的长生命周期任务统一由 `TaskSupervisor` 取消和限时回收。 同一 session 的普通消息由专属有界队列串行处理，不同 session 可以并发；活动 Turn 使用 latest-wins 快照，慢展示端只跳过中间状态，不反压模型。普通出站消息按 `(channel, chat_id)` 分 lane 保序，两条投递路径共享目标写锁。Gateway 的长生命周期任务统一由 `TaskSupervisor` 取消和限时回收。
核心边界： 核心边界：
@ -223,7 +224,8 @@ WebUI/TUI 上传文件默认保存到 `~/.picobot/media/cli_chat`，单文件上
| `bus` | 异步消息队列，承载 inbound、outbound、control 三类消息 | | `bus` | 异步消息队列，承载 inbound、outbound、control 三类消息 |
| `session` | 管理会话生命周期、dialog 操作、上下文、记忆召回、压缩和持久化 | | `session` | 管理会话生命周期、dialog 操作、上下文、记忆召回、压缩和持久化 |
| `agent` | 执行无状态 LLM/tool 循环，处理模型响应和工具调用 | | `agent` | 执行无状态 LLM/tool 循环，处理模型响应和工具调用 |
| `providers` | OpenAI 兼容接口和 Anthropic Messages API 客户端 | | `providers` | OpenAI 兼容接口和 Anthropic Messages API 的原生流解析与回放 |
| `delivery` | 活动 Turn 的展示过滤、latest-wins 节流、终态投递与 TurnSink 生命周期 |
| `tools` | Agent 可调用工具集合 | | `tools` | Agent 可调用工具集合 |
| `storage` | SQLite schema、CRUD、消息和任务持久化 | | `storage` | SQLite schema、CRUD、消息和任务持久化 |
| `scheduler` | 领取定时任务，执行普通/巡检 Agent，并按投递策略记录或发送结果 | | `scheduler` | 领取定时任务，执行普通/巡检 Agent，并按投递策略记录或发送结果 |
@ -241,6 +243,8 @@ WebUI/TUI 上传文件默认保存到 `~/.picobot/media/cli_chat`，单文件上
| `cli_chat` | Ratatui 终端客户端，通过 WebSocket 连接 Gateway | | `cli_chat` | Ratatui 终端客户端，通过 WebSocket 连接 Gateway |
| `feishu` | 飞书/Lark 消息、反应、文件上传下载和媒体引用 | | `feishu` | 飞书/Lark 消息、反应、文件上传下载和媒体引用 |
飞书默认只发送终态结果。设置 `channels.feishu.live_updates=true` 后会创建一张卡片并持续编辑；`live_update_interval_ms` 默认 500ms，运行时限制在 250–5000ms。外部渠道始终不会收到模型 reasoning。
### 会话 ### 会话
Session ID 使用三段式： Session ID 使用三段式：
@ -342,6 +346,8 @@ Skill 是包含 `SKILL.md` 的目录。加载优先级从高到低：
| `memory.timeline_retention_days` | `90` | | `memory.timeline_retention_days` | `90` |
| `mcp.tool_timeout_secs` | `180` | | `mcp.tool_timeout_secs` | `180` |
| `browser.enabled` | `false` | | `browser.enabled` | `false` |
| `channels.feishu.live_updates` | `false` |
| `channels.feishu.live_update_interval_ms` | `500` |
更完整的配置字段说明见 [resources/skills/about-picobot/references/config.md](resources/skills/about-picobot/references/config.md)。 更完整的配置字段说明见 [resources/skills/about-picobot/references/config.md](resources/skills/about-picobot/references/config.md)。
@ -373,7 +379,7 @@ Inbound 消息类型：
| `get_slash_commands` | 无 | | `get_slash_commands` | 无 |
| `ping` | 无 | | `ping` | 无 |
Outbound 消息类型包括 `assistant_response`、`error`、`session_established`、`session_created`、`session_list`、`session_loaded`、`session_history`、`session_renamed`、`session_archived`、`session_deleted`、`history_cleared`、`slash_commands_list`、`pong`、`command_executed` 和 `system_notification`。其中异步 `assistant_response` / `system_notification` 可携带 `session_id`，客户端应避免把迟到结果显示到其他 dialog。 Outbound 消息类型包括活动 Turn 使用的 `turn_updated`，以及 `assistant_response`、`error`、`session_established`、`session_created`、`session_list`、`session_loaded`、`session_history`、`session_renamed`、`session_archived`、`session_deleted`、`history_cleared`、`slash_commands_list`、`pong`、`command_executed` 和 `system_notification`。`turn_updated` 每次携带完整快照和单调 revision，客户端只替换当前 session 的活动 Turn；`assistant_response` 保留给独立完整消息。历史消息可包含 reasoning、turn/iteration、completion status 和结构化工具元数据，但不会暴露 Provider 私有回放状态。
## 测试 ## 测试
@ -402,6 +408,7 @@ src/
channels/ CLI chat 和飞书/Lark 集成 channels/ CLI chat 和飞书/Lark 集成
client/ Ratatui 终端 UI client/ Ratatui 终端 UI
config/ 配置加载、环境变量替换、路径展开 config/ 配置加载、环境变量替换、路径展开
delivery/ 活动 Turn 快照投影、节流与 TurnSink 生命周期
gateway/ Axum HTTP/WebSocket server 和 GatewayState 装配 gateway/ Axum HTTP/WebSocket server 和 GatewayState 装配
mcp/ MCP 客户端连接和工具包装 mcp/ MCP 客户端连接和工具包装
memory/ 记忆管理和记忆类型 memory/ 记忆管理和记忆类型

49
docs/ARCHITECTURE.md
View File

@ -2,6 +2,8 @@
本文档描述 PicoBot 当前实现的运行时边界、数据流、并发模型和演进约束。它面向维护者和后续参与改进的 Agent，是代码架构的主入口；行为细节仍以代码和测试为最终依据。 本文档描述 PicoBot 当前实现的运行时边界、数据流、并发模型和演进约束。它面向维护者和后续参与改进的 Agent，是代码架构的主入口；行为细节仍以代码和测试为最终依据。
流式模型输出、reasoning 展示、活动 Turn 快照和 Channel 实时投递的详细设计与取舍见 [STREAMING_TURN_DESIGN.md](STREAMING_TURN_DESIGN.md)。
## 1. 设计目标 ## 1. 设计目标
PicoBot 是一个单进程、异步、可扩展的个人 AI 助手运行时。核心目标是： PicoBot 是一个单进程、异步、可扩展的个人 AI 助手运行时。核心目标是：
@ -43,6 +45,8 @@ flowchart LR
Agent --> Providers[LLM providers] Agent --> Providers[LLM providers]
Agent --> Tools[ToolRegistry / MCP] Agent --> Tools[ToolRegistry / MCP]
Sessions <--> Storage[(SQLite)] Sessions <--> Storage[(SQLite)]
Sessions -->|TurnSnapshot| Delivery[DeliveryCoordinator]
Delivery -->|TurnSink| Channels
Sessions -->|OutboundMessage| Bus Sessions -->|OutboundMessage| Bus
Scheduler[Scheduler] --> Sessions Scheduler[Scheduler] --> Sessions
Bus --> Dispatcher[OutboundDispatcher] Bus --> Dispatcher[OutboundDispatcher]
@ -61,8 +65,9 @@ flowchart LR
| `channels` | 外部协议适配、权限检查、媒体收发 | 会话选择、LLM 调用 | | `channels` | 外部协议适配、权限检查、媒体收发 | 会话选择、LLM 调用 |
| `bus` | 三条有界异步队列与出站投递协调 | 会话路由、业务状态 | | `bus` | 三条有界异步队列与出站投递协调 | 会话路由、业务状态 |
| `session` | dialog 路由、会话状态、串行工作队列、上下文和持久化协调 | 外部渠道协议 | | `session` | dialog 路由、会话状态、串行工作队列、上下文和持久化协调 | 外部渠道协议 |
| `agent` | 单次无状态模型/工具循环、上下文压缩、子 Agent | 持有 dialog 生命周期 | | `agent` | 单次无状态模型/工具循环、上下文压缩、子 Agent、Turn 语义事件 | 持有 dialog 生命周期 |
| `providers` | 把统一请求映射到模型 API | Session、Bus 或 Channel 感知 | | `providers` | 把统一请求映射为原生模型流，并归一化正文、reasoning、工具和 usage | Session、Bus 或 Channel 感知 |
| `delivery` | 活动 Turn 快照投影、latest-wins 节流、终态重试和 TurnSink 生命周期 | Provider 协议、会话历史、平台 API 细节 |
| `tools` / `mcp` | 工具定义、注册和执行适配 | 隐式修改会话路由 | | `tools` / `mcp` | 工具定义、注册和执行适配 | 隐式修改会话路由 |
| `storage` | SQLite schema、迁移、原子 CRUD | 运行时调度策略 | | `storage` | SQLite schema、迁移、原子 CRUD | 运行时调度策略 |
| `memory` | Knowledge/Timeline 的存取与召回 | 直接驱动消息发送 | | `memory` | Knowledge/Timeline 的存取与召回 | 直接驱动消息发送 |
@ -87,7 +92,9 @@ sequenceDiagram
participant G as Message processor participant G as Message processor
participant S as SessionManager participant S as SessionManager
participant W as Per-session worker participant W as Per-session worker
participant A as AgentLoop participant A as AgentLoop / Provider
participant T as TurnController
participant L as DeliveryCoordinator
participant D as OutboundDispatcher participant D as OutboundDispatcher
C->>B: publish InboundMessage C->>B: publish InboundMessage
@ -95,9 +102,18 @@ sequenceDiagram
G->>S: handle_message G->>S: handle_message
S->>W: try_send AgentTask S->>W: try_send AgentTask
S-->>G: AgentProcessing S-->>G: AgentProcessing
W->>A: process(history) W->>T: start Turn
A-->>W: final response W->>L: subscribe latest snapshots
W->>B: publish OutboundMessage W->>A: process_streaming(history)
A-->>T: reasoning/text/tool events
T-->>L: complete TurnSnapshot
L->>C: TurnSink update (best effort)
A-->>W: emitted messages
W->>W: atomic persistence
W->>T: Completed
T-->>L: terminal snapshot
L->>C: TurnSink finish (bounded retry)
W->>B: independent messages/fallback only
B->>D: consume outbound B->>D: consume outbound
D->>C: Channel::send D->>C: Channel::send
``` ```
@ -108,6 +124,20 @@ sequenceDiagram
- 每个 session 有一条容量为 32 的队列，同一 session 串行处理，不同 session 的 worker 可并发执行。 - 每个 session 有一条容量为 32 的队列，同一 session 串行处理，不同 session 的 worker 可并发执行。
- 队列满时明确拒绝新消息，不允许无界积压。 - 队列满时明确拒绝新消息，不允许无界积压。
- Slash command 不进入 Agent 队列，由 `SessionManager` 直接执行，因此 `/stop` 等控制操作不会排在长模型调用之后。 - Slash command 不进入 Agent 队列，由 `SessionManager` 直接执行，因此 `/stop` 等控制操作不会排在长模型调用之后。
- Session 为每个 Agent 请求创建一个 `TurnController`。Provider 向 AgentLoop 发 delta，AgentLoop 发结构化 TurnEvent，只有 TurnController 能把事件归约为有序 block 和单调 revision 的完整快照。
- Turn 快照经 Tokio `watch` 发布，语义为 latest-wins；慢客户端或慢渠道跳过中间状态，不反压 Provider。终态明确编码在快照中，不依赖 sender 关闭。
- Agent 本轮消息原子持久化成功后才发布 `Completed`。取消或失败若已有可见正文，则保存为 `cancelled`/`interrupted` partial；只有 reasoning 时不创建 assistant 历史。
### 活动 Turn 投递
`DeliveryCoordinator` 与普通出站投递并列：
- `TurnDeliveryService` 根据 Channel 创建本轮独占的 `TurnSink`；sink 私有保存远端消息 ID 和清理资源。
- `PresentationPolicy` 在快照离开 Gateway 核心前过滤内容。TUI/WebUI 展示独立 reasoning 和详细工具状态；外部 Channel 不接收 reasoning，只接收紧凑工具状态；无人值守投递只保留正文。
- `LivePolicy::Snapshot` 按渠道间隔发送最新运行态；`FinalOnly` 忽略运行态，只处理终态。终态绕过节流并只对明确的瞬态错误重试。
- `cli_chat` 将同一 `turn_updated` 快照发给 TUI 和 WebUI。客户端只保留当前 session 中 revision 更新的 `active_turn`，终态随后由持久化历史校准。
- 飞书默认 `FinalOnly`；开启 `live_updates` 后，第一个可见快照创建卡片，后续编辑同一卡片，终态编辑失败则发送完整结果兜底。reaction 清理在 finish、abort 和 Gateway shutdown 中幂等执行。
- DeliveryCoordinator 与 OutboundDispatcher 共享 `(channel, chat_id)` 写锁，避免活动 Turn 终态与独立消息并发写入同一目标。
### 出站投递 ### 出站投递
@ -164,7 +194,7 @@ SessionManager 负责组装会话上下文：系统提示、Skills、召回的 K
- 5 秒 busy timeout。 - 5 秒 busy timeout。
- schema version 迁移。 - schema version 迁移。
持久化范围包括 sessions、messages、memories、task plans/items、scheduled jobs、job runs 和 background tasks。修改 schema 时应： 持久化范围包括 sessions、messages、memories、task plans/items、scheduled jobs、job runs 和 background tasks。消息保存可展示 `reasoning_content`、Provider 私有回放状态、`turn_id`、iteration 和 completion status；私有 Provider 状态不进入 WebSocket/Channel，且只允许回放给同一 Provider。Provider 诊断和失败审计只记录模型、消息数、工具数等请求摘要，不保存正文、reasoning 或签名 payload。修改 schema 时应：
1. 更新集中式 schema/迁移逻辑。 1. 更新集中式 schema/迁移逻辑。
2. 保留已有数据库的升级路径。 2. 保留已有数据库的升级路径。
@ -181,7 +211,7 @@ SessionManager 负责组装会话上下文：系统提示、Skills、召回的 K
## 7. 后台任务与生命周期 ## 7. 后台任务与生命周期
`TaskSupervisor` 是 Gateway 内部后台任务的统一所有者。message processor、outbound dispatcher、scheduler、session workers、outbound lanes、通知消费者和子 Agent 后台任务都应通过它注册。 `TaskSupervisor` 是 Gateway 内部后台任务的统一所有者。message processor、outbound dispatcher、scheduler、session workers、Turn delivery、outbound lanes、通知消费者和子 Agent 后台任务都应通过它注册。
两种注册方式： 两种注册方式：
@ -197,6 +227,8 @@ SessionManager 负责组装会话上下文：系统提示、Skills、召回的 K
WebSocket 每个客户端的 writer task 是连接局部任务，由连接 handler 自己限时回收；它不跨越连接生命周期。 WebSocket 每个客户端的 writer task 是连接局部任务，由连接 handler 自己限时回收；它不跨越连接生命周期。
Turn delivery 使用 `spawn_graceful`。全局取消发生时先停止读取快照，再在共享目标写锁下有界调用 `TurnSink::abort`，使平台 reaction 等私有资源能在 Supervisor 宽限期内清理。
### WebUI 与管理 API ### WebUI 与管理 API
Gateway 在 `/` 提供随二进制编译的 HTML/CSS/JavaScript，不依赖外部 CDN 或前端运行服务。前端源码位于 `webui/`，由 Svelte 5 + Vite 构建，Bits UI 提供无样式的可访问组件原语。`build.rs` 监听前端源码、锁文件和构建配置，增量地将生产资源生成到 Cargo `OUT_DIR`，Rust 再从该目录编译嵌入；前端产物不进入仓库，最终用户使用发布二进制时不需要 Node.js。浏览器聊天继续使用 `/ws` 和 `cli_chat` 渠道，因此复用现有 dialog scope、每会话串行 worker、历史持久化和出站 lane；会话历史帧保留工具调用 ID、名称、参数和工具结果角色，WebUI 在本轮完成后刷新历史并将其渲染为默认折叠的工具卡片；聊天页的 Todo 侧栏默认隐藏，按 session 保存快照和未读状态，并通过结构化 `session_plan`/`plan_updated` 帧刷新，计划变化不会写入聊天历史；斜杠命令补全通过 `get_slash_commands` 获取 Gateway 的实时命令与别名清单，不在前端重复定义；WebUI 不直接调用 Provider 或 SessionManager。 Gateway 在 `/` 提供随二进制编译的 HTML/CSS/JavaScript，不依赖外部 CDN 或前端运行服务。前端源码位于 `webui/`，由 Svelte 5 + Vite 构建，Bits UI 提供无样式的可访问组件原语。`build.rs` 监听前端源码、锁文件和构建配置，增量地将生产资源生成到 Cargo `OUT_DIR`，Rust 再从该目录编译嵌入；前端产物不进入仓库，最终用户使用发布二进制时不需要 Node.js。浏览器聊天继续使用 `/ws` 和 `cli_chat` 渠道，因此复用现有 dialog scope、每会话串行 worker、历史持久化和出站 lane；会话历史帧保留工具调用 ID、名称、参数和工具结果角色，WebUI 在本轮完成后刷新历史并将其渲染为默认折叠的工具卡片；聊天页的 Todo 侧栏默认隐藏，按 session 保存快照和未读状态，并通过结构化 `session_plan`/`plan_updated` 帧刷新，计划变化不会写入聊天历史；斜杠命令补全通过 `get_slash_commands` 获取 Gateway 的实时命令与别名清单，不在前端重复定义；WebUI 不直接调用 Provider 或 SessionManager。
@ -249,6 +281,7 @@ WebUI/TUI 文件字节通过受鉴权的 HTTP 接口流式传输，WebSocket 只
3. 将可重试错误表示为 `ConnectionError`/`SendError`，永久错误使用其他类型。 3. 将可重试错误表示为 `ConnectionError`/`SendError`，永久错误使用其他类型。
4. 为 start/stop 幂等性、取消建连、投递失败和媒体边界增加测试。 4. 为 start/stop 幂等性、取消建连、投递失败和媒体边界增加测试。
5. 不要从 Channel 直接调用 SessionManager 或 Provider。 5. 不要从 Channel 直接调用 SessionManager 或 Provider。
6. 若支持活动 Turn，实现 `live_policy`、`presentation_policy` 和每 Turn 一个实例的 `open_turn`；sink 必须消费完整快照而不是拼接 token，并使 finish/abort 清理幂等。
### 新增 Tool ### 新增 Tool

789
docs/STREAMING_TURN_DESIGN.md Normal file
View File

@ -0,0 +1,789 @@
# 流式 Turn、Reasoning 展示与 Channel 投递设计
> 状态：已实现（2026-07）。
>
> 本文记录 PicoBot 流式模型输出、reasoning 展示、工具过程展示和 Channel 实时投递的设计依据与架构决策。当前运行时总览见 `docs/ARCHITECTURE.md`，具体行为以代码和测试为准。
## 1. 背景
改造前 Provider 只提供一次性 `chat()` 调用。AgentLoop 等待完整响应，将 `content`、`reasoning_content` 和 tool calls 组装为 `ChatMessage`，Session 在 AgentLoop 完成后原子持久化本轮消息，再通过 `OutboundMessage` 发送最终正文。
这个模型具有清晰的持久化语义，但无法表达：
- 正文和 reasoning 的实时增量；
- reasoning、正文、工具调用在一个 Turn 中的自然交错；
- TUI 和 WebUI 对同一个运行中 Turn 的一致展示；
- 飞书等 Channel 通过编辑消息呈现流式效果；
- 不支持编辑的 Channel 自动降级为只发送最终结果；
- 取消、失败、慢消费者和投递失败时的确定行为。
旧 `Channel::send_delta(chat_id, delta)` 没有 Turn 身份、消息身份、reasoning/text 分类、工具边界、终态和取消语义，也绕过出站排序机制，因此已被 `TurnSink` 取代。
## 2. 参考实现结论
本设计综合了 `reference/ryvos`、`reference/zeroclaw`、`reference/hermes-agent` 和 PicoBot 当前实现。
### 2.1 Ryvos
Ryvos 把 thinking 作为正式内容块，并区分 `TextDelta` 与 `ThinkingDelta`。它说明 Provider 层必须结构化解析 reasoning、正文和工具调用，不能依赖最终字符串中的 `<think>` 后处理。
可吸收：
- Provider 流的类型化增量；
- reasoning 与正文分离；
- 工具参数增量组装；
- thinking-only 响应的明确处理；
- reasoning effort/thinking budget 的统一配置概念。
### 2.2 ZeroClaw
ZeroClaw 把 reasoning 作为不透明 Provider 数据保留，用于要求历史回放的模型，同时处理 `reasoning_content`、`reasoning`、内联 `<think>` 和不同 Provider 的回放限制。
可吸收：
- 可展示 reasoning 与 Provider 回放状态分离；
- reasoning 字段别名归一化；
- Provider 专用历史状态不能跨 Provider 发送；
- 内联 think block 必须在 Provider 归一化边界处理；
- reasoning 是否展示与是否回放是两个独立策略。
### 2.3 Hermes
Hermes 新增了 Agent 到 Gateway 的结构化展示事件，并明确规定流事件属于 presentation，而不是 conversation history。它还通过 message segment boundary 处理“工具前正文 → 工具 → 工具后正文”。
可吸收：
- 流事件描述发生的事实，不携带平台发送策略；
- 展示流与持久化历史严格分离；
- 工具边界必须结束当前正文 segment；
- 高频更新需要合并，关键事件发送前需要 flush；
- TUI 的活动 Turn 与已完成 transcript 分离；
- Channel/平台决定如何呈现统一状态。
不直接复制：
- typed events、旧 callbacks 和 TUI 字符串事件并存；
- reasoning 使用独立 callback，没有进入新事件模型；
- 客户端用大型 TurnController 重建服务端状态；
- 一个 GatewayStreamConsumer 同时承担聚合、限流、平台编辑、think 清理、overflow 和 fallback。
### 2.4 PicoBot
PicoBot 已有以下适合保留的不变量：
- 同一 Session 由单 worker 串行处理；
- 不同 Session 并发；
- `worker_generation` 和 `state_version` 防止迟到结果提交；
- 完整 Agent Turn 通过原子持久化接口提交；
- 普通出站消息按 `(channel, chat_id)` 有序投递；
- Channel、Session、Agent、Provider 和 Storage 边界明确。
流式设计不能破坏这些不变量。
## 3. 设计目标
1. OpenAI-compatible 和 Anthropic Provider 支持流式正文、reasoning、tool calls 和 usage。
2. TUI 与 WebUI 使用同一运行态模型显示正文、reasoning、工具进度和取消/失败状态。
3. Channel 可以选择实时更新或只发送最终结果。
4. 支持消息编辑的 Channel 能以同一远端消息呈现流式效果。
5. 慢客户端或慢 Channel 不得反压 Provider 和 AgentLoop，也不得积压大量过时 token。
6. 丢失任意中间更新后，下一次更新必须自动收敛到正确状态。
7. `Completed` 必须表示本轮数据库提交已经成功。
8. reasoning 展示文本与 Provider 回放状态必须隔离。
9. Channel 展示差异不能改变模型上下文和数据库历史。
10. 保持模块数量、事件词汇和状态 owner 尽可能少。
## 4. 非目标
- 不保留旧 WebSocket、TUI、WebUI 或 Channel 流式协议兼容性。
- 不要求每个 Provider 都能返回可展示 reasoning。
- 不把 Provider 的加密/签名 reasoning payload 展示给用户。
- 不逐 token 持久化数据库。
- 不保证重连后恢复尚未完成 Turn 的每一个历史帧。
- 不让所有外部 Channel 默认以多条追加消息模拟流式效果。
- 不把流式展示事件作为可重放的事件溯源日志。
## 5. 核心决策
### 5.1 只有两个权威模型
系统只维护两个跨层权威模型：
- `ConversationMessage`：最终持久化事实，用于会话历史和下轮模型上下文；
- `TurnState`：单次运行的临时展示状态。
Provider 的 SSE chunk 是 Provider 内部输入；Channel 的远端消息 ID 是单个 TurnSink 的私有投递状态。两者都不是全局领域模型。
### 5.2 服务端拥有唯一运行态
Session 侧 `TurnController` 是运行中 Turn 的唯一状态 owner。TUI、WebUI 和 Channel 不根据一串增量自行重建 reasoning、正文、工具和 segment 关系，只渲染服务端发布的 `TurnSnapshot`。
### 5.3 下发幂等快照，不下发可靠 token 流
Provider 到 AgentLoop 使用 delta；TurnController 到展示端使用包含完整当前状态的快照。
每个快照带单调递增 `revision`。消费者只接受 revision 更大的快照，并以新快照整体替换旧状态。因此：
- 丢失中间更新不会损坏内容；
- 慢消费者可以跳过过时状态；
- 更新重试不会重复拼接正文；
- 最终快照能修复暂态渲染；
- 外部 Channel 编辑消息天然获得完整累计内容。
### 5.4 latest-wins，而不是 token 队列
每个活动 Turn 使用 Tokio `watch` 或等价 latest-value primitive 发布 `Arc<TurnSnapshot>`。生产者覆盖旧值，消费者读取最新值。终态显式编码在快照中，不能仅依赖 sender 关闭表达完成。
### 5.5 每个 Channel Turn 使用独立 TurnSink
Channel 为每次 Turn 创建一个 sink。Sink 独占远端消息 ID、编辑状态和平台私有资源，终态后销毁。通用协调器不保存平台消息映射，也不理解飞书卡片 API。
### 5.6 完整消息投递与活动 Turn 展示分离
- `MessageBus` / `OutboundDispatcher`：完整消息、通知、命令结果和需要可靠确认的独立投递；
- `DeliveryCoordinator`：活动 Turn 快照、展示策略、节流和 TurnSink 生命周期。
不把 token 或快照塞入现有 outbound MPSC。
## 6. 数据模型
### 6.1 Provider 私有流
```rust
pub enum ProviderChunk {
Text(String),
Reasoning(String),
ToolCallStart {
index: usize,
id: Option<String>,
name: Option<String>,
},
ToolCallArguments {
index: usize,
delta: String,
},
ProviderState(ProviderReasoningState),
Usage(Usage),
Done(FinishReason),
}
```
`ProviderChunk` 只允许在 `providers` 与 `agent` 模块间使用，不进入 Bus、Session 协议或 Channel。
### 6.2 可展示 reasoning 与回放状态
```rust
pub struct ProviderReasoningState {
pub provider: String,
pub payload: serde_json::Value,
}
pub struct AssistantMessageData {
pub content: String,
pub reasoning: Option<String>,
pub provider_state: Option<ProviderReasoningState>,
pub tool_calls: Vec<ToolCall>,
}
```
约束：
- `reasoning` 可以按展示策略下发；
- `provider_state` 永远不下发给客户端或 Channel；
- `provider_state.provider` 与当前 Provider 不一致时禁止回放；
- 压缩历史时默认不把原始 reasoning 写入 Timeline；
- 日志不得记录完整 reasoning 或 provider payload。
### 6.3 Turn 标识
```rust
pub struct TurnId(pub uuid::Uuid);
pub struct BlockId(pub uuid::Uuid);
```
一次用户输入对应一个 Turn。Turn 开始时预分配最终 assistant `message_id`，使运行态和最终历史能稳定关联。
### 6.4 TurnState
```rust
pub struct TurnState {
pub id: TurnId,
pub session_id: String,
pub message_id: String,
pub revision: u64,
pub status: TurnStatus,
pub phase: TurnPhase,
pub blocks: Vec<TurnBlock>,
pub usage: Option<Usage>,
pub error: Option<String>,
}
pub enum TurnStatus {
Running,
Completed,
Cancelled,
Failed,
}
pub enum TurnPhase {
Queued,
Reasoning,
Responding,
Acting,
Finalizing,
}
```
`TurnPhase` 是展示状态，不等于模型 reasoning：
- 等待首个模型 chunk 时可显示 `Queued`；
- 收到 reasoning delta 时进入 `Reasoning`；
- 收到正文 delta 时进入 `Responding`；
- 执行工具时进入 `Acting`；
- 模型结束、等待持久化时进入 `Finalizing`。
### 6.5 有序 TurnBlock
```rust
pub enum TurnBlock {
Reasoning {
id: BlockId,
iteration: u32,
text: String,
},
Assistant {
id: BlockId,
iteration: u32,
text: String,
},
Tool {
id: String,
iteration: u32,
name: String,
arguments: serde_json::Value,
status: ToolStatus,
preview: Option<String>,
},
}
pub enum ToolStatus {
Running,
Completed,
Failed,
}
```
有序 block 直接表达 reasoning、正文和工具的交错，不再维护多个平行字符串或让客户端猜测工具边界。
### 6.6 Agent 语义事件
```rust
pub enum TurnEvent {
ReasoningDelta {
iteration: u32,
delta: String,
},
TextDelta {
iteration: u32,
delta: String,
},
TextSegmentFinished {
iteration: u32,
},
ToolStarted {
iteration: u32,
call: ToolCall,
},
ToolFinished {
iteration: u32,
call_id: String,
success: bool,
preview: Option<String>,
},
}
```
AgentLoop 只发过程事实。Turn 的 start、finalize、complete、cancel 和 fail 由 Session worker 调用 TurnController，因为 Session 才拥有生命周期、持久化和 stale-state 判断。
## 7. Provider 层
### 7.1 流式优先接口
```rust
#[async_trait]
pub trait LLMProvider: Send + Sync {
async fn stream(
&self,
request: ChatCompletionRequest,
) -> Result<ProviderStream, ProviderError>;
}
```
标题生成、压缩等需要完整响应的代码通过 `collect_provider_stream()` 收集同一实现，避免分别维护 stream 和 non-stream HTTP 路径。
### 7.2 OpenAI-compatible
至少处理：
- `delta.content`；
- `delta.reasoning_content`；
- `delta.reasoning`；
- 同一 payload 同时出现 content 与 reasoning；
- tool call id/name/arguments 分片；
- usage-only final chunk；
- reasoning-only 响应；
- 内联 think tag 被任意 SSE chunk 切分。
`<think>`、`<reasoning>` 等内联标签使用有状态 parser 在 Provider 归一化边界转换为 `ProviderChunk::Reasoning`，不能在 Channel 或客户端重复清理。
### 7.3 Anthropic
至少处理：
- text content block；
- thinking/redacted thinking block；
- signature 或其他回放元数据；
- tool_use block 和 input JSON delta；
- content block start/delta/stop；
- message usage 和 stop reason。
thinking 文本进入 `reasoning`，签名和原始 block 进入 `provider_state`。历史回放必须保持 Provider 要求的块顺序和签名完整性。
### 7.4 reasoning-only
Provider 不擅自把 reasoning 提升为正文。最终轮只有 reasoning、没有正文且没有工具调用时，AgentLoop 保存空正文 assistant 消息及其 reasoning，Turn 正常进入 Completed；交互 UI 仍可显示 reasoning，但不会把它冒充最终答案。
## 8. AgentLoop 与 TurnController
### 8.1 AgentLoop
AgentLoop 消费 `ProviderChunk` 并：
- 累积本轮完整 AssistantMessageData；
- 将展示事实发给 `TurnEmitter`；
- 组装 tool calls；
- Provider 完成当前迭代后执行工具；
- 在工具开始前发出 `TextSegmentFinished`；
- 将完整 assistant/tool messages 加入内部历史；
- 返回最终 `AgentProcessResult`。
AgentLoop 不访问 MessageBus、DeliveryCoordinator 或 Channel。
### 8.2 TurnController
TurnController：
- 是 TurnState 的唯一写入者；
- 把 TurnEvent reduce 为有序 blocks；
- 维护 revision、status 和 phase；
- 发布最新 TurnSnapshot；
- 不执行 Provider、工具、数据库或 Channel I/O。
推荐接口：
```rust
impl TurnController {
pub fn start(... ) -> (Self, TurnEmitter, watch::Receiver<Arc<TurnSnapshot>>);
pub fn begin_finalizing(&mut self);
pub fn complete(&mut self, usage: Option<Usage>);
pub fn cancel(&mut self, reason: Option<String>);
pub fn fail(&mut self, error: String);
}
```
`TurnEmitter` 应轻量、无 Channel 感知，并在 Turn 终止后拒绝新事件。
### 8.3 stale-state
Session worker 在以下位置校验 `worker_generation` 和必要的 `state_version`：
- 创建 Turn 后、调用 Provider 前；
- 发布会产生用户可见变化的快照前；
- 工具批次完成后；
- 最终数据库提交前；
- 发布 Completed 前。
旧 generation 的 Turn 必须进入 Cancelled 或静默终止，禁止继续编辑 Channel 远端消息。
## 9. DeliveryCoordinator
DeliveryCoordinator 是活动 Turn 的唯一展示协调器，职责包括：
1. 订阅 `watch::Receiver<TurnSnapshot>`；
2. 解析当前目标的 PresentationPolicy；
3. 在下发前移除隐藏的 reasoning/tool blocks；
4. 根据 Channel LivePolicy 节流；
5. 为 Turn 创建并持有 TurnSink；
6. 对 Running 快照进行 best-effort 更新；
7. 对终态快照立即 flush；
8. 有界等待 sink 结束；
9. 报告最终投递结果。
### 9.1 合并和背压
- `watch` 自动覆盖过时快照；
- WebSocket/cli_chat 建议最多约 30 FPS；
- 飞书建议从 500ms 更新间隔开始；
- 正在滚动或渲染压力高时，客户端无需向服务端反馈节流，丢弃中间快照即可；
- 终态不等待节流 timer，必须立即发送；
- Running 更新失败不重试旧快照，等待下一最新快照；
- Completed 最终投递使用正常可靠重试语义。
### 9.2 与 OutboundDispatcher 的关系
DeliveryCoordinator 不取代 OutboundDispatcher。
| 组件 | 负责 |
|------|------|
| OutboundDispatcher | 完整独立消息、通知、命令结果、定时投递、可靠重试 |
| DeliveryCoordinator | 活动 Turn 的运行快照、节流、展示过滤和 sink 生命周期 |
两者对同一 `(channel, chat_id)` 的最终写操作必须有统一排序边界。实现时可以复用 per-conversation lane owner，但不能把所有快照排进 lane 的普通 MPSC；lane 应只持有 TurnSink 任务或 latest snapshot receiver。
## 10. Channel 与 TurnSink
### 10.1 接口
删除 `Channel::send_delta`，保留普通 `send`，增加：
```rust
pub enum LivePolicy {
FinalOnly,
Snapshot {
min_interval: Duration,
},
}
#[async_trait]
pub trait Channel: Send + Sync + 'static {
fn live_policy(&self) -> LivePolicy;
async fn open_turn(
&self,
target: TurnTarget,
) -> Result<Box<dyn TurnSink>, ChannelError>;
async fn send(&self, msg: OutboundMessage) -> Result<(), ChannelError>;
}
#[async_trait]
pub trait TurnSink: Send {
async fn update(&mut self, snapshot: &TurnSnapshot) -> Result<(), ChannelError>;
async fn finish(&mut self, snapshot: &TurnSnapshot) -> Result<(), ChannelError>;
async fn abort(&mut self, snapshot: &TurnSnapshot) -> Result<(), ChannelError>;
}
```
`TurnSink` 的每个调用都接收完整、过滤后的快照。Sink 不拼接 token。终态方法保留 `&mut self`，使协调器可以在瞬态错误或超时后重试同一个、仍持有远端消息 ID 的 sink；终态成功或重试耗尽后由协调器销毁 sink。
### 10.2 cli_chat
- `LivePolicy::Snapshot`，默认约 33ms；
- update/finish/abort 都发送统一 `turn_updated` frame；
- WebUI 与 TUI 使用完全相同的 TurnSnapshot；
- Session 不匹配时客户端忽略渲染，但可标记未读；
- 终态后客户端可请求历史作最终校准。
### 10.3 飞书
- 配置关闭实时展示时使用 FinalOnly sink；
- 开启时使用 Snapshot sink；
- 第一个有可见内容的 Running 快照创建卡片；
- 后续快照编辑同一卡片；
- sink 内持有远端 message ID；
- reasoning/tool block 由 PresentationPolicy 决定是否进入卡片；
- 中间编辑失败不影响 Agent；
- finish 做最终编辑，必要时退化为发送一条完整最终消息；
- 卡片长度限制、拆分和平台限流属于 FeishuTurnSink 私有实现。
### 10.4 不支持编辑的 Channel
实现 FinalOnlyTurnSink：忽略 Running 快照，只在 finish 时发送最终投影。默认不通过多条追加消息模拟流式，以免产生无法收回的碎片消息。
## 11. 展示策略
```rust
pub struct PresentationPolicy {
pub live: bool,
pub reasoning: ReasoningVisibility,
pub tools: ToolVisibility,
}
pub enum ReasoningVisibility {
Hidden,
Collapsed,
Expanded,
}
pub enum ToolVisibility {
Hidden,
Compact,
Detailed,
}
```
建议默认值：
- TUI/WebUI：live=true，reasoning=Collapsed，tools=Detailed；
- 外部 Channel：live 由 Channel 配置决定，reasoning=Hidden，tools=Compact；
- Scheduler/无人值守投递：FinalOnly，reasoning=Hidden。
策略由 DeliveryCoordinator 在数据离开 Gateway 核心前应用。Channel 和客户端不能只靠“隐藏 UI”实现 reasoning 保密。
## 12. TUI 与 WebUI
客户端状态简化为：
```text
history 已持久化消息
active_turn 当前 TurnSnapshot（每个 session 最多一个）
```
收到快照时：
```text
if snapshot.revision > active_turn.revision:
active_turn = snapshot
```
渲染规则：
- Reasoning block 显示为折叠或展开区域；
- Assistant block 显示为正文 segment；
- Tool block 显示运行中/成功/失败状态；
- phase 控制 spinner 文案；
- Completed/Cancelled/Failed 显示明确终态；
- 当前 session 之外的快照不进入当前消息列表；
- Completed 后以历史响应替换 active turn，避免展示态与数据库态长期并存。
WebUI 对 Running Markdown 可以按动画帧或快照频率渲染，Completed 时进行最终 sanitize。TUI 原地重绘 active turn，不把每个更新追加为新 transcript 行，也不在用户向上滚动时强制跳到底部。
## 13. WebSocket 协议
Agent 主 Turn 不再通过 `assistant_response` 发送最终正文，活动与终态都使用统一 frame。`assistant_response` 仅保留给不属于活动 Turn 的完整独立消息：
```json
{
"type": "turn_updated",
"snapshot": {
"id": "...",
"session_id": "...",
"message_id": "...",
"revision": 12,
"status": "running",
"phase": "responding",
"blocks": []
}
}
```
Completed、Cancelled 和 Failed 仍使用 `turn_updated`，只改变完整快照的 status。避免为每个生命周期阶段增加一组容易漂移的 frame 类型。
历史协议应返回持久化后的 reasoning、completion_status、turn_id 和 iteration，但不返回 provider_state。
## 14. 持久化和提交顺序
流式过程中不逐 token 写 SQLite。完成顺序固定为：
```text
Provider 完成
→ AgentLoop 组装 emitted_messages
→ Session 校验 generation/state_version
→ 原子写入本轮全部消息
→ TurnController.complete()
→ 发布 Completed 快照
→ TurnSink.finish()
```
因此 `TurnStatus::Completed` 的含义是：数据库已经提交成功，最终展示可以安全收敛到历史。
### 14.1 取消
采用以下语义：
- 没有 Assistant 正文：不持久化 assistant 消息，Turn 标记 Cancelled；
- 已向用户展示部分正文：持久化部分正文并标记 `completion_status=cancelled`；
- 已完成的 assistant/tool/tool-result 链必须保持 Provider 可接受的结构；
- reasoning 可随取消消息保存，但展示仍受 policy 控制；
- 取消后 TurnEmitter 关闭，迟到 delta 被丢弃。
### 14.2 失败
- Provider 在任何可见正文前失败：Turn Failed，错误作为结构化 error 展示，不创建 assistant 历史；
- 已产生部分正文后失败：按 interrupted partial 保存，标记 `completion_status=interrupted`；
- 持久化失败：不得发布 Completed，Turn Failed，并明确告知用户流式预览未保存；
- Running Channel 更新失败不改变 Turn 结果；最终 finish 失败按现有投递错误处理。
## 15. SQLite 迁移
schema v4 为 `messages` 增加：
```text
turn_id TEXT NULL
iteration INTEGER NULL
completion_status TEXT NOT NULL DEFAULT 'completed'
reasoning_content TEXT NULL -- 已存在，语义调整为可展示 reasoning
provider_state TEXT NULL -- JSON，Provider 私有回放状态
```
要求：
- 新库 schema 测试；
- 旧库迁移测试；
- 已有 `reasoning_content` 数据原样保留；
- `provider_state` 解析失败时降级为不回放，不能使历史不可读；
- Session 加载继续修复 tool-call chains；
- 原子提交覆盖完整 Turn 的所有 emitted messages。
当前不新增 `turns` 表。运行中 Turn 只存在内存，历史可通过 messages.turn_id 分组。如果未来要跨 Gateway 重启恢复运行态，再单独设计 durable turn lease/state。
## 16. 生命周期与并发不变量
1. 每个 Session 最多一个活动主 Turn。
2. TurnController 是 TurnState 的唯一写入者。
3. AgentLoop、DeliveryCoordinator 和 TurnSink 不持有 Session mutex 执行慢 I/O。
4. `worker_generation` 变化后，旧 Turn 不得发布新快照或提交消息。
5. Running 快照是 best-effort；终态快照必须显式、完整且有界投递。
6. 慢 sink 只能跳过中间状态，不能阻塞 Provider 或 AgentLoop。
7. Completed 必须晚于数据库成功提交。
8. TurnSink 的生命周期由 DeliveryCoordinator 所有，并通过 TaskSupervisor 回收。
9. Provider stream、节流 timer、Channel 编辑、取消和 shutdown 都必须有硬时间界限。
10. presentation 过滤不能修改 ConversationMessage 或 Agent 历史。
## 17. 模块布局
```text
src/providers/stream.rs
ProviderChunk、ProviderStream、FinishReason、collect helper
src/agent/turn_event.rs
TurnEvent、TurnEmitter
src/session/turn.rs
TurnState、TurnBlock、TurnController、TurnSnapshot
src/delivery/mod.rs
src/delivery/coordinator.rs
src/delivery/policy.rs
watch 订阅、展示过滤、节流、sink 生命周期
src/channels/base.rs
Channel、LivePolicy、TurnSink
src/channels/cli_chat.rs
WebSocketTurnSink
src/channels/feishu.rs
FeishuTurnSink
src/protocol.rs
TurnSnapshot 序列化
```
`observability` 继续记录 Agent/tool 遥测，不承担 UI stream。`MessageBus` 不新增 token/turn 队列。
## 18. 明确拒绝的替代方案
### 18.1 每 token 一个 OutboundMessage
拒绝原因：填满 bounded bus/lane、重试乱序、慢 Channel 反压 Agent、最终消息和暂态更新语义混淆。
### 18.2 端到端 delta 协议
拒绝原因：客户端和每个 Channel 都必须实现累积、去重、segment、取消和丢帧恢复状态机，最终产生多个事实 owner。
### 18.3 客户端自行组合 reasoning/tool/text
拒绝原因：TUI、WebUI 和 Channel 行为会漂移；重连和切 session 时难以恢复；服务端已经拥有全部事实。
### 18.4 把流式事件写入数据库
拒绝原因：消息历史膨胀，事务语义复杂，压缩和 Provider 回放被展示细节污染。
### 18.5 在 Channel 全局保存 turn_id 映射
拒绝原因：owner 和清理边界不清晰。每 Turn 一个 sink 可以让远端消息状态自然随生命周期释放。
### 18.6 一个巨型跨平台 StreamConsumer
拒绝原因：通用合并/策略与平台 API 细节耦合。DeliveryCoordinator 只做统一调度，具体远端编辑由各 TurnSink 自己实现。
## 19. 验证策略
### 19.1 Provider
- SSE 任意字节和 UTF-8 边界切分；
- reasoning/content 同时出现；
- reasoning 与 content 字段别名；
- think tag 跨 chunk；
- 多 tool calls 交错参数 delta；
- usage-only chunk；
- Anthropic thinking signature round-trip；
- 中途断线、超时和取消。
### 19.2 TurnController
- reasoning → text → tool → reasoning → text 的 block 顺序；
- segment boundary；
- parallel tool status；
- revision 严格递增；
- 终态后拒绝新事件；
- reasoning-only、empty response、取消和失败。
建议用 property tests 验证：任意合法 TurnEvent 序列 reduce 后不产生相邻可合并同类 block、重复 tool id 或终态后变更。
### 19.3 DeliveryCoordinator
- 慢 sink 只收到最新快照；
- 终态绕过节流；
- hidden reasoning 在到达 sink 前已经移除；
- Running 更新失败后能由下一快照恢复；
- finish 可靠重试；
- shutdown 有界；
- stale generation 停止更新。
### 19.4 客户端和 Channel
- WebUI/TUI revision 去重；
- session 切换不显示迟到 Turn；
- Completed 后历史校准；
- Markdown 未完成块与最终块；
- 飞书 create/edit/final fallback；
- FinalOnly sink 不发送中间内容；
- 远端长度限制与节流。
### 19.5 Storage
- 新库 schema；
- 旧 schema 迁移；
- provider_state 损坏降级；
- cancelled/interrupted 消息恢复；
- 完整 Turn 原子提交失败不产生部分历史。
## 20. 实施记录
实现按可独立验证的里程碑完成：SQLite 消息语义、Turn 状态机、OpenAI 原生流、Agent/Session 生命周期、DeliveryCoordinator、WebSocket/TUI/WebUI、Anthropic 签名回放、FeishuTurnSink，最后删除过渡适配并同步运行时文档。每个里程碑均保持非流式最终回复可用，且没有为旧增量协议保留双栈。
## 21. 已采用的产品策略
这些选择不改变架构，但决定默认产品行为：
1. `/stop` 后持久化已展示的部分正文并标记 `cancelled`；只有 reasoning 时不创建 assistant 历史。
2. TUI/WebUI 在独立区域展示 reasoning；WebUI 默认折叠，TUI 直接显示。
3. 外部 Channel 默认隐藏 reasoning，工具只显示紧凑状态。
4. reasoning-only 不提升为正文。
5. `provider_state` 随 assistant 消息保留，用于同 Provider 精确回放；不下发客户端或 Channel，损坏时安全忽略。
## 22. 架构验收标准
设计完成实现后，应能用以下陈述准确描述系统：
- Provider 只负责模型协议，AgentLoop 只负责模型/工具语义。
- Session 拥有 Turn 生命周期和最终持久化。
- TurnController 是运行态的唯一事实来源。
- DeliveryCoordinator 只投影展示，不修改历史。
- 每个 Channel Turn 的远端状态只存在于一个 TurnSink。
- 客户端只渲染服务端快照，不重建领域状态。
- 中间快照可以丢，最终状态一定可收敛。
- reasoning 展示、reasoning 回放和“系统正在工作”是三个不同概念。
- Completed 永远意味着数据库已经提交成功。

24
resources/skills/about-picobot/references/architecture.md
View File

@ -7,6 +7,8 @@ Channel → MessageBus.inbound → Gateway processor → SessionManager → per-
↑ │ ↑ │
└── Channel ← OutboundDispatcher ← per-conversation lane ← MessageBus.outbound └── Channel ← OutboundDispatcher ← per-conversation lane ← MessageBus.outbound
AgentLoop → TurnEvent → TurnController → latest TurnSnapshot → DeliveryCoordinator → TurnSink → Channel
WebSocket/Channel → MessageBus.control → Gateway processor → SessionManager (dialog 操作) WebSocket/Channel → MessageBus.control → Gateway processor → SessionManager (dialog 操作)
Scheduler → SessionManager.handle_cron_message → AgentLoop → send_message Scheduler → SessionManager.handle_cron_message → AgentLoop → send_message
``` ```
@ -19,9 +21,10 @@ Scheduler → SessionManager.handle_cron_message → AgentLoop → send_message
| `client` | TUI 聊天客户端 | | `client` | TUI 聊天客户端 |
| `channels` | 外部集成（飞书、CLI），仅收发消息 | | `channels` | 外部集成（飞书、CLI），仅收发消息 |
| `bus` | 有界 inbound/outbound/control 队列；出站 dispatcher 与分目标 lane | | `bus` | 有界 inbound/outbound/control 队列；出站 dispatcher 与分目标 lane |
| `session` | 会话生命周期、dialog 操作、每 session 串行队列、上下文与持久化协调 | | `session` | 会话生命周期、dialog 操作、每 session 串行队列、Turn 状态、上下文与持久化协调 |
| `agent` | LLM 调用循环、工具执行、上下文压缩、媒体处理、子 Agent | | `agent` | LLM 调用循环、工具执行、上下文压缩、媒体处理、子 Agent、Turn 语义事件 |
| `providers` | LLM API 客户端（OpenAI 兼容、Anthropic） | | `providers` | OpenAI/Anthropic 原生流解析，统一正文、reasoning、工具、usage 与私有回放状态 |
| `delivery` | 完整 Turn 快照的展示过滤、latest-wins 节流、终态投递和 TurnSink 生命周期 |
| `tools` | Agent 工具（bash、文件操作、搜索、HTTP、web、browser、memory、delegate 等） | | `tools` | Agent 工具（bash、文件操作、搜索、HTTP、web、browser、memory、delegate 等） |
| `skills` | Skill 加载、管理和 prompt 构建 | | `skills` | Skill 加载、管理和 prompt 构建 |
| `storage` | SQLite 持久化 | | `storage` | SQLite 持久化 |
@ -36,11 +39,14 @@ Scheduler → SessionManager.handle_cron_message → AgentLoop → send_message
## 功能边界 ## 功能边界
- Channels 仅收发消息，不感知 session 或 LLM - Channels 通过 MessageBus 发布入站消息，通过 OutboundDispatcher 或每 Turn 一个的 TurnSink 接收出站写入，不感知 session 或 LLM
- MessageBus 本体持有三条有界队列；出站路由、顺序和重试由 `OutboundDispatcher` 负责 - MessageBus 本体持有三条有界队列；出站路由、顺序和重试由 `OutboundDispatcher` 负责
- SessionManager 拥有 session 状态、dialog 路由、上下文构建和每 session worker，并通过 worker 创建 AgentLoop - SessionManager 拥有 session 状态、dialog 路由、上下文构建和每 session worker，并通过 worker 创建 AgentLoop
- TurnController 是活动 Turn 状态的唯一 owner；Session 在消息原子提交成功后才发布 Completed
- AgentLoop 跨轮无状态，接收已准备的 history 调用 LLM、执行工具并返回一次结果 - AgentLoop 跨轮无状态，接收已准备的 history 调用 LLM、执行工具并返回一次结果
- Providers 是纯 HTTP 客户端，无 bus/session/channel 感知 - Providers 是纯 HTTP 流客户端，无 bus/session/channel 感知；签名 reasoning 状态只回放给匹配 Provider，不下发客户端或 Channel
- DeliveryCoordinator 只投影完整快照，不修改会话历史；慢消费者跳过中间 revision，终态显式、有界投递
- 每个活动 Turn 独占一个 TurnSink；平台 message ID 和 reaction 清理状态只存在于 sink 内
- Tools 接收原始参数，返回字符串结果 - Tools 接收原始参数，返回字符串结果
- MCP 工具在 Gateway 初始化时连接服务器、发现工具，并包装成普通 Tool 注册到 ToolRegistry - MCP 工具在 Gateway 初始化时连接服务器、发现工具，并包装成普通 Tool 注册到 ToolRegistry
- 子 Agent 由 `delegate` 工具创建，复用 provider 配置和按需过滤后的工具集；后台任务结果通过 MessageBus 发回原会话 - 子 Agent 由 `delegate` 工具创建，复用 provider 配置和按需过滤后的工具集；后台任务结果通过 MessageBus 发回原会话
@ -58,6 +64,8 @@ Scheduler → SessionManager.handle_cron_message → AgentLoop → send_message
- `browser` 工具只有在 `browser.enabled=true` 时注册，依赖 Chrome/Chromium 与 WebDriver - `browser` 工具只有在 `browser.enabled=true` 时注册，依赖 Chrome/Chromium 与 WebDriver
- 同一 session 的普通消息串行处理，不同 session 可并发；session 队列容量为 32，满时明确拒绝 - 同一 session 的普通消息串行处理，不同 session 可并发；session 队列容量为 32，满时明确拒绝
- 出站消息按 `(channel, chat_id)` 分 lane 保序；lane 容量为 64，慢目标不阻塞其他目标 - 出站消息按 `(channel, chat_id)` 分 lane 保序；lane 容量为 64，慢目标不阻塞其他目标
- 活动 Turn 与普通出站消息共享 `(channel, chat_id)` 写锁；禁止把 token delta 放入 MessageBus
- `cli_chat` 向 TUI/WebUI 发送统一 `turn_updated` 完整快照；飞书默认 FinalOnly，开启 `live_updates` 后编辑同一卡片
- 长生命周期后台任务由 TaskSupervisor 管理；连接局部任务由其 owner 限时 join 或 abort - 长生命周期后台任务由 TaskSupervisor 管理；连接局部任务由其 owner 限时 join 或 abort
- 外部建连、重试等待和关停 join 必须可取消且有硬超时 - 外部建连、重试等待和关停 join 必须可取消且有硬超时
- 不得记录 API Key、Authorization header 或包含临时凭据的完整连接 URL - 不得记录 API Key、Authorization header 或包含临时凭据的完整连接 URL
@ -138,6 +146,12 @@ Worker 的处理原则：
4. Session 持久化由独立 `persistence_lock` 串行化；批量消息使用原子写入，失败时精确回滚内存后缀。 4. Session 持久化由独立 `persistence_lock` 串行化；批量消息使用原子写入，失败时精确回滚内存后缀。
5. 上下文溢出时按 Provider 返回的真实限制重新压缩并重试。 5. 上下文溢出时按 Provider 返回的真实限制重新压缩并重试。
### 活动 Turn
每个主 Agent 请求会创建一个内存 Turn。Provider delta 经 AgentLoop 转换为 reasoning、正文、工具开始/完成等语义事件，TurnController 归约为有序 block 和单调 revision 的完整快照。TUI/WebUI 使用 `history + active_turn` 渲染，不自行拼接 token；中间帧可丢，下一快照会自动收敛。
展示策略在 Gateway 核心出口应用：交互客户端可显示 reasoning 和详细工具状态；外部渠道隐藏 reasoning、工具仅显示紧凑状态；无人值守投递只保留正文。运行态不逐 token 入库，完成、取消或中断时才原子保存消息及 completion status。
### 会话恢复 ### 会话恢复
从 Storage 恢复 session 时： 从 Storage 恢复 session 时：

4
resources/skills/about-picobot/references/config.md
View File

@ -97,6 +97,10 @@ Gateway WebUI 的“配置”页可以编辑实际加载的配置文件。读取
| `agent` | string | - | 使用的 agent 名称 | | `agent` | string | - | 使用的 agent 名称 |
| `media_dir` | string | ~/.picobot/media/feishu | 配置默认值；Gateway 注册渠道时会覆盖为 `{workspace}/media/feishu` | | `media_dir` | string | ~/.picobot/media/feishu | 配置默认值；Gateway 注册渠道时会覆盖为 `{workspace}/media/feishu` |
| `reaction_emoji` | string | "Typing" | 回复意向表达的表情 | | `reaction_emoji` | string | "Typing" | 回复意向表达的表情 |
| `live_updates` | bool | false | 是否用单张卡片实时编辑活动 Turn；关闭时只发送终态 |
| `live_update_interval_ms` | int | 500 | 卡片更新最小间隔，运行时限制在 250–5000ms |
飞书属于外部渠道：无论是否开启实时卡片，都不会接收模型 reasoning；工具只显示紧凑状态。配置修改需重启 Gateway 生效。
## mcp 字段 ## mcp 字段

12
resources/skills/about-picobot/references/db-schema.md
View File

@ -2,7 +2,7 @@
数据库为 SQLite，默认位于 workspace 下的 `picobot.db`。 数据库为 SQLite，默认位于 workspace 下的 `picobot.db`。
连接启用 WAL、`synchronous=NORMAL`、foreign keys、5 秒 busy timeout，连接池最多 8 个连接。当前 `PRAGMA user_version=1`；启动时会在事务内补齐旧库字段和索引，遇到比程序更新的 schema version 会拒绝启动。 连接启用 WAL、`synchronous=NORMAL`、foreign keys、5 秒 busy timeout，连接池最多 8 个连接。当前 `PRAGMA user_version=4`；启动时会在事务内补齐旧库字段和索引，遇到比程序更新的 schema version 会拒绝启动。
## sessions 表 ## sessions 表
@ -41,7 +41,11 @@
| `tool_calls` | TEXT | 工具调用参数 JSON | | `tool_calls` | TEXT | 工具调用参数 JSON |
| `source` | TEXT | 消息来源（跨会话消息时标记来源 session_id） | | `source` | TEXT | 消息来源（跨会话消息时标记来源 session_id） |
| `created_at` | INTEGER | 创建时间（Unix 毫秒） | | `created_at` | INTEGER | 创建时间（Unix 毫秒） |
| `reasoning_content` | TEXT | provider 返回的推理内容（如有） | | `reasoning_content` | TEXT | 可展示的模型 reasoning（如有） |
| `provider_state` | TEXT | Provider 私有回放状态 JSON；只回放给匹配 Provider，不下发客户端或 Channel |
| `turn_id` | TEXT | 产生该消息的活动 Turn ID |
| `iteration` | INTEGER | Agent 工具循环中的迭代序号 |
| `completion_status` | TEXT | `completed` / `cancelled` / `interrupted`，旧数据默认 completed |
`(session_id, seq)` 有唯一索引，防止并发写入重复序号。删除 session 会通过外键级联删除 messages。 `(session_id, seq)` 有唯一索引，防止并发写入重复序号。删除 session 会通过外键级联删除 messages。
@ -137,9 +141,9 @@ Scheduler 使用原子 `UPDATE ... RETURNING` 领取到期任务。任务结果
| `created_at` | INTEGER | 调用时间 | | `created_at` | INTEGER | 调用时间 |
| `provider` | TEXT | 提供商类型 | | `provider` | TEXT | 提供商类型 |
| `model` | TEXT | 模型名称 | | `model` | TEXT | 模型名称 |
| `request_body` | TEXT | 请求体 JSON | | `request_body` | TEXT | 请求摘要 JSON；旧记录可能是完整请求体 |
| `response_body` | TEXT | 响应体 JSON | | `response_body` | TEXT | 响应体 JSON |
| `error` | TEXT | 错误信息 | | `error` | TEXT | 错误信息 |
| `duration_ms` | INTEGER | 耗时（毫秒） | | `duration_ms` | INTEGER | 耗时（毫秒） |
`request_body`/`response_body` 可能包含用户内容，排障和导出数据库时应按敏感数据处理。 旧数据中的 `request_body`/`response_body` 可能包含用户内容，排障和导出数据库时应按敏感数据处理。新 Provider 请求的 `request_body` 只保存模型、消息数、工具数和 stream 标志等摘要；错误响应仍可能包含服务端回显内容。

4
resources/templates/config.example.json
View File

@ -70,7 +70,9 @@
"allow_from": ["*"], "allow_from": ["*"],
"agent": "default", "agent": "default",
"media_dir": "~/.picobot/media/feishu", "media_dir": "~/.picobot/media/feishu",
"reaction_emoji": "Typing" "reaction_emoji": "Typing",
"live_updates": false,
"live_update_interval_ms": 500
} }
}, },
"memory": { "memory": {

388
src/agent/agent_loop.rs
View File

@ -1,11 +1,15 @@
use crate::agent::context_compressor::estimate_tokens; use crate::agent::context_compressor::estimate_tokens;
use crate::agent::media_handler::MediaHandlerRegistry; use crate::agent::media_handler::MediaHandlerRegistry;
use crate::agent::system_prompt::build_system_prompt; use crate::agent::system_prompt::build_system_prompt;
use crate::agent::turn_event::{AgentTurnContext, TurnEvent};
use crate::bus::message::ContentBlock; use crate::bus::message::ContentBlock;
use crate::bus::{ChatMessage, MediaRef}; use crate::bus::{ChatMessage, MediaRef};
use crate::config::LLMProviderConfig; use crate::config::LLMProviderConfig;
use crate::observability::{Observer, ObserverEvent, ToolExecutionOutcome, truncate_args}; use crate::observability::{Observer, ObserverEvent, ToolExecutionOutcome, truncate_args};
use crate::providers::{ChatCompletionRequest, LLMProvider, Message, ToolCall, create_provider}; use crate::providers::{
ChatCompletionRequest, ChatCompletionResponse, LLMProvider, Message, ProviderChunk,
ProviderResponseAccumulator, ToolCall, create_provider,
};
use crate::tools::ToolRegistry; use crate::tools::ToolRegistry;
use std::collections::VecDeque; use std::collections::VecDeque;
use std::hash::{Hash, Hasher}; use std::hash::{Hash, Hasher};
@ -13,11 +17,14 @@ use std::path::PathBuf;
use std::sync::Arc; use std::sync::Arc;
use std::time::Instant; use std::time::Instant;
use futures_util::StreamExt;
/// Maximum characters in a tool result before truncation. /// Maximum characters in a tool result before truncation.
/// Prevents context overflow from large tool outputs. /// Prevents context overflow from large tool outputs.
const MAX_TOOL_RESULT_CHARS: usize = 16_000; const MAX_TOOL_RESULT_CHARS: usize = 16_000;
/// Minimum characters to keep when truncating /// Minimum characters to keep when truncating
const TRUNCATION_SUFFIX_LEN: usize = 200; const TRUNCATION_SUFFIX_LEN: usize = 200;
const TOOL_PREVIEW_CHARS: usize = 1_000;
enum MediaOrigin<'a> { enum MediaOrigin<'a> {
User, User,
@ -164,6 +171,17 @@ fn truncate_tool_result(output: &str) -> String {
} }
} }
fn tool_result_preview(output: &str) -> String {
if output.len() <= TOOL_PREVIEW_CHARS {
output.to_string()
} else {
format!(
"{}…",
&output[..output.floor_char_boundary(TOOL_PREVIEW_CHARS)]
)
}
}
/// Loop detection result. /// Loop detection result.
#[derive(Debug, Clone, PartialEq, Eq)] #[derive(Debug, Clone, PartialEq, Eq)]
enum LoopDetectionResult { enum LoopDetectionResult {
@ -296,6 +314,32 @@ pub struct AgentProcessResult {
pub final_response: ChatMessage, pub final_response: ChatMessage,
pub emitted_messages: Vec<ChatMessage>, pub emitted_messages: Vec<ChatMessage>,
pub total_tokens: Option<u32>, pub total_tokens: Option<u32>,
pub usage: Option<crate::providers::Usage>,
}
fn merge_usage(total: &mut crate::providers::Usage, next: &crate::providers::Usage) {
total.prompt_tokens = total.prompt_tokens.saturating_add(next.prompt_tokens);
total.completion_tokens = total
.completion_tokens
.saturating_add(next.completion_tokens);
total.total_tokens = total.total_tokens.saturating_add(next.total_tokens);
total.cached_tokens = sum_optional_tokens(total.cached_tokens, next.cached_tokens);
total.cache_read_input_tokens =
sum_optional_tokens(total.cache_read_input_tokens, next.cache_read_input_tokens);
total.cache_creation_input_tokens = sum_optional_tokens(
total.cache_creation_input_tokens,
next.cache_creation_input_tokens,
);
}
fn sum_optional_tokens(left: Option<u32>, right: Option<u32>) -> Option<u32> {
match (left, right) {
(None, None) => None,
(left, right) => Some(
left.unwrap_or_default()
.saturating_add(right.unwrap_or_default()),
),
}
} }
impl AgentLoop { impl AgentLoop {
@ -451,6 +495,60 @@ impl AgentLoop {
&self.tools &self.tools
} }
async fn stream_completion(
&self,
request: ChatCompletionRequest,
iteration: u32,
turn: Option<&AgentTurnContext>,
) -> Result<ChatCompletionResponse, AgentError> {
let mut provider_stream = self.provider.stream(request).await.map_err(|error| {
tracing::error!(error = %error, "LLM request failed");
AgentError::LlmError(error.to_string())
})?;
let mut accumulator = ProviderResponseAccumulator::default();
while let Some(chunk) = provider_stream.next().await {
let chunk = chunk.map_err(|error| {
tracing::error!(error = %error, "LLM stream failed");
AgentError::LlmError(error.to_string())
})?;
if let Some(turn) = turn {
let event = match &chunk {
ProviderChunk::Reasoning(delta) => Some(TurnEvent::ReasoningDelta {
iteration,
delta: delta.clone(),
}),
ProviderChunk::Text(delta) => Some(TurnEvent::TextDelta {
iteration,
delta: delta.clone(),
}),
_ => None,
};
if let Some(event) = event {
turn.emitter.emit(event).map_err(|error| {
AgentError::Other(format!("turn event rejected: {error}"))
})?;
}
}
accumulator.push(chunk);
}
Ok(accumulator.finish())
}
fn annotate_message(
message: &mut ChatMessage,
turn: Option<&AgentTurnContext>,
iteration: u32,
final_response: bool,
) {
message.iteration = Some(iteration);
if let Some(turn) = turn {
message.turn_id = Some(turn.turn_id.clone());
if final_response {
message.id = turn.message_id.clone();
}
}
}
fn chat_message_to_llm_message(&self, m: &ChatMessage, include_media: bool) -> Message { fn chat_message_to_llm_message(&self, m: &ChatMessage, include_media: bool) -> Message {
let content = if m.media_refs.is_empty() || !include_media { let content = if m.media_refs.is_empty() || !include_media {
vec![ContentBlock::text(&m.content)] vec![ContentBlock::text(&m.content)]
@ -473,6 +571,7 @@ impl AgentLoop {
role: m.role.clone(), role: m.role.clone(),
content, content,
reasoning_content: m.reasoning_content.clone(), reasoning_content: m.reasoning_content.clone(),
provider_state: m.provider_state.clone(),
tool_call_id: m.tool_call_id.clone(), tool_call_id: m.tool_call_id.clone(),
name: m.tool_name.clone(), name: m.tool_name.clone(),
tool_calls: m.tool_calls.clone(), tool_calls: m.tool_calls.clone(),
@ -498,8 +597,24 @@ impl AgentLoop {
/// - The LLM returns no more tool calls (final response) /// - The LLM returns no more tool calls (final response)
/// - Maximum iterations are reached /// - Maximum iterations are reached
pub async fn process( pub async fn process(
&self,
messages: Vec<ChatMessage>,
) -> Result<AgentProcessResult, AgentError> {
self.process_inner(messages, None).await
}
pub async fn process_streaming(
&self,
messages: Vec<ChatMessage>,
turn: AgentTurnContext,
) -> Result<AgentProcessResult, AgentError> {
self.process_inner(messages, Some(turn)).await
}
async fn process_inner(
&self, &self,
mut messages: Vec<ChatMessage>, mut messages: Vec<ChatMessage>,
turn: Option<AgentTurnContext>,
) -> Result<AgentProcessResult, AgentError> { ) -> Result<AgentProcessResult, AgentError> {
#[cfg(debug_assertions)] #[cfg(debug_assertions)]
tracing::debug!( tracing::debug!(
@ -522,6 +637,7 @@ impl AgentLoop {
let mut loop_detector = LoopDetector::new(LoopDetectorConfig::default()); let mut loop_detector = LoopDetector::new(LoopDetectorConfig::default());
let mut emitted_messages = Vec::new(); let mut emitted_messages = Vec::new();
let mut accumulated_tokens: u32 = 0; let mut accumulated_tokens: u32 = 0;
let mut accumulated_usage = crate::providers::Usage::default();
for iteration in 0..self.max_iterations { for iteration in 0..self.max_iterations {
#[cfg(debug_assertions)] #[cfg(debug_assertions)]
@ -564,12 +680,14 @@ impl AgentLoop {
}; };
// Call LLM // Call LLM
let response = (*self.provider).chat(request).await.map_err(|e| { let iteration = u32::try_from(iteration)
tracing::error!(error = %e, "LLM request failed"); .map_err(|_| AgentError::Other("tool iteration exceeds u32".to_string()))?;
AgentError::LlmError(e.to_string()) let response = self
})?; .stream_completion(request, iteration, turn.as_ref())
.await?;
accumulated_tokens += response.usage.total_tokens; accumulated_tokens = accumulated_tokens.saturating_add(response.usage.total_tokens);
merge_usage(&mut accumulated_usage, &response.usage);
#[cfg(debug_assertions)] #[cfg(debug_assertions)]
tracing::debug!( tracing::debug!(
@ -583,14 +701,23 @@ impl AgentLoop {
if response.tool_calls.is_empty() { if response.tool_calls.is_empty() {
let mut assistant_message = ChatMessage::assistant(response.content); let mut assistant_message = ChatMessage::assistant(response.content);
assistant_message.reasoning_content = response.reasoning_content; assistant_message.reasoning_content = response.reasoning_content;
assistant_message.provider_state = response.provider_state;
Self::annotate_message(&mut assistant_message, turn.as_ref(), iteration, true);
emitted_messages.push(assistant_message.clone()); emitted_messages.push(assistant_message.clone());
return Ok(AgentProcessResult { return Ok(AgentProcessResult {
final_response: assistant_message, final_response: assistant_message,
emitted_messages, emitted_messages,
total_tokens: Some(accumulated_tokens), total_tokens: Some(accumulated_tokens),
usage: Some(accumulated_usage),
}); });
} }
if let Some(turn) = turn.as_ref() {
turn.emitter
.emit(TurnEvent::TextSegmentFinished { iteration })
.map_err(|error| AgentError::Other(format!("turn event rejected: {error}")))?;
}
// Execute tool calls — log and notify immediately // Execute tool calls — log and notify immediately
{ {
let tools_info: Vec<String> = response let tools_info: Vec<String> = response
@ -614,11 +741,15 @@ impl AgentLoop {
response.tool_calls.clone(), response.tool_calls.clone(),
); );
assistant_message.reasoning_content = response.reasoning_content; assistant_message.reasoning_content = response.reasoning_content;
assistant_message.provider_state = response.provider_state;
Self::annotate_message(&mut assistant_message, turn.as_ref(), iteration, false);
messages.push(assistant_message.clone()); messages.push(assistant_message.clone());
emitted_messages.push(assistant_message); emitted_messages.push(assistant_message);
// Execute tools and add results to messages // Execute tools and add results to messages
let tool_results = self.execute_tools(&response.tool_calls).await; let tool_results = self
.execute_tools(&response.tool_calls, iteration, turn.as_ref())
.await?;
for (tool_call, result) in response.tool_calls.iter().zip(tool_results.iter()) { for (tool_call, result) in response.tool_calls.iter().zip(tool_results.iter()) {
// Log function call with name and arguments // Log function call with name and arguments
@ -644,22 +775,24 @@ impl AgentLoop {
"Loop warning: {}", "Loop warning: {}",
msg msg
); );
let tool_message = ChatMessage::tool_with_media( let mut tool_message = ChatMessage::tool_with_media(
tool_call.id.clone(), tool_call.id.clone(),
tool_call.name.clone(), tool_call.name.clone(),
format!("{}\n\n[上一条结果]\n{}", msg, truncated_output), format!("{}\n\n[上一条结果]\n{}", msg, truncated_output),
result.media_refs.clone(), result.media_refs.clone(),
); );
Self::annotate_message(&mut tool_message, turn.as_ref(), iteration, false);
messages.push(tool_message.clone()); messages.push(tool_message.clone());
emitted_messages.push(tool_message); emitted_messages.push(tool_message);
} }
LoopDetectionResult::Ok => { LoopDetectionResult::Ok => {
let tool_message = ChatMessage::tool_with_media( let mut tool_message = ChatMessage::tool_with_media(
tool_call.id.clone(), tool_call.id.clone(),
tool_call.name.clone(), tool_call.name.clone(),
truncated_output, truncated_output,
result.media_refs.clone(), result.media_refs.clone(),
); );
Self::annotate_message(&mut tool_message, turn.as_ref(), iteration, false);
messages.push(tool_message.clone()); messages.push(tool_message.clone());
emitted_messages.push(tool_message); emitted_messages.push(tool_message);
} }
@ -695,25 +828,56 @@ impl AgentLoop {
tools: None, // No tools in final summary call tools: None, // No tools in final summary call
}; };
match (*self.provider).chat(request).await { let summary_iteration = u32::try_from(self.max_iterations)
.map_err(|_| AgentError::Other("tool iteration exceeds u32".to_string()))?;
match self
.stream_completion(request, summary_iteration, turn.as_ref())
.await
{
Ok(response) => { Ok(response) => {
accumulated_tokens += response.usage.total_tokens; accumulated_tokens = accumulated_tokens.saturating_add(response.usage.total_tokens);
merge_usage(&mut accumulated_usage, &response.usage);
let mut assistant_message = ChatMessage::assistant(response.content); let mut assistant_message = ChatMessage::assistant(response.content);
assistant_message.reasoning_content = response.reasoning_content; assistant_message.reasoning_content = response.reasoning_content;
assistant_message.provider_state = response.provider_state;
Self::annotate_message(
&mut assistant_message,
turn.as_ref(),
summary_iteration,
true,
);
emitted_messages.push(assistant_message.clone()); emitted_messages.push(assistant_message.clone());
Ok(AgentProcessResult { Ok(AgentProcessResult {
final_response: assistant_message, final_response: assistant_message,
emitted_messages, emitted_messages,
total_tokens: Some(accumulated_tokens), total_tokens: Some(accumulated_tokens),
usage: Some(accumulated_usage),
}) })
} }
Err(e) => { Err(e) => {
// Fallback if summary call fails // Fallback if summary call fails
tracing::error!(error = %e, "Failed to get summary from LLM"); tracing::error!(error = %e, "Failed to get summary from LLM");
let final_message = ChatMessage::assistant(format!( let fallback = format!(
"I reached the maximum number of tool call iterations ({}) without completing the task. The work done so far has been lost due to an error. Please try breaking the task into smaller steps.", "I reached the maximum number of tool call iterations ({}) without completing the task. The work done so far has been lost due to an error. Please try breaking the task into smaller steps.",
self.max_iterations self.max_iterations
)); );
if let Some(turn) = turn.as_ref() {
turn.emitter
.emit(TurnEvent::TextSegmentFinished {
iteration: summary_iteration,
})
.and_then(|()| {
turn.emitter.emit(TurnEvent::TextDelta {
iteration: summary_iteration,
delta: fallback.clone(),
})
})
.map_err(|error| {
AgentError::Other(format!("turn event rejected: {error}"))
})?;
}
let mut final_message = ChatMessage::assistant(fallback);
Self::annotate_message(&mut final_message, turn.as_ref(), summary_iteration, true);
emitted_messages.push(final_message.clone()); emitted_messages.push(final_message.clone());
Ok(AgentProcessResult { Ok(AgentProcessResult {
final_response: final_message, final_response: final_message,
@ -723,6 +887,7 @@ impl AgentLoop {
} else { } else {
None None
}, },
usage: (accumulated_usage.total_tokens > 0).then_some(accumulated_usage),
}) })
} }
} }
@ -753,42 +918,76 @@ impl AgentLoop {
} }
/// Execute multiple tool calls, choosing parallel or sequential based on conditions. /// Execute multiple tool calls, choosing parallel or sequential based on conditions.
async fn execute_tools(&self, tool_calls: &[ToolCall]) -> Vec<ToolExecutionOutcome> { async fn execute_tools(
&self,
tool_calls: &[ToolCall],
iteration: u32,
turn: Option<&AgentTurnContext>,
) -> Result<Vec<ToolExecutionOutcome>, AgentError> {
if self.should_execute_in_parallel(tool_calls) { if self.should_execute_in_parallel(tool_calls) {
tracing::debug!("Executing {} tools in parallel", tool_calls.len()); tracing::debug!("Executing {} tools in parallel", tool_calls.len());
self.execute_tools_parallel(tool_calls).await self.execute_tools_parallel(tool_calls, iteration, turn)
.await
} else { } else {
tracing::debug!("Executing {} tools sequentially", tool_calls.len()); tracing::debug!("Executing {} tools sequentially", tool_calls.len());
self.execute_tools_sequential(tool_calls).await self.execute_tools_sequential(tool_calls, iteration, turn)
.await
} }
} }
/// Execute tools in parallel using join_all. /// Execute tools in parallel using join_all.
async fn execute_tools_parallel(&self, tool_calls: &[ToolCall]) -> Vec<ToolExecutionOutcome> { async fn execute_tools_parallel(
&self,
tool_calls: &[ToolCall],
iteration: u32,
turn: Option<&AgentTurnContext>,
) -> Result<Vec<ToolExecutionOutcome>, AgentError> {
let futures: Vec<_> = tool_calls let futures: Vec<_> = tool_calls
.iter() .iter()
.map(|tc| self.execute_one_tool(tc)) .map(|tool_call| self.execute_one_tool(tool_call, iteration, turn))
.collect(); .collect();
futures_util::future::join_all(futures).await futures_util::future::join_all(futures)
.await
.into_iter()
.collect()
} }
/// Execute tools sequentially. /// Execute tools sequentially.
async fn execute_tools_sequential(&self, tool_calls: &[ToolCall]) -> Vec<ToolExecutionOutcome> { async fn execute_tools_sequential(
&self,
tool_calls: &[ToolCall],
iteration: u32,
turn: Option<&AgentTurnContext>,
) -> Result<Vec<ToolExecutionOutcome>, AgentError> {
let mut outcomes = Vec::with_capacity(tool_calls.len()); let mut outcomes = Vec::with_capacity(tool_calls.len());
for tool_call in tool_calls { for tool_call in tool_calls {
outcomes.push(self.execute_one_tool(tool_call).await); outcomes.push(self.execute_one_tool(tool_call, iteration, turn).await?);
} }
outcomes Ok(outcomes)
} }
/// Execute a single tool and return the outcome with event tracking. /// Execute a single tool and return the outcome with event tracking.
async fn execute_one_tool(&self, tool_call: &ToolCall) -> ToolExecutionOutcome { async fn execute_one_tool(
&self,
tool_call: &ToolCall,
iteration: u32,
turn: Option<&AgentTurnContext>,
) -> Result<ToolExecutionOutcome, AgentError> {
let start = Instant::now(); let start = Instant::now();
let tool_name = tool_call.name.clone(); let tool_name = tool_call.name.clone();
if let Some(turn) = turn {
turn.emitter
.emit(TurnEvent::ToolStarted {
iteration,
call: tool_call.clone(),
})
.map_err(|error| AgentError::Other(format!("turn event rejected: {error}")))?;
}
// Record ToolCallStart event // Record ToolCallStart event
if let Some(ref observer) = self.observer { if let Some(ref observer) = self.observer {
observer.record_event(&ObserverEvent::ToolCallStart { observer.record_event(&ObserverEvent::ToolCallStart {
@ -800,6 +999,17 @@ impl AgentLoop {
let result = self.execute_tool_internal(tool_call).await; let result = self.execute_tool_internal(tool_call).await;
let duration = start.elapsed(); let duration = start.elapsed();
if let Some(turn) = turn {
turn.emitter
.emit(TurnEvent::ToolFinished {
iteration,
call_id: tool_call.id.clone(),
success: result.success,
preview: Some(tool_result_preview(&truncate_tool_result(&result.output))),
})
.map_err(|error| AgentError::Other(format!("turn event rejected: {error}")))?;
}
// Record ToolCall event // Record ToolCall event
if let Some(ref observer) = self.observer { if let Some(ref observer) = self.observer {
observer.record_event(&ObserverEvent::ToolCall { observer.record_event(&ObserverEvent::ToolCall {
@ -810,7 +1020,7 @@ impl AgentLoop {
} }
// Apply duration // Apply duration
ToolExecutionOutcome { duration, ..result } Ok(ToolExecutionOutcome { duration, ..result })
} }
/// Internal tool execution without event tracking. /// Internal tool execution without event tracking.
@ -851,13 +1061,115 @@ impl AgentLoop {
mod tests { mod tests {
use super::*; use super::*;
use crate::observability::{MultiObserver, Observer}; use crate::observability::{MultiObserver, Observer};
use crate::providers::{ChatCompletionResponse, Usage}; use crate::providers::{
ChatCompletionResponse, FinishReason, ProviderChunk, ProviderStream, Usage,
};
use crate::session::{TurnBlock, TurnController};
use crate::tools::FileReadTool; use crate::tools::FileReadTool;
struct TestObserver { struct TestObserver {
events: std::sync::Mutex<Vec<ObserverEvent>>, events: std::sync::Mutex<Vec<ObserverEvent>>,
} }
struct StreamingTextProvider;
#[async_trait::async_trait]
impl LLMProvider for StreamingTextProvider {
async fn stream(
&self,
_request: ChatCompletionRequest,
) -> Result<ProviderStream, crate::providers::DynProviderError> {
let chunks = vec![
ProviderChunk::Metadata {
id: "response".into(),
model: "streaming-test".into(),
},
ProviderChunk::Reasoning("because ".into()),
ProviderChunk::Reasoning("facts".into()),
ProviderChunk::Text("hello ".into()),
ProviderChunk::Text("world".into()),
ProviderChunk::ProviderState(crate::bus::ProviderReasoningState {
provider: "test".into(),
payload: serde_json::json!({"opaque":"state"}),
}),
ProviderChunk::Usage(Usage {
prompt_tokens: 2,
completion_tokens: 3,
total_tokens: 5,
..Usage::default()
}),
ProviderChunk::Done(FinishReason::Stop),
];
Ok(Box::pin(futures_util::stream::iter(
chunks.into_iter().map(Ok),
)))
}
fn ptype(&self) -> &str {
"test"
}
fn name(&self) -> &str {
"streaming-test"
}
fn model_id(&self) -> &str {
"streaming-test"
}
}
#[tokio::test]
async fn process_streaming_emits_turn_blocks_and_stamps_durable_message() {
let agent = AgentLoop::with_provider(
Arc::new(StreamingTextProvider),
1,
"streaming-test".into(),
PathBuf::from("."),
Vec::new(),
);
let (controller, emitter, _) = TurnController::start("session", "assistant-id");
let initial = controller.snapshot();
let context =
AgentTurnContext::new(initial.id.0.clone(), initial.message_id.clone(), emitter);
let result = agent
.process_streaming(vec![ChatMessage::user("hi")], context)
.await
.unwrap();
assert_eq!(result.final_response.id, "assistant-id");
assert_eq!(result.final_response.content, "hello world");
assert_eq!(
result.final_response.reasoning_content.as_deref(),
Some("because facts")
);
assert_eq!(result.final_response.turn_id, Some(initial.id.0.clone()));
assert_eq!(result.final_response.iteration, Some(0));
assert_eq!(
result
.final_response
.provider_state
.as_ref()
.map(|state| state.provider.as_str()),
Some("test")
);
assert_eq!(
result.usage.as_ref().map(|usage| usage.total_tokens),
Some(5)
);
let snapshot = controller.snapshot();
assert_eq!(snapshot.blocks.len(), 2);
assert!(matches!(
&snapshot.blocks[0],
TurnBlock::Reasoning { text, .. } if text == "because facts"
));
assert!(matches!(
&snapshot.blocks[1],
TurnBlock::Assistant { text, .. } if text == "hello world"
));
}
impl TestObserver { impl TestObserver {
fn new() -> Self { fn new() -> Self {
Self { Self {
@ -899,16 +1211,16 @@ mod tests {
#[async_trait::async_trait] #[async_trait::async_trait]
impl LLMProvider for ToolMediaProvider { impl LLMProvider for ToolMediaProvider {
async fn chat( async fn stream(
&self, &self,
request: ChatCompletionRequest, request: ChatCompletionRequest,
) -> Result<ChatCompletionResponse, Box<dyn std::error::Error + Send + Sync>> { ) -> Result<crate::providers::ProviderStream, crate::providers::DynProviderError> {
let call_number = { let call_number = {
let mut requests = self.requests.lock().unwrap(); let mut requests = self.requests.lock().unwrap();
requests.push(request); requests.push(request);
requests.len() requests.len()
}; };
Ok(ChatCompletionResponse { let response = ChatCompletionResponse {
id: format!("response-{call_number}"), id: format!("response-{call_number}"),
model: "vision-test".to_string(), model: "vision-test".to_string(),
content: if call_number == 1 { content: if call_number == 1 {
@ -917,6 +1229,7 @@ mod tests {
"image seen".to_string() "image seen".to_string()
}, },
reasoning_content: None, reasoning_content: None,
provider_state: None,
tool_calls: if call_number == 1 { tool_calls: if call_number == 1 {
vec![ToolCall { vec![ToolCall {
id: "call-image".to_string(), id: "call-image".to_string(),
@ -934,7 +1247,8 @@ mod tests {
cache_read_input_tokens: None, cache_read_input_tokens: None,
cache_creation_input_tokens: None, cache_creation_input_tokens: None,
}, },
}) };
Ok(crate::providers::provider_stream_for_test(response))
} }
fn ptype(&self) -> &str { fn ptype(&self) -> &str {
@ -971,8 +1285,13 @@ mod tests {
vec!["text".to_string(), "image".to_string()], vec!["text".to_string(), "image".to_string()],
); );
let (controller, emitter, _) = TurnController::start("session", "assistant-id");
let turn = controller.snapshot();
let result = agent let result = agent
.process(vec![ChatMessage::user("inspect the image")]) .process_streaming(
vec![ChatMessage::user("inspect the image")],
AgentTurnContext::new(turn.id.0.clone(), turn.message_id.clone(), emitter),
)
.await .await
.unwrap(); .unwrap();
@ -997,6 +1316,14 @@ mod tests {
.iter() .iter()
.any(|media| media.media_type == "image") .any(|media| media.media_type == "image")
})); }));
assert!(controller.snapshot().blocks.iter().any(|block| matches!(
block,
TurnBlock::Tool {
id,
status: crate::session::ToolStatus::Completed,
..
} if id == "call-image"
)));
} }
#[test] #[test]
@ -1031,6 +1358,7 @@ mod tests {
role: chat_message.role.clone(), role: chat_message.role.clone(),
content, content,
reasoning_content: None, reasoning_content: None,
provider_state: None,
tool_call_id: chat_message.tool_call_id.clone(), tool_call_id: chat_message.tool_call_id.clone(),
name: chat_message.tool_name.clone(), name: chat_message.tool_name.clone(),
tool_calls: chat_message.tool_calls.clone(), tool_calls: chat_message.tool_calls.clone(),

17
src/agent/context_compressor.rs
View File

@ -669,11 +669,11 @@ mod tests {
#[async_trait] #[async_trait]
impl LLMProvider for MockProvider { impl LLMProvider for MockProvider {
async fn chat( async fn stream(
&self, &self,
_request: ChatCompletionRequest, _request: ChatCompletionRequest,
) -> Result<ChatCompletionResponse, Box<dyn std::error::Error + Send + Sync>> { ) -> Result<crate::providers::ProviderStream, crate::providers::DynProviderError> {
panic!("MockProvider.chat() called - not expected in test") panic!("MockProvider.stream() called - not expected in test")
} }
fn ptype(&self) -> &str { fn ptype(&self) -> &str {
@ -699,15 +699,17 @@ mod tests {
#[async_trait] #[async_trait]
impl LLMProvider for MockSummarizer { impl LLMProvider for MockSummarizer {
async fn chat( async fn stream(
&self, &self,
_request: ChatCompletionRequest, _request: ChatCompletionRequest,
) -> Result<ChatCompletionResponse, Box<dyn std::error::Error + Send + Sync>> { ) -> Result<crate::providers::ProviderStream, crate::providers::DynProviderError> {
Ok(ChatCompletionResponse { Ok(crate::providers::provider_stream_for_test(
ChatCompletionResponse {
id: "mock".into(), id: "mock".into(),
model: "mock".into(), model: "mock".into(),
content: "[summarized]".into(), content: "[summarized]".into(),
reasoning_content: None, reasoning_content: None,
provider_state: None,
tool_calls: vec![], tool_calls: vec![],
usage: Usage { usage: Usage {
prompt_tokens: 0, prompt_tokens: 0,
@ -717,7 +719,8 @@ mod tests {
cache_read_input_tokens: None, cache_read_input_tokens: None,
cache_creation_input_tokens: None, cache_creation_input_tokens: None,
}, },
}) },
))
} }
fn ptype(&self) -> &str { fn ptype(&self) -> &str {

2
src/agent/mod.rs
View File

@ -3,6 +3,7 @@ pub mod context_compressor;
pub mod media_handler; pub mod media_handler;
pub mod sub_agent; pub mod sub_agent;
pub mod system_prompt; pub mod system_prompt;
pub mod turn_event;
pub use agent_loop::{AgentError, AgentLoop, AgentProcessResult}; pub use agent_loop::{AgentError, AgentLoop, AgentProcessResult};
pub use context_compressor::{ContextCompressor, estimate_tokens}; pub use context_compressor::{ContextCompressor, estimate_tokens};
@ -14,3 +15,4 @@ pub use system_prompt::{
PromptContext, PromptSection, SystemPromptBuilder, build_sub_agent_system_prompt, PromptContext, PromptSection, SystemPromptBuilder, build_sub_agent_system_prompt,
build_system_prompt, build_system_prompt,
}; };
pub use turn_event::{AgentTurnContext, TurnEmitError, TurnEmitter, TurnEvent};

107
src/agent/turn_event.rs Normal file
View File

@ -0,0 +1,107 @@
use std::sync::{Arc, Mutex};
use thiserror::Error;
use crate::providers::ToolCall;
/// Presentation facts emitted while AgentLoop processes one model turn.
///
/// Events contain no persistence or channel-delivery decisions. Session owns
/// the lifecycle around these facts and reduces them into authoritative state.
#[derive(Debug, Clone)]
pub enum TurnEvent {
ReasoningDelta {
iteration: u32,
delta: String,
},
TextDelta {
iteration: u32,
delta: String,
},
TextSegmentFinished {
iteration: u32,
},
ToolStarted {
iteration: u32,
call: ToolCall,
},
ToolFinished {
iteration: u32,
call_id: String,
success: bool,
preview: Option<String>,
},
}
#[derive(Debug, Clone, PartialEq, Eq, Error)]
pub enum TurnEmitError {
#[error("turn is no longer active")]
Inactive,
#[error("tool call {0} already exists in this turn")]
DuplicateTool(String),
#[error("tool call {0} does not exist in this turn")]
UnknownTool(String),
}
type EmitFn = dyn Fn(TurnEvent) -> Result<(), TurnEmitError> + Send + Sync;
/// Cheap cloneable handle used by AgentLoop to report presentation facts.
#[derive(Clone)]
pub struct TurnEmitter {
emit: Arc<EmitFn>,
enabled: Arc<Mutex<bool>>,
}
/// Session-owned identity and emitter for one AgentLoop execution.
#[derive(Clone)]
pub struct AgentTurnContext {
pub turn_id: String,
pub message_id: String,
pub emitter: TurnEmitter,
}
impl AgentTurnContext {
pub fn new(
turn_id: impl Into<String>,
message_id: impl Into<String>,
emitter: TurnEmitter,
) -> Self {
Self {
turn_id: turn_id.into(),
message_id: message_id.into(),
emitter,
}
}
}
impl TurnEmitter {
pub(crate) fn new<F>(emit: F) -> Self
where
F: Fn(TurnEvent) -> Result<(), TurnEmitError> + Send + Sync + 'static,
{
Self {
emit: Arc::new(emit),
enabled: Arc::new(Mutex::new(true)),
}
}
pub fn emit(&self, event: TurnEvent) -> Result<(), TurnEmitError> {
let enabled = self
.enabled
.lock()
.unwrap_or_else(|poisoned| poisoned.into_inner());
if !*enabled {
return Ok(());
}
(self.emit)(event)
}
/// Stop forwarding new presentation facts without changing durable or
/// terminal Turn status. Session uses this before invalidating a worker.
pub fn deactivate(&self) {
*self
.enabled
.lock()
.unwrap_or_else(|poisoned| poisoned.into_inner()) = false;
}
}

62
src/bus/dispatcher.rs
View File

@ -7,6 +7,7 @@ use tokio::sync::mpsc;
use crate::bus::{MessageBus, OutboundMessage}; use crate::bus::{MessageBus, OutboundMessage};
use crate::channels::ChannelManager; use crate::channels::ChannelManager;
use crate::channels::base::{Channel, ChannelError}; use crate::channels::base::{Channel, ChannelError};
use crate::delivery::ConversationWriteLocks;
use crate::task_supervisor::TaskSupervisor; use crate::task_supervisor::TaskSupervisor;
const LANE_CAPACITY: usize = 64; const LANE_CAPACITY: usize = 64;
@ -20,6 +21,7 @@ pub struct OutboundDispatcher {
bus: Arc<MessageBus>, bus: Arc<MessageBus>,
channel_manager: ChannelManager, channel_manager: ChannelManager,
task_supervisor: TaskSupervisor, task_supervisor: TaskSupervisor,
write_locks: ConversationWriteLocks,
} }
impl OutboundDispatcher { impl OutboundDispatcher {
@ -27,11 +29,13 @@ impl OutboundDispatcher {
bus: Arc<MessageBus>, bus: Arc<MessageBus>,
channel_manager: ChannelManager, channel_manager: ChannelManager,
task_supervisor: TaskSupervisor, task_supervisor: TaskSupervisor,
write_locks: ConversationWriteLocks,
) -> Self { ) -> Self {
Self { Self {
bus, bus,
channel_manager, channel_manager,
task_supervisor, task_supervisor,
write_locks,
} }
} }
@ -120,6 +124,7 @@ impl OutboundDispatcher {
channel_name: String, channel_name: String,
chat_id: String, chat_id: String,
) -> bool { ) -> bool {
let target_lock = self.write_locks.for_target(&channel_name, &chat_id);
self.task_supervisor.spawn( self.task_supervisor.spawn(
format!("outbound-lane:{channel_name}:{chat_id}"), format!("outbound-lane:{channel_name}:{chat_id}"),
async move { async move {
@ -128,7 +133,7 @@ impl OutboundDispatcher {
Ok(Some(msg)) => msg, Ok(Some(msg)) => msg,
Ok(None) | Err(_) => break, Ok(None) | Err(_) => break,
}; };
let result = Self::send_with_retry(&*channel, &msg).await; let result = Self::send_with_retry(&*channel, &msg, &target_lock).await;
if let Err(error) = &result { if let Err(error) = &result {
tracing::error!( tracing::error!(
channel = %channel_name, channel = %channel_name,
@ -146,7 +151,9 @@ impl OutboundDispatcher {
async fn send_with_retry( async fn send_with_retry(
channel: &dyn Channel, channel: &dyn Channel,
msg: &OutboundMessage, msg: &OutboundMessage,
target_lock: &tokio::sync::Mutex<()>,
) -> Result<(), ChannelError> { ) -> Result<(), ChannelError> {
let _guard = target_lock.lock().await;
const DELAYS: &[u64] = &[1, 2, 4]; const DELAYS: &[u64] = &[1, 2, 4];
for (attempt, &delay) in DELAYS.iter().enumerate() { for (attempt, &delay) in DELAYS.iter().enumerate() {
@ -263,7 +270,12 @@ mod tests {
manager.register_channel("recording", channel.clone()).await; manager.register_channel("recording", channel.clone()).await;
let supervisor = TaskSupervisor::new(); let supervisor = TaskSupervisor::new();
let dispatcher = OutboundDispatcher::new(bus.clone(), manager, supervisor.clone()); let dispatcher = OutboundDispatcher::new(
bus.clone(),
manager,
supervisor.clone(),
ConversationWriteLocks::default(),
);
let task = tokio::spawn(async move { dispatcher.run().await }); let task = tokio::spawn(async move { dispatcher.run().await });
bus.publish_outbound(outbound("slow", "slow-1")) bus.publish_outbound(outbound("slow", "slow-1"))
.await .await
@ -301,7 +313,12 @@ mod tests {
bus.clone(), bus.clone(),
); );
let supervisor = TaskSupervisor::new(); let supervisor = TaskSupervisor::new();
let dispatcher = OutboundDispatcher::new(bus.clone(), manager, supervisor.clone()); let dispatcher = OutboundDispatcher::new(
bus.clone(),
manager,
supervisor.clone(),
ConversationWriteLocks::default(),
);
let task = tokio::spawn(async move { dispatcher.run().await }); let task = tokio::spawn(async move { dispatcher.run().await });
let mut message = outbound("missing", "not delivered"); let mut message = outbound("missing", "not delivered");
@ -326,7 +343,12 @@ mod tests {
}); });
manager.register_channel("recording", channel.clone()).await; manager.register_channel("recording", channel.clone()).await;
let supervisor = TaskSupervisor::new(); let supervisor = TaskSupervisor::new();
let dispatcher = OutboundDispatcher::new(bus.clone(), manager, supervisor.clone()); let dispatcher = OutboundDispatcher::new(
bus.clone(),
manager,
supervisor.clone(),
ConversationWriteLocks::default(),
);
let task = tokio::spawn(async move { dispatcher.run().await }); let task = tokio::spawn(async move { dispatcher.run().await });
bus.deliver_outbound(outbound("confirmed", "delivered")) bus.deliver_outbound(outbound("confirmed", "delivered"))
@ -344,11 +366,41 @@ mod tests {
attempts: AtomicUsize::new(0), attempts: AtomicUsize::new(0),
}; };
let error = OutboundDispatcher::send_with_retry(&channel, &outbound("invalid", "message")) let target_lock = tokio::sync::Mutex::new(());
let error = OutboundDispatcher::send_with_retry(
&channel,
&outbound("invalid", "message"),
&target_lock,
)
.await .await
.unwrap_err(); .unwrap_err();
assert!(matches!(error, ChannelError::Other(_))); assert!(matches!(error, ChannelError::Other(_)));
assert_eq!(channel.attempts.load(Ordering::SeqCst), 1); assert_eq!(channel.attempts.load(Ordering::SeqCst), 1);
} }
#[tokio::test]
async fn shared_write_lock_orders_dispatcher_with_live_turn_writes() {
let channel = RecordingChannel {
sent: Mutex::new(Vec::new()),
notify: Notify::new(),
};
let write_locks = ConversationWriteLocks::default();
let target_lock = write_locks.for_target("recording", "same-chat");
let live_write = target_lock.lock().await;
let message = outbound("same-chat", "after-live-update");
let send = OutboundDispatcher::send_with_retry(&channel, &message, &target_lock);
tokio::pin!(send);
assert!(
tokio::time::timeout(Duration::from_millis(10), &mut send)
.await
.is_err()
);
assert!(channel.sent.lock().await.is_empty());
drop(live_write);
send.await.unwrap();
assert_eq!(channel.sent.lock().await.as_slice(), &["after-live-update"]);
}
} }

105
src/bus/message.rs
View File

@ -3,6 +3,52 @@ use std::collections::HashMap;
use crate::providers::ToolCall; use crate::providers::ToolCall;
/// Provider-private state required to faithfully replay an assistant message.
///
/// This is durable conversation data, but it is never presentation data. UI and
/// channel projections must not serialize it to end users.
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub struct ProviderReasoningState {
pub provider: String,
pub payload: serde_json::Value,
}
impl ProviderReasoningState {
/// Decode persisted provider state without making conversation history
/// unreadable when an old or damaged payload is encountered.
pub fn from_json_lossy(value: &str) -> Option<Self> {
serde_json::from_str(value).ok()
}
}
/// Describes whether a persisted message represents a complete model result.
#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum CompletionStatus {
#[default]
Completed,
Cancelled,
Interrupted,
}
impl CompletionStatus {
pub fn as_str(self) -> &'static str {
match self {
Self::Completed => "completed",
Self::Cancelled => "cancelled",
Self::Interrupted => "interrupted",
}
}
pub fn from_storage(value: &str) -> Self {
match value {
"cancelled" => Self::Cancelled,
"interrupted" => Self::Interrupted,
_ => Self::Completed,
}
}
}
// ============================================================================ // ============================================================================
// ContentBlock - Multimodal content representation (OpenAI-style) // ContentBlock - Multimodal content representation (OpenAI-style)
// ============================================================================ // ============================================================================
@ -85,6 +131,15 @@ pub struct ChatMessage {
pub role: String, pub role: String,
pub content: String, pub content: String,
pub reasoning_content: Option<String>, pub reasoning_content: Option<String>,
/// Opaque state used only when replaying history to the same provider.
#[serde(skip_serializing_if = "Option::is_none")]
pub provider_state: Option<ProviderReasoningState>,
#[serde(skip_serializing_if = "Option::is_none")]
pub turn_id: Option<String>,
#[serde(skip_serializing_if = "Option::is_none")]
pub iteration: Option<u32>,
#[serde(default)]
pub completion_status: CompletionStatus,
pub media_refs: Vec<MediaRef>, pub media_refs: Vec<MediaRef>,
pub timestamp: i64, pub timestamp: i64,
#[serde(skip_serializing_if = "Option::is_none")] #[serde(skip_serializing_if = "Option::is_none")]
@ -124,6 +179,10 @@ impl ChatMessage {
role: "user".to_string(), role: "user".to_string(),
content: content.into(), content: content.into(),
reasoning_content: None, reasoning_content: None,
provider_state: None,
turn_id: None,
iteration: None,
completion_status: CompletionStatus::Completed,
media_refs: Vec::new(), media_refs: Vec::new(),
timestamp: current_timestamp(), timestamp: current_timestamp(),
tool_call_id: None, tool_call_id: None,
@ -139,6 +198,10 @@ impl ChatMessage {
role: "user".to_string(), role: "user".to_string(),
content: content.into(), content: content.into(),
reasoning_content: None, reasoning_content: None,
provider_state: None,
turn_id: None,
iteration: None,
completion_status: CompletionStatus::Completed,
media_refs, media_refs,
timestamp: current_timestamp(), timestamp: current_timestamp(),
tool_call_id: None, tool_call_id: None,
@ -154,6 +217,10 @@ impl ChatMessage {
role: "assistant".to_string(), role: "assistant".to_string(),
content: content.into(), content: content.into(),
reasoning_content: None, reasoning_content: None,
provider_state: None,
turn_id: None,
iteration: None,
completion_status: CompletionStatus::Completed,
media_refs: Vec::new(), media_refs: Vec::new(),
timestamp: current_timestamp(), timestamp: current_timestamp(),
tool_call_id: None, tool_call_id: None,
@ -172,6 +239,10 @@ impl ChatMessage {
role: "assistant".to_string(), role: "assistant".to_string(),
content: content.into(), content: content.into(),
reasoning_content: None, reasoning_content: None,
provider_state: None,
turn_id: None,
iteration: None,
completion_status: CompletionStatus::Completed,
media_refs: Vec::new(), media_refs: Vec::new(),
timestamp: current_timestamp(), timestamp: current_timestamp(),
tool_call_id: None, tool_call_id: None,
@ -187,6 +258,10 @@ impl ChatMessage {
role: "assistant".to_string(), role: "assistant".to_string(),
content: content.into(), content: content.into(),
reasoning_content: None, reasoning_content: None,
provider_state: None,
turn_id: None,
iteration: None,
completion_status: CompletionStatus::Completed,
media_refs: Vec::new(), media_refs: Vec::new(),
timestamp: current_timestamp(), timestamp: current_timestamp(),
tool_call_id: None, tool_call_id: None,
@ -202,6 +277,10 @@ impl ChatMessage {
role: "system".to_string(), role: "system".to_string(),
content: content.into(), content: content.into(),
reasoning_content: None, reasoning_content: None,
provider_state: None,
turn_id: None,
iteration: None,
completion_status: CompletionStatus::Completed,
media_refs: Vec::new(), media_refs: Vec::new(),
timestamp: current_timestamp(), timestamp: current_timestamp(),
tool_call_id: None, tool_call_id: None,
@ -230,6 +309,10 @@ impl ChatMessage {
role: "tool".to_string(), role: "tool".to_string(),
content: content.into(), content: content.into(),
reasoning_content: None, reasoning_content: None,
provider_state: None,
turn_id: None,
iteration: None,
completion_status: CompletionStatus::Completed,
media_refs, media_refs,
timestamp: current_timestamp(), timestamp: current_timestamp(),
tool_call_id: Some(tool_call_id.into()), tool_call_id: Some(tool_call_id.into()),
@ -245,6 +328,10 @@ impl ChatMessage {
role: "user".to_string(), role: "user".to_string(),
content: content.into(), content: content.into(),
reasoning_content: None, reasoning_content: None,
provider_state: None,
turn_id: None,
iteration: None,
completion_status: CompletionStatus::Completed,
media_refs: Vec::new(), media_refs: Vec::new(),
timestamp: current_timestamp(), timestamp: current_timestamp(),
tool_call_id: None, tool_call_id: None,
@ -255,6 +342,24 @@ impl ChatMessage {
} }
} }
#[cfg(test)]
mod conversation_message_tests {
use super::*;
#[test]
fn damaged_provider_state_is_ignored() {
assert!(ProviderReasoningState::from_json_lossy("not-json").is_none());
}
#[test]
fn unknown_completion_status_is_backward_compatible() {
assert_eq!(
CompletionStatus::from_storage("future-status"),
CompletionStatus::Completed
);
}
}
// ============================================================================ // ============================================================================
// InboundMessage - Message from Channel to Bus (user input) // InboundMessage - Message from Channel to Bus (user input)
// ============================================================================ // ============================================================================

4
src/bus/mod.rs
View File

@ -3,8 +3,8 @@ pub mod message;
pub use dispatcher::OutboundDispatcher; pub use dispatcher::OutboundDispatcher;
pub use message::{ pub use message::{
ChatMessage, ContentBlock, ControlMessage, InboundMessage, MediaItem, MediaRef, MessageSource, ChatMessage, CompletionStatus, ContentBlock, ControlMessage, InboundMessage, MediaItem,
OutboundMessage, SourceKind, MediaRef, MessageSource, OutboundMessage, ProviderReasoningState, SourceKind,
}; };
use std::sync::Arc; use std::sync::Arc;

48
src/channels/base.rs
View File

@ -1,7 +1,33 @@
use async_trait::async_trait; use async_trait::async_trait;
use std::collections::HashMap;
use std::sync::Arc; use std::sync::Arc;
use std::time::Duration;
use crate::bus::{BusError, InboundMessage, MessageBus, OutboundMessage}; use crate::bus::{BusError, InboundMessage, MessageBus, OutboundMessage};
use crate::delivery::PresentationPolicy;
use crate::session::TurnSnapshot;
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum LivePolicy {
FinalOnly,
Snapshot { min_interval: Duration },
}
#[derive(Debug, Clone)]
pub struct TurnTarget {
pub channel: String,
pub chat_id: String,
pub session_id: String,
pub reply_to: Option<String>,
pub metadata: HashMap<String, String>,
}
#[async_trait]
pub trait TurnSink: Send {
async fn update(&mut self, snapshot: &TurnSnapshot) -> Result<(), ChannelError>;
async fn finish(&mut self, snapshot: &TurnSnapshot) -> Result<(), ChannelError>;
async fn abort(&mut self, snapshot: &TurnSnapshot) -> Result<(), ChannelError>;
}
#[derive(Debug)] #[derive(Debug)]
pub enum ChannelError { pub enum ChannelError {
@ -49,16 +75,24 @@ pub trait Channel: Send + Sync + 'static {
/// Stop the channel /// Stop the channel
async fn stop(&self) -> Result<(), ChannelError>; async fn stop(&self) -> Result<(), ChannelError>;
fn live_policy(&self) -> LivePolicy {
LivePolicy::FinalOnly
}
fn presentation_policy(&self) -> PresentationPolicy {
PresentationPolicy::external(matches!(self.live_policy(), LivePolicy::Snapshot { .. }))
}
async fn open_turn(&self, _target: TurnTarget) -> Result<Box<dyn TurnSink>, ChannelError> {
Err(ChannelError::Other(format!(
"channel {} does not support turn delivery",
self.name()
)))
}
/// Send a message to the channel (called by OutboundDispatcher) /// Send a message to the channel (called by OutboundDispatcher)
async fn send(&self, msg: OutboundMessage) -> Result<(), ChannelError>; async fn send(&self, msg: OutboundMessage) -> Result<(), ChannelError>;
/// Send a streaming delta (optional, for channels that support it)
async fn send_delta(&self, chat_id: &str, delta: &str) -> Result<(), ChannelError> {
let _ = chat_id;
let _ = delta;
Ok(())
}
/// Check if a sender is allowed to use this channel /// Check if a sender is allowed to use this channel
fn is_allowed(&self, _sender_id: &str) -> bool { fn is_allowed(&self, _sender_id: &str) -> bool {
true true

124
src/channels/cli_chat.rs
View File

@ -1,6 +1,7 @@
use async_trait::async_trait; use async_trait::async_trait;
use std::collections::HashMap; use std::collections::HashMap;
use std::sync::Arc; use std::sync::Arc;
use std::time::Duration;
use tokio::sync::{Mutex, mpsc}; use tokio::sync::{Mutex, mpsc};
use crate::bus::{ControlMessage, InboundMessage, MessageBus, OutboundMessage}; use crate::bus::{ControlMessage, InboundMessage, MessageBus, OutboundMessage};
@ -8,9 +9,10 @@ use crate::gateway::uploads::UploadRegistry;
use crate::protocol::{ use crate::protocol::{
HistoryMessage, MessageAttachment, SlashCommandInfo, WsInbound, WsOutbound, parse_inbound, HistoryMessage, MessageAttachment, SlashCommandInfo, WsInbound, WsOutbound, parse_inbound,
}; };
use crate::session::TurnSnapshot;
use crate::session::{SessionCommand, SessionEvent, UnifiedSessionId}; use crate::session::{SessionCommand, SessionEvent, UnifiedSessionId};
use super::base::{Channel, ChannelError}; use super::base::{Channel, ChannelError, LivePolicy, TurnSink, TurnTarget};
// ============================================================================ // ============================================================================
// Client - Connected CLI client // Client - Connected CLI client
@ -34,7 +36,7 @@ impl Client {
pub struct CliChatChannel { pub struct CliChatChannel {
bus: std::sync::Mutex<Option<Arc<MessageBus>>>, bus: std::sync::Mutex<Option<Arc<MessageBus>>>,
clients: Mutex<HashMap<String, Arc<Client>>>, clients: Arc<Mutex<HashMap<String, Arc<Client>>>>,
uploads: UploadRegistry, uploads: UploadRegistry,
} }
@ -52,7 +54,7 @@ impl CliChatChannel {
pub fn with_upload_registry(uploads: UploadRegistry) -> Self { pub fn with_upload_registry(uploads: UploadRegistry) -> Self {
Self { Self {
bus: std::sync::Mutex::new(None), bus: std::sync::Mutex::new(None),
clients: Mutex::new(HashMap::new()), clients: Arc::new(Mutex::new(HashMap::new())),
uploads, uploads,
} }
} }
@ -451,6 +453,8 @@ impl CliChatChannel {
seq: message.seq, seq: message.seq,
role: message.role, role: message.role,
content: message.content, content: message.content,
reasoning_content: message.reasoning_content,
completion_status: message.completion_status,
created_at: message.created_at, created_at: message.created_at,
tool_call_id: message.tool_call_id, tool_call_id: message.tool_call_id,
tool_name: message.tool_name, tool_name: message.tool_name,
@ -803,6 +807,54 @@ impl CliChatChannel {
} }
} }
struct CliChatTurnSink {
clients: Arc<Mutex<HashMap<String, Arc<Client>>>>,
chat_id: String,
}
impl CliChatTurnSink {
async fn publish(&self, snapshot: &TurnSnapshot) {
let client = self.clients.lock().await.get(&self.chat_id).cloned();
let Some(client) = client else {
return;
};
if client
.sender
.send(WsOutbound::TurnUpdated {
snapshot: snapshot.clone(),
})
.await
.is_err()
{
let mut clients = self.clients.lock().await;
if clients
.get(&self.chat_id)
.is_some_and(|registered| Arc::ptr_eq(registered, &client))
{
clients.remove(&self.chat_id);
}
}
}
}
#[async_trait]
impl TurnSink for CliChatTurnSink {
async fn update(&mut self, snapshot: &TurnSnapshot) -> Result<(), ChannelError> {
self.publish(snapshot).await;
Ok(())
}
async fn finish(&mut self, snapshot: &TurnSnapshot) -> Result<(), ChannelError> {
self.publish(snapshot).await;
Ok(())
}
async fn abort(&mut self, snapshot: &TurnSnapshot) -> Result<(), ChannelError> {
self.publish(snapshot).await;
Ok(())
}
}
#[async_trait] #[async_trait]
impl Channel for CliChatChannel { impl Channel for CliChatChannel {
fn name(&self) -> &str { fn name(&self) -> &str {
@ -824,6 +876,23 @@ impl Channel for CliChatChannel {
Ok(()) Ok(())
} }
fn live_policy(&self) -> LivePolicy {
LivePolicy::Snapshot {
min_interval: Duration::from_millis(33),
}
}
fn presentation_policy(&self) -> crate::delivery::PresentationPolicy {
crate::delivery::PresentationPolicy::interactive()
}
async fn open_turn(&self, target: TurnTarget) -> Result<Box<dyn TurnSink>, ChannelError> {
Ok(Box::new(CliChatTurnSink {
clients: self.clients.clone(),
chat_id: target.chat_id,
}))
}
async fn send(&self, msg: OutboundMessage) -> Result<(), ChannelError> { async fn send(&self, msg: OutboundMessage) -> Result<(), ChannelError> {
let client = self.clients.lock().await.get(&msg.chat_id).cloned(); let client = self.clients.lock().await.get(&msg.chat_id).cloned();
let Some(client) = client else { let Some(client) = client else {
@ -1026,4 +1095,53 @@ mod tests {
other => panic!("unexpected outbound: {other:?}"), other => panic!("unexpected outbound: {other:?}"),
} }
} }
#[tokio::test]
async fn turn_sink_sends_the_same_snapshot_shape_for_running_and_terminal_states() {
let channel = CliChatChannel::new();
let (sender, mut receiver) = mpsc::channel(2);
let client = Arc::new(Client {
sender,
chat_id: "client".into(),
current_session_id: Mutex::new(None),
});
channel.clients.lock().await.insert("client".into(), client);
let mut sink = channel
.open_turn(TurnTarget {
channel: "cli_chat".into(),
chat_id: "client".into(),
session_id: "cli_chat:client:dialog".into(),
reply_to: None,
metadata: HashMap::new(),
})
.await
.unwrap();
let (controller, emitter, _) =
crate::session::TurnController::start("cli_chat:client:dialog", "message");
emitter
.emit(crate::agent::TurnEvent::TextDelta {
iteration: 0,
delta: "stream".into(),
})
.unwrap();
let running = controller.snapshot();
sink.update(&running).await.unwrap();
controller.complete(None);
let completed = controller.snapshot();
sink.finish(&completed).await.unwrap();
match receiver.recv().await.unwrap() {
WsOutbound::TurnUpdated { snapshot } => {
assert_eq!(snapshot.status, crate::session::TurnStatus::Running);
}
other => panic!("unexpected outbound: {other:?}"),
}
match receiver.recv().await.unwrap() {
WsOutbound::TurnUpdated { snapshot } => {
assert_eq!(snapshot.status, crate::session::TurnStatus::Completed);
assert!(snapshot.revision > running.revision);
}
other => panic!("unexpected outbound: {other:?}"),
}
}
} }

494
src/channels/feishu.rs
View File

@ -6,15 +6,15 @@ use std::time::{Duration, Instant};
use async_trait::async_trait; use async_trait::async_trait;
use futures_util::{SinkExt, StreamExt}; use futures_util::{SinkExt, StreamExt};
use prost::{Message as ProstMessage, bytes::Bytes}; use prost::{Message as ProstMessage, bytes::Bytes};
use regex::Regex;
use serde::Deserialize; use serde::Deserialize;
use tokio::sync::{Mutex, RwLock}; use tokio::sync::{Mutex, RwLock};
use tokio::task::JoinHandle; use tokio::task::JoinHandle;
use tokio_util::sync::CancellationToken; use tokio_util::sync::CancellationToken;
use crate::bus::{MediaItem, MessageBus, OutboundMessage}; use crate::bus::{MediaItem, MessageBus, OutboundMessage};
use crate::channels::base::{Channel, ChannelError}; use crate::channels::base::{Channel, ChannelError, LivePolicy, TurnSink, TurnTarget};
use crate::config::FeishuChannelConfig; use crate::config::FeishuChannelConfig;
use crate::session::{ToolStatus, TurnBlock, TurnSnapshot, TurnStatus};
const FEISHU_API_BASE: &str = "https://open.feishu.cn/open-apis"; const FEISHU_API_BASE: &str = "https://open.feishu.cn/open-apis";
const FEISHU_WS_BASE: &str = "https://open.feishu.cn"; const FEISHU_WS_BASE: &str = "https://open.feishu.cn";
@ -1807,14 +1807,6 @@ async fn resolve_unique_path(dir: &Path, filename: &str) -> std::path::PathBuf {
} }
impl FeishuChannel { impl FeishuChannel {
fn strip_thinking_tags(content: &str) -> String {
use std::sync::LazyLock;
static THINK_RE: LazyLock<Regex> =
LazyLock::new(|| Regex::new(r"(?s)<think>.*?</think>").unwrap());
let stripped = THINK_RE.replace_all(content, "");
stripped.trim().to_string()
}
/// Build a Card JSON 2.0 interactive card with a single markdown element. /// Build a Card JSON 2.0 interactive card with a single markdown element.
fn build_card_content(markdown: &str) -> String { fn build_card_content(markdown: &str) -> String {
serde_json::json!({ serde_json::json!({
@ -1849,7 +1841,7 @@ impl FeishuChannel {
break; break;
} }
let end = start + Self::CARD_MARKDOWN_MAX_BYTES; let end = text.floor_char_boundary(start + Self::CARD_MARKDOWN_MAX_BYTES);
let search_region = &text[start..end]; let search_region = &text[start..end];
let split_at = search_region let split_at = search_region
.rfind('\n') .rfind('\n')
@ -1886,7 +1878,7 @@ impl FeishuChannel {
receive_id: &str, receive_id: &str,
receive_id_type: &str, receive_id_type: &str,
card_content: &str, card_content: &str,
) -> Result<(), ChannelError> { ) -> Result<String, ChannelError> {
let token = self.get_tenant_access_token().await?; let token = self.get_tenant_access_token().await?;
let resp = self let resp = self
@ -1912,6 +1904,12 @@ impl FeishuChannel {
struct SendResp { struct SendResp {
code: i32, code: i32,
msg: String, msg: String,
data: Option<SendData>,
}
#[derive(Deserialize)]
struct SendData {
message_id: String,
} }
let send_resp: SendResp = resp.json().await.map_err(|e| { let send_resp: SendResp = resp.json().await.map_err(|e| {
@ -1925,10 +1923,244 @@ impl FeishuChannel {
))); )));
} }
send_resp
.data
.map(|data| data.message_id)
.filter(|message_id| !message_id.is_empty())
.ok_or_else(|| ChannelError::Other("Feishu send response has no message_id".into()))
}
async fn update_interactive_card(
&self,
message_id: &str,
card_content: &str,
) -> Result<(), ChannelError> {
let token = self.get_tenant_access_token().await?;
let card: serde_json::Value = serde_json::from_str(card_content)
.map_err(|error| ChannelError::Other(format!("Invalid card JSON: {error}")))?;
let response = self
.http_client
.patch(format!("{}/im/v1/messages/{}", FEISHU_API_BASE, message_id))
.header("Content-Type", "application/json")
.header("Authorization", format!("Bearer {}", token))
.json(&serde_json::json!({ "card": card }))
.send()
.await
.map_err(|error| {
ChannelError::ConnectionError(format!("Update card HTTP error: {error}"))
})?;
#[derive(Deserialize)]
struct UpdateResp {
code: i32,
msg: String,
}
let result: UpdateResp = response.json().await.map_err(|error| {
ChannelError::Other(format!("Parse update card response error: {error}"))
})?;
if result.code != 0 {
return Err(ChannelError::Other(format!(
"Update card failed: code={} msg={}",
result.code, result.msg
)));
}
Ok(()) Ok(())
} }
} }
#[async_trait]
trait FeishuTurnApi: Send {
async fn create_card(&mut self, markdown: &str) -> Result<String, ChannelError>;
async fn update_card(&mut self, message_id: &str, markdown: &str) -> Result<(), ChannelError>;
async fn cleanup(&mut self);
}
struct FeishuTurnBackend {
channel: FeishuChannel,
receive_id: String,
receive_id_type: &'static str,
metadata: HashMap<String, String>,
}
#[async_trait]
impl FeishuTurnApi for FeishuTurnBackend {
async fn create_card(&mut self, markdown: &str) -> Result<String, ChannelError> {
let card = FeishuChannel::build_card_content(markdown);
self.channel
.send_interactive_card(&self.receive_id, self.receive_id_type, &card)
.await
}
async fn update_card(&mut self, message_id: &str, markdown: &str) -> Result<(), ChannelError> {
let card = FeishuChannel::build_card_content(markdown);
self.channel
.update_interactive_card(message_id, &card)
.await
}
async fn cleanup(&mut self) {
self.channel
.remove_reaction_from_metadata(&self.metadata)
.await;
}
}
struct FeishuTurnSink {
api: Box<dyn FeishuTurnApi>,
message_id: Option<String>,
cleaned_up: bool,
}
impl FeishuTurnSink {
fn new(api: Box<dyn FeishuTurnApi>) -> Self {
Self {
api,
message_id: None,
cleaned_up: false,
}
}
async fn cleanup(&mut self) {
if !self.cleaned_up {
self.api.cleanup().await;
self.cleaned_up = true;
}
}
async fn send_chunks(&mut self, chunks: &[String]) -> Result<(), ChannelError> {
for chunk in chunks {
self.api.create_card(chunk).await?;
}
Ok(())
}
async fn finish_snapshot(&mut self, snapshot: &TurnSnapshot) -> Result<(), ChannelError> {
let markdown = render_feishu_turn(snapshot);
let chunks = if markdown.is_empty() {
Vec::new()
} else {
FeishuChannel::split_markdown_chunks(&markdown)
};
let result = if chunks.is_empty() {
Ok(())
} else if let Some(message_id) = self.message_id.clone() {
match self.api.update_card(&message_id, &chunks[0]).await {
Ok(()) => self.send_chunks(&chunks[1..]).await,
Err(error) => {
tracing::warn!(error = %error, "Final Feishu card update failed; sending complete fallback");
self.send_chunks(&chunks).await
}
}
} else {
self.send_chunks(&chunks).await
};
self.cleanup().await;
result
}
}
#[async_trait]
impl TurnSink for FeishuTurnSink {
async fn update(&mut self, snapshot: &TurnSnapshot) -> Result<(), ChannelError> {
let markdown = render_feishu_turn(snapshot);
if markdown.is_empty() {
return Ok(());
}
let live_markdown = truncate_feishu_live_markdown(&markdown);
if let Some(message_id) = self.message_id.clone() {
self.api.update_card(&message_id, &live_markdown).await
} else {
self.message_id = Some(self.api.create_card(&live_markdown).await?);
Ok(())
}
}
async fn finish(&mut self, snapshot: &TurnSnapshot) -> Result<(), ChannelError> {
self.finish_snapshot(snapshot).await
}
async fn abort(&mut self, snapshot: &TurnSnapshot) -> Result<(), ChannelError> {
self.finish_snapshot(snapshot).await
}
}
fn render_feishu_turn(snapshot: &TurnSnapshot) -> String {
let mut sections = Vec::new();
for block in &snapshot.blocks {
match block {
TurnBlock::Reasoning { text, .. } if !text.trim().is_empty() => {
sections.push(format!(
"> **思考过程**\n> {}",
text.trim().replace('\n', "\n> ")
));
}
TurnBlock::Assistant { text, .. } if !text.trim().is_empty() => {
sections.push(text.trim().to_string());
}
TurnBlock::Tool {
name,
status,
preview,
..
} => {
let status = match status {
ToolStatus::Running => "执行中",
ToolStatus::Completed => "已完成",
ToolStatus::Failed => "失败",
};
let mut section = format!("> 🔧 **{name}** · {status}");
if let Some(preview) = preview.as_deref().filter(|value| !value.trim().is_empty()) {
section.push_str("\n> ");
section.push_str(&preview.trim().replace('\n', "\n> "));
}
sections.push(section);
}
_ => {}
}
}
if sections.is_empty() {
if snapshot.status == TurnStatus::Failed {
sections.push(format!(
"⚠️ 回复失败：{}",
snapshot.error.as_deref().unwrap_or("未知错误")
));
} else if snapshot.status == TurnStatus::Cancelled {
sections.push("已停止生成。".to_string());
} else {
return String::new();
}
}
let status = match snapshot.status {
TurnStatus::Running => Some(match snapshot.phase {
crate::session::TurnPhase::Queued => "排队中",
crate::session::TurnPhase::Reasoning => "思考中",
crate::session::TurnPhase::Responding => "生成中",
crate::session::TurnPhase::Acting => "调用工具中",
crate::session::TurnPhase::Finalizing => "收尾中",
}),
TurnStatus::Cancelled => Some("已停止"),
TurnStatus::Failed => Some("失败"),
TurnStatus::Completed => None,
};
if let Some(status) = status {
sections.push(format!("_{status}_"));
}
sections.join("\n\n")
}
fn truncate_feishu_live_markdown(markdown: &str) -> String {
if markdown.len() <= FeishuChannel::CARD_MARKDOWN_MAX_BYTES {
return markdown.to_string();
}
const SUFFIX: &str = "\n\n_内容仍在生成，已暂时截断…_";
let limit = FeishuChannel::CARD_MARKDOWN_MAX_BYTES.saturating_sub(SUFFIX.len());
let boundary = markdown.floor_char_boundary(limit);
format!("{}{SUFFIX}", &markdown[..boundary])
}
#[async_trait] #[async_trait]
impl Channel for FeishuChannel { impl Channel for FeishuChannel {
fn name(&self) -> &str { fn name(&self) -> &str {
@ -2043,11 +2275,37 @@ impl Channel for FeishuChannel {
self.running.try_read().map(|r| *r).unwrap_or(false) self.running.try_read().map(|r| *r).unwrap_or(false)
} }
async fn send(&self, msg: OutboundMessage) -> Result<(), ChannelError> { fn live_policy(&self) -> LivePolicy {
let msg = OutboundMessage { if self.config.live_updates {
content: Self::strip_thinking_tags(&msg.content), LivePolicy::Snapshot {
..msg min_interval: Duration::from_millis(
self.config.live_update_interval_ms.clamp(250, 5_000),
),
}
} else {
LivePolicy::FinalOnly
}
}
fn presentation_policy(&self) -> crate::delivery::PresentationPolicy {
crate::delivery::PresentationPolicy::external(self.config.live_updates)
}
async fn open_turn(&self, target: TurnTarget) -> Result<Box<dyn TurnSink>, ChannelError> {
let (receive_id, receive_id_type) = if target.chat_id.starts_with("oc_") {
(target.chat_id, "chat_id")
} else {
(target.reply_to.unwrap_or(target.chat_id), "open_id")
}; };
Ok(Box::new(FeishuTurnSink::new(Box::new(FeishuTurnBackend {
channel: self.clone(),
receive_id,
receive_id_type,
metadata: target.metadata,
}))))
}
async fn send(&self, msg: OutboundMessage) -> Result<(), ChannelError> {
let receive_id = if msg.chat_id.starts_with("oc_") { let receive_id = if msg.chat_id.starts_with("oc_") {
&msg.chat_id &msg.chat_id
} else { } else {
@ -2252,6 +2510,53 @@ impl Channel for FeishuChannel {
#[cfg(test)] #[cfg(test)]
mod tests { mod tests {
use super::*; use super::*;
use crate::agent::TurnEvent;
use crate::delivery::{PresentationPolicy, project_snapshot};
#[derive(Default)]
struct MockTurnState {
created: Vec<String>,
updated: Vec<(String, String)>,
cleanups: usize,
fail_updates: usize,
}
struct MockTurnApi {
state: Arc<Mutex<MockTurnState>>,
}
#[async_trait]
impl FeishuTurnApi for MockTurnApi {
async fn create_card(&mut self, markdown: &str) -> Result<String, ChannelError> {
let mut state = self.state.lock().await;
state.created.push(markdown.to_string());
Ok(format!("card-{}", state.created.len()))
}
async fn update_card(
&mut self,
message_id: &str,
markdown: &str,
) -> Result<(), ChannelError> {
let mut state = self.state.lock().await;
if state.fail_updates > 0 {
state.fail_updates -= 1;
return Err(ChannelError::Other("card can no longer be edited".into()));
}
state
.updated
.push((message_id.to_string(), markdown.to_string()));
Ok(())
}
async fn cleanup(&mut self) {
self.state.lock().await.cleanups += 1;
}
}
fn mock_sink(state: Arc<Mutex<MockTurnState>>) -> FeishuTurnSink {
FeishuTurnSink::new(Box::new(MockTurnApi { state }))
}
fn test_channel() -> FeishuChannel { fn test_channel() -> FeishuChannel {
FeishuChannel::new( FeishuChannel::new(
@ -2263,12 +2568,169 @@ mod tests {
agent: String::new(), agent: String::new(),
media_dir: String::new(), media_dir: String::new(),
reaction_emoji: "THUMBSUP".to_string(), reaction_emoji: "THUMBSUP".to_string(),
live_updates: false,
live_update_interval_ms: 500,
}, },
Path::new("/tmp"), Path::new("/tmp"),
) )
.expect("test channel should be valid") .expect("test channel should be valid")
} }
#[tokio::test]
async fn turn_sink_creates_once_updates_same_card_and_cleans_up_at_finish() {
let state = Arc::new(Mutex::new(MockTurnState::default()));
let mut sink = mock_sink(state.clone());
let (controller, emitter, _) =
crate::session::TurnController::start("feishu:chat:dialog", "message");
emitter
.emit(TurnEvent::TextDelta {
iteration: 0,
delta: "hello".into(),
})
.unwrap();
sink.update(&project_snapshot(
&controller.snapshot(),
PresentationPolicy::external(true),
))
.await
.unwrap();
emitter
.emit(TurnEvent::TextDelta {
iteration: 0,
delta: " world".into(),
})
.unwrap();
sink.update(&project_snapshot(
&controller.snapshot(),
PresentationPolicy::external(true),
))
.await
.unwrap();
controller.complete(None);
sink.finish(&project_snapshot(
&controller.snapshot(),
PresentationPolicy::external(true),
))
.await
.unwrap();
let state = state.lock().await;
assert_eq!(state.created.len(), 1);
assert_eq!(state.updated.len(), 2);
assert!(state.updated.iter().all(|(id, _)| id == "card-1"));
assert!(state.updated.last().unwrap().1.contains("hello world"));
assert_eq!(state.cleanups, 1);
}
#[tokio::test]
async fn final_update_failure_sends_complete_fallback_and_cleanup_is_idempotent() {
let state = Arc::new(Mutex::new(MockTurnState::default()));
let mut sink = mock_sink(state.clone());
let (controller, emitter, _) =
crate::session::TurnController::start("feishu:chat:dialog", "message");
emitter
.emit(TurnEvent::TextDelta {
iteration: 0,
delta: "partial".into(),
})
.unwrap();
sink.update(&controller.snapshot()).await.unwrap();
state.lock().await.fail_updates = 1;
emitter
.emit(TurnEvent::TextDelta {
iteration: 0,
delta: " final".into(),
})
.unwrap();
controller.complete(None);
sink.finish(&controller.snapshot()).await.unwrap();
sink.finish(&controller.snapshot()).await.unwrap();
let state = state.lock().await;
assert_eq!(state.created.len(), 2);
assert!(state.created[1].contains("partial final"));
assert_eq!(state.cleanups, 1);
}
#[tokio::test]
async fn final_only_sink_sends_no_fragments_and_abort_without_text_is_visible() {
let state = Arc::new(Mutex::new(MockTurnState::default()));
let mut sink = mock_sink(state.clone());
let (controller, _emitter, _) =
crate::session::TurnController::start("feishu:chat:dialog", "message");
controller.fail("provider unavailable");
sink.abort(&controller.snapshot()).await.unwrap();
let state = state.lock().await;
assert_eq!(state.created.len(), 1);
assert!(state.created[0].contains("provider unavailable"));
assert_eq!(state.updated.len(), 0);
assert_eq!(state.cleanups, 1);
}
#[test]
fn external_projection_removes_reasoning_before_feishu_rendering() {
let (controller, emitter, _) =
crate::session::TurnController::start("feishu:chat:dialog", "message");
emitter
.emit(TurnEvent::ReasoningDelta {
iteration: 0,
delta: "private".into(),
})
.unwrap();
emitter
.emit(TurnEvent::TextDelta {
iteration: 0,
delta: "public".into(),
})
.unwrap();
let projected =
project_snapshot(&controller.snapshot(), PresentationPolicy::external(true));
let markdown = render_feishu_turn(&projected);
assert!(markdown.contains("public"));
assert!(!markdown.contains("private"));
}
#[test]
fn live_card_truncation_preserves_utf8_and_payload_limit() {
let markdown = "你".repeat(FeishuChannel::CARD_MARKDOWN_MAX_BYTES);
let truncated = truncate_feishu_live_markdown(&markdown);
assert!(truncated.len() <= FeishuChannel::CARD_MARKDOWN_MAX_BYTES);
assert!(truncated.ends_with("_内容仍在生成，已暂时截断…_"));
}
#[test]
fn final_card_chunking_preserves_long_utf8_content() {
let markdown = "你".repeat(FeishuChannel::CARD_MARKDOWN_MAX_BYTES);
let chunks = FeishuChannel::split_markdown_chunks(&markdown);
assert!(chunks.len() > 1);
assert!(
chunks
.iter()
.all(|chunk| chunk.len() <= FeishuChannel::CARD_MARKDOWN_MAX_BYTES)
);
assert_eq!(chunks.concat(), markdown);
}
#[test]
fn live_policy_uses_configured_bounded_interval() {
let mut channel = test_channel();
assert_eq!(channel.live_policy(), LivePolicy::FinalOnly);
channel.config.live_updates = true;
channel.config.live_update_interval_ms = 10;
assert_eq!(
channel.live_policy(),
LivePolicy::Snapshot {
min_interval: Duration::from_millis(250)
}
);
}
#[tokio::test] #[tokio::test]
async fn stop_aborts_connection_task_that_ignores_cancellation() { async fn stop_aborts_connection_task_that_ignores_cancellation() {
let channel = test_channel(); let channel = test_channel();

2
src/channels/mod.rs
View File

@ -4,7 +4,7 @@ pub mod feishu;
pub mod manager; pub mod manager;
pub mod slash_command; pub mod slash_command;
pub use base::{Channel, ChannelError}; pub use base::{Channel, ChannelError, LivePolicy, TurnSink, TurnTarget};
pub use cli_chat::CliChatChannel; pub use cli_chat::CliChatChannel;
pub use feishu::FeishuChannel; pub use feishu::FeishuChannel;
pub use manager::ChannelManager; pub use manager::ChannelManager;

22
src/client/mod.rs
View File

@ -253,6 +253,28 @@ async fn run_app(
async fn handle_ws_message(app: &mut App, outbound: WsOutbound) { async fn handle_ws_message(app: &mut App, outbound: WsOutbound) {
match outbound { match outbound {
WsOutbound::TurnUpdated { snapshot } => {
let terminal = snapshot.status != crate::session::TurnStatus::Running;
let session_id = snapshot.session_id.clone();
if terminal {
app.pending_responses = app.pending_responses.saturating_sub(1);
}
if app.apply_turn_snapshot(snapshot) {
if terminal {
app.status_message = None;
if app.current_session_id.as_deref() == Some(&session_id) {
request_history(app, session_id).await;
} else {
app.status_message = Some("另一个会话已完成响应".to_string());
}
request_session_list(app).await;
} else {
app.status_message = Some("正在生成回复…".to_string());
}
} else if terminal && app.current_session_id.as_deref() != Some(&session_id) {
app.status_message = Some("另一个会话已完成响应".to_string());
}
}
WsOutbound::AssistantResponse { WsOutbound::AssistantResponse {
id, id,
content, content,

100
src/client/tui/app.rs
View File

@ -1,6 +1,7 @@
use crate::protocol::{ use crate::protocol::{
HistoryMessage, MessageAttachment, SessionSummary, SlashCommandInfo, UploadDescriptor, HistoryMessage, MessageAttachment, SessionSummary, SlashCommandInfo, UploadDescriptor,
}; };
use crate::session::{TurnSnapshot, TurnStatus};
use std::collections::VecDeque; use std::collections::VecDeque;
use tokio_tungstenite::tungstenite::Message; use tokio_tungstenite::tungstenite::Message;
@ -19,6 +20,8 @@ pub struct ChatMessage {
pub id: String, pub id: String,
pub role: MessageRole, pub role: MessageRole,
pub content: String, pub content: String,
pub reasoning_content: Option<String>,
pub completion_status: crate::bus::CompletionStatus,
pub attachments: Vec<MessageAttachment>, pub attachments: Vec<MessageAttachment>,
} }
@ -63,6 +66,7 @@ pub struct App {
pub selected_session: usize, pub selected_session: usize,
pub show_archived: bool, pub show_archived: bool,
pub messages: VecDeque<ChatMessage>, pub messages: VecDeque<ChatMessage>,
pub active_turn: Option<TurnSnapshot>,
pub input: String, pub input: String,
/// UTF-8 byte offset. It is always maintained at a character boundary. /// UTF-8 byte offset. It is always maintained at a character boundary.
pub input_cursor_pos: usize, pub input_cursor_pos: usize,
@ -95,6 +99,7 @@ impl App {
selected_session: 0, selected_session: 0,
show_archived: false, show_archived: false,
messages: VecDeque::new(), messages: VecDeque::new(),
active_turn: None,
input: String::new(), input: String::new(),
input_cursor_pos: 0, input_cursor_pos: 0,
focus: Focus::Input, focus: Focus::Input,
@ -132,6 +137,8 @@ impl App {
id, id,
role, role,
content, content,
reasoning_content: None,
completion_status: crate::bus::CompletionStatus::Completed,
attachments, attachments,
}); });
while self.messages.len() > MAX_MESSAGES { while self.messages.len() > MAX_MESSAGES {
@ -144,6 +151,10 @@ impl App {
if self.current_session_id.as_deref() != Some(session_id) { if self.current_session_id.as_deref() != Some(session_id) {
return; return;
} }
let calibrates_terminal = self.active_turn.as_ref().is_some_and(|turn| {
turn.status != TurnStatus::Running
&& messages.iter().any(|message| message.id == turn.message_id)
});
self.messages = messages self.messages = messages
.into_iter() .into_iter()
.filter_map(|message| { .filter_map(|message| {
@ -157,6 +168,8 @@ impl App {
id: message.id, id: message.id,
role, role,
content: message.content, content: message.content,
reasoning_content: message.reasoning_content,
completion_status: message.completion_status,
attachments: message.attachments, attachments: message.attachments,
}) })
}) })
@ -166,6 +179,9 @@ impl App {
} }
self.chat_scroll_from_bottom = 0; self.chat_scroll_from_bottom = 0;
self.status_message = None; self.status_message = None;
if calibrates_terminal {
self.active_turn = None;
}
} }
pub fn set_sessions(&mut self, sessions: Vec<SessionSummary>) { pub fn set_sessions(&mut self, sessions: Vec<SessionSummary>) {
@ -185,6 +201,7 @@ impl App {
if self.current_session_id != session_id { if self.current_session_id != session_id {
self.current_session_id = session_id; self.current_session_id = session_id;
self.messages.clear(); self.messages.clear();
self.active_turn = None;
self.pending_uploads.clear(); self.pending_uploads.clear();
self.chat_scroll_from_bottom = 0; self.chat_scroll_from_bottom = 0;
} }
@ -198,6 +215,21 @@ impl App {
} }
} }
pub fn apply_turn_snapshot(&mut self, snapshot: TurnSnapshot) -> bool {
if self.current_session_id.as_deref() != Some(&snapshot.session_id) {
return false;
}
if let Some(current) = &self.active_turn
&& current.id == snapshot.id
&& current.revision >= snapshot.revision
{
return false;
}
self.active_turn = Some(snapshot);
self.chat_scroll_from_bottom = 0;
true
}
pub fn current_title(&self) -> &str { pub fn current_title(&self) -> &str {
self.current_session_id self.current_session_id
.as_ref() .as_ref()
@ -383,6 +415,21 @@ fn next_boundary(value: &str, offset: usize) -> Option<usize> {
#[cfg(test)] #[cfg(test)]
mod tests { mod tests {
use super::*; use super::*;
use crate::session::{TurnId, TurnPhase, TurnState};
fn turn(revision: u64, status: TurnStatus) -> TurnSnapshot {
TurnState {
id: TurnId("turn".into()),
session_id: "current".into(),
message_id: "message".into(),
revision,
status,
phase: TurnPhase::Responding,
blocks: Vec::new(),
usage: None,
error: None,
}
}
#[test] #[test]
fn unicode_cursor_edits_only_at_character_boundaries() { fn unicode_cursor_edits_only_at_character_boundaries() {
@ -404,4 +451,57 @@ mod tests {
app.set_history("old", Vec::new()); app.set_history("old", Vec::new());
assert_eq!(app.messages.len(), 1); assert_eq!(app.messages.len(), 1);
} }
#[test]
fn active_turn_ignores_stale_revisions_and_other_sessions() {
let mut app = App::new();
app.set_current_session(Some("current".into()));
assert!(app.apply_turn_snapshot(turn(2, TurnStatus::Running)));
assert!(!app.apply_turn_snapshot(turn(1, TurnStatus::Running)));
let mut other = turn(3, TurnStatus::Running);
other.session_id = "other".into();
assert!(!app.apply_turn_snapshot(other));
assert_eq!(app.active_turn.as_ref().unwrap().revision, 2);
}
#[test]
fn terminal_turn_remains_visible_until_history_calibrates_it() {
let mut app = App::new();
app.set_current_session(Some("current".into()));
app.apply_turn_snapshot(turn(3, TurnStatus::Completed));
assert!(app.active_turn.is_some());
app.set_history(
"current",
vec![HistoryMessage {
id: "message".into(),
seq: 1,
role: "assistant".into(),
content: "done".into(),
reasoning_content: None,
completion_status: crate::bus::CompletionStatus::Completed,
created_at: 1,
tool_call_id: None,
tool_name: None,
tool_calls: None,
attachments: Vec::new(),
}],
);
assert!(app.active_turn.is_none());
}
#[test]
fn failed_turn_without_a_durable_message_remains_visible_after_history_refresh() {
let mut app = App::new();
app.set_current_session(Some("current".into()));
app.apply_turn_snapshot(turn(3, TurnStatus::Failed));
app.set_history("current", Vec::new());
assert_eq!(
app.active_turn.as_ref().map(|turn| turn.status),
Some(TurnStatus::Failed)
);
}
} }

98
src/client/tui/components/chat_history.rs
View File

@ -1,4 +1,5 @@
use crate::client::tui::app::{App, MessageRole}; use crate::client::tui::app::{App, MessageRole};
use crate::session::{ToolStatus, TurnBlock, TurnPhase, TurnStatus};
use ratatui::{ use ratatui::{
Frame, Frame,
layout::Rect, layout::Rect,
@ -23,6 +24,15 @@ pub fn render(frame: &mut Frame, area: Rect, app: &App) {
label, label,
Style::default().fg(color).add_modifier(Modifier::BOLD), Style::default().fg(color).add_modifier(Modifier::BOLD),
))); )));
if let Some(reasoning) = &message.reasoning_content {
lines.push(Line::from(Span::styled(
"思考过程",
Style::default()
.fg(Color::DarkGray)
.add_modifier(Modifier::ITALIC),
)));
push_wrapped(&mut lines, reasoning, content_width, Color::DarkGray);
}
for source_line in message.content.lines() { for source_line in message.content.lines() {
let wrapped = textwrap::wrap(source_line, content_width); let wrapped = textwrap::wrap(source_line, content_width);
if wrapped.is_empty() { if wrapped.is_empty() {
@ -35,6 +45,12 @@ pub fn render(frame: &mut Frame, area: Rect, app: &App) {
); );
} }
} }
if message.completion_status != crate::bus::CompletionStatus::Completed {
lines.push(Line::from(Span::styled(
format!("[{}]", message.completion_status.as_str()),
Style::default().fg(Color::Yellow),
)));
}
for attachment in &message.attachments { for attachment in &message.attachments {
lines.push(Line::from(Span::styled( lines.push(Line::from(Span::styled(
format!(" [附件 {}] {}", attachment.index + 1, attachment.name), format!(" [附件 {}] {}", attachment.index + 1, attachment.name),
@ -43,7 +59,74 @@ pub fn render(frame: &mut Frame, area: Rect, app: &App) {
} }
lines.push(Line::from("")); lines.push(Line::from(""));
} }
if app.pending_responses > 0 { if let Some(turn) = &app.active_turn {
lines.push(Line::from(Span::styled(
"PicoBot",
Style::default()
.fg(Color::Green)
.add_modifier(Modifier::BOLD),
)));
for block in &turn.blocks {
match block {
TurnBlock::Reasoning { text, .. } => {
lines.push(Line::from(Span::styled(
"思考过程",
Style::default()
.fg(Color::DarkGray)
.add_modifier(Modifier::ITALIC),
)));
push_wrapped(&mut lines, text, content_width, Color::DarkGray);
}
TurnBlock::Assistant { text, .. } => {
push_wrapped(&mut lines, text, content_width, Color::Reset);
}
TurnBlock::Tool {
name,
status,
preview,
..
} => {
let status = match status {
ToolStatus::Running => "执行中",
ToolStatus::Completed => "已完成",
ToolStatus::Failed => "失败",
};
lines.push(Line::from(Span::styled(
format!("工具 · {name} · {status}"),
Style::default().fg(Color::Magenta),
)));
if let Some(preview) = preview {
push_wrapped(&mut lines, preview, content_width, Color::DarkGray);
}
}
}
}
let phase = match turn.phase {
TurnPhase::Queued => "排队中",
TurnPhase::Reasoning => "思考中",
TurnPhase::Responding => "生成中",
TurnPhase::Acting => "调用工具中",
TurnPhase::Finalizing => "收尾中",
};
let status = match turn.status {
TurnStatus::Running => phase,
TurnStatus::Completed => "已完成",
TurnStatus::Cancelled => "已停止",
TurnStatus::Failed => "失败",
};
lines.push(Line::from(Span::styled(
format!("● {status}"),
Style::default().fg(if turn.status == TurnStatus::Failed {
Color::Red
} else {
Color::Cyan
}),
)));
if let Some(error) = &turn.error {
push_wrapped(&mut lines, error, content_width, Color::Red);
}
lines.push(Line::from(""));
} else if app.pending_responses > 0 {
lines.push(Line::from(Span::styled( lines.push(Line::from(Span::styled(
"● 正在思考…", "● 正在思考…",
Style::default().fg(Color::Cyan), Style::default().fg(Color::Cyan),
@ -66,3 +149,16 @@ pub fn render(frame: &mut Frame, area: Rect, app: &App) {
area, area,
); );
} }
fn push_wrapped(lines: &mut Vec<Line<'static>>, text: &str, width: usize, color: Color) {
for source_line in text.lines() {
let wrapped = textwrap::wrap(source_line, width);
if wrapped.is_empty() {
lines.push(Line::from(""));
} else {
lines.extend(wrapped.into_iter().map(|line| {
Line::from(Span::styled(line.into_owned(), Style::default().fg(color)))
}));
}
}
}

34
src/client/tui/ui.rs
View File

@ -189,4 +189,38 @@ mod tests {
); );
terminal.draw(|frame| render_ui(frame, &app)).unwrap(); terminal.draw(|frame| render_ui(frame, &app)).unwrap();
} }
#[test]
fn active_reasoning_text_and_tool_snapshot_renders_without_panic() {
let backend = TestBackend::new(96, 24);
let mut terminal = Terminal::new(backend).unwrap();
let mut app = App::new();
app.set_current_session(Some("session".into()));
let (controller, emitter, _) = crate::session::TurnController::start("session", "message");
emitter
.emit(crate::agent::TurnEvent::ReasoningDelta {
iteration: 0,
delta: "先检查状态".into(),
})
.unwrap();
emitter
.emit(crate::agent::TurnEvent::TextDelta {
iteration: 0,
delta: "正在处理".into(),
})
.unwrap();
emitter
.emit(crate::agent::TurnEvent::ToolStarted {
iteration: 0,
call: crate::providers::ToolCall {
id: "call".into(),
name: "bash".into(),
arguments: serde_json::json!({"cmd": "pwd"}),
},
})
.unwrap();
app.apply_turn_snapshot((*controller.snapshot()).clone());
terminal.draw(|frame| render_ui(frame, &app)).unwrap();
}
} }

9
src/config/mod.rs
View File

@ -78,6 +78,11 @@ pub struct FeishuChannelConfig {
/// Emoji type for message reactions (e.g. "THUMBSUP", "OK", "EYES"). /// Emoji type for message reactions (e.g. "THUMBSUP", "OK", "EYES").
#[serde(default = "default_reaction_emoji")] #[serde(default = "default_reaction_emoji")]
pub reaction_emoji: String, pub reaction_emoji: String,
/// Edit one card with latest Turn snapshots instead of sending only the final result.
#[serde(default)]
pub live_updates: bool,
#[serde(default = "default_feishu_live_update_interval_ms")]
pub live_update_interval_ms: u64,
} }
fn default_allow_from() -> Vec<String> { fn default_allow_from() -> Vec<String> {
@ -95,6 +100,10 @@ fn default_reaction_emoji() -> String {
"Typing".to_string() "Typing".to_string()
} }
fn default_feishu_live_update_interval_ms() -> u64 {
500
}
#[derive(Debug, Clone, Deserialize, Serialize)] #[derive(Debug, Clone, Deserialize, Serialize)]
pub struct ProviderConfig { pub struct ProviderConfig {
#[serde(rename = "type")] #[serde(rename = "type")]

676
src/delivery/coordinator.rs Normal file
View File

@ -0,0 +1,676 @@
use std::collections::HashMap;
use std::sync::{Arc, Mutex, Weak};
use std::time::Duration;
use tokio::sync::{Mutex as AsyncMutex, oneshot, watch};
use tokio::time::{Instant, sleep_until, timeout};
use crate::channels::{Channel, ChannelError, LivePolicy, TurnSink, TurnTarget};
use crate::delivery::{PresentationPolicy, project_snapshot};
use crate::session::{TurnSnapshot, TurnStatus};
use crate::task_supervisor::TaskSupervisor;
const SINK_CALL_TIMEOUT: Duration = Duration::from_secs(30);
const FINAL_RETRY_DELAYS: &[Duration] = &[
Duration::from_secs(1),
Duration::from_secs(2),
Duration::from_secs(4),
];
#[derive(Debug)]
pub enum DeliveryError {
ChannelNotFound(String),
OpenFailed(ChannelError),
SnapshotStreamClosed,
SupervisorStopping,
FinalTimedOut,
FinalFailed(ChannelError),
}
impl std::fmt::Display for DeliveryError {
fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
match self {
Self::ChannelNotFound(channel) => write!(formatter, "channel not found: {channel}"),
Self::OpenFailed(error) => write!(formatter, "failed to open turn sink: {error}"),
Self::SnapshotStreamClosed => {
formatter.write_str("turn snapshot stream closed before a terminal state")
}
Self::SupervisorStopping => {
formatter.write_str("cannot start turn delivery while Gateway is stopping")
}
Self::FinalTimedOut => formatter.write_str("final turn delivery timed out"),
Self::FinalFailed(error) => write!(formatter, "final turn delivery failed: {error}"),
}
}
}
impl std::error::Error for DeliveryError {}
/// Shared ordering boundary for writes to one `(channel, chat_id)` target.
///
/// The registry stores weak references so inactive conversations disappear
/// without a cleanup task. Callers hold the returned lock only around one
/// external write, never for the lifetime of a Turn.
#[derive(Clone, Default)]
pub struct ConversationWriteLocks {
locks: Arc<Mutex<HashMap<String, Weak<AsyncMutex<()>>>>>,
}
impl ConversationWriteLocks {
pub fn for_target(&self, channel: &str, chat_id: &str) -> Arc<AsyncMutex<()>> {
let key = format!("{channel}\0{chat_id}");
let mut locks = self
.locks
.lock()
.unwrap_or_else(|poisoned| poisoned.into_inner());
if let Some(existing) = locks.get(&key).and_then(Weak::upgrade) {
return existing;
}
let lock = Arc::new(AsyncMutex::new(()));
locks.insert(key, Arc::downgrade(&lock));
lock
}
}
#[derive(Clone)]
pub struct DeliveryCoordinator {
write_locks: ConversationWriteLocks,
sink_call_timeout: Duration,
final_retry_delays: Arc<[Duration]>,
}
pub(crate) struct SinkRoute {
pub channel: String,
pub chat_id: String,
pub live_policy: LivePolicy,
pub presentation: PresentationPolicy,
}
impl DeliveryCoordinator {
pub fn new(write_locks: ConversationWriteLocks) -> Self {
Self {
write_locks,
sink_call_timeout: SINK_CALL_TIMEOUT,
final_retry_delays: FINAL_RETRY_DELAYS.into(),
}
}
#[cfg(test)]
fn for_test(
sink_call_timeout: Duration,
final_retry_delays: impl Into<Arc<[Duration]>>,
) -> Self {
Self {
write_locks: ConversationWriteLocks::default(),
sink_call_timeout,
final_retry_delays: final_retry_delays.into(),
}
}
pub fn write_locks(&self) -> ConversationWriteLocks {
self.write_locks.clone()
}
pub async fn open_and_deliver(
&self,
channel: Arc<dyn Channel + Send + Sync>,
target: TurnTarget,
presentation: PresentationPolicy,
snapshots: watch::Receiver<Arc<TurnSnapshot>>,
) -> Result<(), DeliveryError> {
let live_policy = channel.live_policy();
let sink = channel
.open_turn(target.clone())
.await
.map_err(DeliveryError::OpenFailed)?;
self.deliver(
&target.channel,
&target.chat_id,
live_policy,
presentation,
snapshots,
sink,
)
.await
}
/// Start one sink lifecycle under the Gateway's task owner and return a
/// bounded completion report to the caller.
pub fn spawn(
&self,
supervisor: &TaskSupervisor,
channel: Arc<dyn Channel + Send + Sync>,
target: TurnTarget,
presentation: PresentationPolicy,
snapshots: watch::Receiver<Arc<TurnSnapshot>>,
) -> Result<oneshot::Receiver<Result<(), DeliveryError>>, DeliveryError> {
let (result_tx, result_rx) = oneshot::channel();
let coordinator = self.clone();
let task_name = format!("turn-delivery:{}:{}", target.channel, target.chat_id);
let spawned = supervisor.spawn(task_name, async move {
let result = coordinator
.open_and_deliver(channel, target, presentation, snapshots)
.await;
if let Err(error) = &result {
tracing::error!(error = %error, "Turn delivery failed");
}
let _ = result_tx.send(result);
});
if !spawned {
return Err(DeliveryError::SupervisorStopping);
}
Ok(result_rx)
}
pub(crate) fn spawn_sink(
&self,
supervisor: &TaskSupervisor,
route: SinkRoute,
mut snapshots: watch::Receiver<Arc<TurnSnapshot>>,
mut sink: Box<dyn TurnSink>,
) -> Result<oneshot::Receiver<Result<(), DeliveryError>>, DeliveryError> {
let SinkRoute {
channel,
chat_id,
live_policy,
presentation,
} = route;
let (result_tx, result_rx) = oneshot::channel();
let coordinator = self.clone();
let task_name = format!("turn-delivery:{channel}:{chat_id}");
let cancellation = supervisor.cancellation_token();
let shutdown_snapshot = snapshots.clone();
let spawned = supervisor.spawn_graceful(task_name, async move {
let mut delivery = Box::pin(coordinator.deliver_sink(
&channel,
&chat_id,
live_policy,
presentation,
&mut snapshots,
&mut *sink,
));
let result = tokio::select! {
result = &mut delivery => result,
() = cancellation.cancelled() => {
drop(delivery);
let snapshot = shutdown_snapshot.borrow().clone();
let projected = project_snapshot(&snapshot, presentation);
coordinator
.abort_for_shutdown(&channel, &chat_id, &mut *sink, &projected)
.await
}
};
if let Err(error) = &result {
tracing::error!(channel, chat_id, error = %error, "Turn delivery failed");
}
let _ = result_tx.send(result);
});
if !spawned {
return Err(DeliveryError::SupervisorStopping);
}
Ok(result_rx)
}
pub async fn deliver(
&self,
channel: &str,
chat_id: &str,
live_policy: LivePolicy,
presentation: PresentationPolicy,
mut snapshots: watch::Receiver<Arc<TurnSnapshot>>,
mut sink: Box<dyn TurnSink>,
) -> Result<(), DeliveryError> {
self.deliver_sink(
channel,
chat_id,
live_policy,
presentation,
&mut snapshots,
&mut *sink,
)
.await
}
async fn deliver_sink(
&self,
channel: &str,
chat_id: &str,
live_policy: LivePolicy,
presentation: PresentationPolicy,
snapshots: &mut watch::Receiver<Arc<TurnSnapshot>>,
sink: &mut dyn TurnSink,
) -> Result<(), DeliveryError> {
let target_lock = self.write_locks.for_target(channel, chat_id);
let min_interval = match live_policy {
LivePolicy::FinalOnly => None,
LivePolicy::Snapshot { min_interval } if presentation.live => Some(min_interval),
LivePolicy::Snapshot { .. } => None,
};
let mut next_update_at = Instant::now();
loop {
let snapshot = snapshots.borrow_and_update().clone();
if snapshot.status != TurnStatus::Running {
let projected = project_snapshot(&snapshot, presentation);
return self.deliver_terminal(&target_lock, sink, &projected).await;
}
if let Some(interval) = min_interval {
while Instant::now() < next_update_at {
tokio::select! {
changed = snapshots.changed() => {
changed.map_err(|_| DeliveryError::SnapshotStreamClosed)?;
let latest = snapshots.borrow_and_update().clone();
if latest.status != TurnStatus::Running {
let projected = project_snapshot(&latest, presentation);
return self.deliver_terminal(&target_lock, sink, &projected).await;
}
}
() = sleep_until(next_update_at) => break,
}
}
let latest = snapshots.borrow_and_update().clone();
if latest.status != TurnStatus::Running {
let projected = project_snapshot(&latest, presentation);
return self.deliver_terminal(&target_lock, sink, &projected).await;
}
let projected = project_snapshot(&latest, presentation);
let _guard = target_lock.lock().await;
match timeout(self.sink_call_timeout, sink.update(&projected)).await {
Ok(Ok(())) => {}
Ok(Err(error)) => {
tracing::warn!(error = %error, revision = projected.revision, "Live turn update failed; waiting for a newer snapshot");
}
Err(_) => {
tracing::warn!(
revision = projected.revision,
"Live turn update timed out; waiting for a newer snapshot"
);
}
}
next_update_at = Instant::now() + interval;
}
snapshots
.changed()
.await
.map_err(|_| DeliveryError::SnapshotStreamClosed)?;
}
}
async fn abort_for_shutdown(
&self,
channel: &str,
chat_id: &str,
sink: &mut dyn TurnSink,
snapshot: &TurnSnapshot,
) -> Result<(), DeliveryError> {
let target_lock = self.write_locks.for_target(channel, chat_id);
let _guard = target_lock.lock().await;
match timeout(self.sink_call_timeout, sink.abort(snapshot)).await {
Ok(Ok(())) => Ok(()),
Ok(Err(error)) => Err(DeliveryError::FinalFailed(error)),
Err(_) => Err(DeliveryError::FinalTimedOut),
}
}
async fn deliver_terminal(
&self,
target_lock: &Arc<AsyncMutex<()>>,
sink: &mut dyn TurnSink,
snapshot: &TurnSnapshot,
) -> Result<(), DeliveryError> {
let attempts = self.final_retry_delays.len() + 1;
for attempt in 0..attempts {
let _guard = target_lock.lock().await;
let result = if snapshot.status == TurnStatus::Completed {
timeout(self.sink_call_timeout, sink.finish(snapshot)).await
} else {
timeout(self.sink_call_timeout, sink.abort(snapshot)).await
};
drop(_guard);
match result {
Ok(Ok(())) => return Ok(()),
Ok(Err(error))
if error.is_transient() && attempt < self.final_retry_delays.len() =>
{
sleep_until(Instant::now() + self.final_retry_delays[attempt]).await;
}
Ok(Err(error)) => return Err(DeliveryError::FinalFailed(error)),
Err(_) if attempt < self.final_retry_delays.len() => {
sleep_until(Instant::now() + self.final_retry_delays[attempt]).await;
}
Err(_) => return Err(DeliveryError::FinalTimedOut),
}
}
unreachable!()
}
}
#[cfg(test)]
mod tests {
use super::*;
use async_trait::async_trait;
use std::sync::atomic::{AtomicUsize, Ordering};
use tokio::sync::{Mutex as TokioMutex, Notify};
use crate::agent::TurnEvent;
use crate::bus::{MessageBus, OutboundMessage};
use crate::channels::{Channel, TurnSink};
use crate::session::{TurnBlock, TurnController};
#[derive(Default)]
struct SinkState {
updates: TokioMutex<Vec<TurnSnapshot>>,
terminal: TokioMutex<Vec<TurnSnapshot>>,
update_started: Notify,
release_update: Notify,
block_first_update: bool,
fail_updates: AtomicUsize,
fail_finish: AtomicUsize,
}
struct RecordingSink(Arc<SinkState>);
struct SinkChannel {
state: Arc<SinkState>,
opened: AtomicUsize,
}
#[async_trait]
impl Channel for SinkChannel {
fn name(&self) -> &str {
"sink-channel"
}
fn is_running(&self) -> bool {
true
}
async fn start(&self, _bus: Arc<MessageBus>) -> Result<(), ChannelError> {
Ok(())
}
async fn stop(&self) -> Result<(), ChannelError> {
Ok(())
}
fn live_policy(&self) -> LivePolicy {
LivePolicy::FinalOnly
}
async fn open_turn(&self, _target: TurnTarget) -> Result<Box<dyn TurnSink>, ChannelError> {
self.opened.fetch_add(1, Ordering::SeqCst);
Ok(sink(self.state.clone()))
}
async fn send(&self, _msg: OutboundMessage) -> Result<(), ChannelError> {
Ok(())
}
}
#[async_trait]
impl TurnSink for RecordingSink {
async fn update(&mut self, snapshot: &TurnSnapshot) -> Result<(), ChannelError> {
self.0.update_started.notify_waiters();
if self.0.block_first_update && self.0.updates.lock().await.is_empty() {
self.0.release_update.notified().await;
}
if self
.0
.fail_updates
.fetch_update(Ordering::SeqCst, Ordering::SeqCst, |remaining| {
remaining.checked_sub(1)
})
.is_ok()
{
return Err(ChannelError::SendError("update".into()));
}
self.0.updates.lock().await.push(snapshot.clone());
Ok(())
}
async fn finish(&mut self, snapshot: &TurnSnapshot) -> Result<(), ChannelError> {
if self
.0
.fail_finish
.fetch_update(Ordering::SeqCst, Ordering::SeqCst, |remaining| {
remaining.checked_sub(1)
})
.is_ok()
{
return Err(ChannelError::SendError("finish".into()));
}
self.0.terminal.lock().await.push(snapshot.clone());
Ok(())
}
async fn abort(&mut self, snapshot: &TurnSnapshot) -> Result<(), ChannelError> {
self.0.terminal.lock().await.push(snapshot.clone());
Ok(())
}
}
fn sink(state: Arc<SinkState>) -> Box<dyn TurnSink> {
Box::new(RecordingSink(state))
}
#[tokio::test]
async fn slow_sink_observes_latest_snapshot_and_terminal_bypasses_throttle() {
let state = Arc::new(SinkState {
block_first_update: true,
..SinkState::default()
});
let (controller, emitter, receiver) = TurnController::start("session", "message");
let coordinator = DeliveryCoordinator::for_test(Duration::from_secs(30), []);
let task = tokio::spawn({
let state = state.clone();
async move {
coordinator
.deliver(
"cli_chat",
"chat",
LivePolicy::Snapshot {
min_interval: Duration::from_secs(10),
},
PresentationPolicy::interactive(),
receiver,
sink(state),
)
.await
}
});
state.update_started.notified().await;
emitter
.emit(TurnEvent::TextDelta {
iteration: 0,
delta: "a".into(),
})
.unwrap();
emitter
.emit(TurnEvent::TextDelta {
iteration: 0,
delta: "b".into(),
})
.unwrap();
state.release_update.notify_waiters();
tokio::task::yield_now().await;
controller.complete(None);
assert!(task.await.unwrap().is_ok());
let terminal = state.terminal.lock().await;
assert_eq!(terminal.len(), 1);
assert_eq!(terminal[0].status, TurnStatus::Completed);
assert!(
matches!(&terminal[0].blocks[0], TurnBlock::Assistant { text, .. } if text == "ab")
);
}
#[tokio::test]
async fn hidden_reasoning_is_removed_before_sink_and_failed_update_recovers() {
let state = Arc::new(SinkState {
fail_updates: AtomicUsize::new(1),
..SinkState::default()
});
let (controller, emitter, receiver) = TurnController::start("session", "message");
let coordinator = DeliveryCoordinator::for_test(Duration::from_secs(30), []);
let task = tokio::spawn({
let state = state.clone();
async move {
coordinator
.deliver(
"feishu",
"chat",
LivePolicy::Snapshot {
min_interval: Duration::ZERO,
},
PresentationPolicy::external(true),
receiver,
sink(state),
)
.await
}
});
emitter
.emit(TurnEvent::ReasoningDelta {
iteration: 0,
delta: "secret".into(),
})
.unwrap();
tokio::task::yield_now().await;
emitter
.emit(TurnEvent::TextDelta {
iteration: 0,
delta: "public".into(),
})
.unwrap();
tokio::task::yield_now().await;
controller.complete(None);
assert!(task.await.unwrap().is_ok());
let terminal = state.terminal.lock().await;
assert!(
terminal[0]
.blocks
.iter()
.all(|block| !matches!(block, TurnBlock::Reasoning { .. }))
);
}
#[tokio::test]
async fn final_only_skips_updates_and_retries_transient_finish() {
let state = Arc::new(SinkState {
fail_finish: AtomicUsize::new(2),
..SinkState::default()
});
let (controller, emitter, receiver) = TurnController::start("session", "message");
let coordinator = DeliveryCoordinator::for_test(
Duration::from_secs(30),
[Duration::from_millis(1), Duration::from_millis(2)],
);
let task = tokio::spawn({
let state = state.clone();
async move {
coordinator
.deliver(
"channel",
"chat",
LivePolicy::FinalOnly,
PresentationPolicy::unattended(),
receiver,
sink(state),
)
.await
}
});
emitter
.emit(TurnEvent::TextDelta {
iteration: 0,
delta: "done".into(),
})
.unwrap();
controller.complete(None);
assert!(task.await.unwrap().is_ok());
assert!(state.updates.lock().await.is_empty());
assert_eq!(state.terminal.lock().await.len(), 1);
}
#[tokio::test]
async fn open_and_deliver_owns_sink_creation_and_terminal_lifecycle() {
let state = Arc::new(SinkState::default());
let channel = Arc::new(SinkChannel {
state: state.clone(),
opened: AtomicUsize::new(0),
});
let (controller, emitter, receiver) = TurnController::start("session", "message");
let coordinator = DeliveryCoordinator::for_test(Duration::from_secs(30), []);
let target = TurnTarget {
channel: "sink-channel".into(),
chat_id: "chat".into(),
session_id: "session".into(),
reply_to: None,
metadata: HashMap::new(),
};
let task = tokio::spawn({
let channel = channel.clone();
async move {
coordinator
.open_and_deliver(channel, target, PresentationPolicy::unattended(), receiver)
.await
}
});
emitter
.emit(TurnEvent::TextDelta {
iteration: 0,
delta: "done".into(),
})
.unwrap();
controller.complete(None);
assert!(task.await.unwrap().is_ok());
assert_eq!(channel.opened.load(Ordering::SeqCst), 1);
assert_eq!(state.terminal.lock().await.len(), 1);
}
#[tokio::test]
async fn supervisor_shutdown_aborts_sink_and_waits_for_cleanup() {
let state = Arc::new(SinkState::default());
let (_controller, emitter, receiver) = TurnController::start("session", "message");
emitter
.emit(TurnEvent::TextDelta {
iteration: 0,
delta: "partial".into(),
})
.unwrap();
let supervisor = TaskSupervisor::new();
let coordinator = DeliveryCoordinator::for_test(Duration::from_secs(1), []);
let result = coordinator
.spawn_sink(
&supervisor,
SinkRoute {
channel: "channel".into(),
chat_id: "chat".into(),
live_policy: LivePolicy::FinalOnly,
presentation: PresentationPolicy::unattended(),
},
receiver,
sink(state.clone()),
)
.unwrap();
tokio::task::yield_now().await;
supervisor.shutdown(Duration::from_secs(1)).await;
assert!(result.await.unwrap().is_ok());
let terminal = state.terminal.lock().await;
assert_eq!(terminal.len(), 1);
assert_eq!(terminal[0].status, TurnStatus::Running);
}
}

7
src/delivery/mod.rs Normal file
View File

@ -0,0 +1,7 @@
mod coordinator;
mod policy;
mod service;
pub use coordinator::{ConversationWriteLocks, DeliveryCoordinator, DeliveryError};
pub use policy::{PresentationPolicy, ReasoningVisibility, ToolVisibility, project_snapshot};
pub use service::TurnDeliveryService;

138
src/delivery/policy.rs Normal file
View File

@ -0,0 +1,138 @@
use crate::session::{TurnBlock, TurnSnapshot};
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum ReasoningVisibility {
Hidden,
Collapsed,
Expanded,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum ToolVisibility {
Hidden,
Compact,
Detailed,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct PresentationPolicy {
pub live: bool,
pub reasoning: ReasoningVisibility,
pub tools: ToolVisibility,
}
impl PresentationPolicy {
pub const fn interactive() -> Self {
Self {
live: true,
reasoning: ReasoningVisibility::Collapsed,
tools: ToolVisibility::Detailed,
}
}
pub const fn external(live: bool) -> Self {
Self {
live,
reasoning: ReasoningVisibility::Hidden,
tools: ToolVisibility::Compact,
}
}
pub const fn unattended() -> Self {
Self {
live: false,
reasoning: ReasoningVisibility::Hidden,
tools: ToolVisibility::Hidden,
}
}
}
/// Produce the immutable view that is allowed to leave the Gateway core.
///
/// Presentation filtering deliberately clones the snapshot. Conversation
/// history and the authoritative TurnController state remain untouched.
pub fn project_snapshot(snapshot: &TurnSnapshot, policy: PresentationPolicy) -> TurnSnapshot {
let mut projected = snapshot.clone();
projected.blocks.retain_mut(|block| match block {
TurnBlock::Reasoning { .. } => policy.reasoning != ReasoningVisibility::Hidden,
TurnBlock::Assistant { .. } => true,
TurnBlock::Tool {
arguments, preview, ..
} => match policy.tools {
ToolVisibility::Hidden => false,
ToolVisibility::Compact => {
*arguments = serde_json::Value::Null;
*preview = None;
true
}
ToolVisibility::Detailed => true,
},
});
projected
}
#[cfg(test)]
mod tests {
use super::*;
use crate::session::{BlockId, ToolStatus, TurnId, TurnPhase, TurnState, TurnStatus};
fn snapshot() -> TurnSnapshot {
TurnState {
id: TurnId("turn".into()),
session_id: "session".into(),
message_id: "message".into(),
revision: 3,
status: TurnStatus::Running,
phase: TurnPhase::Acting,
blocks: vec![
TurnBlock::Reasoning {
id: BlockId("reasoning".into()),
iteration: 0,
text: "private chain".into(),
},
TurnBlock::Assistant {
id: BlockId("text".into()),
iteration: 0,
text: "visible".into(),
},
TurnBlock::Tool {
id: "tool".into(),
iteration: 0,
name: "bash".into(),
arguments: serde_json::json!({"token": "secret"}),
status: ToolStatus::Completed,
preview: Some("sensitive output".into()),
},
],
usage: None,
error: None,
}
}
#[test]
fn external_projection_removes_reasoning_and_tool_details_without_mutating_source() {
let source = snapshot();
let projected = project_snapshot(&source, PresentationPolicy::external(true));
assert_eq!(projected.blocks.len(), 2);
assert!(matches!(projected.blocks[0], TurnBlock::Assistant { .. }));
assert!(matches!(
&projected.blocks[1],
TurnBlock::Tool {
arguments: serde_json::Value::Null,
preview: None,
..
}
));
assert_eq!(source.blocks.len(), 3);
assert!(matches!(source.blocks[0], TurnBlock::Reasoning { .. }));
}
#[test]
fn unattended_projection_keeps_only_assistant_blocks() {
let projected = project_snapshot(&snapshot(), PresentationPolicy::unattended());
assert_eq!(projected.blocks.len(), 1);
assert!(matches!(projected.blocks[0], TurnBlock::Assistant { .. }));
}
}

64
src/delivery/service.rs Normal file
View File

@ -0,0 +1,64 @@
use std::sync::Arc;
use tokio::sync::watch;
use crate::channels::{ChannelManager, TurnTarget};
use crate::delivery::{DeliveryCoordinator, DeliveryError};
use crate::session::TurnSnapshot;
use crate::task_supervisor::TaskSupervisor;
use super::coordinator::SinkRoute;
/// Gateway-owned facade that resolves a target Channel and starts exactly one
/// TurnSink lifecycle. Session workers depend on this abstraction rather than
/// on Channel implementations or WebSocket/Feishu protocols.
#[derive(Clone)]
pub struct TurnDeliveryService {
coordinator: DeliveryCoordinator,
channels: ChannelManager,
supervisor: TaskSupervisor,
}
impl TurnDeliveryService {
pub fn new(
coordinator: DeliveryCoordinator,
channels: ChannelManager,
supervisor: TaskSupervisor,
) -> Self {
Self {
coordinator,
channels,
supervisor,
}
}
pub async fn start(
&self,
target: TurnTarget,
snapshots: watch::Receiver<Arc<TurnSnapshot>>,
) -> Result<(), DeliveryError> {
let channel = self
.channels
.get_channel(&target.channel)
.await
.ok_or_else(|| DeliveryError::ChannelNotFound(target.channel.clone()))?;
let live_policy = channel.live_policy();
let presentation = channel.presentation_policy();
let sink = channel
.open_turn(target.clone())
.await
.map_err(DeliveryError::OpenFailed)?;
let _result = self.coordinator.spawn_sink(
&self.supervisor,
SinkRoute {
channel: target.channel,
chat_id: target.chat_id,
live_policy,
presentation,
},
snapshots,
sink,
)?;
Ok(())
}
}

35
src/gateway/mod.rs
View File

@ -12,11 +12,12 @@ use crate::bus::{ControlMessage, MessageBus, OutboundDispatcher};
use crate::channels::base::ChannelError; use crate::channels::base::ChannelError;
use crate::channels::{ChannelManager, CliChatChannel}; use crate::channels::{ChannelManager, CliChatChannel};
use crate::config::{Config, ensure_workspace_dir, expand_path}; use crate::config::{Config, ensure_workspace_dir, expand_path};
use crate::delivery::{ConversationWriteLocks, DeliveryCoordinator, TurnDeliveryService};
use crate::logging; use crate::logging;
use crate::mcp; use crate::mcp;
use crate::memory::MemoryManager; use crate::memory::MemoryManager;
use crate::scheduler::Scheduler; use crate::scheduler::Scheduler;
use crate::session::SessionManager; use crate::session::{SessionManager, SessionManagerServices};
use crate::task_supervisor::TaskSupervisor; use crate::task_supervisor::TaskSupervisor;
pub struct GatewayState { pub struct GatewayState {
@ -27,6 +28,7 @@ pub struct GatewayState {
pub channel_manager: ChannelManager, pub channel_manager: ChannelManager,
pub storage: Arc<crate::storage::Storage>, pub storage: Arc<crate::storage::Storage>,
pub task_supervisor: TaskSupervisor, pub task_supervisor: TaskSupervisor,
pub delivery_coordinator: DeliveryCoordinator,
pub connection_shutdown: tokio_util::sync::CancellationToken, pub connection_shutdown: tokio_util::sync::CancellationToken,
pub auth: auth::AuthManager, pub auth: auth::AuthManager,
pub uploads: uploads::UploadRegistry, pub uploads: uploads::UploadRegistry,
@ -44,6 +46,7 @@ impl GatewayState {
config_path: std::path::PathBuf, config_path: std::path::PathBuf,
) -> Result<Self, Box<dyn std::error::Error>> { ) -> Result<Self, Box<dyn std::error::Error>> {
let task_supervisor = TaskSupervisor::new(); let task_supervisor = TaskSupervisor::new();
let delivery_coordinator = DeliveryCoordinator::new(ConversationWriteLocks::default());
let connection_shutdown = tokio_util::sync::CancellationToken::new(); let connection_shutdown = tokio_util::sync::CancellationToken::new();
let auth = auth::AuthManager::load( let auth = auth::AuthManager::load(
config.gateway.require_pairing, config.gateway.require_pairing,
@ -109,6 +112,20 @@ impl GatewayState {
// Create MessageBus first (shared by SessionManager and ChannelManager) // Create MessageBus first (shared by SessionManager and ChannelManager)
let bus = MessageBus::new(100); let bus = MessageBus::new(100);
// Channels are resolved by TurnDeliveryService, while Session workers
// depend only on that protocol-neutral delivery facade.
let cli_chat_channel = Arc::new(CliChatChannel::with_upload_registry(uploads.clone()));
let channel_manager = ChannelManager::with_bus(cli_chat_channel, bus.clone());
channel_manager
.init(&config, workspace_path.clone())
.await
.map_err(|e| format!("Failed to init channels: {}", e))?;
let turn_delivery = TurnDeliveryService::new(
delivery_coordinator.clone(),
channel_manager.clone(),
task_supervisor.clone(),
);
let browser_config = if config.browser.enabled { let browser_config = if config.browser.enabled {
Some(config.browser.clone()) Some(config.browser.clone())
} else { } else {
@ -119,22 +136,17 @@ impl GatewayState {
let session_manager = SessionManager::new( let session_manager = SessionManager::new(
provider_config.clone(), provider_config.clone(),
storage.clone(), storage.clone(),
SessionManagerServices::new(
bus.clone(), bus.clone(),
memory_manager, memory_manager,
task_supervisor.clone(),
turn_delivery,
),
browser_config, browser_config,
config.gateway.max_concurrent_background_tasks, config.gateway.max_concurrent_background_tasks,
task_supervisor.clone(),
)?; )?;
let session_manager = Arc::new(session_manager); let session_manager = Arc::new(session_manager);
// Create ChannelManager and init channels
let cli_chat_channel = Arc::new(CliChatChannel::with_upload_registry(uploads.clone()));
let channel_manager = ChannelManager::with_bus(cli_chat_channel, bus);
channel_manager
.init(&config, workspace_path.clone())
.await
.map_err(|e| format!("Failed to init channels: {}", e))?;
// Register send_message tool with available channel names // Register send_message tool with available channel names
let available_channels = channel_manager.list_channel_names().await; let available_channels = channel_manager.list_channel_names().await;
let valid_channels = available_channels.clone(); let valid_channels = available_channels.clone();
@ -209,6 +221,7 @@ impl GatewayState {
channel_manager, channel_manager,
storage, storage,
task_supervisor, task_supervisor,
delivery_coordinator,
connection_shutdown, connection_shutdown,
auth, auth,
uploads, uploads,
@ -282,6 +295,7 @@ impl GatewayState {
&inbound.chat_id, &inbound.chat_id,
&inbound.content, &inbound.content,
inbound.media, inbound.media,
inbound.forwarded_metadata.clone(),
).await { ).await {
Ok(crate::session::session::HandleResult::AgentResponse(content)) => { Ok(crate::session::session::HandleResult::AgentResponse(content)) => {
let outbound = crate::bus::OutboundMessage { let outbound = crate::bus::OutboundMessage {
@ -342,6 +356,7 @@ impl GatewayState {
bus_for_outbound, bus_for_outbound,
self.channel_manager.clone(), self.channel_manager.clone(),
self.task_supervisor.clone(), self.task_supervisor.clone(),
self.delivery_coordinator.write_locks(),
); );
self.task_supervisor self.task_supervisor

10
src/gateway/ws.rs
View File

@ -56,10 +56,12 @@ async fn handle_socket(
let _ = sender let _ = sender
.send(WsOutbound::SessionEstablished { .send(WsOutbound::SessionEstablished {
session_id: session_id.clone(), session_id: session_id.clone(),
capabilities: if state.uploads.enabled() { capabilities: {
vec!["file_transfer_v1".to_string()] let mut capabilities = vec!["turn_snapshots_v1".to_string()];
} else { if state.uploads.enabled() {
Vec::new() capabilities.push("file_transfer_v1".to_string());
}
capabilities
}, },
}) })
.await; .await;

1
src/lib.rs
View File

@ -3,6 +3,7 @@ pub mod bus;
pub mod channels; pub mod channels;
pub mod client; pub mod client;
pub mod config; pub mod config;
pub mod delivery;
pub mod gateway; pub mod gateway;
pub mod logging; pub mod logging;
pub mod mcp; pub mod mcp;

56
src/protocol.rs
View File

@ -62,6 +62,10 @@ pub struct HistoryMessage {
pub seq: i64, pub seq: i64,
pub role: String, pub role: String,
pub content: String, pub content: String,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub reasoning_content: Option<String>,
#[serde(default)]
pub completion_status: crate::bus::CompletionStatus,
pub created_at: i64, pub created_at: i64,
#[serde(default, skip_serializing_if = "Option::is_none")] #[serde(default, skip_serializing_if = "Option::is_none")]
pub tool_call_id: Option<String>, pub tool_call_id: Option<String>,
@ -140,6 +144,10 @@ pub enum WsInbound {
#[derive(Debug, Clone, Serialize, Deserialize)] #[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(tag = "type")] #[serde(tag = "type")]
pub enum WsOutbound { pub enum WsOutbound {
#[serde(rename = "turn_updated")]
TurnUpdated {
snapshot: crate::session::TurnSnapshot,
},
#[serde(rename = "assistant_response")] #[serde(rename = "assistant_response")]
AssistantResponse { AssistantResponse {
id: String, id: String,
@ -222,3 +230,51 @@ pub fn serialize_inbound(msg: &WsInbound) -> Result<String, serde_json::Error> {
pub fn serialize_outbound(msg: &WsOutbound) -> Result<String, serde_json::Error> { pub fn serialize_outbound(msg: &WsOutbound) -> Result<String, serde_json::Error> {
serde_json::to_string(msg) serde_json::to_string(msg)
} }
#[cfg(test)]
mod tests {
use super::*;
use crate::session::{TurnId, TurnPhase, TurnState, TurnStatus};
#[test]
fn turn_updated_serializes_as_one_complete_snapshot_frame() {
let frame = WsOutbound::TurnUpdated {
snapshot: TurnState {
id: TurnId("turn-1".into()),
session_id: "cli_chat:client:dialog".into(),
message_id: "message-1".into(),
revision: 7,
status: TurnStatus::Running,
phase: TurnPhase::Responding,
blocks: Vec::new(),
usage: None,
error: None,
},
};
let json = serialize_outbound(&frame).unwrap();
let value: serde_json::Value = serde_json::from_str(&json).unwrap();
assert_eq!(value["type"], "turn_updated");
assert_eq!(value["snapshot"]["revision"], 7);
assert_eq!(value["snapshot"]["status"], "running");
}
#[test]
fn history_defaults_new_reasoning_fields_for_old_frames() {
let message: HistoryMessage = serde_json::from_value(serde_json::json!({
"id": "message",
"seq": 1,
"role": "assistant",
"content": "answer",
"created_at": 1,
"attachments": []
}))
.unwrap();
assert_eq!(message.reasoning_content, None);
assert_eq!(
message.completion_status,
crate::bus::CompletionStatus::Completed
);
}
}

614
src/providers/anthropic.rs
View File

@ -1,12 +1,19 @@
use async_trait::async_trait; use async_trait::async_trait;
use futures_util::stream;
use reqwest::Client; use reqwest::Client;
use serde::{Deserialize, Serialize}; use serde::Serialize;
use std::collections::HashMap; use serde_json::Value;
use std::collections::{BTreeMap, HashMap, VecDeque};
use std::time::Duration; use std::time::Duration;
use thiserror::Error;
use super::stream::SseFramer;
use super::traits::Usage; use super::traits::Usage;
use super::{ChatCompletionRequest, ChatCompletionResponse, LLMProvider, Message, Tool, ToolCall}; use super::{
use crate::bus::message::ContentBlock; ChatCompletionRequest, DynProviderError, FinishReason, LLMProvider, Message, ProviderChunk,
ProviderStream, Tool,
};
use crate::bus::{ProviderReasoningState, message::ContentBlock};
use crate::storage::Storage; use crate::storage::Storage;
use std::sync::Arc; use std::sync::Arc;
@ -127,6 +134,7 @@ struct AnthropicRequest {
messages: Vec<AnthropicMessage>, messages: Vec<AnthropicMessage>,
max_tokens: u32, max_tokens: u32,
temperature: Option<f32>, temperature: Option<f32>,
stream: bool,
#[serde(skip_serializing_if = "Option::is_none")] #[serde(skip_serializing_if = "Option::is_none")]
tools: Option<Vec<AnthropicTool>>, tools: Option<Vec<AnthropicTool>>,
#[serde(flatten)] #[serde(flatten)]
@ -154,6 +162,8 @@ fn convert_messages(messages: &[Message]) -> Vec<AnthropicMessage> {
"tool_use_id": tool_call_id, "tool_use_id": tool_call_id,
"content": convert_content_blocks(&message.content, false), "content": convert_content_blocks(&message.content, false),
})] })]
} else if let Some(native) = native_anthropic_content(message) {
native
} else { } else {
let mut blocks = convert_content_blocks(&message.content, message.role == "system"); let mut blocks = convert_content_blocks(&message.content, message.role == "system");
if let Some(tool_calls) = message if let Some(tool_calls) = message
@ -177,6 +187,26 @@ fn convert_messages(messages: &[Message]) -> Vec<AnthropicMessage> {
.collect() .collect()
} }
fn native_anthropic_content(message: &Message) -> Option<Vec<Value>> {
if message.role != "assistant" {
return None;
}
let state = message.provider_state.as_ref()?;
if state.provider != "anthropic" {
return None;
}
let blocks = state.payload.get("content")?.as_array()?;
blocks
.iter()
.all(|block| {
block
.as_object()
.and_then(|value| value.get("type"))
.is_some()
})
.then(|| blocks.clone())
}
#[derive(Serialize)] #[derive(Serialize)]
struct AnthropicTool { struct AnthropicTool {
name: String, name: String,
@ -186,56 +216,298 @@ struct AnthropicTool {
cache_control: Option<CacheControl>, cache_control: Option<CacheControl>,
} }
#[derive(Deserialize)] #[derive(Debug, Error)]
struct AnthropicResponse { enum AnthropicStreamError {
id: Option<String>, #[error("invalid UTF-8 in Anthropic SSE event: {0}")]
model: Option<String>, Utf8(#[from] std::string::FromUtf8Error),
#[serde(default)] #[error("invalid Anthropic SSE payload: {0}")]
content: Vec<AnthropicContent>, Json(#[from] serde_json::Error),
#[serde(default)] #[error("Anthropic stream error: {0}")]
usage: Option<AnthropicUsage>, Api(String),
#[error("Anthropic stream ended without message_stop")]
MissingFinish,
} }
#[derive(Deserialize)] #[derive(Default)]
#[serde(tag = "type", rename_all = "snake_case")] struct AnthropicSseDecoder {
enum AnthropicContent { framer: SseFramer,
Text { blocks: BTreeMap<usize, Value>,
#[serde(alias = "content")] tool_json: HashMap<usize, String>,
text: String, usage: Usage,
}, finish_reason: Option<FinishReason>,
Thinking { done_emitted: bool,
#[serde(alias = "content")]
thinking: String,
},
#[serde(rename = "tool_use")]
ToolUse {
id: String,
name: String,
#[serde(alias = "arguments")]
input: serde_json::Value,
},
#[serde(other)]
Unknown,
} }
#[derive(Deserialize)] impl AnthropicSseDecoder {
struct AnthropicUsage { fn push(&mut self, bytes: &[u8]) -> Result<Vec<ProviderChunk>, AnthropicStreamError> {
#[serde(default)] let frames = self.framer.push(bytes)?;
input_tokens: u32, self.decode_frames(frames)
#[serde(default)] }
output_tokens: u32,
#[serde(default)] fn finish(&mut self) -> Result<Vec<ProviderChunk>, AnthropicStreamError> {
cache_read_input_tokens: Option<u32>, let frames = self.framer.finish()?;
#[serde(default)] let chunks = self.decode_frames(frames)?;
cache_creation_input_tokens: Option<u32>, if !self.done_emitted {
return Err(AnthropicStreamError::MissingFinish);
}
Ok(chunks)
}
fn decode_frames(
&mut self,
frames: Vec<String>,
) -> Result<Vec<ProviderChunk>, AnthropicStreamError> {
let mut chunks = Vec::new();
for data in frames {
let payload: Value = serde_json::from_str(&data)?;
match payload.get("type").and_then(Value::as_str) {
Some("message_start") => self.message_start(&payload, &mut chunks),
Some("content_block_start") => self.block_start(&payload, &mut chunks),
Some("content_block_delta") => self.block_delta(&payload, &mut chunks),
Some("content_block_stop") => self.block_stop(&payload),
Some("message_delta") => self.message_delta(&payload, &mut chunks),
Some("message_stop") => self.message_stop(&mut chunks),
Some("error") => {
let message = payload
.pointer("/error/message")
.and_then(Value::as_str)
.unwrap_or("unknown streaming error");
return Err(AnthropicStreamError::Api(message.to_string()));
}
Some("ping") | None | Some(_) => {}
}
}
Ok(chunks)
}
fn message_start(&mut self, payload: &Value, chunks: &mut Vec<ProviderChunk>) {
let message = payload.get("message").unwrap_or(&Value::Null);
let id = message
.get("id")
.and_then(Value::as_str)
.unwrap_or_default();
let model = message
.get("model")
.and_then(Value::as_str)
.unwrap_or_default();
chunks.push(ProviderChunk::Metadata {
id: id.to_string(),
model: model.to_string(),
});
if let Some(usage) = message.get("usage") {
update_anthropic_usage(&mut self.usage, usage);
chunks.push(ProviderChunk::Usage(self.usage.clone()));
}
}
fn block_start(&mut self, payload: &Value, chunks: &mut Vec<ProviderChunk>) {
let Some(index) = event_index(payload) else {
return;
};
let Some(block) = payload.get("content_block") else {
return;
};
self.blocks.insert(index, block.clone());
match block.get("type").and_then(Value::as_str) {
Some("text") => {
if let Some(text) = block.get("text").and_then(Value::as_str)
&& !text.is_empty()
{
chunks.push(ProviderChunk::Text(text.to_string()));
}
}
Some("thinking") => {
if let Some(thinking) = block.get("thinking").and_then(Value::as_str)
&& !thinking.is_empty()
{
chunks.push(ProviderChunk::Reasoning(thinking.to_string()));
}
}
Some("tool_use") => {
chunks.push(ProviderChunk::ToolCallStart {
index,
id: block.get("id").and_then(Value::as_str).map(str::to_string),
name: block
.get("name")
.and_then(Value::as_str)
.map(str::to_string),
});
let input = block.get("input").cloned().unwrap_or(Value::Null);
if !input.is_null() && input != serde_json::json!({}) {
chunks.push(ProviderChunk::ToolCallArguments {
index,
delta: input.to_string(),
});
}
}
_ => {}
}
}
fn block_delta(&mut self, payload: &Value, chunks: &mut Vec<ProviderChunk>) {
let Some(index) = event_index(payload) else {
return;
};
let Some(delta) = payload.get("delta") else {
return;
};
match delta.get("type").and_then(Value::as_str) {
Some("text_delta") => {
if let Some(text) = delta.get("text").and_then(Value::as_str) {
append_block_string(&mut self.blocks, index, "text", text);
if !text.is_empty() {
chunks.push(ProviderChunk::Text(text.to_string()));
}
}
}
Some("thinking_delta") => {
if let Some(thinking) = delta.get("thinking").and_then(Value::as_str) {
append_block_string(&mut self.blocks, index, "thinking", thinking);
if !thinking.is_empty() {
chunks.push(ProviderChunk::Reasoning(thinking.to_string()));
}
}
}
Some("signature_delta") => {
if let Some(signature) = delta.get("signature").and_then(Value::as_str) {
append_block_string(&mut self.blocks, index, "signature", signature);
}
}
Some("input_json_delta") => {
if let Some(partial) = delta.get("partial_json").and_then(Value::as_str) {
self.tool_json.entry(index).or_default().push_str(partial);
if !partial.is_empty() {
chunks.push(ProviderChunk::ToolCallArguments {
index,
delta: partial.to_string(),
});
}
}
}
_ => {}
}
}
fn block_stop(&mut self, payload: &Value) {
let Some(index) = event_index(payload) else {
return;
};
let Some(json) = self.tool_json.remove(&index) else {
return;
};
let input = serde_json::from_str(&json).unwrap_or(Value::Null);
if let Some(block) = self.blocks.get_mut(&index)
&& let Some(object) = block.as_object_mut()
{
object.insert("input".to_string(), input);
}
}
fn message_delta(&mut self, payload: &Value, chunks: &mut Vec<ProviderChunk>) {
if let Some(reason) = payload
.pointer("/delta/stop_reason")
.and_then(Value::as_str)
{
self.finish_reason = Some(FinishReason::from_provider(reason));
}
if let Some(usage) = payload.get("usage") {
update_anthropic_usage(&mut self.usage, usage);
chunks.push(ProviderChunk::Usage(self.usage.clone()));
}
}
fn message_stop(&mut self, chunks: &mut Vec<ProviderChunk>) {
if self.done_emitted {
return;
}
let content = self.blocks.values().cloned().collect::<Vec<_>>();
chunks.push(ProviderChunk::ProviderState(ProviderReasoningState {
provider: "anthropic".to_string(),
payload: serde_json::json!({ "version": 1, "content": content }),
}));
chunks.push(ProviderChunk::Done(
self.finish_reason.clone().unwrap_or(FinishReason::Stop),
));
self.done_emitted = true;
}
}
fn event_index(payload: &Value) -> Option<usize> {
payload
.get("index")
.and_then(Value::as_u64)
.and_then(|value| usize::try_from(value).ok())
}
fn append_block_string(blocks: &mut BTreeMap<usize, Value>, index: usize, key: &str, delta: &str) {
let Some(object) = blocks.get_mut(&index).and_then(Value::as_object_mut) else {
return;
};
let value = object
.entry(key.to_string())
.or_insert_with(|| Value::String(String::new()));
if let Some(current) = value.as_str() {
*value = Value::String(format!("{current}{delta}"));
}
}
fn update_anthropic_usage(usage: &mut Usage, value: &Value) {
if let Some(input) = json_u32(value, "input_tokens") {
usage.prompt_tokens = input;
}
if let Some(output) = json_u32(value, "output_tokens") {
usage.completion_tokens = output;
}
if let Some(cache_read) = json_u32(value, "cache_read_input_tokens") {
usage.cached_tokens = Some(cache_read);
usage.cache_read_input_tokens = Some(cache_read);
}
if let Some(cache_creation) = json_u32(value, "cache_creation_input_tokens") {
usage.cache_creation_input_tokens = Some(cache_creation);
}
usage.total_tokens = usage.prompt_tokens.saturating_add(usage.completion_tokens);
}
fn json_u32(value: &Value, key: &str) -> Option<u32> {
value
.get(key)
.and_then(Value::as_u64)
.and_then(|number| u32::try_from(number).ok())
}
struct AnthropicHttpStream {
response: reqwest::Response,
decoder: AnthropicSseDecoder,
pending: VecDeque<ProviderChunk>,
reached_eof: bool,
}
async fn next_anthropic_chunk(
mut state: AnthropicHttpStream,
) -> Result<Option<(ProviderChunk, AnthropicHttpStream)>, DynProviderError> {
loop {
if let Some(chunk) = state.pending.pop_front() {
return Ok(Some((chunk, state)));
}
if state.reached_eof {
return Ok(None);
}
match state.response.chunk().await? {
Some(bytes) => state.pending.extend(state.decoder.push(&bytes)?),
None => {
state.pending.extend(state.decoder.finish()?);
state.reached_eof = true;
}
}
}
} }
#[async_trait] #[async_trait]
impl LLMProvider for AnthropicProvider { impl LLMProvider for AnthropicProvider {
async fn chat( async fn stream(
&self, &self,
request: ChatCompletionRequest, request: ChatCompletionRequest,
) -> Result<ChatCompletionResponse, Box<dyn std::error::Error + Send + Sync>> { ) -> Result<ProviderStream, DynProviderError> {
let start = std::time::Instant::now(); let start = std::time::Instant::now();
let url = format!("{}/v1/messages", self.base_url); let url = format!("{}/v1/messages", self.base_url);
let max_tokens = request.max_tokens.or(self.max_tokens).unwrap_or(1024); let max_tokens = request.max_tokens.or(self.max_tokens).unwrap_or(1024);
@ -257,6 +529,7 @@ impl LLMProvider for AnthropicProvider {
messages: convert_messages(&request.messages), messages: convert_messages(&request.messages),
max_tokens, max_tokens,
temperature: request.temperature.or(self.temperature), temperature: request.temperature.or(self.temperature),
stream: true,
tools, tools,
extra: self.model_extra.clone(), extra: self.model_extra.clone(),
}; };
@ -272,8 +545,16 @@ impl LLMProvider for AnthropicProvider {
req_builder = req_builder.header(key.as_str(), value.as_str()); req_builder = req_builder.header(key.as_str(), value.as_str());
} }
let req_body_str = serde_json::to_string_pretty(&body).unwrap_or_default(); let request_summary = super::stream::diagnostic_request_summary(
tracing::debug!(req_body = %req_body_str, "LLM request"); &self.model_id,
body.messages.len(),
body.tools.as_ref().map_or(0, Vec::len),
);
tracing::debug!(
message_count = body.messages.len(),
tool_count = body.tools.as_ref().map_or(0, Vec::len),
"Anthropic streaming request"
);
let resp = req_builder.json(&body).send().await.inspect_err(|e| { let resp = req_builder.json(&body).send().await.inspect_err(|e| {
let is_timeout = e.is_timeout(); let is_timeout = e.is_timeout();
@ -289,10 +570,8 @@ impl LLMProvider for AnthropicProvider {
})?; })?;
let status = resp.status(); let status = resp.status();
let body_text = resp.text().await?;
tracing::debug!(status = %status, resp_body = %body_text, "LLM response");
if !status.is_success() { if !status.is_success() {
let body_text = resp.text().await?;
let error_msg = serde_json::from_str::<serde_json::Value>(&body_text) let error_msg = serde_json::from_str::<serde_json::Value>(&body_text)
.ok() .ok()
.and_then(|v| { .and_then(|v| {
@ -315,7 +594,7 @@ impl LLMProvider for AnthropicProvider {
.append_llm_call( .append_llm_call(
&self.name, &self.name,
&self.model_id, &self.model_id,
&req_body_str, &request_summary,
Some(&body_text), Some(&body_text),
Some(&error_msg), Some(&error_msg),
start.elapsed().as_millis() as u64, start.elapsed().as_millis() as u64,
@ -324,110 +603,16 @@ impl LLMProvider for AnthropicProvider {
} }
return Err(format!("API error ({}): {}", status.as_u16(), error_msg).into()); return Err(format!("API error ({}): {}", status.as_u16(), error_msg).into());
} }
tracing::debug!(status = %status, "Anthropic streaming response started");
let anthropic_resp: AnthropicResponse = match serde_json::from_str(&body_text) { Ok(Box::pin(stream::try_unfold(
Ok(response) => response, AnthropicHttpStream {
Err(e) => { response: resp,
let err_msg = format!("decode error: {} | body: {}", e, &body_text); decoder: AnthropicSseDecoder::default(),
if let Some(ref storage) = self.storage { pending: VecDeque::new(),
let dur = start.elapsed().as_millis() as u64; reached_eof: false,
if let Err(error) = storage
.append_llm_call(
&self.name,
&self.model_id,
&req_body_str,
Some(&body_text),
Some(&err_msg),
dur,
)
.await
{
tracing::warn!("failed to persist LLM call (decode error): {}", error);
}
}
return Err(err_msg.into());
}
};
let mut content = String::new();
let mut reasoning = None;
let mut tool_calls = Vec::new();
for c in &anthropic_resp.content {
match c {
AnthropicContent::Text { text } => {
if !text.is_empty() {
if !content.is_empty() {
content.push('\n');
}
content.push_str(text);
}
}
AnthropicContent::Thinking { thinking } => {
reasoning = Some(thinking.clone());
}
AnthropicContent::Unknown => {}
AnthropicContent::ToolUse { id, name, input } => {
tool_calls.push(ToolCall {
id: id.clone(),
name: name.clone(),
arguments: input.clone(),
});
}
}
}
let response = ChatCompletionResponse {
id: anthropic_resp.id.unwrap_or_default(),
model: anthropic_resp.model.unwrap_or_default(),
content,
reasoning_content: reasoning,
tool_calls,
usage: Usage {
prompt_tokens: anthropic_resp
.usage
.as_ref()
.map(|u| u.input_tokens)
.unwrap_or(0),
completion_tokens: anthropic_resp
.usage
.as_ref()
.map(|u| u.output_tokens)
.unwrap_or(0),
total_tokens: anthropic_resp
.usage
.as_ref()
.map(|u| u.input_tokens + u.output_tokens)
.unwrap_or(0),
cached_tokens: anthropic_resp
.usage
.as_ref()
.and_then(|u| u.cache_read_input_tokens),
cache_read_input_tokens: anthropic_resp
.usage
.as_ref()
.and_then(|u| u.cache_read_input_tokens),
cache_creation_input_tokens: anthropic_resp
.usage
.as_ref()
.and_then(|u| u.cache_creation_input_tokens),
}, },
}; next_anthropic_chunk,
)))
if let Some(ref storage) = self.storage {
let _ = storage
.append_llm_call(
&self.name,
&self.model_id,
&req_body_str,
Some(&body_text),
None,
start.elapsed().as_millis() as u64,
)
.await;
}
Ok(response)
} }
fn ptype(&self) -> &str { fn ptype(&self) -> &str {
@ -446,6 +631,7 @@ impl LLMProvider for AnthropicProvider {
#[cfg(test)] #[cfg(test)]
mod tests { mod tests {
use super::*; use super::*;
use crate::providers::ProviderResponseAccumulator;
use serde_json::json; use serde_json::json;
#[test] #[test]
@ -487,6 +673,7 @@ mod tests {
ContentBlock::image_url("data:image/png;base64,AAAA"), ContentBlock::image_url("data:image/png;base64,AAAA"),
], ],
reasoning_content: None, reasoning_content: None,
provider_state: None,
tool_call_id: Some("call_1".to_string()), tool_call_id: Some("call_1".to_string()),
name: Some("file_read".to_string()), name: Some("file_read".to_string()),
tool_calls: None, tool_calls: None,
@ -502,4 +689,125 @@ mod tests {
assert_eq!(result["content"][1]["source"]["media_type"], "image/png"); assert_eq!(result["content"][1]["source"]["media_type"], "image/png");
assert_eq!(result["content"][1]["source"]["data"], "AAAA"); assert_eq!(result["content"][1]["source"]["data"], "AAAA");
} }
#[test]
fn native_stream_decodes_thinking_signature_tools_usage_and_replay_state() {
let events = [
json!({"type":"message_start","message":{"id":"msg_1","model":"claude-test","usage":{"input_tokens":11,"output_tokens":0,"cache_read_input_tokens":3}}}),
json!({"type":"content_block_start","index":0,"content_block":{"type":"thinking","thinking":"","signature":""}}),
json!({"type":"content_block_delta","index":0,"delta":{"type":"thinking_delta","thinking":"check "}}),
json!({"type":"content_block_delta","index":0,"delta":{"type":"thinking_delta","thinking":"facts"}}),
json!({"type":"content_block_delta","index":0,"delta":{"type":"signature_delta","signature":"sig=="}}),
json!({"type":"content_block_stop","index":0}),
json!({"type":"content_block_start","index":1,"content_block":{"type":"tool_use","id":"tool_1","name":"lookup","input":{}}}),
json!({"type":"content_block_delta","index":1,"delta":{"type":"input_json_delta","partial_json":"{\"q\":"}}),
json!({"type":"content_block_delta","index":1,"delta":{"type":"input_json_delta","partial_json":"\"rust\"}"}}),
json!({"type":"content_block_stop","index":1}),
json!({"type":"content_block_start","index":2,"content_block":{"type":"text","text":""}}),
json!({"type":"content_block_delta","index":2,"delta":{"type":"text_delta","text":"answer"}}),
json!({"type":"content_block_stop","index":2}),
json!({"type":"message_delta","delta":{"stop_reason":"tool_use","stop_sequence":null},"usage":{"output_tokens":7}}),
json!({"type":"message_stop"}),
];
let wire = events
.iter()
.map(|event| format!("event: ignored\ndata: {event}\n\n"))
.collect::<String>();
let mut decoder = AnthropicSseDecoder::default();
let mut accumulator = ProviderResponseAccumulator::default();
for bytes in wire.as_bytes().chunks(7) {
for chunk in decoder.push(bytes).unwrap() {
accumulator.push(chunk);
}
}
for chunk in decoder.finish().unwrap() {
accumulator.push(chunk);
}
let response = accumulator.finish();
assert_eq!(response.id, "msg_1");
assert_eq!(response.model, "claude-test");
assert_eq!(response.reasoning_content.as_deref(), Some("check facts"));
assert_eq!(response.content, "answer");
assert_eq!(response.tool_calls.len(), 1);
assert_eq!(response.tool_calls[0].id, "tool_1");
assert_eq!(response.tool_calls[0].arguments, json!({"q":"rust"}));
assert_eq!(response.usage.prompt_tokens, 11);
assert_eq!(response.usage.completion_tokens, 7);
assert_eq!(response.usage.total_tokens, 18);
assert_eq!(response.usage.cache_read_input_tokens, Some(3));
let state = response.provider_state.unwrap();
assert_eq!(state.provider, "anthropic");
assert_eq!(state.payload["content"][0]["thinking"], "check facts");
assert_eq!(state.payload["content"][0]["signature"], "sig==");
assert_eq!(state.payload["content"][1]["input"], json!({"q":"rust"}));
assert_eq!(state.payload["content"][2]["text"], "answer");
}
#[test]
fn matching_provider_state_replays_native_blocks_without_generic_duplicates() {
let native = json!([
{"type":"thinking","thinking":"signed thought","signature":"sig=="},
{"type":"tool_use","id":"tool_1","name":"lookup","input":{"q":"rust"}}
]);
let message = Message {
role: "assistant".into(),
content: vec![ContentBlock::text("generic text must not be appended")],
reasoning_content: Some("display copy".into()),
provider_state: Some(ProviderReasoningState {
provider: "anthropic".into(),
payload: json!({"version":1,"content":native}),
}),
tool_call_id: None,
name: None,
tool_calls: Some(vec![crate::providers::ToolCall {
id: "duplicate".into(),
name: "duplicate".into(),
arguments: json!({}),
}]),
};
let converted = convert_messages(&[message]);
assert_eq!(converted[0].content.len(), 2);
assert_eq!(converted[0].content[0]["signature"], "sig==");
assert_eq!(converted[0].content[1]["id"], "tool_1");
}
#[test]
fn foreign_provider_state_is_not_replayed_to_anthropic() {
let message = Message {
role: "assistant".into(),
content: vec![ContentBlock::text("answer")],
reasoning_content: Some("unsigned display reasoning".into()),
provider_state: Some(ProviderReasoningState {
provider: "openai".into(),
payload: json!({"private":"state"}),
}),
tool_call_id: None,
name: None,
tool_calls: None,
};
let converted = convert_messages(&[message]);
assert_eq!(
converted[0].content,
vec![json!({"type":"text","text":"answer"})]
);
}
#[test]
fn stream_requires_message_stop() {
let mut decoder = AnthropicSseDecoder::default();
decoder
.push(b"data: {\"type\":\"message_start\",\"message\":{}}\n\n")
.unwrap();
assert!(matches!(
decoder.finish(),
Err(AnthropicStreamError::MissingFinish)
));
}
} }

7
src/providers/mod.rs
View File

@ -1,11 +1,18 @@
pub mod anthropic; pub mod anthropic;
pub mod openai; pub mod openai;
pub mod stream;
pub mod traits; pub mod traits;
pub use self::anthropic::AnthropicProvider; pub use self::anthropic::AnthropicProvider;
pub use self::openai::OpenAIProvider; pub use self::openai::OpenAIProvider;
use crate::config::LLMProviderConfig; use crate::config::LLMProviderConfig;
#[cfg(test)]
pub use stream::provider_stream_for_test;
pub use stream::{
DynProviderError, FinishReason, ProviderChunk, ProviderResponseAccumulator, ProviderStream,
ProviderStreamItem, collect_provider_stream,
};
pub use traits::{ pub use traits::{
ChatCompletionRequest, ChatCompletionResponse, LLMProvider, Message, Tool, ToolCall, ChatCompletionRequest, ChatCompletionResponse, LLMProvider, Message, Tool, ToolCall,
ToolFunction, Usage, ToolFunction, Usage,

662
src/providers/openai.rs
View File

@ -1,12 +1,16 @@
use async_trait::async_trait; use async_trait::async_trait;
use futures_util::stream;
use reqwest::Client; use reqwest::Client;
use serde::Deserialize;
use serde_json::{Value, json}; use serde_json::{Value, json};
use std::collections::HashMap; use std::collections::{HashMap, VecDeque};
use std::time::Duration; use std::time::Duration;
use thiserror::Error;
use super::traits::Usage; use super::stream::SseFramer;
use super::{ChatCompletionRequest, ChatCompletionResponse, LLMProvider, Message, ToolCall}; use super::{
ChatCompletionRequest, DynProviderError, FinishReason, LLMProvider, Message, ProviderChunk,
ProviderStream, Usage,
};
use crate::bus::message::ContentBlock; use crate::bus::message::ContentBlock;
use crate::storage::Storage; use crate::storage::Storage;
use std::sync::Arc; use std::sync::Arc;
@ -208,82 +212,348 @@ impl OpenAIProvider {
body body
} }
fn build_stream_request_body(&self, request: &ChatCompletionRequest) -> Value {
let mut body = self.build_request_body(request);
body["stream"] = Value::Bool(true);
body["stream_options"] = json!({ "include_usage": true });
body
}
} }
#[derive(Deserialize)] #[derive(Debug, Error)]
struct OpenAIResponse { enum OpenAIStreamError {
id: String, #[error("invalid UTF-8 in SSE event: {0}")]
model: String, Utf8(#[from] std::string::FromUtf8Error),
choices: Vec<OpenAIChoice>, #[error("invalid OpenAI-compatible SSE payload: {0}")]
#[serde(default)] Json(#[from] serde_json::Error),
usage: OpenAIUsage, #[error("OpenAI-compatible stream ended without a finish marker")]
MissingFinish,
} }
#[derive(Deserialize)] #[derive(Clone, Copy, PartialEq, Eq)]
struct OpenAIChoice { enum InlineMode {
message: OpenAIMessage, Text,
Reasoning,
} }
fn null_or_missing_tool_calls<'de, D>(deserializer: D) -> Result<Vec<OpenAIToolCall>, D::Error> struct InlineReasoningParser {
where mode: InlineMode,
D: serde::Deserializer<'de>, pending: String,
{
Ok(Option::<Vec<OpenAIToolCall>>::deserialize(deserializer)?.unwrap_or_default())
} }
#[derive(Deserialize)] impl Default for InlineReasoningParser {
struct OpenAIMessage { fn default() -> Self {
#[serde(default)] Self {
content: Option<String>, mode: InlineMode::Text,
#[serde(default)] pending: String::new(),
reasoning_content: Option<String>, }
#[serde(default, deserialize_with = "null_or_missing_tool_calls")] }
tool_calls: Vec<OpenAIToolCall>,
} }
#[derive(Deserialize)] impl InlineReasoningParser {
struct OpenAIToolCall { const TAGS: [(&'static str, InlineMode); 4] = [
id: String, ("<think>", InlineMode::Reasoning),
#[serde(rename = "function")] ("<reasoning>", InlineMode::Reasoning),
function: OAIFunction, ("</think>", InlineMode::Text),
("</reasoning>", InlineMode::Text),
];
fn push(&mut self, delta: &str) -> Vec<ProviderChunk> {
self.pending.push_str(delta);
self.drain(false)
}
fn finish(&mut self) -> Vec<ProviderChunk> {
self.drain(true)
}
fn drain(&mut self, finish: bool) -> Vec<ProviderChunk> {
let mut chunks = Vec::new();
loop {
let next_tag = Self::TAGS
.iter()
.filter_map(|(tag, mode)| self.pending.find(tag).map(|index| (index, *tag, *mode)))
.min_by_key(|(index, _, _)| *index);
if let Some((index, tag, mode)) = next_tag {
let text = self.pending[..index].to_string();
self.emit_text(text, &mut chunks);
self.pending.drain(..index + tag.len());
self.mode = mode;
continue;
}
let retained = if finish {
0
} else {
longest_tag_prefix_suffix(&self.pending, &Self::TAGS)
};
let emit_len = self.pending.len() - retained;
if emit_len > 0 {
let text = self.pending[..emit_len].to_string();
self.pending.drain(..emit_len);
self.emit_text(text, &mut chunks);
}
break;
}
chunks
}
fn emit_text(&self, text: String, chunks: &mut Vec<ProviderChunk>) {
if text.is_empty() {
return;
}
chunks.push(match self.mode {
InlineMode::Text => ProviderChunk::Text(text),
InlineMode::Reasoning => ProviderChunk::Reasoning(text),
});
}
} }
#[derive(Deserialize)] fn longest_tag_prefix_suffix(value: &str, tags: &[(&str, InlineMode)]) -> usize {
struct OAIFunction { let mut best = 0;
name: String, for boundary in value
arguments: String, .char_indices()
.map(|(index, _)| index)
.chain([value.len()])
{
let suffix = &value[boundary..];
if tags.iter().any(|(tag, _)| tag.starts_with(suffix)) {
best = best.max(suffix.len());
}
}
best
} }
#[derive(Deserialize, Default)] #[derive(Default)]
struct OpenAIUsage { struct PartialStreamTool {
#[serde(default)] id: Option<String>,
prompt_tokens: u32, name: Option<String>,
#[serde(default)] started: bool,
completion_tokens: u32,
#[serde(default)]
total_tokens: u32,
#[serde(default)]
cached_tokens: Option<u32>,
#[serde(default)]
prompt_tokens_details: Option<OpenAIPromptTokensDetails>,
} }
#[derive(Deserialize, Default)] #[derive(Default)]
struct OpenAIPromptTokensDetails { struct OpenAISseDecoder {
#[serde(default)] framer: SseFramer,
cached_tokens: Option<u32>, inline_reasoning: InlineReasoningParser,
tools: HashMap<usize, PartialStreamTool>,
metadata_emitted: bool,
done_emitted: bool,
}
impl OpenAISseDecoder {
fn push(&mut self, bytes: &[u8]) -> Result<Vec<ProviderChunk>, OpenAIStreamError> {
let frames = self.framer.push(bytes)?;
self.decode_frames(frames)
}
fn finish(&mut self) -> Result<Vec<ProviderChunk>, OpenAIStreamError> {
let frames = self.framer.finish()?;
let mut chunks = self.decode_frames(frames)?;
chunks.extend(self.inline_reasoning.finish());
if !self.done_emitted {
return Err(OpenAIStreamError::MissingFinish);
}
Ok(chunks)
}
fn decode_frames(
&mut self,
frames: Vec<String>,
) -> Result<Vec<ProviderChunk>, OpenAIStreamError> {
let mut chunks = Vec::new();
for data in frames {
if data == "[DONE]" {
chunks.extend(self.inline_reasoning.finish());
if !self.done_emitted {
chunks.push(ProviderChunk::Done(FinishReason::Stop));
self.done_emitted = true;
}
continue;
}
let payload: Value = serde_json::from_str(&data)?;
if !self.metadata_emitted {
let id = payload
.get("id")
.and_then(Value::as_str)
.unwrap_or_default();
let model = payload
.get("model")
.and_then(Value::as_str)
.unwrap_or_default();
if !id.is_empty() || !model.is_empty() {
chunks.push(ProviderChunk::Metadata {
id: id.to_string(),
model: model.to_string(),
});
self.metadata_emitted = true;
}
}
if let Some(usage) = payload.get("usage").filter(|value| !value.is_null()) {
chunks.push(ProviderChunk::Usage(parse_openai_usage(usage)));
}
let Some(choice) = payload
.get("choices")
.and_then(Value::as_array)
.and_then(|choices| choices.first())
else {
continue;
};
if let Some(delta) = choice.get("delta") {
if let Some(reasoning) = delta
.get("reasoning_content")
.and_then(Value::as_str)
.or_else(|| delta.get("reasoning").and_then(Value::as_str))
.filter(|value| !value.is_empty())
{
chunks.push(ProviderChunk::Reasoning(reasoning.to_string()));
}
if let Some(content) = delta
.get("content")
.and_then(Value::as_str)
.filter(|value| !value.is_empty())
{
chunks.extend(self.inline_reasoning.push(content));
}
if let Some(tool_calls) = delta.get("tool_calls").and_then(Value::as_array) {
self.decode_tool_calls(tool_calls, &mut chunks);
}
}
if let Some(reason) = choice.get("finish_reason").and_then(Value::as_str) {
chunks.extend(self.inline_reasoning.finish());
self.flush_unstarted_tools(&mut chunks);
if !self.done_emitted {
chunks.push(ProviderChunk::Done(FinishReason::from_provider(reason)));
self.done_emitted = true;
}
}
}
Ok(chunks)
}
fn decode_tool_calls(&mut self, calls: &[Value], chunks: &mut Vec<ProviderChunk>) {
for (fallback_index, call) in calls.iter().enumerate() {
let index = call
.get("index")
.and_then(Value::as_u64)
.and_then(|value| usize::try_from(value).ok())
.unwrap_or(fallback_index);
let tool = self.tools.entry(index).or_default();
if let Some(id) = call.get("id").and_then(Value::as_str) {
tool.id = Some(id.to_string());
}
let function = call.get("function");
if let Some(name) = function
.and_then(|value| value.get("name"))
.and_then(Value::as_str)
.filter(|value| !value.is_empty())
{
tool.name = Some(name.to_string());
}
if !tool.started && tool.name.is_some() {
chunks.push(ProviderChunk::ToolCallStart {
index,
id: tool.id.clone(),
name: tool.name.clone(),
});
tool.started = true;
}
if let Some(arguments) = function
.and_then(|value| value.get("arguments"))
.and_then(Value::as_str)
.filter(|value| !value.is_empty())
{
chunks.push(ProviderChunk::ToolCallArguments {
index,
delta: arguments.to_string(),
});
}
}
}
fn flush_unstarted_tools(&mut self, chunks: &mut Vec<ProviderChunk>) {
let mut indexes = self.tools.keys().copied().collect::<Vec<_>>();
indexes.sort_unstable();
for index in indexes {
let tool = self
.tools
.get_mut(&index)
.expect("tool index came from map");
if !tool.started {
chunks.push(ProviderChunk::ToolCallStart {
index,
id: tool.id.clone(),
name: tool.name.clone(),
});
tool.started = true;
}
}
}
}
fn parse_openai_usage(value: &Value) -> Usage {
let direct_cached = value.get("cached_tokens").and_then(Value::as_u64);
let nested_cached = value
.get("prompt_tokens_details")
.and_then(|details| details.get("cached_tokens"))
.and_then(Value::as_u64);
Usage {
prompt_tokens: json_u32(value, "prompt_tokens"),
completion_tokens: json_u32(value, "completion_tokens"),
total_tokens: json_u32(value, "total_tokens"),
cached_tokens: nested_cached
.or(direct_cached)
.and_then(|tokens| u32::try_from(tokens).ok()),
cache_read_input_tokens: None,
cache_creation_input_tokens: None,
}
}
fn json_u32(value: &Value, key: &str) -> u32 {
value
.get(key)
.and_then(Value::as_u64)
.and_then(|number| u32::try_from(number).ok())
.unwrap_or_default()
}
struct OpenAIHttpStream {
response: reqwest::Response,
decoder: OpenAISseDecoder,
pending: VecDeque<ProviderChunk>,
reached_eof: bool,
}
async fn next_openai_chunk(
mut state: OpenAIHttpStream,
) -> Result<Option<(ProviderChunk, OpenAIHttpStream)>, DynProviderError> {
loop {
if let Some(chunk) = state.pending.pop_front() {
return Ok(Some((chunk, state)));
}
if state.reached_eof {
return Ok(None);
}
match state.response.chunk().await? {
Some(bytes) => state.pending.extend(state.decoder.push(&bytes)?),
None => {
state.pending.extend(state.decoder.finish()?);
state.reached_eof = true;
}
}
}
} }
#[async_trait] #[async_trait]
impl LLMProvider for OpenAIProvider { impl LLMProvider for OpenAIProvider {
async fn chat( async fn stream(
&self, &self,
request: ChatCompletionRequest, request: ChatCompletionRequest,
) -> Result<ChatCompletionResponse, Box<dyn std::error::Error + Send + Sync>> { ) -> Result<ProviderStream, DynProviderError> {
let start = std::time::Instant::now(); let start = std::time::Instant::now();
let url = format!("{}/chat/completions", self.base_url); let url = format!("{}/chat/completions", self.base_url);
let body = self.build_request_body(&request); let body = self.build_stream_request_body(&request);
// Debug: Log LLM request summary (only in debug builds) // Debug: Log LLM request summary (only in debug builds)
#[cfg(debug_assertions)] #[cfg(debug_assertions)]
@ -322,8 +592,11 @@ impl LLMProvider for OpenAIProvider {
req_builder = req_builder.header(key.as_str(), value.as_str()); req_builder = req_builder.header(key.as_str(), value.as_str());
} }
let req_body_str = serde_json::to_string_pretty(&body).unwrap_or_default(); let request_summary = super::stream::diagnostic_request_summary(
tracing::debug!(req_body = %req_body_str, "LLM request"); &self.model_id,
body["messages"].as_array().map_or(0, Vec::len),
body["tools"].as_array().map_or(0, Vec::len),
);
let resp = req_builder.json(&body).send().await.inspect_err(|e| { let resp = req_builder.json(&body).send().await.inspect_err(|e| {
let is_timeout = e.is_timeout(); let is_timeout = e.is_timeout();
@ -339,10 +612,8 @@ impl LLMProvider for OpenAIProvider {
})?; })?;
let status = resp.status(); let status = resp.status();
let text = resp.text().await?;
tracing::debug!(status = %status, resp_body = %text, "LLM response");
if !status.is_success() { if !status.is_success() {
let text = resp.text().await?;
let error = format!("API error {}: {}", status, text); let error = format!("API error {}: {}", status, text);
tracing::error!( tracing::error!(
provider = %self.name, provider = %self.name,
@ -357,7 +628,7 @@ impl LLMProvider for OpenAIProvider {
.append_llm_call( .append_llm_call(
&self.name, &self.name,
&self.model_id, &self.model_id,
&req_body_str, &request_summary,
Some(&text), Some(&text),
Some(&error), Some(&error),
start.elapsed().as_millis() as u64, start.elapsed().as_millis() as u64,
@ -369,93 +640,13 @@ impl LLMProvider for OpenAIProvider {
return Err(error.into()); return Err(error.into());
} }
let openai_resp: OpenAIResponse = match serde_json::from_str(&text) { let state = OpenAIHttpStream {
Ok(response) => response, response: resp,
Err(e) => { decoder: OpenAISseDecoder::default(),
let err_msg = format!("decode error: {} | body: {}", e, &text); pending: VecDeque::new(),
if let Some(ref storage) = self.storage { reached_eof: false,
let dur = start.elapsed().as_millis() as u64;
if let Err(error) = storage
.append_llm_call(
&self.name,
&self.model_id,
&req_body_str,
Some(&text),
Some(&err_msg),
dur,
)
.await
{
tracing::warn!("failed to persist LLM call (decode error): {}", error);
}
}
return Err(err_msg.into());
}
}; };
Ok(Box::pin(stream::try_unfold(state, next_openai_chunk)))
let first_choice = openai_resp
.choices
.into_iter()
.next()
.ok_or("no choices in response")?;
let content = first_choice
.message
.content
.as_ref()
.unwrap_or(&String::new())
.clone();
let tool_calls: Vec<ToolCall> = first_choice
.message
.tool_calls
.iter()
.map(|tc| ToolCall {
id: tc.id.clone(),
name: tc.function.name.clone(),
arguments: serde_json::from_str(&tc.function.arguments)
.unwrap_or(serde_json::Value::Null),
})
.collect();
let usage = openai_resp.usage;
let nested_cached_tokens = usage
.prompt_tokens_details
.as_ref()
.and_then(|d| d.cached_tokens);
let cached_tokens = nested_cached_tokens.or(usage.cached_tokens);
let response = ChatCompletionResponse {
id: openai_resp.id,
model: openai_resp.model,
content,
reasoning_content: first_choice.message.reasoning_content,
tool_calls,
usage: Usage {
prompt_tokens: usage.prompt_tokens,
completion_tokens: usage.completion_tokens,
total_tokens: usage.total_tokens,
cached_tokens,
cache_read_input_tokens: None,
cache_creation_input_tokens: None,
},
};
if let Some(ref storage) = self.storage
&& let Err(e) = storage
.append_llm_call(
&self.name,
&self.model_id,
&req_body_str,
Some(&text),
None,
start.elapsed().as_millis() as u64,
)
.await
{
tracing::warn!("failed to persist LLM call: {}", e);
}
Ok(response)
} }
fn ptype(&self) -> &str { fn ptype(&self) -> &str {
@ -474,7 +665,7 @@ impl LLMProvider for OpenAIProvider {
#[cfg(test)] #[cfg(test)]
mod tests { mod tests {
use super::*; use super::*;
use crate::providers::Message; use crate::providers::{Message, ToolCall};
#[test] #[test]
fn test_build_request_body_includes_assistant_tool_calls() { fn test_build_request_body_includes_assistant_tool_calls() {
@ -494,6 +685,7 @@ mod tests {
role: "assistant".to_string(), role: "assistant".to_string(),
content: vec![ContentBlock::text("calling tool")], content: vec![ContentBlock::text("calling tool")],
reasoning_content: None, reasoning_content: None,
provider_state: None,
tool_call_id: None, tool_call_id: None,
name: None, name: None,
tool_calls: Some(vec![ToolCall { tool_calls: Some(vec![ToolCall {
@ -519,6 +711,10 @@ mod tests {
tool_calls[0]["function"]["arguments"], tool_calls[0]["function"]["arguments"],
"{\"expression\":\"1+1\"}" "{\"expression\":\"1+1\"}"
); );
let stream_body = provider.build_stream_request_body(&request);
assert_eq!(stream_body["stream"], true);
assert_eq!(stream_body["stream_options"]["include_usage"], true);
} }
#[test] #[test]
@ -532,6 +728,7 @@ mod tests {
ContentBlock::image_url("data:image/png;base64,AAAA"), ContentBlock::image_url("data:image/png;base64,AAAA"),
], ],
reasoning_content: None, reasoning_content: None,
provider_state: None,
tool_call_id: Some("call_2".to_string()), tool_call_id: Some("call_2".to_string()),
name: Some("file_read".to_string()), name: Some("file_read".to_string()),
tool_calls: None, tool_calls: None,
@ -552,76 +749,113 @@ mod tests {
assert_eq!(converted[1]["content"], "second image"); assert_eq!(converted[1]["content"], "second image");
} }
#[test] #[tokio::test]
fn test_decode_response_accepts_null_tool_calls() { async fn sse_decoder_handles_byte_boundaries_reasoning_content_and_usage() {
let text = r#"{ let input = concat!(
"id": "d21abaa6552741949e2aba76bde59359", "data: {\"id\":\"r1\",\"model\":\"m1\",\"choices\":[{\"delta\":{\"reasoning_content\":\"why\",\"content\":\"你好\"},\"finish_reason\":null}]}\r\n\r\n",
"choices": [{ "data: {\"choices\":[{\"delta\":{},\"finish_reason\":\"stop\"}]}\n\n",
"finish_reason": "stop", "data: {\"choices\":[],\"usage\":{\"prompt_tokens\":10,\"completion_tokens\":5,\"total_tokens\":15,\"prompt_tokens_details\":{\"cached_tokens\":3}}}\n\n",
"index": 0, "data: [DONE]\n\n"
"message": {
"content": "你好！",
"role": "assistant",
"tool_calls": null,
"reasoning_content": "The user sent a greeting."
}
}],
"created": 1781622889,
"model": "mimo-v2.5",
"object": "chat.completion",
"usage": {
"completion_tokens": 65,
"prompt_tokens": 11741,
"total_tokens": 11806,
"completion_tokens_details": {"reasoning_tokens": 40},
"prompt_tokens_details": {}
}
}"#;
let response: OpenAIResponse = serde_json::from_str(text).unwrap();
let message = &response.choices[0].message;
assert_eq!(message.content.as_deref(), Some("你好！"));
assert_eq!(
message.reasoning_content.as_deref(),
Some("The user sent a greeting.")
); );
assert!(message.tool_calls.is_empty()); let mut decoder = OpenAISseDecoder::default();
assert_eq!(response.usage.total_tokens, 11806); let mut chunks = Vec::new();
for byte in input.as_bytes() {
chunks.extend(decoder.push(std::slice::from_ref(byte)).unwrap());
}
chunks.extend(decoder.finish().unwrap());
let response = crate::providers::collect_provider_stream(Box::pin(
futures_util::stream::iter(chunks.into_iter().map(Ok)),
))
.await
.unwrap();
assert_eq!(response.id, "r1");
assert_eq!(response.model, "m1");
assert_eq!(response.reasoning_content.as_deref(), Some("why"));
assert_eq!(response.content, "你好");
assert_eq!(response.usage.total_tokens, 15);
assert_eq!(response.usage.cached_tokens, Some(3));
} }
#[test] #[test]
fn test_decode_response_exposes_cached_tokens() { fn inline_reasoning_tags_may_span_sse_chunks() {
let text = r#"{ let input = concat!(
"id": "d21abaa6552741949e2aba76bde59359", "data: {\"choices\":[{\"delta\":{\"content\":\"<thi\"},\"finish_reason\":null}]}\n\n",
"choices": [{ "data: {\"choices\":[{\"delta\":{\"content\":\"nk>secret</th\"},\"finish_reason\":null}]}\n\n",
"finish_reason": "stop", "data: {\"choices\":[{\"delta\":{\"content\":\"ink>answer\"},\"finish_reason\":\"stop\"}]}\n\n",
"index": 0, "data: [DONE]\n\n"
"message": {
"content": "你好！",
"role": "assistant",
"tool_calls": null
}
}],
"created": 1781622889,
"model": "mimo-v2.5",
"object": "chat.completion",
"usage": {
"completion_tokens": 65,
"prompt_tokens": 11741,
"total_tokens": 11806,
"prompt_tokens_details": {"cached_tokens": 1200}
}
}"#;
let response: OpenAIResponse = serde_json::from_str(text).unwrap();
assert_eq!(
response
.usage
.prompt_tokens_details
.as_ref()
.and_then(|d| d.cached_tokens),
Some(1200)
); );
let mut decoder = OpenAISseDecoder::default();
let chunks = decoder.push(input.as_bytes()).unwrap();
assert_eq!(
chunks
.iter()
.filter_map(|chunk| match chunk {
ProviderChunk::Reasoning(value) => Some(value.as_str()),
_ => None,
})
.collect::<String>(),
"secret"
);
assert_eq!(
chunks
.iter()
.filter_map(|chunk| match chunk {
ProviderChunk::Text(value) => Some(value.as_str()),
_ => None,
})
.collect::<String>(),
"answer"
);
}
#[test]
fn reasoning_alias_is_used_when_reasoning_content_is_null() {
let input = concat!(
"data: {\"choices\":[{\"delta\":{\"reasoning_content\":null,\"reasoning\":\"alias\"},\"finish_reason\":\"stop\"}]}\n\n",
"data: [DONE]\n\n"
);
let mut decoder = OpenAISseDecoder::default();
let chunks = decoder.push(input.as_bytes()).unwrap();
assert!(
chunks
.iter()
.any(|chunk| matches!(chunk, ProviderChunk::Reasoning(value) if value == "alias"))
);
}
#[tokio::test]
async fn tool_call_arguments_are_assembled_across_sse_events() {
let input = concat!(
"data: {\"choices\":[{\"delta\":{\"tool_calls\":[{\"index\":0,\"id\":\"call_1\",\"function\":{\"name\":\"calculator\",\"arguments\":\"{\\\"expression\\\":\"}}]},\"finish_reason\":null}]}\n\n",
"data: {\"choices\":[{\"delta\":{\"tool_calls\":[{\"index\":0,\"function\":{\"arguments\":\"\\\"1+1\\\"}\"}}]},\"finish_reason\":\"tool_calls\"}]}\n\n",
"data: [DONE]\n\n"
);
let mut decoder = OpenAISseDecoder::default();
let chunks = decoder.push(input.as_bytes()).unwrap();
let response = crate::providers::collect_provider_stream(Box::pin(
futures_util::stream::iter(chunks.into_iter().map(Ok)),
))
.await
.unwrap();
assert_eq!(response.tool_calls.len(), 1);
assert_eq!(response.tool_calls[0].id, "call_1");
assert_eq!(response.tool_calls[0].name, "calculator");
assert_eq!(
response.tool_calls[0].arguments,
serde_json::json!({"expression":"1+1"})
);
}
#[test]
fn missing_finish_marker_is_an_error() {
let mut decoder = OpenAISseDecoder::default();
decoder
.push(b"data: {\"choices\":[{\"delta\":{\"content\":\"partial\"}}]}\n\n")
.unwrap();
assert!(matches!(
decoder.finish(),
Err(OpenAIStreamError::MissingFinish)
));
} }
} }

334
src/providers/stream.rs Normal file
View File

@ -0,0 +1,334 @@
use std::collections::BTreeMap;
use std::error::Error;
use std::pin::Pin;
use futures_util::{Stream, StreamExt};
use serde::{Deserialize, Serialize};
use crate::bus::ProviderReasoningState;
use super::{ChatCompletionResponse, ToolCall, Usage};
pub type DynProviderError = Box<dyn Error + Send + Sync>;
pub type ProviderStreamItem = Result<ProviderChunk, DynProviderError>;
pub type ProviderStream = Pin<Box<dyn Stream<Item = ProviderStreamItem> + Send>>;
pub(crate) fn diagnostic_request_summary(
model: &str,
message_count: usize,
tool_count: usize,
) -> String {
serde_json::json!({
"model": model,
"message_count": message_count,
"tool_count": tool_count,
"stream": true,
})
.to_string()
}
/// Incremental framing shared by SSE-based providers. It accepts arbitrary
/// byte/UTF-8 boundaries and returns only joined `data:` payloads.
#[derive(Default)]
pub(crate) struct SseFramer {
buffer: Vec<u8>,
}
impl SseFramer {
pub(crate) fn push(&mut self, bytes: &[u8]) -> Result<Vec<String>, std::string::FromUtf8Error> {
self.buffer.extend_from_slice(bytes);
self.drain_frames(false)
}
pub(crate) fn finish(&mut self) -> Result<Vec<String>, std::string::FromUtf8Error> {
self.drain_frames(true)
}
fn drain_frames(&mut self, finish: bool) -> Result<Vec<String>, std::string::FromUtf8Error> {
let mut frames = Vec::new();
loop {
let Some((position, delimiter_len)) = find_sse_delimiter(&self.buffer) else {
break;
};
let frame = self.buffer.drain(..position).collect::<Vec<_>>();
self.buffer.drain(..delimiter_len);
if let Some(data) = sse_data(frame)? {
frames.push(data);
}
}
if finish && !self.buffer.is_empty() {
let frame = std::mem::take(&mut self.buffer);
if let Some(data) = sse_data(frame)? {
frames.push(data);
}
}
Ok(frames)
}
}
fn find_sse_delimiter(buffer: &[u8]) -> Option<(usize, usize)> {
let lf = buffer.windows(2).position(|window| window == b"\n\n");
let crlf = buffer.windows(4).position(|window| window == b"\r\n\r\n");
match (lf, crlf) {
(Some(left), Some(right)) if left <= right => Some((left, 2)),
(Some(_), Some(right)) => Some((right, 4)),
(Some(position), None) => Some((position, 2)),
(None, Some(position)) => Some((position, 4)),
(None, None) => None,
}
}
fn sse_data(frame: Vec<u8>) -> Result<Option<String>, std::string::FromUtf8Error> {
let frame = String::from_utf8(frame)?;
let data = frame
.lines()
.filter_map(|line| {
line.strip_prefix("data:")
.map(|value| value.strip_prefix(' ').unwrap_or(value))
})
.collect::<Vec<_>>()
.join("\n");
Ok((!data.is_empty()).then_some(data))
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum FinishReason {
Stop,
ToolCalls,
Length,
ContentFilter,
Other(String),
}
impl FinishReason {
pub fn from_provider(value: &str) -> Self {
match value {
"stop" | "end_turn" | "stop_sequence" => Self::Stop,
"tool_calls" | "tool_use" => Self::ToolCalls,
"length" | "max_tokens" => Self::Length,
"content_filter" => Self::ContentFilter,
other => Self::Other(other.to_string()),
}
}
}
#[derive(Debug, Clone, PartialEq)]
pub enum ProviderChunk {
Metadata {
id: String,
model: String,
},
Text(String),
Reasoning(String),
ToolCallStart {
index: usize,
id: Option<String>,
name: Option<String>,
},
ToolCallArguments {
index: usize,
delta: String,
},
ProviderState(ProviderReasoningState),
Usage(Usage),
Done(FinishReason),
}
#[derive(Default)]
struct PartialToolCall {
id: Option<String>,
name: Option<String>,
arguments: String,
}
#[derive(Default)]
pub struct ProviderResponseAccumulator {
id: String,
model: String,
content: String,
reasoning_content: String,
provider_state: Option<ProviderReasoningState>,
usage: Usage,
tool_calls: BTreeMap<usize, PartialToolCall>,
}
impl ProviderResponseAccumulator {
pub fn push(&mut self, chunk: ProviderChunk) {
match chunk {
ProviderChunk::Metadata {
id: response_id,
model: response_model,
} => {
self.id = response_id;
self.model = response_model;
}
ProviderChunk::Text(delta) => self.content.push_str(&delta),
ProviderChunk::Reasoning(delta) => self.reasoning_content.push_str(&delta),
ProviderChunk::ToolCallStart {
index,
id: call_id,
name,
} => {
let partial = self.tool_calls.entry(index).or_default();
if call_id.is_some() {
partial.id = call_id;
}
if name.is_some() {
partial.name = name;
}
}
ProviderChunk::ToolCallArguments { index, delta } => {
self.tool_calls
.entry(index)
.or_default()
.arguments
.push_str(&delta);
}
ProviderChunk::ProviderState(state) => self.provider_state = Some(state),
ProviderChunk::Usage(value) => self.usage = value,
ProviderChunk::Done(_) => {}
}
}
pub fn finish(self) -> ChatCompletionResponse {
let tool_calls = self
.tool_calls
.into_iter()
.map(|(index, partial)| ToolCall {
id: partial.id.unwrap_or_else(|| format!("tool_call_{index}")),
name: partial.name.unwrap_or_default(),
arguments: serde_json::from_str(&partial.arguments)
.unwrap_or(serde_json::Value::Null),
})
.collect();
ChatCompletionResponse {
id: self.id,
model: self.model,
content: self.content,
reasoning_content: (!self.reasoning_content.is_empty())
.then_some(self.reasoning_content),
provider_state: self.provider_state,
tool_calls,
usage: self.usage,
}
}
}
pub async fn collect_provider_stream(
mut provider_stream: ProviderStream,
) -> Result<ChatCompletionResponse, DynProviderError> {
let mut accumulator = ProviderResponseAccumulator::default();
while let Some(chunk) = provider_stream.next().await {
accumulator.push(chunk?);
}
Ok(accumulator.finish())
}
#[cfg(test)]
pub fn provider_stream_for_test(response: ChatCompletionResponse) -> ProviderStream {
let finish_reason = if response.tool_calls.is_empty() {
FinishReason::Stop
} else {
FinishReason::ToolCalls
};
let mut chunks = vec![ProviderChunk::Metadata {
id: response.id,
model: response.model,
}];
if let Some(reasoning) = response.reasoning_content {
chunks.push(ProviderChunk::Reasoning(reasoning));
}
if !response.content.is_empty() {
chunks.push(ProviderChunk::Text(response.content));
}
for (index, call) in response.tool_calls.into_iter().enumerate() {
chunks.push(ProviderChunk::ToolCallStart {
index,
id: Some(call.id),
name: Some(call.name),
});
chunks.push(ProviderChunk::ToolCallArguments {
index,
delta: serde_json::to_string(&call.arguments).unwrap_or_else(|_| "null".to_string()),
});
}
if let Some(state) = response.provider_state {
chunks.push(ProviderChunk::ProviderState(state));
}
chunks.push(ProviderChunk::Usage(response.usage));
chunks.push(ProviderChunk::Done(finish_reason));
Box::pin(futures_util::stream::iter(chunks.into_iter().map(Ok)))
}
#[cfg(test)]
mod tests {
use super::*;
use futures_util::stream;
#[test]
fn diagnostic_summary_contains_counts_without_message_content() {
let summary = diagnostic_request_summary("model", 3, 2);
assert_eq!(
serde_json::from_str::<serde_json::Value>(&summary).unwrap(),
serde_json::json!({
"model": "model",
"message_count": 3,
"tool_count": 2,
"stream": true
})
);
assert!(!summary.contains("reasoning"));
assert!(!summary.contains("signature"));
}
#[tokio::test]
async fn collector_assembles_interleaved_tool_argument_fragments() {
let chunks = vec![
ProviderChunk::Metadata {
id: "response".into(),
model: "model".into(),
},
ProviderChunk::Reasoning("why".into()),
ProviderChunk::Text("answer".into()),
ProviderChunk::ToolCallStart {
index: 1,
id: Some("second".into()),
name: Some("b".into()),
},
ProviderChunk::ToolCallArguments {
index: 1,
delta: "{\"n\":".into(),
},
ProviderChunk::ToolCallStart {
index: 0,
id: Some("first".into()),
name: Some("a".into()),
},
ProviderChunk::ToolCallArguments {
index: 0,
delta: "{}".into(),
},
ProviderChunk::ToolCallArguments {
index: 1,
delta: "2}".into(),
},
ProviderChunk::Done(FinishReason::ToolCalls),
];
let response = collect_provider_stream(Box::pin(stream::iter(chunks.into_iter().map(Ok))))
.await
.unwrap();
assert_eq!(response.content, "answer");
assert_eq!(response.reasoning_content.as_deref(), Some("why"));
assert_eq!(response.tool_calls.len(), 2);
assert_eq!(response.tool_calls[0].id, "first");
assert_eq!(
response.tool_calls[1].arguments,
serde_json::json!({"n": 2})
);
}
}

21
src/providers/traits.rs
View File

@ -2,12 +2,17 @@ use crate::bus::message::ContentBlock;
use async_trait::async_trait; use async_trait::async_trait;
use serde::{Deserialize, Serialize}; use serde::{Deserialize, Serialize};
use super::stream::{DynProviderError, ProviderStream, collect_provider_stream};
#[derive(Debug, Clone, Serialize, Deserialize)] #[derive(Debug, Clone, Serialize, Deserialize)]
pub struct Message { pub struct Message {
pub role: String, pub role: String,
pub content: Vec<ContentBlock>, pub content: Vec<ContentBlock>,
#[serde(skip_serializing_if = "Option::is_none")] #[serde(skip_serializing_if = "Option::is_none")]
pub reasoning_content: Option<String>, pub reasoning_content: Option<String>,
/// Opaque state replayed only by the provider that produced it.
#[serde(skip)]
pub provider_state: Option<crate::bus::ProviderReasoningState>,
#[serde(skip_serializing_if = "Option::is_none")] #[serde(skip_serializing_if = "Option::is_none")]
pub tool_call_id: Option<String>, pub tool_call_id: Option<String>,
#[serde(skip_serializing_if = "Option::is_none")] #[serde(skip_serializing_if = "Option::is_none")]
@ -22,6 +27,7 @@ impl Message {
role: "user".to_string(), role: "user".to_string(),
content: vec![ContentBlock::text(content)], content: vec![ContentBlock::text(content)],
reasoning_content: None, reasoning_content: None,
provider_state: None,
tool_call_id: None, tool_call_id: None,
name: None, name: None,
tool_calls: None, tool_calls: None,
@ -33,6 +39,7 @@ impl Message {
role: "assistant".to_string(), role: "assistant".to_string(),
content: vec![ContentBlock::text(content)], content: vec![ContentBlock::text(content)],
reasoning_content: None, reasoning_content: None,
provider_state: None,
tool_call_id: None, tool_call_id: None,
name: None, name: None,
tool_calls: None, tool_calls: None,
@ -44,6 +51,7 @@ impl Message {
role: "system".to_string(), role: "system".to_string(),
content: vec![ContentBlock::text(content)], content: vec![ContentBlock::text(content)],
reasoning_content: None, reasoning_content: None,
provider_state: None,
tool_call_id: None, tool_call_id: None,
name: None, name: None,
tool_calls: None, tool_calls: None,
@ -59,6 +67,7 @@ impl Message {
role: "tool".to_string(), role: "tool".to_string(),
content: vec![ContentBlock::text(content)], content: vec![ContentBlock::text(content)],
reasoning_content: None, reasoning_content: None,
provider_state: None,
tool_call_id: Some(tool_call_id.into()), tool_call_id: Some(tool_call_id.into()),
name: Some(tool_name.into()), name: Some(tool_name.into()),
tool_calls: None, tool_calls: None,
@ -101,11 +110,12 @@ pub struct ChatCompletionResponse {
pub model: String, pub model: String,
pub content: String, pub content: String,
pub reasoning_content: Option<String>, pub reasoning_content: Option<String>,
pub provider_state: Option<crate::bus::ProviderReasoningState>,
pub tool_calls: Vec<ToolCall>, pub tool_calls: Vec<ToolCall>,
pub usage: Usage, pub usage: Usage,
} }
#[derive(Debug, Clone, Serialize, Deserialize)] #[derive(Debug, Clone, Default, PartialEq, Eq, Serialize, Deserialize)]
pub struct Usage { pub struct Usage {
pub prompt_tokens: u32, pub prompt_tokens: u32,
pub completion_tokens: u32, pub completion_tokens: u32,
@ -120,10 +130,17 @@ pub struct Usage {
#[async_trait] #[async_trait]
pub trait LLMProvider: Send + Sync { pub trait LLMProvider: Send + Sync {
async fn stream(
&self,
request: ChatCompletionRequest,
) -> Result<ProviderStream, DynProviderError>;
async fn chat( async fn chat(
&self, &self,
request: ChatCompletionRequest, request: ChatCompletionRequest,
) -> Result<ChatCompletionResponse, Box<dyn std::error::Error + Send + Sync>>; ) -> Result<ChatCompletionResponse, DynProviderError> {
collect_provider_stream(self.stream(request).await?).await
}
fn ptype(&self) -> &str; fn ptype(&self) -> &str;

7
src/session/mod.rs
View File

@ -7,9 +7,14 @@ mod persistence;
#[allow(clippy::module_inception)] #[allow(clippy::module_inception)]
pub mod session; pub mod session;
pub mod session_id; pub mod session_id;
pub mod turn;
pub use commands::SessionCommand; pub use commands::SessionCommand;
pub use error::SessionError; pub use error::SessionError;
pub use events::{DialogInfo, SessionEvent}; pub use events::{DialogInfo, SessionEvent};
pub use session::{SLASH_COMMANDS, Session, SessionManager, SlashCommand}; pub use session::{SLASH_COMMANDS, Session, SessionManager, SessionManagerServices, SlashCommand};
pub use session_id::UnifiedSessionId; pub use session_id::UnifiedSessionId;
pub use turn::{
BlockId, ToolStatus, TurnBlock, TurnController, TurnId, TurnPhase, TurnSnapshot, TurnState,
TurnStatus,
};

73
src/session/persistence.rs
View File

@ -1,10 +1,12 @@
use std::sync::Arc; use std::sync::Arc;
use std::{fmt::Display, future::Future};
use tokio::sync::Mutex; use tokio::sync::Mutex;
use super::session::{MessagePersistSnapshot, Session}; use super::session::{MessagePersistSnapshot, Session};
use crate::bus::ChatMessage; use crate::bus::ChatMessage;
use crate::storage::StorageError; use crate::storage::StorageError;
use crate::{providers::Usage, session::TurnController};
async fn persist_added_messages( async fn persist_added_messages(
snapshots: Vec<Option<MessagePersistSnapshot>>, snapshots: Vec<Option<MessagePersistSnapshot>>,
@ -63,3 +65,74 @@ pub(super) async fn append_persisted_messages(
} }
Ok(()) Ok(())
} }
/// Publish `Completed` only after the supplied durable write succeeds.
///
/// Keeping this ordering in one helper makes the user-visible terminal status
/// impossible to publish optimistically before SQLite commits.
pub(super) async fn finalize_turn_after_persistence<F, T, E>(
controller: &TurnController,
usage: Option<Usage>,
persistence: F,
) -> Result<T, E>
where
F: Future<Output = Result<T, E>>,
E: Display,
{
controller.begin_finalizing();
match persistence.await {
Ok(value) => {
controller.complete(usage);
Ok(value)
}
Err(error) => {
controller.fail(format!("failed to persist turn: {error}"));
Err(error)
}
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::session::{TurnController, TurnStatus};
#[tokio::test]
async fn completed_is_published_only_after_persistence_succeeds() {
let (controller, _emitter, receiver) = TurnController::start("session", "message");
let persisted = Arc::new(std::sync::atomic::AtomicBool::new(false));
let persisted_in_future = persisted.clone();
let result: Result<(), String> =
finalize_turn_after_persistence(&controller, None, async move {
assert_eq!(receiver.borrow().status, TurnStatus::Running);
assert_eq!(
receiver.borrow().phase,
crate::session::TurnPhase::Finalizing
);
persisted_in_future.store(true, std::sync::atomic::Ordering::SeqCst);
Ok(())
})
.await;
assert!(result.is_ok());
assert!(persisted.load(std::sync::atomic::Ordering::SeqCst));
assert_eq!(controller.snapshot().status, TurnStatus::Completed);
}
#[tokio::test]
async fn persistence_failure_never_publishes_completed() {
let (controller, _emitter, _receiver) = TurnController::start("session", "message");
let result: Result<(), &str> =
finalize_turn_after_persistence(&controller, None, async { Err("database down") })
.await;
assert_eq!(result, Err("database down"));
let snapshot = controller.snapshot();
assert_eq!(snapshot.status, TurnStatus::Failed);
assert_eq!(
snapshot.error.as_deref(),
Some("failed to persist turn: database down")
);
}
}

401
src/session/session.rs
View File

@ -3,8 +3,11 @@ use std::sync::Arc;
use tokio::sync::{Mutex, mpsc, oneshot}; use tokio::sync::{Mutex, mpsc, oneshot};
use super::persistence::append_persisted_messages; use super::persistence::{append_persisted_messages, finalize_turn_after_persistence};
use crate::bus::{ChatMessage, MediaItem, MediaRef, MessageSource, OutboundMessage, SourceKind}; use super::turn::{TurnBlock, TurnController, TurnSnapshot};
use crate::bus::{
ChatMessage, CompletionStatus, MediaItem, MediaRef, MessageSource, OutboundMessage, SourceKind,
};
use crate::mcp::get_mcp_status; use crate::mcp::get_mcp_status;
use crate::storage::{Storage, StorageError}; use crate::storage::{Storage, StorageError};
use std::sync::Arc as StdArc; use std::sync::Arc as StdArc;
@ -22,6 +25,15 @@ fn outbound_session_metadata(session_id: &str) -> HashMap<String, String> {
HashMap::from([("_session_id".to_string(), session_id.to_string())]) HashMap::from([("_session_id".to_string(), session_id.to_string())])
} }
fn outbound_turn_metadata(
session_id: &str,
forwarded: &HashMap<String, String>,
) -> HashMap<String, String> {
let mut metadata = forwarded.clone();
metadata.insert("_session_id".to_string(), session_id.to_string());
metadata
}
tokio::task_local! { tokio::task_local! {
pub(super) static CURRENT_SOURCE_SESSION: Option<String>; pub(super) static CURRENT_SOURCE_SESSION: Option<String>;
} }
@ -37,10 +49,11 @@ pub enum HandleResult {
} }
use crate::agent::context_compressor::ContextCompressionConfig; use crate::agent::context_compressor::ContextCompressionConfig;
use crate::agent::system_prompt::{build_runtime_context, build_system_prompt}; use crate::agent::system_prompt::{build_runtime_context, build_system_prompt};
use crate::agent::{AgentError, AgentLoop, ContextCompressor}; use crate::agent::{AgentError, AgentLoop, AgentTurnContext, ContextCompressor, TurnEmitter};
use crate::channels::slash_command::parse_slash_command; use crate::channels::slash_command::parse_slash_command;
use crate::config::BrowserConfig; use crate::config::BrowserConfig;
use crate::config::LLMProviderConfig; use crate::config::LLMProviderConfig;
use crate::delivery::TurnDeliveryService;
/// Check if an LLM error message indicates a context window overflow. /// Check if an LLM error message indicates a context window overflow.
fn is_context_overflow_error(msg: &str) -> bool { fn is_context_overflow_error(msg: &str) -> bool {
@ -53,6 +66,139 @@ fn is_context_overflow_error(msg: &str) -> bool {
|| lower.contains("prompt is too long") || lower.contains("prompt is too long")
|| lower.contains("input is too long") || lower.contains("input is too long")
} }
fn partial_assistant_message(
snapshot: &TurnSnapshot,
completion_status: CompletionStatus,
) -> Option<ChatMessage> {
let mut assistant_segments = Vec::new();
let mut reasoning_segments = Vec::new();
let mut last_iteration = None;
for block in &snapshot.blocks {
match block {
TurnBlock::Assistant {
iteration, text, ..
} if !text.is_empty() => {
assistant_segments.push(text.as_str());
last_iteration = Some(*iteration);
}
TurnBlock::Reasoning { text, .. } if !text.is_empty() => {
reasoning_segments.push(text.as_str());
}
_ => {}
}
}
if assistant_segments.is_empty() {
return None;
}
let mut message = ChatMessage::assistant(assistant_segments.join("\n\n"));
message.id = snapshot.message_id.clone();
message.turn_id = Some(snapshot.id.0.clone());
message.iteration = last_iteration;
message.completion_status = completion_status;
message.reasoning_content =
(!reasoning_segments.is_empty()).then(|| reasoning_segments.join("\n\n"));
Some(message)
}
async fn fail_turn_with_partial(
controller: &TurnController,
session: &Arc<Mutex<Session>>,
error: String,
) {
let snapshot = controller.snapshot();
let partial = partial_assistant_message(&snapshot, CompletionStatus::Interrupted);
if let Some(partial) = partial {
controller.begin_finalizing();
if let Err(persistence_error) = append_persisted_messages(session, vec![partial]).await {
controller.fail(format!(
"{error}; failed to persist interrupted turn: {persistence_error}"
));
return;
}
}
controller.fail(error);
}
#[cfg(test)]
mod cancelled_partial_tests {
use super::*;
use crate::agent::TurnEvent;
#[test]
fn turn_metadata_preserves_channel_cleanup_fields() {
let forwarded = HashMap::from([
("feishu.message_id".to_string(), "message-1".to_string()),
("feishu.reaction_id".to_string(), "reaction-1".to_string()),
("_session_id".to_string(), "stale".to_string()),
]);
let metadata = outbound_turn_metadata("session-1", &forwarded);
assert_eq!(
metadata.get("_session_id").map(String::as_str),
Some("session-1")
);
assert_eq!(
metadata.get("feishu.message_id").map(String::as_str),
Some("message-1")
);
assert_eq!(
metadata.get("feishu.reaction_id").map(String::as_str),
Some("reaction-1")
);
}
#[test]
fn visible_partial_text_becomes_cancelled_persisted_message() {
let (controller, emitter, _) = TurnController::start("session", "message-id");
emitter
.emit(TurnEvent::ReasoningDelta {
iteration: 0,
delta: "reason".into(),
})
.unwrap();
emitter
.emit(TurnEvent::TextDelta {
iteration: 0,
delta: "first".into(),
})
.unwrap();
emitter
.emit(TurnEvent::TextSegmentFinished { iteration: 0 })
.unwrap();
emitter
.emit(TurnEvent::TextDelta {
iteration: 1,
delta: "second".into(),
})
.unwrap();
let message =
partial_assistant_message(&controller.snapshot(), CompletionStatus::Cancelled).unwrap();
assert_eq!(message.id, "message-id");
assert_eq!(message.content, "first\n\nsecond");
assert_eq!(message.reasoning_content.as_deref(), Some("reason"));
assert_eq!(message.iteration, Some(1));
assert_eq!(message.completion_status, CompletionStatus::Cancelled);
}
#[test]
fn reasoning_only_cancel_does_not_create_assistant_history() {
let (controller, emitter, _) = TurnController::start("session", "message-id");
emitter
.emit(TurnEvent::ReasoningDelta {
iteration: 0,
delta: "private".into(),
})
.unwrap();
assert!(
partial_assistant_message(&controller.snapshot(), CompletionStatus::Cancelled,)
.is_none()
);
}
}
use crate::bus::MessageBus; use crate::bus::MessageBus;
use crate::providers::{LLMProvider, create_provider}; use crate::providers::{LLMProvider, create_provider};
use crate::session::events::DialogInfo; use crate::session::events::DialogInfo;
@ -93,6 +239,7 @@ pub struct Session {
agent_tx: Option<mpsc::Sender<AgentTask>>, agent_tx: Option<mpsc::Sender<AgentTask>>,
/// Cancel signal for the currently executing agent task /// Cancel signal for the currently executing agent task
current_cancel: Option<oneshot::Sender<()>>, current_cancel: Option<oneshot::Sender<()>>,
active_turn_emitter: Option<ActiveTurnEmitter>,
/// Monotonic counter to detect stale workers /// Monotonic counter to detect stale workers
worker_generation: u64, worker_generation: u64,
/// Monotonic counter for in-memory session mutations. /// Monotonic counter for in-memory session mutations.
@ -108,12 +255,18 @@ pub struct Session {
pub(super) persistence_lock: Arc<Mutex<()>>, pub(super) persistence_lock: Arc<Mutex<()>>,
} }
struct ActiveTurnEmitter {
turn_id: String,
emitter: TurnEmitter,
}
/// A task to be processed by the per-session agent worker /// A task to be processed by the per-session agent worker
struct AgentTask { struct AgentTask {
channel: String, channel: String,
chat_id: String, chat_id: String,
content: String, content: String,
media: Vec<MediaItem>, media: Vec<MediaItem>,
forwarded_metadata: HashMap<String, String>,
} }
#[derive(Clone)] #[derive(Clone)]
@ -123,6 +276,7 @@ struct AgentWorkerDeps {
work_manager: Arc<crate::work::WorkManager>, work_manager: Arc<crate::work::WorkManager>,
skills_loader: Arc<SkillsLoader>, skills_loader: Arc<SkillsLoader>,
task_supervisor: crate::task_supervisor::TaskSupervisor, task_supervisor: crate::task_supervisor::TaskSupervisor,
turn_delivery: TurnDeliveryService,
} }
impl Session { impl Session {
@ -178,6 +332,7 @@ impl Session {
memory_manager, memory_manager,
agent_tx: None, agent_tx: None,
current_cancel: None, current_cancel: None,
active_turn_emitter: None,
worker_generation: 0, worker_generation: 0,
state_version: 0, state_version: 0,
persistence_lock: Arc::new(Mutex::new(())), persistence_lock: Arc::new(Mutex::new(())),
@ -258,6 +413,12 @@ impl Session {
role: m.role, role: m.role,
content: m.content, content: m.content,
reasoning_content: m.reasoning_content, reasoning_content: m.reasoning_content,
provider_state: m.provider_state.and_then(|state| {
crate::bus::ProviderReasoningState::from_json_lossy(&state)
}),
turn_id: m.turn_id,
iteration: m.iteration.and_then(|value| u32::try_from(value).ok()),
completion_status: m.completion_status,
media_refs: m media_refs: m
.media_refs .media_refs
.map(|refs| serde_json::from_str(&refs).unwrap_or_default()) .map(|refs| serde_json::from_str(&refs).unwrap_or_default())
@ -293,6 +454,12 @@ impl Session {
role: m.role, role: m.role,
content: m.content, content: m.content,
reasoning_content: m.reasoning_content, reasoning_content: m.reasoning_content,
provider_state: m.provider_state.and_then(|state| {
crate::bus::ProviderReasoningState::from_json_lossy(&state)
}),
turn_id: m.turn_id,
iteration: m.iteration.and_then(|value| u32::try_from(value).ok()),
completion_status: m.completion_status,
media_refs: m media_refs: m
.media_refs .media_refs
.map(|refs| serde_json::from_str(&refs).unwrap_or_default()) .map(|refs| serde_json::from_str(&refs).unwrap_or_default())
@ -354,6 +521,7 @@ impl Session {
memory_manager, memory_manager,
agent_tx: None, agent_tx: None,
current_cancel: None, current_cancel: None,
active_turn_emitter: None,
worker_generation: 0, worker_generation: 0,
state_version: 0, state_version: 0,
persistence_lock: Arc::new(Mutex::new(())), persistence_lock: Arc::new(Mutex::new(())),
@ -386,6 +554,13 @@ impl Session {
role: message.role.clone(), role: message.role.clone(),
content: message.content.clone(), content: message.content.clone(),
reasoning_content: message.reasoning_content.clone(), reasoning_content: message.reasoning_content.clone(),
provider_state: message
.provider_state
.as_ref()
.and_then(|state| serde_json::to_string(state).ok()),
turn_id: message.turn_id.clone(),
iteration: message.iteration.map(i64::from),
completion_status: message.completion_status,
media_refs: if message.media_refs.is_empty() { media_refs: if message.media_refs.is_empty() {
None None
} else { } else {
@ -897,6 +1072,31 @@ pub struct SessionManager {
work_manager: Arc<crate::work::WorkManager>, work_manager: Arc<crate::work::WorkManager>,
sub_agent_manager: Arc<crate::agent::SubAgentManager>, sub_agent_manager: Arc<crate::agent::SubAgentManager>,
task_supervisor: crate::task_supervisor::TaskSupervisor, task_supervisor: crate::task_supervisor::TaskSupervisor,
turn_delivery: TurnDeliveryService,
}
/// Gateway-owned runtime services shared by all Session workers.
pub struct SessionManagerServices {
bus: Arc<MessageBus>,
memory_manager: Arc<crate::memory::MemoryManager>,
task_supervisor: crate::task_supervisor::TaskSupervisor,
turn_delivery: TurnDeliveryService,
}
impl SessionManagerServices {
pub fn new(
bus: Arc<MessageBus>,
memory_manager: Arc<crate::memory::MemoryManager>,
task_supervisor: crate::task_supervisor::TaskSupervisor,
turn_delivery: TurnDeliveryService,
) -> Self {
Self {
bus,
memory_manager,
task_supervisor,
turn_delivery,
}
}
} }
struct SessionManagerInner { struct SessionManagerInner {
@ -1010,18 +1210,23 @@ impl SessionManager {
work_manager: self.work_manager.clone(), work_manager: self.work_manager.clone(),
skills_loader: self.skills_loader.clone(), skills_loader: self.skills_loader.clone(),
task_supervisor: self.task_supervisor.clone(), task_supervisor: self.task_supervisor.clone(),
turn_delivery: self.turn_delivery.clone(),
} }
} }
pub fn new( pub fn new(
provider_config: LLMProviderConfig, provider_config: LLMProviderConfig,
storage: Arc<Storage>, storage: Arc<Storage>,
bus: Arc<MessageBus>, services: SessionManagerServices,
memory_manager: Arc<crate::memory::MemoryManager>,
browser_config: Option<BrowserConfig>, browser_config: Option<BrowserConfig>,
max_concurrent_background_tasks: usize, max_concurrent_background_tasks: usize,
task_supervisor: crate::task_supervisor::TaskSupervisor,
) -> Result<Self, AgentError> { ) -> Result<Self, AgentError> {
let SessionManagerServices {
bus,
memory_manager,
task_supervisor,
turn_delivery,
} = services;
let mut skills_loader = SkillsLoader::new(); let mut skills_loader = SkillsLoader::new();
skills_loader.load_skills(); skills_loader.load_skills();
skills_loader.set_workspace_skills_dir(provider_config.workspace_dir.clone()); skills_loader.set_workspace_skills_dir(provider_config.workspace_dir.clone());
@ -1108,6 +1313,7 @@ impl SessionManager {
work_manager, work_manager,
sub_agent_manager, sub_agent_manager,
task_supervisor, task_supervisor,
turn_delivery,
}) })
} }
@ -1454,6 +1660,9 @@ impl SessionManager {
if guard.current_cancel.take().is_some() { if guard.current_cancel.take().is_some() {
msgs.push("当前任务已发送停止信号。".to_string()); msgs.push("当前任务已发送停止信号。".to_string());
} }
if let Some(active_turn) = guard.active_turn_emitter.take() {
active_turn.emitter.deactivate();
}
if guard.agent_tx.take().is_some() { if guard.agent_tx.take().is_some() {
msgs.push("消息队列已清空。".to_string()); msgs.push("消息队列已清空。".to_string());
} }
@ -2001,6 +2210,7 @@ impl SessionManager {
chat_id: &str, chat_id: &str,
content: &str, content: &str,
media: Vec<MediaItem>, media: Vec<MediaItem>,
forwarded_metadata: HashMap<String, String>,
) -> Result<HandleResult, AgentError> { ) -> Result<HandleResult, AgentError> {
let unified_id = self.resolve_dialog_id(channel, chat_id).await?; let unified_id = self.resolve_dialog_id(channel, chat_id).await?;
tracing::debug!(unified_id = %unified_id, "handle_message resolved unified_id"); tracing::debug!(unified_id = %unified_id, "handle_message resolved unified_id");
@ -2038,6 +2248,7 @@ impl SessionManager {
chat_id: chat_id.to_string(), chat_id: chat_id.to_string(),
content: content.to_string(), content: content.to_string(),
media, media,
forwarded_metadata,
}; };
let session_clone = session.clone(); let session_clone = session.clone();
let unified_str = unified_id.to_string(); let unified_str = unified_id.to_string();
@ -2167,6 +2378,7 @@ fn spawn_agent_worker(
work_manager, work_manager,
skills_loader, skills_loader,
task_supervisor, task_supervisor,
turn_delivery,
} = deps; } = deps;
let worker_supervisor = task_supervisor.clone(); let worker_supervisor = task_supervisor.clone();
task_supervisor.spawn(format!("session-worker:{unified_str}"), async move { task_supervisor.spawn(format!("session-worker:{unified_str}"), async move {
@ -2175,6 +2387,7 @@ fn spawn_agent_worker(
'tasks: while let Some(task) = task_rx.recv().await { 'tasks: while let Some(task) = task_rx.recv().await {
let task_chan = task.channel.clone(); let task_chan = task.channel.clone();
let task_cid = task.chat_id.clone(); let task_cid = task.chat_id.clone();
let task_metadata = task.forwarded_metadata.clone();
let notification_session_id = unified_str.clone(); let notification_session_id = unified_str.clone();
let (notify_tx, mut notify_rx) = mpsc::unbounded_channel(); let (notify_tx, mut notify_rx) = mpsc::unbounded_channel();
@ -2231,7 +2444,7 @@ fn spawn_agent_worker(
content: "Failed to save your message, please try again.".to_string(), content: "Failed to save your message, please try again.".to_string(),
reply_to: None, reply_to: None,
media: vec![], media: vec![],
metadata: outbound_session_metadata(&unified_str), metadata: outbound_turn_metadata(&unified_str, &task_metadata),
delivery: None, delivery: None,
}; };
let _ = bus.publish_outbound(err_outbound).await; let _ = bus.publish_outbound(err_outbound).await;
@ -2258,7 +2471,7 @@ fn spawn_agent_worker(
.to_string(), .to_string(),
reply_to: None, reply_to: None,
media: vec![], media: vec![],
metadata: outbound_session_metadata(&unified_str), metadata: outbound_turn_metadata(&unified_str, &task_metadata),
delivery: None, delivery: None,
}; };
let _ = bus.publish_outbound(err_outbound).await; let _ = bus.publish_outbound(err_outbound).await;
@ -2372,12 +2585,64 @@ fn spawn_agent_worker(
Session::append_runtime_context_to_user_message(last_msg, &runtime_context); Session::append_runtime_context_to_user_message(last_msg, &runtime_context);
} }
let (turn_controller, turn_emitter, turn_receiver) = TurnController::start(
unified_str.clone(),
uuid::Uuid::new_v4().to_string(),
);
let initial_turn = turn_controller.snapshot();
let active_turn_id = initial_turn.id.0.clone();
let live_delivery_started = match turn_delivery
.start(
crate::channels::TurnTarget {
channel: task_chan.clone(),
chat_id: task_cid.clone(),
session_id: unified_str.clone(),
reply_to: None,
metadata: outbound_turn_metadata(&unified_str, &task_metadata),
},
turn_receiver,
)
.await
{
Ok(()) => true,
Err(error) => {
tracing::debug!(
channel = %task_chan,
error = %error,
"Live turn delivery unavailable; using ordinary final delivery"
);
false
}
};
{
let mut guard = session.lock().await;
if guard.worker_generation != worker_gen || guard.state_version != base_version {
turn_emitter.deactivate();
turn_controller.cancel(Some(
"session changed before model execution".to_string(),
));
guard.current_cancel = None;
continue 'tasks;
}
guard.active_turn_emitter = Some(ActiveTurnEmitter {
turn_id: initial_turn.id.0.clone(),
emitter: turn_emitter.clone(),
});
}
let agent_turn = AgentTurnContext::new(
initial_turn.id.0.clone(),
initial_turn.message_id.clone(),
turn_emitter,
);
// Phase 2 + 3: LLM call with cancellation // Phase 2 + 3: LLM call with cancellation
let session2 = session.clone(); let session2 = session.clone();
let bus2 = bus.clone(); let bus2 = bus.clone();
let chan2 = task_chan.clone(); let chan2 = task_chan.clone();
let cid2 = task_cid.clone(); let cid2 = task_cid.clone();
let unified_str2 = unified_str.clone(); let unified_str2 = unified_str.clone();
let task_metadata2 = task_metadata.clone();
let turn_lifecycle = &turn_controller;
let process_future = async move { let process_future = async move {
let response_session_id = unified_str2.clone(); let response_session_id = unified_str2.clone();
let process_result = crate::agent::sub_agent::DELEGATE_CONTEXT.scope( let process_result = crate::agent::sub_agent::DELEGATE_CONTEXT.scope(
@ -2386,7 +2651,7 @@ fn spawn_agent_worker(
channel: chan2.clone(), channel: chan2.clone(),
chat_id: cid2.clone(), chat_id: cid2.clone(),
}, },
agent.process(history_out.clone()), agent.process_streaming(history_out.clone(), agent_turn.clone()),
).await; ).await;
let result = match process_result { let result = match process_result {
Ok(r) => r, Ok(r) => r,
@ -2416,6 +2681,12 @@ fn spawn_agent_worker(
Ok(r) => r, Ok(r) => r,
Err(e) => { Err(e) => {
tracing::error!(error = %e, "Retry compression failed"); tracing::error!(error = %e, "Retry compression failed");
fail_turn_with_partial(
turn_lifecycle,
&session2,
format!("context overflow handling failed: {e}"),
)
.await;
let err_outbound = OutboundMessage { let err_outbound = OutboundMessage {
channel: chan2, channel: chan2,
chat_id: cid2, chat_id: cid2,
@ -2423,12 +2694,15 @@ fn spawn_agent_worker(
.to_string(), .to_string(),
reply_to: None, reply_to: None,
media: vec![], media: vec![],
metadata: outbound_session_metadata( metadata: outbound_turn_metadata(
&response_session_id, &response_session_id,
&task_metadata2,
), ),
delivery: None, delivery: None,
}; };
if !live_delivery_started {
let _ = bus2.publish_outbound(err_outbound).await; let _ = bus2.publish_outbound(err_outbound).await;
}
return; return;
} }
}; };
@ -2440,6 +2714,10 @@ fn spawn_agent_worker(
session_id = %guard.id, session_id = %guard.id,
"Session changed while retry-compressing after context overflow" "Session changed while retry-compressing after context overflow"
); );
turn_lifecycle.cancel(Some(
"session changed during context overflow recovery"
.to_string(),
));
return; return;
} }
guard.compressor.set_context_window(new_window); guard.compressor.set_context_window(new_window);
@ -2471,54 +2749,99 @@ fn spawn_agent_worker(
retry retry
}; };
match agent.process(retry_history).await { match agent
.process_streaming(retry_history, agent_turn.clone())
.await
{
Ok(r) => r, Ok(r) => r,
Err(e) => { Err(e) => {
tracing::error!( tracing::error!(
error = %e, error = %e,
"Agent retry after overflow failed" "Agent retry after overflow failed"
); );
fail_turn_with_partial(
turn_lifecycle,
&session2,
e.to_string(),
)
.await;
let err_outbound = OutboundMessage { let err_outbound = OutboundMessage {
channel: chan2, channel: chan2,
chat_id: cid2, chat_id: cid2,
content: format!("Processing error: {}", e), content: format!("Processing error: {}", e),
reply_to: None, reply_to: None,
media: vec![], media: vec![],
metadata: outbound_session_metadata(&response_session_id), metadata: outbound_turn_metadata(
&response_session_id,
&task_metadata2,
),
delivery: None, delivery: None,
}; };
if !live_delivery_started {
let _ = bus2.publish_outbound(err_outbound).await; let _ = bus2.publish_outbound(err_outbound).await;
}
return; return;
} }
} }
} }
Err(e) => { Err(e) => {
tracing::error!(error = %e, "Agent processing error"); tracing::error!(error = %e, "Agent processing error");
fail_turn_with_partial(
turn_lifecycle,
&session2,
e.to_string(),
)
.await;
let err_outbound = OutboundMessage { let err_outbound = OutboundMessage {
channel: chan2, channel: chan2,
chat_id: cid2, chat_id: cid2,
content: format!("Processing error: {}", e), content: format!("Processing error: {}", e),
reply_to: None, reply_to: None,
media: vec![], media: vec![],
metadata: outbound_session_metadata(&response_session_id), metadata: outbound_turn_metadata(
&response_session_id,
&task_metadata2,
),
delivery: None, delivery: None,
}; };
if !live_delivery_started {
let _ = bus2.publish_outbound(err_outbound).await; let _ = bus2.publish_outbound(err_outbound).await;
}
return; return;
} }
}; };
let response_content = result.final_response.content; let response_content = result.final_response.content;
let total_tokens = result.total_tokens; let total_tokens = result.total_tokens;
let response = let usage = result.usage;
if let Err(e) = append_persisted_messages(&session2, result.emitted_messages).await { {
tracing::error!(error = %e, "Failed to atomically persist agent turn"); let guard = session2.lock().await;
None if guard.worker_generation != worker_gen
} else { || guard.state_version != base_version
{
turn_lifecycle.cancel(Some(
"session changed before turn commit".to_string(),
));
return;
}
}
let response = match finalize_turn_after_persistence(
turn_lifecycle,
usage,
append_persisted_messages(&session2, result.emitted_messages),
)
.await
{
Ok(()) => {
let mut guard = session2.lock().await; let mut guard = session2.lock().await;
let sent_count = guard.messages.len(); let sent_count = guard.messages.len();
guard.compressor.set_last_api_info(sent_count, total_tokens); guard.compressor.set_last_api_info(sent_count, total_tokens);
Some(response_content) Some(response_content)
}
Err(e) => {
tracing::error!(error = %e, "Failed to atomically persist agent turn");
None
}
}; };
let Some(response) = response else { let Some(response) = response else {
@ -2529,10 +2852,15 @@ fn spawn_agent_worker(
.to_string(), .to_string(),
reply_to: None, reply_to: None,
media: vec![], media: vec![],
metadata: outbound_session_metadata(&response_session_id), metadata: outbound_turn_metadata(
&response_session_id,
&task_metadata2,
),
delivery: None, delivery: None,
}; };
if !live_delivery_started {
let _ = bus2.publish_outbound(err_outbound).await; let _ = bus2.publish_outbound(err_outbound).await;
}
return; return;
}; };
@ -2540,27 +2868,60 @@ fn spawn_agent_worker(
tracing::warn!("failed to generate title: {}", e); tracing::warn!("failed to generate title: {}", e);
} }
if !live_delivery_started {
let outbound = OutboundMessage { let outbound = OutboundMessage {
channel: chan2, channel: chan2,
chat_id: cid2, chat_id: cid2,
content: response, content: response,
reply_to: None, reply_to: None,
media: vec![], media: vec![],
metadata: outbound_session_metadata(&response_session_id), metadata: outbound_turn_metadata(
&response_session_id,
&task_metadata2,
),
delivery: None, delivery: None,
}; };
let _ = bus2.publish_outbound(outbound).await; let _ = bus2.publish_outbound(outbound).await;
}
}; };
tokio::select! { tokio::select! {
() = process_future => {} () = process_future => {}
_ = cancel_rx => { _ = cancel_rx => {
// cancelled — current_cancel already taken by /stop // cancelled — current_cancel already taken by /stop
let snapshot = turn_controller.snapshot();
if let Some(partial) = partial_assistant_message(
&snapshot,
CompletionStatus::Cancelled,
) {
turn_controller.begin_finalizing();
match append_persisted_messages(&session, vec![partial]).await {
Ok(()) => {
turn_controller.cancel(Some("stopped by user".to_string()));
}
Err(error) => {
tracing::error!(error = %error, "Failed to persist cancelled partial turn");
turn_controller.fail(format!(
"failed to persist cancelled turn: {error}"
));
}
}
} else {
turn_controller.cancel(Some("stopped by user".to_string()));
}
} }
} }
// Clean up // Clean up
let mut guard = session.lock().await; let mut guard = session.lock().await;
if guard
.active_turn_emitter
.as_ref()
.is_some_and(|active| active.turn_id == active_turn_id)
&& let Some(active) = guard.active_turn_emitter.take()
{
active.emitter.deactivate();
}
if guard.worker_generation == worker_gen { if guard.worker_generation == worker_gen {
guard.current_cancel = None; guard.current_cancel = None;
} }

606
src/session/turn.rs Normal file
View File

@ -0,0 +1,606 @@
use std::sync::{Arc, Mutex, MutexGuard, Weak};
use serde::{Deserialize, Serialize};
use tokio::sync::watch;
use crate::agent::{TurnEmitError, TurnEmitter, TurnEvent};
use crate::providers::Usage;
#[derive(Debug, Clone, PartialEq, Eq, Hash, Serialize, Deserialize)]
#[serde(transparent)]
pub struct TurnId(pub String);
impl TurnId {
pub fn new() -> Self {
Self(uuid::Uuid::new_v4().to_string())
}
}
impl Default for TurnId {
fn default() -> Self {
Self::new()
}
}
#[derive(Debug, Clone, PartialEq, Eq, Hash, Serialize, Deserialize)]
#[serde(transparent)]
pub struct BlockId(pub String);
impl BlockId {
fn new() -> Self {
Self(uuid::Uuid::new_v4().to_string())
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum TurnStatus {
Running,
Completed,
Cancelled,
Failed,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum TurnPhase {
Queued,
Reasoning,
Responding,
Acting,
Finalizing,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum ToolStatus {
Running,
Completed,
Failed,
}
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
#[serde(tag = "type", rename_all = "snake_case")]
pub enum TurnBlock {
Reasoning {
id: BlockId,
iteration: u32,
text: String,
},
Assistant {
id: BlockId,
iteration: u32,
text: String,
},
Tool {
id: String,
iteration: u32,
name: String,
arguments: serde_json::Value,
status: ToolStatus,
preview: Option<String>,
},
}
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub struct TurnState {
pub id: TurnId,
pub session_id: String,
pub message_id: String,
pub revision: u64,
pub status: TurnStatus,
pub phase: TurnPhase,
pub blocks: Vec<TurnBlock>,
pub usage: Option<Usage>,
pub error: Option<String>,
}
pub type TurnSnapshot = TurnState;
struct TurnControllerInner {
state: TurnState,
snapshots: watch::Sender<Arc<TurnSnapshot>>,
text_segment_open: bool,
}
impl TurnControllerInner {
fn emit(&mut self, event: TurnEvent) -> Result<(), TurnEmitError> {
if self.state.status != TurnStatus::Running {
return Err(TurnEmitError::Inactive);
}
let changed = match event {
TurnEvent::ReasoningDelta { iteration, delta } => {
if delta.is_empty() {
false
} else {
self.state.phase = TurnPhase::Reasoning;
match self.state.blocks.last_mut() {
Some(TurnBlock::Reasoning {
iteration: current,
text,
..
}) if *current == iteration => text.push_str(&delta),
_ => self.state.blocks.push(TurnBlock::Reasoning {
id: BlockId::new(),
iteration,
text: delta,
}),
}
true
}
}
TurnEvent::TextDelta { iteration, delta } => {
if delta.is_empty() {
false
} else {
self.state.phase = TurnPhase::Responding;
if self.text_segment_open {
match self.state.blocks.last_mut() {
Some(TurnBlock::Assistant {
iteration: current,
text,
..
}) if *current == iteration => text.push_str(&delta),
_ => {
self.push_text_block(iteration, delta);
}
}
} else {
self.push_text_block(iteration, delta);
}
self.text_segment_open = true;
true
}
}
TurnEvent::TextSegmentFinished { .. } => {
self.text_segment_open = false;
false
}
TurnEvent::ToolStarted { iteration, call } => {
if self
.state
.blocks
.iter()
.any(|block| matches!(block, TurnBlock::Tool { id, .. } if id == &call.id))
{
return Err(TurnEmitError::DuplicateTool(call.id));
}
self.text_segment_open = false;
self.state.phase = TurnPhase::Acting;
self.state.blocks.push(TurnBlock::Tool {
id: call.id,
iteration,
name: call.name,
arguments: call.arguments,
status: ToolStatus::Running,
preview: None,
});
true
}
TurnEvent::ToolFinished {
iteration,
call_id,
success,
preview,
} => {
let Some(TurnBlock::Tool {
status,
preview: current_preview,
..
}) = self.state.blocks.iter_mut().find(|block| {
matches!(block, TurnBlock::Tool { id, iteration: current, .. } if id == &call_id && *current == iteration)
})
else {
return Err(TurnEmitError::UnknownTool(call_id));
};
*status = if success {
ToolStatus::Completed
} else {
ToolStatus::Failed
};
*current_preview = preview;
true
}
};
if changed {
self.publish();
}
Ok(())
}
fn push_text_block(&mut self, iteration: u32, text: String) {
self.state.blocks.push(TurnBlock::Assistant {
id: BlockId::new(),
iteration,
text,
});
}
fn publish(&mut self) {
self.state.revision = self.state.revision.wrapping_add(1);
self.snapshots.send_replace(Arc::new(self.state.clone()));
}
fn transition_terminal(
&mut self,
status: TurnStatus,
usage: Option<Usage>,
error: Option<String>,
) -> bool {
if self.state.status != TurnStatus::Running {
return false;
}
self.text_segment_open = false;
self.state.status = status;
self.state.phase = TurnPhase::Finalizing;
self.state.usage = usage;
self.state.error = error;
self.publish();
true
}
}
/// The sole writer for one running turn's presentation state.
///
/// Mutation is synchronous and bounded to a small in-memory reduction. This
/// lets AgentLoop emit facts without creating an unbounded token queue or a
/// reducer background task.
pub struct TurnController {
inner: Arc<Mutex<TurnControllerInner>>,
}
impl TurnController {
pub fn start(
session_id: impl Into<String>,
message_id: impl Into<String>,
) -> (Self, TurnEmitter, watch::Receiver<Arc<TurnSnapshot>>) {
let initial = TurnState {
id: TurnId::new(),
session_id: session_id.into(),
message_id: message_id.into(),
revision: 0,
status: TurnStatus::Running,
phase: TurnPhase::Queued,
blocks: Vec::new(),
usage: None,
error: None,
};
let (snapshots, receiver) = watch::channel(Arc::new(initial.clone()));
let inner = Arc::new(Mutex::new(TurnControllerInner {
state: initial,
snapshots,
text_segment_open: false,
}));
let weak: Weak<Mutex<TurnControllerInner>> = Arc::downgrade(&inner);
let emitter = TurnEmitter::new(move |event| {
let Some(inner) = weak.upgrade() else {
return Err(TurnEmitError::Inactive);
};
lock_unpoisoned(&inner).emit(event)
});
(Self { inner }, emitter, receiver)
}
pub fn snapshot(&self) -> Arc<TurnSnapshot> {
Arc::new(lock_unpoisoned(&self.inner).state.clone())
}
pub fn begin_finalizing(&self) -> bool {
let mut inner = lock_unpoisoned(&self.inner);
if inner.state.status != TurnStatus::Running || inner.state.phase == TurnPhase::Finalizing {
return false;
}
inner.text_segment_open = false;
inner.state.phase = TurnPhase::Finalizing;
inner.publish();
true
}
pub fn complete(&self, usage: Option<Usage>) -> bool {
lock_unpoisoned(&self.inner).transition_terminal(TurnStatus::Completed, usage, None)
}
pub fn cancel(&self, reason: Option<String>) -> bool {
lock_unpoisoned(&self.inner).transition_terminal(TurnStatus::Cancelled, None, reason)
}
pub fn fail(&self, error: impl Into<String>) -> bool {
lock_unpoisoned(&self.inner).transition_terminal(
TurnStatus::Failed,
None,
Some(error.into()),
)
}
}
fn lock_unpoisoned<T>(mutex: &Mutex<T>) -> MutexGuard<'_, T> {
mutex
.lock()
.unwrap_or_else(|poisoned| poisoned.into_inner())
}
#[cfg(test)]
mod tests {
use super::*;
use crate::providers::ToolCall;
fn start() -> (
TurnController,
TurnEmitter,
watch::Receiver<Arc<TurnSnapshot>>,
) {
TurnController::start("cli:test:dialog", "assistant-message")
}
#[test]
fn ordered_blocks_preserve_reasoning_text_tool_and_iterations() {
let (controller, emitter, _) = start();
emitter
.emit(TurnEvent::ReasoningDelta {
iteration: 0,
delta: "plan ".into(),
})
.unwrap();
emitter
.emit(TurnEvent::ReasoningDelta {
iteration: 0,
delta: "step".into(),
})
.unwrap();
emitter
.emit(TurnEvent::TextDelta {
iteration: 0,
delta: "checking".into(),
})
.unwrap();
emitter
.emit(TurnEvent::ToolStarted {
iteration: 0,
call: ToolCall {
id: "call-1".into(),
name: "bash".into(),
arguments: serde_json::json!({"cmd":"pwd"}),
},
})
.unwrap();
emitter
.emit(TurnEvent::ToolFinished {
iteration: 0,
call_id: "call-1".into(),
success: true,
preview: Some("/tmp".into()),
})
.unwrap();
emitter
.emit(TurnEvent::ReasoningDelta {
iteration: 1,
delta: "done thinking".into(),
})
.unwrap();
emitter
.emit(TurnEvent::TextDelta {
iteration: 1,
delta: "final".into(),
})
.unwrap();
let snapshot = controller.snapshot();
assert_eq!(snapshot.revision, 7);
assert_eq!(snapshot.phase, TurnPhase::Responding);
assert_eq!(snapshot.blocks.len(), 5);
assert!(matches!(
&snapshot.blocks[0],
TurnBlock::Reasoning { iteration: 0, text, .. } if text == "plan step"
));
assert!(matches!(
&snapshot.blocks[1],
TurnBlock::Assistant { iteration: 0, text, .. } if text == "checking"
));
assert!(matches!(
&snapshot.blocks[2],
TurnBlock::Tool {
id,
status: ToolStatus::Completed,
preview: Some(preview),
..
} if id == "call-1" && preview == "/tmp"
));
assert!(matches!(
&snapshot.blocks[3],
TurnBlock::Reasoning { iteration: 1, .. }
));
assert!(matches!(
&snapshot.blocks[4],
TurnBlock::Assistant { iteration: 1, .. }
));
}
#[test]
fn explicit_segment_boundary_prevents_text_coalescing() {
let (controller, emitter, _) = start();
emitter
.emit(TurnEvent::TextDelta {
iteration: 0,
delta: "before".into(),
})
.unwrap();
emitter
.emit(TurnEvent::TextSegmentFinished { iteration: 0 })
.unwrap();
emitter
.emit(TurnEvent::TextDelta {
iteration: 0,
delta: "after".into(),
})
.unwrap();
let snapshot = controller.snapshot();
assert_eq!(snapshot.revision, 2);
assert_eq!(snapshot.blocks.len(), 2);
assert!(
matches!(&snapshot.blocks[0], TurnBlock::Assistant { text, .. } if text == "before")
);
assert!(
matches!(&snapshot.blocks[1], TurnBlock::Assistant { text, .. } if text == "after")
);
}
#[test]
fn revisions_are_monotonic_and_watch_is_latest_wins() {
let (_controller, emitter, receiver) = start();
for delta in ["a", "b", "c"] {
emitter
.emit(TurnEvent::TextDelta {
iteration: 0,
delta: delta.into(),
})
.unwrap();
}
let latest = receiver.borrow().clone();
assert_eq!(latest.revision, 3);
assert!(matches!(&latest.blocks[0], TurnBlock::Assistant { text, .. } if text == "abc"));
}
#[test]
fn terminal_state_rejects_late_events_and_is_idempotent() {
let (controller, emitter, receiver) = start();
emitter
.emit(TurnEvent::TextDelta {
iteration: 0,
delta: "saved".into(),
})
.unwrap();
assert!(controller.begin_finalizing());
assert!(controller.complete(None));
assert!(!controller.complete(None));
assert_eq!(
emitter.emit(TurnEvent::TextDelta {
iteration: 0,
delta: "late".into(),
}),
Err(TurnEmitError::Inactive)
);
let latest = receiver.borrow().clone();
assert_eq!(latest.status, TurnStatus::Completed);
assert_eq!(latest.phase, TurnPhase::Finalizing);
assert_eq!(latest.revision, 3);
assert!(matches!(&latest.blocks[0], TurnBlock::Assistant { text, .. } if text == "saved"));
}
#[test]
fn invalid_tool_transitions_do_not_publish() {
let (_controller, emitter, receiver) = start();
let unknown = emitter.emit(TurnEvent::ToolFinished {
iteration: 0,
call_id: "missing".into(),
success: false,
preview: None,
});
assert_eq!(unknown, Err(TurnEmitError::UnknownTool("missing".into())));
assert_eq!(receiver.borrow().revision, 0);
let call = ToolCall {
id: "same".into(),
name: "bash".into(),
arguments: serde_json::json!({}),
};
emitter
.emit(TurnEvent::ToolStarted {
iteration: 0,
call: call.clone(),
})
.unwrap();
assert_eq!(
emitter.emit(TurnEvent::ToolStarted { iteration: 0, call }),
Err(TurnEmitError::DuplicateTool("same".into()))
);
assert_eq!(receiver.borrow().revision, 1);
}
#[test]
fn parallel_tools_update_independently() {
let (controller, emitter, _) = start();
for id in ["first", "second"] {
emitter
.emit(TurnEvent::ToolStarted {
iteration: 0,
call: ToolCall {
id: id.into(),
name: "bash".into(),
arguments: serde_json::json!({"cmd": id}),
},
})
.unwrap();
}
emitter
.emit(TurnEvent::ToolFinished {
iteration: 0,
call_id: "second".into(),
success: false,
preview: Some("failed".into()),
})
.unwrap();
let snapshot = controller.snapshot();
assert!(matches!(
&snapshot.blocks[0],
TurnBlock::Tool { id, status: ToolStatus::Running, .. } if id == "first"
));
assert!(matches!(
&snapshot.blocks[1],
TurnBlock::Tool { id, status: ToolStatus::Failed, .. } if id == "second"
));
}
#[test]
fn empty_deltas_are_noops_and_cancel_reason_is_terminal() {
let (controller, emitter, receiver) = start();
emitter
.emit(TurnEvent::ReasoningDelta {
iteration: 0,
delta: String::new(),
})
.unwrap();
emitter
.emit(TurnEvent::TextDelta {
iteration: 0,
delta: String::new(),
})
.unwrap();
assert_eq!(receiver.borrow().revision, 0);
assert!(controller.cancel(Some("stopped by user".into())));
let snapshot = controller.snapshot();
assert_eq!(snapshot.status, TurnStatus::Cancelled);
assert_eq!(snapshot.error.as_deref(), Some("stopped by user"));
assert_eq!(snapshot.revision, 1);
}
#[test]
fn failure_is_published_as_structured_terminal_state() {
let (controller, _emitter, _) = start();
assert!(controller.fail("provider disconnected"));
let snapshot = controller.snapshot();
assert_eq!(snapshot.status, TurnStatus::Failed);
assert_eq!(snapshot.error.as_deref(), Some("provider disconnected"));
assert_eq!(snapshot.phase, TurnPhase::Finalizing);
}
#[test]
fn deactivated_emitter_drops_events_before_stale_worker_reduction() {
let (controller, emitter, _) = start();
emitter.deactivate();
emitter
.emit(TurnEvent::TextDelta {
iteration: 0,
delta: "late".into(),
})
.unwrap();
assert_eq!(controller.snapshot().revision, 0);
assert!(controller.snapshot().blocks.is_empty());
}
}

6
src/storage/message.rs
View File

@ -1,5 +1,7 @@
use serde::{Deserialize, Serialize}; use serde::{Deserialize, Serialize};
use crate::bus::CompletionStatus;
#[derive(Debug, Clone, Serialize, Deserialize)] #[derive(Debug, Clone, Serialize, Deserialize)]
pub struct MessageMeta { pub struct MessageMeta {
pub id: String, pub id: String,
@ -8,6 +10,10 @@ pub struct MessageMeta {
pub role: String, pub role: String,
pub content: String, pub content: String,
pub reasoning_content: Option<String>, pub reasoning_content: Option<String>,
pub provider_state: Option<String>,
pub turn_id: Option<String>,
pub iteration: Option<i64>,
pub completion_status: CompletionStatus,
pub media_refs: Option<String>, pub media_refs: Option<String>,
pub tool_call_id: Option<String>, pub tool_call_id: Option<String>,
pub tool_name: Option<String>, pub tool_name: Option<String>,

344
src/storage/mod.rs
View File

@ -9,15 +9,21 @@ pub use background_task::BackgroundTask;
pub use error::StorageError; pub use error::StorageError;
pub use scheduler::{DeliveryPolicy, JobKind, JobRun, ScheduledJob}; pub use scheduler::{DeliveryPolicy, JobKind, JobRun, ScheduledJob};
use sqlx::sqlite::{SqliteConnectOptions, SqliteJournalMode, SqlitePoolOptions, SqliteSynchronous}; use sqlx::sqlite::{
SqliteConnectOptions, SqliteJournalMode, SqlitePoolOptions, SqliteRow, SqliteSynchronous,
};
use sqlx::{Pool, Row, Sqlite}; use sqlx::{Pool, Row, Sqlite};
use std::path::Path; use std::path::Path;
use tokio::time::{Duration, sleep}; use tokio::time::{Duration, sleep};
const SCHEMA_VERSION: i64 = 3; const SCHEMA_VERSION: i64 = 4;
const INSERT_MESSAGE_SQL: &str = r#" const INSERT_MESSAGE_SQL: &str = r#"
INSERT INTO messages (id, session_id, seq, role, content, reasoning_content, media_refs, tool_call_id, tool_name, tool_calls, source, created_at) INSERT INTO messages (
VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) id, session_id, seq, role, content, reasoning_content, provider_state,
turn_id, iteration, completion_status, media_refs, tool_call_id,
tool_name, tool_calls, source, created_at
)
VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
"#; "#;
fn insert_message_query<'a>( fn insert_message_query<'a>(
@ -31,6 +37,10 @@ fn insert_message_query<'a>(
.bind(&msg.role) .bind(&msg.role)
.bind(&msg.content) .bind(&msg.content)
.bind(&msg.reasoning_content) .bind(&msg.reasoning_content)
.bind(&msg.provider_state)
.bind(&msg.turn_id)
.bind(msg.iteration)
.bind(msg.completion_status.as_str())
.bind(&msg.media_refs) .bind(&msg.media_refs)
.bind(&msg.tool_call_id) .bind(&msg.tool_call_id)
.bind(&msg.tool_name) .bind(&msg.tool_name)
@ -39,6 +49,28 @@ fn insert_message_query<'a>(
.bind(msg.created_at) .bind(msg.created_at)
} }
fn message_meta_from_row(row: SqliteRow) -> crate::storage::message::MessageMeta {
let completion_status: String = row.get("completion_status");
crate::storage::message::MessageMeta {
id: row.get("id"),
session_id: row.get("session_id"),
seq: row.get("seq"),
role: row.get("role"),
content: row.get("content"),
reasoning_content: row.get("reasoning_content"),
provider_state: row.get("provider_state"),
turn_id: row.get("turn_id"),
iteration: row.get("iteration"),
completion_status: crate::bus::CompletionStatus::from_storage(&completion_status),
media_refs: row.get("media_refs"),
tool_call_id: row.get("tool_call_id"),
tool_name: row.get("tool_name"),
tool_calls: row.get("tool_calls"),
source: row.get("source"),
created_at: row.get("created_at"),
}
}
pub struct Storage { pub struct Storage {
pub(crate) pool: Pool<Sqlite>, pub(crate) pool: Pool<Sqlite>,
} }
@ -111,6 +143,10 @@ impl Storage {
tool_calls TEXT, tool_calls TEXT,
source TEXT, source TEXT,
reasoning_content TEXT, reasoning_content TEXT,
provider_state TEXT,
turn_id TEXT,
iteration INTEGER,
completion_status TEXT NOT NULL DEFAULT 'completed',
created_at INTEGER NOT NULL, created_at INTEGER NOT NULL,
FOREIGN KEY (session_id) REFERENCES sessions(id) ON DELETE CASCADE FOREIGN KEY (session_id) REFERENCES sessions(id) ON DELETE CASCADE
) )
@ -351,6 +387,14 @@ impl Storage {
for (table, column, definition) in [ for (table, column, definition) in [
("messages", "source", "source TEXT"), ("messages", "source", "source TEXT"),
("messages", "reasoning_content", "reasoning_content TEXT"), ("messages", "reasoning_content", "reasoning_content TEXT"),
("messages", "provider_state", "provider_state TEXT"),
("messages", "turn_id", "turn_id TEXT"),
("messages", "iteration", "iteration INTEGER"),
(
"messages",
"completion_status",
"completion_status TEXT NOT NULL DEFAULT 'completed'",
),
("sessions", "archived_at", "archived_at INTEGER"), ("sessions", "archived_at", "archived_at INTEGER"),
( (
"sessions", "sessions",
@ -829,7 +873,9 @@ impl Storage {
) -> Result<Vec<crate::storage::message::MessageMeta>, StorageError> { ) -> Result<Vec<crate::storage::message::MessageMeta>, StorageError> {
let rows = sqlx::query( let rows = sqlx::query(
r#" r#"
SELECT id, session_id, seq, role, content, reasoning_content, media_refs, tool_call_id, tool_name, tool_calls, source, created_at SELECT id, session_id, seq, role, content, reasoning_content, provider_state,
turn_id, iteration, completion_status, media_refs, tool_call_id,
tool_name, tool_calls, source, created_at
FROM messages FROM messages
WHERE session_id = ? AND seq >= ? WHERE session_id = ? AND seq >= ?
ORDER BY seq ASC ORDER BY seq ASC
@ -840,23 +886,7 @@ impl Storage {
.fetch_all(self.pool()) .fetch_all(self.pool())
.await?; .await?;
Ok(rows Ok(rows.into_iter().map(message_meta_from_row).collect())
.into_iter()
.map(|row| crate::storage::message::MessageMeta {
id: row.get("id"),
session_id: row.get("session_id"),
seq: row.get("seq"),
role: row.get("role"),
content: row.get("content"),
reasoning_content: row.get("reasoning_content"),
media_refs: row.get("media_refs"),
tool_call_id: row.get("tool_call_id"),
tool_name: row.get("tool_name"),
tool_calls: row.get("tool_calls"),
source: row.get("source"),
created_at: row.get("created_at"),
})
.collect())
} }
pub async fn get_message( pub async fn get_message(
@ -866,8 +896,9 @@ impl Storage {
) -> Result<Option<crate::storage::message::MessageMeta>, StorageError> { ) -> Result<Option<crate::storage::message::MessageMeta>, StorageError> {
let row = sqlx::query( let row = sqlx::query(
r#" r#"
SELECT id, session_id, seq, role, content, reasoning_content, media_refs, SELECT id, session_id, seq, role, content, reasoning_content, provider_state,
tool_call_id, tool_name, tool_calls, source, created_at turn_id, iteration, completion_status, media_refs, tool_call_id,
tool_name, tool_calls, source, created_at
FROM messages FROM messages
WHERE session_id = ? AND id = ? WHERE session_id = ? AND id = ?
"#, "#,
@ -877,20 +908,7 @@ impl Storage {
.fetch_optional(self.pool()) .fetch_optional(self.pool())
.await?; .await?;
Ok(row.map(|row| crate::storage::message::MessageMeta { Ok(row.map(message_meta_from_row))
id: row.get("id"),
session_id: row.get("session_id"),
seq: row.get("seq"),
role: row.get("role"),
content: row.get("content"),
reasoning_content: row.get("reasoning_content"),
media_refs: row.get("media_refs"),
tool_call_id: row.get("tool_call_id"),
tool_name: row.get("tool_name"),
tool_calls: row.get("tool_calls"),
source: row.get("source"),
created_at: row.get("created_at"),
}))
} }
pub async fn get_max_message_seq(&self, session_id: &str) -> Result<i64, StorageError> { pub async fn get_max_message_seq(&self, session_id: &str) -> Result<i64, StorageError> {
@ -912,11 +930,13 @@ impl Storage {
let limit = limit.clamp(1, 2_000); let limit = limit.clamp(1, 2_000);
let rows = sqlx::query( let rows = sqlx::query(
r#" r#"
SELECT id, session_id, seq, role, content, reasoning_content, media_refs, SELECT id, session_id, seq, role, content, reasoning_content, provider_state,
tool_call_id, tool_name, tool_calls, source, created_at turn_id, iteration, completion_status, media_refs, tool_call_id,
tool_name, tool_calls, source, created_at
FROM ( FROM (
SELECT id, session_id, seq, role, content, reasoning_content, media_refs, SELECT id, session_id, seq, role, content, reasoning_content, provider_state,
tool_call_id, tool_name, tool_calls, source, created_at turn_id, iteration, completion_status, media_refs, tool_call_id,
tool_name, tool_calls, source, created_at
FROM messages FROM messages
WHERE session_id = ? WHERE session_id = ?
ORDER BY seq DESC ORDER BY seq DESC
@ -930,23 +950,7 @@ impl Storage {
.fetch_all(self.pool()) .fetch_all(self.pool())
.await?; .await?;
Ok(rows Ok(rows.into_iter().map(message_meta_from_row).collect())
.into_iter()
.map(|row| crate::storage::message::MessageMeta {
id: row.get("id"),
session_id: row.get("session_id"),
seq: row.get("seq"),
role: row.get("role"),
content: row.get("content"),
reasoning_content: row.get("reasoning_content"),
media_refs: row.get("media_refs"),
tool_call_id: row.get("tool_call_id"),
tool_name: row.get("tool_name"),
tool_calls: row.get("tool_calls"),
source: row.get("source"),
created_at: row.get("created_at"),
})
.collect())
} }
pub async fn load_messages_after_timestamp( pub async fn load_messages_after_timestamp(
@ -956,7 +960,9 @@ impl Storage {
) -> Result<Vec<crate::storage::message::MessageMeta>, StorageError> { ) -> Result<Vec<crate::storage::message::MessageMeta>, StorageError> {
let rows = sqlx::query( let rows = sqlx::query(
r#" r#"
SELECT id, session_id, seq, role, content, reasoning_content, media_refs, tool_call_id, tool_name, tool_calls, source, created_at SELECT id, session_id, seq, role, content, reasoning_content, provider_state,
turn_id, iteration, completion_status, media_refs, tool_call_id,
tool_name, tool_calls, source, created_at
FROM messages FROM messages
WHERE session_id = ? AND created_at > ? WHERE session_id = ? AND created_at > ?
ORDER BY seq ASC ORDER BY seq ASC
@ -967,23 +973,7 @@ impl Storage {
.fetch_all(self.pool()) .fetch_all(self.pool())
.await?; .await?;
Ok(rows Ok(rows.into_iter().map(message_meta_from_row).collect())
.into_iter()
.map(|row| crate::storage::message::MessageMeta {
id: row.get("id"),
session_id: row.get("session_id"),
seq: row.get("seq"),
role: row.get("role"),
content: row.get("content"),
reasoning_content: row.get("reasoning_content"),
media_refs: row.get("media_refs"),
tool_call_id: row.get("tool_call_id"),
tool_name: row.get("tool_name"),
tool_calls: row.get("tool_calls"),
source: row.get("source"),
created_at: row.get("created_at"),
})
.collect())
} }
pub async fn query_sessions_range( pub async fn query_sessions_range(
@ -1040,7 +1030,9 @@ impl Storage {
) -> Result<Vec<crate::storage::message::MessageMeta>, StorageError> { ) -> Result<Vec<crate::storage::message::MessageMeta>, StorageError> {
let rows = sqlx::query( let rows = sqlx::query(
r#" r#"
SELECT id, session_id, seq, role, content, reasoning_content, media_refs, tool_call_id, tool_name, tool_calls, source, created_at SELECT id, session_id, seq, role, content, reasoning_content, provider_state,
turn_id, iteration, completion_status, media_refs, tool_call_id,
tool_name, tool_calls, source, created_at
FROM messages FROM messages
WHERE session_id = ? WHERE session_id = ?
ORDER BY seq DESC ORDER BY seq DESC
@ -1052,23 +1044,7 @@ impl Storage {
.fetch_all(self.pool()) .fetch_all(self.pool())
.await?; .await?;
let mut messages: Vec<_> = rows let mut messages: Vec<_> = rows.into_iter().map(message_meta_from_row).collect();
.into_iter()
.map(|row| crate::storage::message::MessageMeta {
id: row.get("id"),
session_id: row.get("session_id"),
seq: row.get("seq"),
role: row.get("role"),
content: row.get("content"),
reasoning_content: row.get("reasoning_content"),
media_refs: row.get("media_refs"),
tool_call_id: row.get("tool_call_id"),
tool_name: row.get("tool_name"),
tool_calls: row.get("tool_calls"),
source: row.get("source"),
created_at: row.get("created_at"),
})
.collect();
messages.reverse(); messages.reverse();
Ok(messages) Ok(messages)
} }
@ -1095,7 +1071,9 @@ impl Storage {
); );
let select_sql = format!( let select_sql = format!(
r#" r#"
SELECT id, session_id, seq, role, content, reasoning_content, media_refs, tool_call_id, tool_name, tool_calls, source, created_at SELECT id, session_id, seq, role, content, reasoning_content, provider_state,
turn_id, iteration, completion_status, media_refs, tool_call_id,
tool_name, tool_calls, source, created_at
FROM messages FROM messages
WHERE session_id = ?{} WHERE session_id = ?{}
ORDER BY seq ASC ORDER BY seq ASC
@ -1127,23 +1105,7 @@ impl Storage {
.fetch_all(self.pool()) .fetch_all(self.pool())
.await?; .await?;
let messages: Vec<_> = rows let messages: Vec<_> = rows.into_iter().map(message_meta_from_row).collect();
.into_iter()
.map(|row| crate::storage::message::MessageMeta {
id: row.get("id"),
session_id: row.get("session_id"),
seq: row.get("seq"),
role: row.get("role"),
content: row.get("content"),
reasoning_content: row.get("reasoning_content"),
media_refs: row.get("media_refs"),
tool_call_id: row.get("tool_call_id"),
tool_name: row.get("tool_name"),
tool_calls: row.get("tool_calls"),
source: row.get("source"),
created_at: row.get("created_at"),
})
.collect();
Ok((messages, total)) Ok((messages, total))
} }
@ -1616,7 +1578,17 @@ mod tests {
let storage = Storage::new(&db_path).await.unwrap(); let storage = Storage::new(&db_path).await.unwrap();
for (table, expected) in [ for (table, expected) in [
("messages", vec!["source", "reasoning_content"]), (
"messages",
vec![
"source",
"reasoning_content",
"provider_state",
"turn_id",
"iteration",
"completion_status",
],
),
( (
"sessions", "sessions",
vec![ vec![
@ -1660,6 +1632,81 @@ mod tests {
} }
} }
#[tokio::test]
async fn v3_migration_preserves_existing_reasoning_and_defaults_completion() {
let dir = tempfile::tempdir().unwrap();
let db_path = dir.path().join("v3.db");
let pool = SqlitePoolOptions::new()
.connect_with(
SqliteConnectOptions::new()
.filename(&db_path)
.create_if_missing(true),
)
.await
.unwrap();
sqlx::query(
r#"
CREATE TABLE sessions (
id TEXT PRIMARY KEY, channel TEXT NOT NULL, chat_id TEXT NOT NULL,
dialog_id TEXT NOT NULL, title TEXT NOT NULL DEFAULT 'new',
created_at INTEGER NOT NULL, last_active_at INTEGER NOT NULL,
message_count INTEGER DEFAULT 0, routing_info TEXT, archived_at INTEGER,
deleted_at INTEGER, last_consolidated_at INTEGER,
last_compressed_message_at INTEGER,
UNIQUE(channel, chat_id, dialog_id)
)
"#,
)
.execute(&pool)
.await
.unwrap();
sqlx::query(
r#"
CREATE TABLE messages (
id TEXT PRIMARY KEY, session_id TEXT NOT NULL, seq INTEGER NOT NULL,
role TEXT NOT NULL, content TEXT NOT NULL, reasoning_content TEXT,
media_refs TEXT, tool_call_id TEXT, tool_name TEXT, tool_calls TEXT,
source TEXT, created_at INTEGER NOT NULL
)
"#,
)
.execute(&pool)
.await
.unwrap();
sqlx::query(
"INSERT INTO sessions (id, channel, chat_id, dialog_id, created_at, last_active_at) VALUES ('cli:c:d', 'cli', 'c', 'd', 1, 1)",
)
.execute(&pool)
.await
.unwrap();
sqlx::query(
"INSERT INTO messages (id, session_id, seq, role, content, reasoning_content, created_at) VALUES ('m1', 'cli:c:d', 1, 'assistant', 'answer', 'existing reasoning', 1)",
)
.execute(&pool)
.await
.unwrap();
sqlx::query("PRAGMA user_version = 3")
.execute(&pool)
.await
.unwrap();
drop(pool);
let storage = Storage::new(&db_path).await.unwrap();
let messages = storage.load_messages("cli:c:d", 0).await.unwrap();
assert_eq!(messages.len(), 1);
assert_eq!(
messages[0].reasoning_content.as_deref(),
Some("existing reasoning")
);
assert_eq!(messages[0].provider_state, None);
assert_eq!(messages[0].turn_id, None);
assert_eq!(messages[0].iteration, None);
assert_eq!(
messages[0].completion_status,
crate::bus::CompletionStatus::Completed
);
}
#[tokio::test] #[tokio::test]
async fn test_upsert_and_get_session() { async fn test_upsert_and_get_session() {
let (storage, _dir) = create_test_storage().await; let (storage, _dir) = create_test_storage().await;
@ -1784,6 +1831,10 @@ mod tests {
role: "user".to_string(), role: "user".to_string(),
content: "你好".to_string(), content: "你好".to_string(),
reasoning_content: None, reasoning_content: None,
provider_state: None,
turn_id: None,
iteration: None,
completion_status: crate::bus::CompletionStatus::Completed,
media_refs: None, media_refs: None,
tool_call_id: None, tool_call_id: None,
tool_name: None, tool_name: None,
@ -1821,6 +1872,67 @@ mod tests {
assert_eq!(recent[1].seq, 5); assert_eq!(recent[1].seq, 5);
} }
#[tokio::test]
async fn streaming_message_metadata_round_trips() {
let (storage, _dir) = create_test_storage().await;
let session_meta = crate::storage::session::SessionMeta {
id: "cli_chat:stream:dialog1".to_string(),
channel: "cli_chat".to_string(),
chat_id: "stream".to_string(),
dialog_id: "dialog1".to_string(),
title: "Stream metadata".to_string(),
created_at: 1000,
last_active_at: 1000,
message_count: 0,
routing_info: None,
archived_at: None,
deleted_at: None,
last_consolidated_at: None,
last_compressed_message_at: None,
};
storage.upsert_session(&session_meta).await.unwrap();
let message = crate::storage::message::MessageMeta {
id: "assistant-1".to_string(),
session_id: session_meta.id.clone(),
seq: 1,
role: "assistant".to_string(),
content: "partial".to_string(),
reasoning_content: Some("visible reasoning".to_string()),
provider_state: Some(
serde_json::json!({"provider":"anthropic","payload":{"signature":"opaque"}})
.to_string(),
),
turn_id: Some("turn-1".to_string()),
iteration: Some(2),
completion_status: crate::bus::CompletionStatus::Interrupted,
media_refs: None,
tool_call_id: None,
tool_name: None,
tool_calls: None,
source: None,
created_at: 1001,
};
storage
.append_message(&session_meta.id, &message)
.await
.unwrap();
let loaded = storage.load_messages(&session_meta.id, 0).await.unwrap();
assert_eq!(loaded.len(), 1);
assert_eq!(
loaded[0].reasoning_content.as_deref(),
Some("visible reasoning")
);
assert_eq!(loaded[0].provider_state, message.provider_state);
assert_eq!(loaded[0].turn_id.as_deref(), Some("turn-1"));
assert_eq!(loaded[0].iteration, Some(2));
assert_eq!(
loaded[0].completion_status,
crate::bus::CompletionStatus::Interrupted
);
}
#[tokio::test] #[tokio::test]
async fn test_persist_message_batch_is_atomic() { async fn test_persist_message_batch_is_atomic() {
let (storage, _dir) = create_test_storage().await; let (storage, _dir) = create_test_storage().await;
@ -1848,6 +1960,10 @@ mod tests {
role: "assistant".to_string(), role: "assistant".to_string(),
content: "must roll back".to_string(), content: "must roll back".to_string(),
reasoning_content: None, reasoning_content: None,
provider_state: None,
turn_id: None,
iteration: None,
completion_status: crate::bus::CompletionStatus::Completed,
media_refs: None, media_refs: None,
tool_call_id: None, tool_call_id: None,
tool_name: None, tool_name: None,

12
src/tools/chat_manager.rs
View File

@ -358,6 +358,10 @@ mod tests {
}, },
content: format!("消息内容 {}", i), content: format!("消息内容 {}", i),
reasoning_content: None, reasoning_content: None,
provider_state: None,
turn_id: None,
iteration: None,
completion_status: crate::bus::CompletionStatus::Completed,
media_refs: None, media_refs: None,
tool_call_id: None, tool_call_id: None,
tool_name: None, tool_name: None,
@ -420,6 +424,10 @@ mod tests {
}, },
content: format!("消息内容 {}", i), content: format!("消息内容 {}", i),
reasoning_content: None, reasoning_content: None,
provider_state: None,
turn_id: None,
iteration: None,
completion_status: crate::bus::CompletionStatus::Completed,
media_refs: None, media_refs: None,
tool_call_id: None, tool_call_id: None,
tool_name: None, tool_name: None,
@ -476,6 +484,10 @@ mod tests {
role: "user".to_string(), role: "user".to_string(),
content: format!("消息内容 {}", i), content: format!("消息内容 {}", i),
reasoning_content: None, reasoning_content: None,
provider_state: None,
turn_id: None,
iteration: None,
completion_status: crate::bus::CompletionStatus::Completed,
media_refs: None, media_refs: None,
tool_call_id: None, tool_call_id: None,
tool_name: None, tool_name: None,

5
tests/test_request_format.rs
View File

@ -135,6 +135,8 @@ fn test_bounded_session_history_protocol() {
seq: 1, seq: 1,
role: "user".to_string(), role: "user".to_string(),
content: "你好".to_string(), content: "你好".to_string(),
reasoning_content: None,
completion_status: picobot::bus::CompletionStatus::Completed,
created_at: 123, created_at: 123,
tool_call_id: None, tool_call_id: None,
tool_name: None, tool_name: None,
@ -182,6 +184,8 @@ fn test_session_history_preserves_tool_call_metadata() {
seq: 2, seq: 2,
role: "assistant".to_string(), role: "assistant".to_string(),
content: String::new(), content: String::new(),
reasoning_content: Some("checking".to_string()),
completion_status: picobot::bus::CompletionStatus::Completed,
created_at: 124, created_at: 124,
tool_call_id: None, tool_call_id: None,
tool_name: None, tool_name: None,
@ -197,6 +201,7 @@ fn test_session_history_preserves_tool_call_metadata() {
let json = serde_json::to_string(&outbound).unwrap(); let json = serde_json::to_string(&outbound).unwrap();
assert!(json.contains(r#""tool_calls""#)); assert!(json.contains(r#""tool_calls""#));
assert!(json.contains(r#""read_file""#)); assert!(json.contains(r#""read_file""#));
assert!(json.contains(r#""reasoning_content":"checking""#));
let decoded: WsOutbound = serde_json::from_str(&json).unwrap(); let decoded: WsOutbound = serde_json::from_str(&json).unwrap();
match decoded { match decoded {
WsOutbound::SessionHistory { messages, .. } => { WsOutbound::SessionHistory { messages, .. } => {

45
webui/src/lib/TurnView.svelte Normal file
View File

@ -0,0 +1,45 @@
<script>
import Markdown from "./Markdown.svelte";
let { turn } = $props();
function phaseLabel(value) {
return ({ queued: "排队中", reasoning: "思考中", responding: "生成中", acting: "调用工具中", finalizing: "收尾中" })[value] || value;
}
function statusLabel(value) {
return ({ running: phaseLabel(turn.phase), completed: "已完成", cancelled: "已停止", failed: "失败" })[value] || value;
}
function toolStatus(value) {
return ({ running: "执行中", completed: "已完成", failed: "失败" })[value] || value;
}
</script>
<div class="message assistant active-turn">
<div class="avatar">P</div>
<div class="message-content">
{#each turn.blocks as block (block.id)}
{#if block.type === "reasoning"}
<details class="reasoning-block">
<summary><span class="pulse"></span>思考过程</summary>
<div class="reasoning-content"><Markdown content={block.text} /></div>
</details>
{:else if block.type === "assistant"}
<div class="bubble streaming"><Markdown content={block.text} /></div>
{:else if block.type === "tool"}
<details class="live-tool">
<summary><span>⌘</span><strong>{block.name}</strong><small>{toolStatus(block.status)}</small></summary>
<div class="live-tool-details">
{#if block.arguments !== null}<pre>{JSON.stringify(block.arguments, null, 2)}</pre>{/if}
{#if block.preview}<div class="tool-result"><Markdown content={block.preview} /></div>{/if}
</div>
</details>
{/if}
{/each}
<div class:failed={turn.status === "failed"} class="turn-status">
{#if turn.status === "running"}<span class="pulse"></span>{/if}{statusLabel(turn.status)}
</div>
{#if turn.error}<div class="turn-error">{turn.error}</div>{/if}
</div>
</div>

36
webui/src/pages/ChatPage.svelte
View File

@ -4,6 +4,7 @@
import { clientId, formatTime, randomId } from "../lib/api.js"; import { clientId, formatTime, randomId } from "../lib/api.js";
import Markdown from "../lib/Markdown.svelte"; import Markdown from "../lib/Markdown.svelte";
import ToolCallCard from "../lib/ToolCallCard.svelte"; import ToolCallCard from "../lib/ToolCallCard.svelte";
import TurnView from "../lib/TurnView.svelte";
let { notify } = $props(); let { notify } = $props();
let socket = $state(null); let socket = $state(null);
@ -17,6 +18,7 @@
let selectedCommand = $state(0); let selectedCommand = $state(0);
let commandMenuDismissed = $state(false); let commandMenuDismissed = $state(false);
let thinking = $state(false); let thinking = $state(false);
let activeTurn = $state(null);
let pendingUploads = $state([]); let pendingUploads = $state([]);
let fileInput; let fileInput;
let plansBySession = $state({}); let plansBySession = $state({});
@ -75,7 +77,7 @@
function handleFrame(frame) { function handleFrame(frame) {
switch (frame.type) { switch (frame.type) {
case "session_established": currentId = frame.session_id; break; case "session_established": currentId = frame.session_id; activeTurn = null; break;
case "session_list": case "session_list":
sessions = frame.sessions || []; sessions = frame.sessions || [];
if (frame.current_session_id) currentId = frame.current_session_id; if (frame.current_session_id) currentId = frame.current_session_id;
@ -84,16 +86,19 @@
case "session_created": case "session_created":
currentId = frame.session_id; currentId = frame.session_id;
messages = []; messages = [];
activeTurn = null;
send({ type: "list_sessions", include_archived: false }); send({ type: "list_sessions", include_archived: false });
break; break;
case "session_loaded": case "session_loaded":
currentId = frame.session_id; currentId = frame.session_id;
activeTurn = null;
send({ type: "get_session_history", session_id: currentId, limit: 1000 }); send({ type: "get_session_history", session_id: currentId, limit: 1000 });
send({ type: "get_session_plan", session_id: currentId }); send({ type: "get_session_plan", session_id: currentId });
break; break;
case "session_history": case "session_history":
if (frame.session_id === currentId) { if (frame.session_id === currentId) {
messages = frame.messages || []; messages = frame.messages || [];
if (activeTurn?.status !== "running" && messages.some((message) => message.id === activeTurn?.message_id)) activeTurn = null;
scrollToBottom(); scrollToBottom();
} }
break; break;
@ -133,6 +138,19 @@
} }
send({ type: "list_sessions", include_archived: false }); send({ type: "list_sessions", include_archived: false });
break; break;
case "turn_updated": {
const next = frame.snapshot;
if (!next || next.session_id !== currentId) break;
if (activeTurn?.id === next.id && activeTurn.revision >= next.revision) break;
activeTurn = next;
thinking = next.status === "running";
scrollToBottom();
if (next.status !== "running") {
send({ type: "get_session_history", session_id: currentId, limit: 1000 });
send({ type: "list_sessions", include_archived: false });
}
break;
}
case "system_notification": case "system_notification":
if (!frame.session_id || frame.session_id === currentId) appendMessage("assistant", frame.content); if (!frame.session_id || frame.session_id === currentId) appendMessage("assistant", frame.content);
break; break;
@ -156,6 +174,8 @@
currentId = id; currentId = id;
clearPendingUploads(); clearPendingUploads();
messages = []; messages = [];
activeTurn = null;
thinking = false;
todoOpen = Boolean(unseenPlanSessions[id]); todoOpen = Boolean(unseenPlanSessions[id]);
unseenPlanSessions[id] = false; unseenPlanSessions[id] = false;
send({ type: "load_session", session_id: id }); send({ type: "load_session", session_id: id });
@ -375,7 +395,7 @@
</div> </div>
</div> </div>
<div class="messages" bind:this={messageBox}> <div class="messages" bind:this={messageBox}>
{#if messages.length === 0} {#if messages.length === 0 && !activeTurn}
<div class="empty"><div class="empty-logo">P</div><h2>今天想做些什么？</h2><p>消息与 CLI 客户端使用同一套会话、记忆和工具能力。</p></div> <div class="empty"><div class="empty-logo">P</div><h2>今天想做些什么？</h2><p>消息与 CLI 客户端使用同一套会话、记忆和工具能力。</p></div>
{/if} {/if}
{#each messages as message (message.id)} {#each messages as message (message.id)}
@ -383,7 +403,16 @@
<div class:user={message.role === "user"} class:assistant={message.role !== "user"} class:has-tools={message.tool_calls?.length} class="message"> <div class:user={message.role === "user"} class:assistant={message.role !== "user"} class:has-tools={message.tool_calls?.length} class="message">
<div class="avatar">{message.role === "user" ? "你" : "P"}</div> <div class="avatar">{message.role === "user" ? "你" : "P"}</div>
<div class="message-content"> <div class="message-content">
{#if message.reasoning_content}
<details class="reasoning-block historical">
<summary>思考过程</summary>
<div class="reasoning-content"><Markdown content={message.reasoning_content} /></div>
</details>
{/if}
{#if message.content}<div class="bubble"><Markdown content={message.content} /></div>{/if} {#if message.content}<div class="bubble"><Markdown content={message.content} /></div>{/if}
{#if message.completion_status && message.completion_status !== "completed"}
<small class="completion-status">{message.completion_status === "cancelled" ? "已停止" : "回复中断"}</small>
{/if}
{#if message.attachments?.length} {#if message.attachments?.length}
<div class="message-attachments"> <div class="message-attachments">
{#each message.attachments as attachment (`${message.id}:${attachment.index}`)} {#each message.attachments as attachment (`${message.id}:${attachment.index}`)}
@ -408,7 +437,8 @@
</div> </div>
{/if} {/if}
{/each} {/each}
{#if thinking}<div class="message assistant typing"><div class="avatar">P</div><div class="bubble"><span class="pulse"></span>正在思考…</div></div>{/if} {#if activeTurn}<TurnView turn={activeTurn} />
{:else if thinking}<div class="message assistant typing"><div class="avatar">P</div><div class="bubble"><span class="pulse"></span>正在思考…</div></div>{/if}
</div> </div>
<form class="composer" onsubmit={(event) => { event.preventDefault(); submit(); }} ondragover={(event) => event.preventDefault()} ondrop={dropFiles}> <form class="composer" onsubmit={(event) => { event.preventDefault(); submit(); }} ondragover={(event) => event.preventDefault()} ondrop={dropFiles}>
{#if commandSuggestions.length} {#if commandSuggestions.length}

15
webui/src/styles.css
View File

@ -150,6 +150,21 @@ main { min-width: 0; height: 100vh; display: flex; flex-direction: column; }
.attachment-card strong { font-size: 12px; }.attachment-card small { margin-top: 3px; color: var(--muted); font-size: 9px; } .attachment-card strong { font-size: 12px; }.attachment-card small { margin-top: 3px; color: var(--muted); font-size: 9px; }
.attachment-card a { width: 28px; height: 28px; display: grid; place-items: center; border: 1px solid var(--line); border-radius: 7px; color: var(--accent); text-decoration: none; } .attachment-card a { width: 28px; height: 28px; display: grid; place-items: center; border: 1px solid var(--line); border-radius: 7px; color: var(--accent); text-decoration: none; }
.typing .bubble { color: var(--muted); } .typing .bubble { color: var(--muted); }
.active-turn .message-content { width: min(82%, 760px); }
.streaming { border-color: var(--accent-border); }
.reasoning-block, .live-tool { width: min(100%, 680px); overflow: hidden; border: 1px solid var(--line); border-radius: 10px; background: var(--panel); }
.reasoning-block summary, .live-tool summary { padding: 9px 12px; color: var(--muted); font-size: 11px; cursor: pointer; user-select: none; }
.reasoning-block summary::marker, .live-tool summary::marker { color: var(--accent); }
.reasoning-content { padding: 10px 13px; border-top: 1px solid var(--line); color: var(--muted); font-size: 12px; }
.reasoning-block.historical { border-style: dashed; }
.live-tool summary { display: grid; grid-template-columns: 24px 1fr auto; align-items: center; gap: 8px; }
.live-tool summary strong { color: var(--text); font: 600 12px/1.4 ui-monospace, SFMono-Regular, Consolas, monospace; }
.live-tool summary small { color: var(--muted); }
.live-tool-details { display: grid; gap: 8px; padding: 10px 12px; border-top: 1px solid var(--line); }
.live-tool-details pre { max-height: 260px; margin: 0; padding: 10px; overflow: auto; border-radius: 8px; color: var(--text-soft); background: var(--code-bg); font: 11px/1.5 ui-monospace, SFMono-Regular, Consolas, monospace; white-space: pre-wrap; }
.turn-status, .completion-status { color: var(--muted); font-size: 10px; }
.turn-status.failed, .turn-error { color: var(--danger); }
.turn-error { padding: 8px 10px; border-radius: 8px; background: var(--danger-soft); font-size: 11px; }
.pulse { display: inline-block; width: 6px; height: 6px; margin-right: 8px; border-radius: 50%; background: var(--accent); animation: pulse 1.1s infinite; } .pulse { display: inline-block; width: 6px; height: 6px; margin-right: 8px; border-radius: 50%; background: var(--accent); animation: pulse 1.1s infinite; }
@keyframes pulse { 50% { opacity: .25; transform: scale(.8); } } @keyframes pulse { 50% { opacity: .25; transform: scale(.8); } }
.markdown-body { min-width: 0; } .markdown-body { min-width: 0; }