diff --git a/AGENTS.md b/AGENTS.md index 7bbb946..b84ae3d 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -47,6 +47,8 @@ Channel → MessageBus.inbound → Gateway processor → SessionManager → per- ↑ │ └── 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) 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()` | | `channels` | External integrations (Feishu, CLI chat) | `ChannelManager`, `Channel` trait | | `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` | -| `agent` | LLM call loop, tool execution, context compression | `AgentLoop` | -| `providers` | LLM API clients (OpenAI-compatible, Anthropic) | `LLMProvider` trait, factory `create_provider()` | +| `session` | Conversation lifecycle, dialog operations, per-session serialization, Turn state, persistence coordination | `SessionManager`, `Session`, `TurnController` | +| `agent` | LLM call loop, tool execution, context compression, semantic Turn events | `AgentLoop`, `TurnEvent` | +| `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 | | `skills` | Skills loading, management, and prompt building | `SkillsLoader`, `Skill` | | `storage` | SQLite persistence for sessions and messages | `Storage`, `SessionMeta`, `MessageMeta` | @@ -75,9 +78,11 @@ Scheduler → SessionManager scheduled execution → AgentLoop → Scheduler del ### 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 - **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 - **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 @@ -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 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 +- **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 ### Concurrency and Lifecycle Invariants - 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 +- 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 - 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 diff --git a/README.md b/README.md index 36953c4..565a8d5 100644 --- a/README.md +++ b/README.md @@ -9,9 +9,10 @@ PicoBot 是一个用 Rust 编写的个人 AI 助手运行时。它在本地启 ## 适合做什么 - 在终端里和本地 AI 助手持续对话。 -- 在浏览器中聊天，并查看日志、任务和记忆，修改运行配置与助手档案。 +- 在 TUI 或浏览器中实时查看正文、思考过程和工具执行状态，并在完成后收敛到持久化历史。 +- 在浏览器中查看日志、任务和记忆，修改运行配置与助手档案。 - 复杂任务可创建 session 级 Todo 计划，把不同子项并行委托给多个子 Agent；聊天页侧栏实时显示进度。 -- 将同一套 Agent 能力接入飞书/Lark。 +- 将同一套 Agent 能力接入飞书/Lark，并可选用单张卡片实时更新回复。 - 让 Agent 使用本地文件、Shell、搜索、HTTP、浏览器、MCP 工具完成任务。 - 把长期偏好、事实和历史摘要存成可检索记忆。 - 用 Cron 定时执行任务，并把结果发回目标渠道。 @@ -143,7 +144,7 @@ picobot pair WebUI 随二进制嵌入，不需要 Node.js、npm 或单独部署静态文件，提供： -- 在线聊天、会话创建/切换、历史回放、Markdown 消息、可折叠工具调用卡片，以及基于 Gateway 实时命令清单的 `/` 斜杠命令补全。 +- 在线聊天、会话创建/切换、历史回放、流式 Markdown、独立思考区、实时工具状态、可折叠历史工具调用卡片，以及基于 Gateway 实时命令清单的 `/` 斜杠命令补全。 - 文件选择、拖放和剪贴板图片上传；消息中的附件可预览或下载。附件按服务端路径引用，原文件移动或删除后历史附件可能不可用。 - 可持久化的浅色/深色主题，首次访问时跟随系统偏好。 - Cron 定时任务、最近运行记录和后台子任务状态。 @@ -189,7 +190,7 @@ picobot service uninstall 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-消息与控制数据流)。 -同一 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 三类消息 | | `session` | 管理会话生命周期、dialog 操作、上下文、记忆召回、压缩和持久化 | | `agent` | 执行无状态 LLM/tool 循环，处理模型响应和工具调用 | -| `providers` | OpenAI 兼容接口和 Anthropic Messages API 客户端 | +| `providers` | OpenAI 兼容接口和 Anthropic Messages API 的原生流解析与回放 | +| `delivery` | 活动 Turn 的展示过滤、latest-wins 节流、终态投递与 TurnSink 生命周期 | | `tools` | Agent 可调用工具集合 | | `storage` | SQLite schema、CRUD、消息和任务持久化 | | `scheduler` | 领取定时任务，执行普通/巡检 Agent，并按投递策略记录或发送结果 | @@ -241,6 +243,8 @@ WebUI/TUI 上传文件默认保存到 `~/.picobot/media/cli_chat`，单文件上 | `cli_chat` | Ratatui 终端客户端，通过 WebSocket 连接 Gateway | | `feishu` | 飞书/Lark 消息、反应、文件上传下载和媒体引用 | +飞书默认只发送终态结果。设置 `channels.feishu.live_updates=true` 后会创建一张卡片并持续编辑；`live_update_interval_ms` 默认 500ms，运行时限制在 250–5000ms。外部渠道始终不会收到模型 reasoning。 + ### 会话 Session ID 使用三段式： @@ -342,6 +346,8 @@ Skill 是包含 `SKILL.md` 的目录。加载优先级从高到低： | `memory.timeline_retention_days` | `90` | | `mcp.tool_timeout_secs` | `180` | | `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)。 @@ -373,7 +379,7 @@ Inbound 消息类型： | `get_slash_commands` | 无 | | `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 集成 client/ Ratatui 终端 UI config/ 配置加载、环境变量替换、路径展开 + delivery/ 活动 Turn 快照投影、节流与 TurnSink 生命周期 gateway/ Axum HTTP/WebSocket server 和 GatewayState 装配 mcp/ MCP 客户端连接和工具包装 memory/ 记忆管理和记忆类型 diff --git a/docs/ARCHITECTURE.md b/docs/ARCHITECTURE.md index 9f7cc43..c0cc9a7 100644 --- a/docs/ARCHITECTURE.md +++ b/docs/ARCHITECTURE.md @@ -2,6 +2,8 @@ 本文档描述 PicoBot 当前实现的运行时边界、数据流、并发模型和演进约束。它面向维护者和后续参与改进的 Agent，是代码架构的主入口；行为细节仍以代码和测试为最终依据。 +流式模型输出、reasoning 展示、活动 Turn 快照和 Channel 实时投递的详细设计与取舍见 [STREAMING_TURN_DESIGN.md](STREAMING_TURN_DESIGN.md)。 + ## 1. 设计目标 PicoBot 是一个单进程、异步、可扩展的个人 AI 助手运行时。核心目标是： @@ -43,6 +45,8 @@ flowchart LR Agent --> Providers[LLM providers] Agent --> Tools[ToolRegistry / MCP] Sessions <--> Storage[(SQLite)] + Sessions -->|TurnSnapshot| Delivery[DeliveryCoordinator] + Delivery -->|TurnSink| Channels Sessions -->|OutboundMessage| Bus Scheduler[Scheduler] --> Sessions Bus --> Dispatcher[OutboundDispatcher] @@ -61,8 +65,9 @@ flowchart LR | `channels` | 外部协议适配、权限检查、媒体收发 | 会话选择、LLM 调用 | | `bus` | 三条有界异步队列与出站投递协调 | 会话路由、业务状态 | | `session` | dialog 路由、会话状态、串行工作队列、上下文和持久化协调 | 外部渠道协议 | -| `agent` | 单次无状态模型/工具循环、上下文压缩、子 Agent | 持有 dialog 生命周期 | -| `providers` | 把统一请求映射到模型 API | Session、Bus 或 Channel 感知 | +| `agent` | 单次无状态模型/工具循环、上下文压缩、子 Agent、Turn 语义事件 | 持有 dialog 生命周期 | +| `providers` | 把统一请求映射为原生模型流，并归一化正文、reasoning、工具和 usage | Session、Bus 或 Channel 感知 | +| `delivery` | 活动 Turn 快照投影、latest-wins 节流、终态重试和 TurnSink 生命周期 | Provider 协议、会话历史、平台 API 细节 | | `tools` / `mcp` | 工具定义、注册和执行适配 | 隐式修改会话路由 | | `storage` | SQLite schema、迁移、原子 CRUD | 运行时调度策略 | | `memory` | Knowledge/Timeline 的存取与召回 | 直接驱动消息发送 | @@ -87,7 +92,9 @@ sequenceDiagram participant G as Message processor participant S as SessionManager 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 C->>B: publish InboundMessage @@ -95,9 +102,18 @@ sequenceDiagram G->>S: handle_message S->>W: try_send AgentTask S-->>G: AgentProcessing - W->>A: process(history) - A-->>W: final response - W->>B: publish OutboundMessage + W->>T: start Turn + W->>L: subscribe latest snapshots + 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 D->>C: Channel::send ``` @@ -108,6 +124,20 @@ sequenceDiagram - 每个 session 有一条容量为 32 的队列，同一 session 串行处理，不同 session 的 worker 可并发执行。 - 队列满时明确拒绝新消息，不允许无界积压。 - 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。 - 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/迁移逻辑。 2. 保留已有数据库的升级路径。 @@ -181,7 +211,7 @@ SessionManager 负责组装会话上下文：系统提示、Skills、召回的 K ## 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 自己限时回收；它不跨越连接生命周期。 +Turn delivery 使用 `spawn_graceful`。全局取消发生时先停止读取快照，再在共享目标写锁下有界调用 `TurnSink::abort`，使平台 reaction 等私有资源能在 Supervisor 宽限期内清理。 + ### 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。 @@ -249,6 +281,7 @@ WebUI/TUI 文件字节通过受鉴权的 HTTP 接口流式传输，WebSocket 只 3. 将可重试错误表示为 `ConnectionError`/`SendError`，永久错误使用其他类型。 4. 为 start/stop 幂等性、取消建连、投递失败和媒体边界增加测试。 5. 不要从 Channel 直接调用 SessionManager 或 Provider。 +6. 若支持活动 Turn，实现 `live_policy`、`presentation_policy` 和每 Turn 一个实例的 `open_turn`；sink 必须消费完整快照而不是拼接 token，并使 finish/abort 清理幂等。 ### 新增 Tool diff --git a/docs/STREAMING_TURN_DESIGN.md b/docs/STREAMING_TURN_DESIGN.md new file mode 100644 index 0000000..81a700f --- /dev/null +++ b/docs/STREAMING_TURN_DESIGN.md @@ -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、正文和工具调用，不能依赖最终字符串中的 `` 后处理。 + +可吸收： + +- Provider 流的类型化增量； +- reasoning 与正文分离； +- 工具参数增量组装； +- thinking-only 响应的明确处理； +- reasoning effort/thinking budget 的统一配置概念。 + +### 2.2 ZeroClaw + +ZeroClaw 把 reasoning 作为不透明 Provider 数据保留，用于要求历史回放的模型，同时处理 `reasoning_content`、`reasoning`、内联 `` 和不同 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`。生产者覆盖旧值，消费者读取最新值。终态显式编码在快照中，不能仅依赖 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, + name: Option, + }, + 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, + pub provider_state: Option, + pub tool_calls: Vec, +} +``` + +约束： + +- `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, + pub usage: Option, + pub error: Option, +} + +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, + }, +} + +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, + }, +} +``` + +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; +} +``` + +标题生成、压缩等需要完整响应的代码通过 `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 切分。 + +``、`` 等内联标签使用有状态 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>); + pub fn begin_finalizing(&mut self); + pub fn complete(&mut self, usage: Option); + pub fn cancel(&mut self, reason: Option); + 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`； +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, 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 永远意味着数据库已经提交成功。 diff --git a/resources/skills/about-picobot/references/architecture.md b/resources/skills/about-picobot/references/architecture.md index 646fa36..001cb1b 100644 --- a/resources/skills/about-picobot/references/architecture.md +++ b/resources/skills/about-picobot/references/architecture.md @@ -7,6 +7,8 @@ Channel → MessageBus.inbound → Gateway processor → SessionManager → per- ↑ │ └── Channel ← OutboundDispatcher ← per-conversation lane ← MessageBus.outbound +AgentLoop → TurnEvent → TurnController → latest TurnSnapshot → DeliveryCoordinator → TurnSink → Channel + WebSocket/Channel → MessageBus.control → Gateway processor → SessionManager (dialog 操作) Scheduler → SessionManager.handle_cron_message → AgentLoop → send_message ``` @@ -19,9 +21,10 @@ Scheduler → SessionManager.handle_cron_message → AgentLoop → send_message | `client` | TUI 聊天客户端 | | `channels` | 外部集成（飞书、CLI），仅收发消息 | | `bus` | 有界 inbound/outbound/control 队列；出站 dispatcher 与分目标 lane | -| `session` | 会话生命周期、dialog 操作、每 session 串行队列、上下文与持久化协调 | -| `agent` | LLM 调用循环、工具执行、上下文压缩、媒体处理、子 Agent | -| `providers` | LLM API 客户端（OpenAI 兼容、Anthropic） | +| `session` | 会话生命周期、dialog 操作、每 session 串行队列、Turn 状态、上下文与持久化协调 | +| `agent` | LLM 调用循环、工具执行、上下文压缩、媒体处理、子 Agent、Turn 语义事件 | +| `providers` | OpenAI/Anthropic 原生流解析，统一正文、reasoning、工具、usage 与私有回放状态 | +| `delivery` | 完整 Turn 快照的展示过滤、latest-wins 节流、终态投递和 TurnSink 生命周期 | | `tools` | Agent 工具（bash、文件操作、搜索、HTTP、web、browser、memory、delegate 等） | | `skills` | Skill 加载、管理和 prompt 构建 | | `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` 负责 - SessionManager 拥有 session 状态、dialog 路由、上下文构建和每 session worker，并通过 worker 创建 AgentLoop +- TurnController 是活动 Turn 状态的唯一 owner；Session 在消息原子提交成功后才发布 Completed - 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 接收原始参数，返回字符串结果 - MCP 工具在 Gateway 初始化时连接服务器、发现工具，并包装成普通 Tool 注册到 ToolRegistry - 子 Agent 由 `delegate` 工具创建，复用 provider 配置和按需过滤后的工具集；后台任务结果通过 MessageBus 发回原会话 @@ -58,6 +64,8 @@ Scheduler → SessionManager.handle_cron_message → AgentLoop → send_message - `browser` 工具只有在 `browser.enabled=true` 时注册，依赖 Chrome/Chromium 与 WebDriver - 同一 session 的普通消息串行处理，不同 session 可并发；session 队列容量为 32，满时明确拒绝 - 出站消息按 `(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 - 外部建连、重试等待和关停 join 必须可取消且有硬超时 - 不得记录 API Key、Authorization header 或包含临时凭据的完整连接 URL @@ -138,6 +146,12 @@ Worker 的处理原则： 4. Session 持久化由独立 `persistence_lock` 串行化；批量消息使用原子写入，失败时精确回滚内存后缀。 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 时： diff --git a/resources/skills/about-picobot/references/config.md b/resources/skills/about-picobot/references/config.md index b4cbee4..8da7608 100644 --- a/resources/skills/about-picobot/references/config.md +++ b/resources/skills/about-picobot/references/config.md @@ -97,6 +97,10 @@ Gateway WebUI 的“配置”页可以编辑实际加载的配置文件。读取 | `agent` | string | - | 使用的 agent 名称 | | `media_dir` | string | ~/.picobot/media/feishu | 配置默认值；Gateway 注册渠道时会覆盖为 `{workspace}/media/feishu` | | `reaction_emoji` | string | "Typing" | 回复意向表达的表情 | +| `live_updates` | bool | false | 是否用单张卡片实时编辑活动 Turn；关闭时只发送终态 | +| `live_update_interval_ms` | int | 500 | 卡片更新最小间隔，运行时限制在 250–5000ms | + +飞书属于外部渠道：无论是否开启实时卡片，都不会接收模型 reasoning；工具只显示紧凑状态。配置修改需重启 Gateway 生效。 ## mcp 字段 diff --git a/resources/skills/about-picobot/references/db-schema.md b/resources/skills/about-picobot/references/db-schema.md index 57d6b1f..7887098 100644 --- a/resources/skills/about-picobot/references/db-schema.md +++ b/resources/skills/about-picobot/references/db-schema.md @@ -2,7 +2,7 @@ 数据库为 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 表 @@ -41,7 +41,11 @@ | `tool_calls` | TEXT | 工具调用参数 JSON | | `source` | TEXT | 消息来源（跨会话消息时标记来源 session_id） | | `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。 @@ -137,9 +141,9 @@ Scheduler 使用原子 `UPDATE ... RETURNING` 领取到期任务。任务结果 | `created_at` | INTEGER | 调用时间 | | `provider` | TEXT | 提供商类型 | | `model` | TEXT | 模型名称 | -| `request_body` | TEXT | 请求体 JSON | +| `request_body` | TEXT | 请求摘要 JSON；旧记录可能是完整请求体 | | `response_body` | TEXT | 响应体 JSON | | `error` | TEXT | 错误信息 | | `duration_ms` | INTEGER | 耗时（毫秒） | -`request_body`/`response_body` 可能包含用户内容，排障和导出数据库时应按敏感数据处理。 +旧数据中的 `request_body`/`response_body` 可能包含用户内容，排障和导出数据库时应按敏感数据处理。新 Provider 请求的 `request_body` 只保存模型、消息数、工具数和 stream 标志等摘要；错误响应仍可能包含服务端回显内容。 diff --git a/resources/templates/config.example.json b/resources/templates/config.example.json index fb18857..bedc9ad 100644 --- a/resources/templates/config.example.json +++ b/resources/templates/config.example.json @@ -70,7 +70,9 @@ "allow_from": ["*"], "agent": "default", "media_dir": "~/.picobot/media/feishu", - "reaction_emoji": "Typing" + "reaction_emoji": "Typing", + "live_updates": false, + "live_update_interval_ms": 500 } }, "memory": { diff --git a/src/agent/agent_loop.rs b/src/agent/agent_loop.rs index 57a020c..4519115 100644 --- a/src/agent/agent_loop.rs +++ b/src/agent/agent_loop.rs @@ -1,11 +1,15 @@ use crate::agent::context_compressor::estimate_tokens; use crate::agent::media_handler::MediaHandlerRegistry; use crate::agent::system_prompt::build_system_prompt; +use crate::agent::turn_event::{AgentTurnContext, TurnEvent}; use crate::bus::message::ContentBlock; use crate::bus::{ChatMessage, MediaRef}; use crate::config::LLMProviderConfig; 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 std::collections::VecDeque; use std::hash::{Hash, Hasher}; @@ -13,11 +17,14 @@ use std::path::PathBuf; use std::sync::Arc; use std::time::Instant; +use futures_util::StreamExt; + /// Maximum characters in a tool result before truncation. /// Prevents context overflow from large tool outputs. const MAX_TOOL_RESULT_CHARS: usize = 16_000; /// Minimum characters to keep when truncating const TRUNCATION_SUFFIX_LEN: usize = 200; +const TOOL_PREVIEW_CHARS: usize = 1_000; enum MediaOrigin<'a> { 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. #[derive(Debug, Clone, PartialEq, Eq)] enum LoopDetectionResult { @@ -296,6 +314,32 @@ pub struct AgentProcessResult { pub final_response: ChatMessage, pub emitted_messages: Vec, pub total_tokens: Option, + pub usage: Option, +} + +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, right: Option) -> Option { + match (left, right) { + (None, None) => None, + (left, right) => Some( + left.unwrap_or_default() + .saturating_add(right.unwrap_or_default()), + ), + } } impl AgentLoop { @@ -451,6 +495,60 @@ impl AgentLoop { &self.tools } + async fn stream_completion( + &self, + request: ChatCompletionRequest, + iteration: u32, + turn: Option<&AgentTurnContext>, + ) -> Result { + 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 { let content = if m.media_refs.is_empty() || !include_media { vec![ContentBlock::text(&m.content)] @@ -473,6 +571,7 @@ impl AgentLoop { role: m.role.clone(), content, reasoning_content: m.reasoning_content.clone(), + provider_state: m.provider_state.clone(), tool_call_id: m.tool_call_id.clone(), name: m.tool_name.clone(), tool_calls: m.tool_calls.clone(), @@ -498,8 +597,24 @@ impl AgentLoop { /// - The LLM returns no more tool calls (final response) /// - Maximum iterations are reached pub async fn process( + &self, + messages: Vec, + ) -> Result { + self.process_inner(messages, None).await + } + + pub async fn process_streaming( + &self, + messages: Vec, + turn: AgentTurnContext, + ) -> Result { + self.process_inner(messages, Some(turn)).await + } + + async fn process_inner( &self, mut messages: Vec, + turn: Option, ) -> Result { #[cfg(debug_assertions)] tracing::debug!( @@ -522,6 +637,7 @@ impl AgentLoop { let mut loop_detector = LoopDetector::new(LoopDetectorConfig::default()); let mut emitted_messages = Vec::new(); let mut accumulated_tokens: u32 = 0; + let mut accumulated_usage = crate::providers::Usage::default(); for iteration in 0..self.max_iterations { #[cfg(debug_assertions)] @@ -564,12 +680,14 @@ impl AgentLoop { }; // Call LLM - let response = (*self.provider).chat(request).await.map_err(|e| { - tracing::error!(error = %e, "LLM request failed"); - AgentError::LlmError(e.to_string()) - })?; + let iteration = u32::try_from(iteration) + .map_err(|_| AgentError::Other("tool iteration exceeds u32".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)] tracing::debug!( @@ -583,14 +701,23 @@ impl AgentLoop { if response.tool_calls.is_empty() { let mut assistant_message = ChatMessage::assistant(response.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()); return Ok(AgentProcessResult { final_response: assistant_message, emitted_messages, 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 { let tools_info: Vec = response @@ -614,11 +741,15 @@ impl AgentLoop { response.tool_calls.clone(), ); 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()); emitted_messages.push(assistant_message); // 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()) { // Log function call with name and arguments @@ -644,22 +775,24 @@ impl AgentLoop { "Loop warning: {}", msg ); - let tool_message = ChatMessage::tool_with_media( + let mut tool_message = ChatMessage::tool_with_media( tool_call.id.clone(), tool_call.name.clone(), format!("{}\n\n[上一条结果]\n{}", msg, truncated_output), result.media_refs.clone(), ); + Self::annotate_message(&mut tool_message, turn.as_ref(), iteration, false); messages.push(tool_message.clone()); emitted_messages.push(tool_message); } LoopDetectionResult::Ok => { - let tool_message = ChatMessage::tool_with_media( + let mut tool_message = ChatMessage::tool_with_media( tool_call.id.clone(), tool_call.name.clone(), truncated_output, result.media_refs.clone(), ); + Self::annotate_message(&mut tool_message, turn.as_ref(), iteration, false); messages.push(tool_message.clone()); emitted_messages.push(tool_message); } @@ -695,25 +828,56 @@ impl AgentLoop { 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) => { - 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); 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()); Ok(AgentProcessResult { final_response: assistant_message, emitted_messages, total_tokens: Some(accumulated_tokens), + usage: Some(accumulated_usage), }) } Err(e) => { // Fallback if summary call fails 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.", 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()); Ok(AgentProcessResult { final_response: final_message, @@ -723,6 +887,7 @@ impl AgentLoop { } else { 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. - async fn execute_tools(&self, tool_calls: &[ToolCall]) -> Vec { + async fn execute_tools( + &self, + tool_calls: &[ToolCall], + iteration: u32, + turn: Option<&AgentTurnContext>, + ) -> Result, AgentError> { if self.should_execute_in_parallel(tool_calls) { 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 { 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. - async fn execute_tools_parallel(&self, tool_calls: &[ToolCall]) -> Vec { + async fn execute_tools_parallel( + &self, + tool_calls: &[ToolCall], + iteration: u32, + turn: Option<&AgentTurnContext>, + ) -> Result, AgentError> { let futures: Vec<_> = tool_calls .iter() - .map(|tc| self.execute_one_tool(tc)) + .map(|tool_call| self.execute_one_tool(tool_call, iteration, turn)) .collect(); - futures_util::future::join_all(futures).await + futures_util::future::join_all(futures) + .await + .into_iter() + .collect() } /// Execute tools sequentially. - async fn execute_tools_sequential(&self, tool_calls: &[ToolCall]) -> Vec { + async fn execute_tools_sequential( + &self, + tool_calls: &[ToolCall], + iteration: u32, + turn: Option<&AgentTurnContext>, + ) -> Result, AgentError> { let mut outcomes = Vec::with_capacity(tool_calls.len()); 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. - 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 { let start = Instant::now(); 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 if let Some(ref observer) = self.observer { observer.record_event(&ObserverEvent::ToolCallStart { @@ -800,6 +999,17 @@ impl AgentLoop { let result = self.execute_tool_internal(tool_call).await; 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 if let Some(ref observer) = self.observer { observer.record_event(&ObserverEvent::ToolCall { @@ -810,7 +1020,7 @@ impl AgentLoop { } // Apply duration - ToolExecutionOutcome { duration, ..result } + Ok(ToolExecutionOutcome { duration, ..result }) } /// Internal tool execution without event tracking. @@ -851,13 +1061,115 @@ impl AgentLoop { mod tests { use super::*; 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; struct TestObserver { events: std::sync::Mutex>, } + struct StreamingTextProvider; + + #[async_trait::async_trait] + impl LLMProvider for StreamingTextProvider { + async fn stream( + &self, + _request: ChatCompletionRequest, + ) -> Result { + 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 { fn new() -> Self { Self { @@ -899,16 +1211,16 @@ mod tests { #[async_trait::async_trait] impl LLMProvider for ToolMediaProvider { - async fn chat( + async fn stream( &self, request: ChatCompletionRequest, - ) -> Result> { + ) -> Result { let call_number = { let mut requests = self.requests.lock().unwrap(); requests.push(request); requests.len() }; - Ok(ChatCompletionResponse { + let response = ChatCompletionResponse { id: format!("response-{call_number}"), model: "vision-test".to_string(), content: if call_number == 1 { @@ -917,6 +1229,7 @@ mod tests { "image seen".to_string() }, reasoning_content: None, + provider_state: None, tool_calls: if call_number == 1 { vec![ToolCall { id: "call-image".to_string(), @@ -934,7 +1247,8 @@ mod tests { cache_read_input_tokens: None, cache_creation_input_tokens: None, }, - }) + }; + Ok(crate::providers::provider_stream_for_test(response)) } fn ptype(&self) -> &str { @@ -971,8 +1285,13 @@ mod tests { vec!["text".to_string(), "image".to_string()], ); + let (controller, emitter, _) = TurnController::start("session", "assistant-id"); + let turn = controller.snapshot(); 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 .unwrap(); @@ -997,6 +1316,14 @@ mod tests { .iter() .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] @@ -1031,6 +1358,7 @@ mod tests { role: chat_message.role.clone(), content, reasoning_content: None, + provider_state: None, tool_call_id: chat_message.tool_call_id.clone(), name: chat_message.tool_name.clone(), tool_calls: chat_message.tool_calls.clone(), diff --git a/src/agent/context_compressor.rs b/src/agent/context_compressor.rs index eb0eddf..14ba089 100644 --- a/src/agent/context_compressor.rs +++ b/src/agent/context_compressor.rs @@ -669,11 +669,11 @@ mod tests { #[async_trait] impl LLMProvider for MockProvider { - async fn chat( + async fn stream( &self, _request: ChatCompletionRequest, - ) -> Result> { - panic!("MockProvider.chat() called - not expected in test") + ) -> Result { + panic!("MockProvider.stream() called - not expected in test") } fn ptype(&self) -> &str { @@ -699,25 +699,28 @@ mod tests { #[async_trait] impl LLMProvider for MockSummarizer { - async fn chat( + async fn stream( &self, _request: ChatCompletionRequest, - ) -> Result> { - Ok(ChatCompletionResponse { - id: "mock".into(), - model: "mock".into(), - content: "[summarized]".into(), - reasoning_content: None, - tool_calls: vec![], - usage: Usage { - prompt_tokens: 0, - completion_tokens: 0, - total_tokens: 0, - cached_tokens: None, - cache_read_input_tokens: None, - cache_creation_input_tokens: None, + ) -> Result { + Ok(crate::providers::provider_stream_for_test( + ChatCompletionResponse { + id: "mock".into(), + model: "mock".into(), + content: "[summarized]".into(), + reasoning_content: None, + provider_state: None, + tool_calls: vec![], + usage: Usage { + prompt_tokens: 0, + completion_tokens: 0, + total_tokens: 0, + cached_tokens: None, + cache_read_input_tokens: None, + cache_creation_input_tokens: None, + }, }, - }) + )) } fn ptype(&self) -> &str { diff --git a/src/agent/mod.rs b/src/agent/mod.rs index 64591be..dc75e49 100644 --- a/src/agent/mod.rs +++ b/src/agent/mod.rs @@ -3,6 +3,7 @@ pub mod context_compressor; pub mod media_handler; pub mod sub_agent; pub mod system_prompt; +pub mod turn_event; pub use agent_loop::{AgentError, AgentLoop, AgentProcessResult}; pub use context_compressor::{ContextCompressor, estimate_tokens}; @@ -14,3 +15,4 @@ pub use system_prompt::{ PromptContext, PromptSection, SystemPromptBuilder, build_sub_agent_system_prompt, build_system_prompt, }; +pub use turn_event::{AgentTurnContext, TurnEmitError, TurnEmitter, TurnEvent}; diff --git a/src/agent/turn_event.rs b/src/agent/turn_event.rs new file mode 100644 index 0000000..da10f92 --- /dev/null +++ b/src/agent/turn_event.rs @@ -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, + }, +} + +#[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, + enabled: Arc>, +} + +/// 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, + message_id: impl Into, + emitter: TurnEmitter, + ) -> Self { + Self { + turn_id: turn_id.into(), + message_id: message_id.into(), + emitter, + } + } +} + +impl TurnEmitter { + pub(crate) fn new(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; + } +} diff --git a/src/bus/dispatcher.rs b/src/bus/dispatcher.rs index 35927c9..edb830e 100644 --- a/src/bus/dispatcher.rs +++ b/src/bus/dispatcher.rs @@ -7,6 +7,7 @@ use tokio::sync::mpsc; use crate::bus::{MessageBus, OutboundMessage}; use crate::channels::ChannelManager; use crate::channels::base::{Channel, ChannelError}; +use crate::delivery::ConversationWriteLocks; use crate::task_supervisor::TaskSupervisor; const LANE_CAPACITY: usize = 64; @@ -20,6 +21,7 @@ pub struct OutboundDispatcher { bus: Arc, channel_manager: ChannelManager, task_supervisor: TaskSupervisor, + write_locks: ConversationWriteLocks, } impl OutboundDispatcher { @@ -27,11 +29,13 @@ impl OutboundDispatcher { bus: Arc, channel_manager: ChannelManager, task_supervisor: TaskSupervisor, + write_locks: ConversationWriteLocks, ) -> Self { Self { bus, channel_manager, task_supervisor, + write_locks, } } @@ -120,6 +124,7 @@ impl OutboundDispatcher { channel_name: String, chat_id: String, ) -> bool { + let target_lock = self.write_locks.for_target(&channel_name, &chat_id); self.task_supervisor.spawn( format!("outbound-lane:{channel_name}:{chat_id}"), async move { @@ -128,7 +133,7 @@ impl OutboundDispatcher { Ok(Some(msg)) => msg, 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 { tracing::error!( channel = %channel_name, @@ -146,7 +151,9 @@ impl OutboundDispatcher { async fn send_with_retry( channel: &dyn Channel, msg: &OutboundMessage, + target_lock: &tokio::sync::Mutex<()>, ) -> Result<(), ChannelError> { + let _guard = target_lock.lock().await; const DELAYS: &[u64] = &[1, 2, 4]; for (attempt, &delay) in DELAYS.iter().enumerate() { @@ -263,7 +270,12 @@ mod tests { manager.register_channel("recording", channel.clone()).await; 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 }); bus.publish_outbound(outbound("slow", "slow-1")) .await @@ -301,7 +313,12 @@ mod tests { bus.clone(), ); 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 mut message = outbound("missing", "not delivered"); @@ -326,7 +343,12 @@ mod tests { }); manager.register_channel("recording", channel.clone()).await; 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 }); bus.deliver_outbound(outbound("confirmed", "delivered")) @@ -344,11 +366,41 @@ mod tests { attempts: AtomicUsize::new(0), }; - let error = OutboundDispatcher::send_with_retry(&channel, &outbound("invalid", "message")) - .await - .unwrap_err(); + let target_lock = tokio::sync::Mutex::new(()); + let error = OutboundDispatcher::send_with_retry( + &channel, + &outbound("invalid", "message"), + &target_lock, + ) + .await + .unwrap_err(); assert!(matches!(error, ChannelError::Other(_))); 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"]); + } } diff --git a/src/bus/message.rs b/src/bus/message.rs index 56ae5ef..21c5659 100644 --- a/src/bus/message.rs +++ b/src/bus/message.rs @@ -3,6 +3,52 @@ use std::collections::HashMap; 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 { + 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) // ============================================================================ @@ -85,6 +131,15 @@ pub struct ChatMessage { pub role: String, pub content: String, pub reasoning_content: Option, + /// Opaque state used only when replaying history to the same provider. + #[serde(skip_serializing_if = "Option::is_none")] + pub provider_state: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub turn_id: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub iteration: Option, + #[serde(default)] + pub completion_status: CompletionStatus, pub media_refs: Vec, pub timestamp: i64, #[serde(skip_serializing_if = "Option::is_none")] @@ -124,6 +179,10 @@ impl ChatMessage { role: "user".to_string(), content: content.into(), reasoning_content: None, + provider_state: None, + turn_id: None, + iteration: None, + completion_status: CompletionStatus::Completed, media_refs: Vec::new(), timestamp: current_timestamp(), tool_call_id: None, @@ -139,6 +198,10 @@ impl ChatMessage { role: "user".to_string(), content: content.into(), reasoning_content: None, + provider_state: None, + turn_id: None, + iteration: None, + completion_status: CompletionStatus::Completed, media_refs, timestamp: current_timestamp(), tool_call_id: None, @@ -154,6 +217,10 @@ impl ChatMessage { role: "assistant".to_string(), content: content.into(), reasoning_content: None, + provider_state: None, + turn_id: None, + iteration: None, + completion_status: CompletionStatus::Completed, media_refs: Vec::new(), timestamp: current_timestamp(), tool_call_id: None, @@ -172,6 +239,10 @@ impl ChatMessage { role: "assistant".to_string(), content: content.into(), reasoning_content: None, + provider_state: None, + turn_id: None, + iteration: None, + completion_status: CompletionStatus::Completed, media_refs: Vec::new(), timestamp: current_timestamp(), tool_call_id: None, @@ -187,6 +258,10 @@ impl ChatMessage { role: "assistant".to_string(), content: content.into(), reasoning_content: None, + provider_state: None, + turn_id: None, + iteration: None, + completion_status: CompletionStatus::Completed, media_refs: Vec::new(), timestamp: current_timestamp(), tool_call_id: None, @@ -202,6 +277,10 @@ impl ChatMessage { role: "system".to_string(), content: content.into(), reasoning_content: None, + provider_state: None, + turn_id: None, + iteration: None, + completion_status: CompletionStatus::Completed, media_refs: Vec::new(), timestamp: current_timestamp(), tool_call_id: None, @@ -230,6 +309,10 @@ impl ChatMessage { role: "tool".to_string(), content: content.into(), reasoning_content: None, + provider_state: None, + turn_id: None, + iteration: None, + completion_status: CompletionStatus::Completed, media_refs, timestamp: current_timestamp(), tool_call_id: Some(tool_call_id.into()), @@ -245,6 +328,10 @@ impl ChatMessage { role: "user".to_string(), content: content.into(), reasoning_content: None, + provider_state: None, + turn_id: None, + iteration: None, + completion_status: CompletionStatus::Completed, media_refs: Vec::new(), timestamp: current_timestamp(), 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) // ============================================================================ diff --git a/src/bus/mod.rs b/src/bus/mod.rs index bde9eb4..f54fd1d 100644 --- a/src/bus/mod.rs +++ b/src/bus/mod.rs @@ -3,8 +3,8 @@ pub mod message; pub use dispatcher::OutboundDispatcher; pub use message::{ - ChatMessage, ContentBlock, ControlMessage, InboundMessage, MediaItem, MediaRef, MessageSource, - OutboundMessage, SourceKind, + ChatMessage, CompletionStatus, ContentBlock, ControlMessage, InboundMessage, MediaItem, + MediaRef, MessageSource, OutboundMessage, ProviderReasoningState, SourceKind, }; use std::sync::Arc; diff --git a/src/channels/base.rs b/src/channels/base.rs index 73e164b..aa460c6 100644 --- a/src/channels/base.rs +++ b/src/channels/base.rs @@ -1,7 +1,33 @@ use async_trait::async_trait; +use std::collections::HashMap; use std::sync::Arc; +use std::time::Duration; 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, + pub metadata: HashMap, +} + +#[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)] pub enum ChannelError { @@ -49,16 +75,24 @@ pub trait Channel: Send + Sync + 'static { /// Stop the channel 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, ChannelError> { + Err(ChannelError::Other(format!( + "channel {} does not support turn delivery", + self.name() + ))) + } + /// Send a message to the channel (called by OutboundDispatcher) 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 fn is_allowed(&self, _sender_id: &str) -> bool { true diff --git a/src/channels/cli_chat.rs b/src/channels/cli_chat.rs index 031cf4e..337eb31 100644 --- a/src/channels/cli_chat.rs +++ b/src/channels/cli_chat.rs @@ -1,6 +1,7 @@ use async_trait::async_trait; use std::collections::HashMap; use std::sync::Arc; +use std::time::Duration; use tokio::sync::{Mutex, mpsc}; use crate::bus::{ControlMessage, InboundMessage, MessageBus, OutboundMessage}; @@ -8,9 +9,10 @@ use crate::gateway::uploads::UploadRegistry; use crate::protocol::{ HistoryMessage, MessageAttachment, SlashCommandInfo, WsInbound, WsOutbound, parse_inbound, }; +use crate::session::TurnSnapshot; use crate::session::{SessionCommand, SessionEvent, UnifiedSessionId}; -use super::base::{Channel, ChannelError}; +use super::base::{Channel, ChannelError, LivePolicy, TurnSink, TurnTarget}; // ============================================================================ // Client - Connected CLI client @@ -34,7 +36,7 @@ impl Client { pub struct CliChatChannel { bus: std::sync::Mutex>>, - clients: Mutex>>, + clients: Arc>>>, uploads: UploadRegistry, } @@ -52,7 +54,7 @@ impl CliChatChannel { pub fn with_upload_registry(uploads: UploadRegistry) -> Self { Self { bus: std::sync::Mutex::new(None), - clients: Mutex::new(HashMap::new()), + clients: Arc::new(Mutex::new(HashMap::new())), uploads, } } @@ -451,6 +453,8 @@ impl CliChatChannel { seq: message.seq, role: message.role, content: message.content, + reasoning_content: message.reasoning_content, + completion_status: message.completion_status, created_at: message.created_at, tool_call_id: message.tool_call_id, tool_name: message.tool_name, @@ -803,6 +807,54 @@ impl CliChatChannel { } } +struct CliChatTurnSink { + clients: Arc>>>, + 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] impl Channel for CliChatChannel { fn name(&self) -> &str { @@ -824,6 +876,23 @@ impl Channel for CliChatChannel { 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, ChannelError> { + Ok(Box::new(CliChatTurnSink { + clients: self.clients.clone(), + chat_id: target.chat_id, + })) + } + async fn send(&self, msg: OutboundMessage) -> Result<(), ChannelError> { let client = self.clients.lock().await.get(&msg.chat_id).cloned(); let Some(client) = client else { @@ -1026,4 +1095,53 @@ mod tests { 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:?}"), + } + } } diff --git a/src/channels/feishu.rs b/src/channels/feishu.rs index 9f6b354..30d1d02 100644 --- a/src/channels/feishu.rs +++ b/src/channels/feishu.rs @@ -6,15 +6,15 @@ use std::time::{Duration, Instant}; use async_trait::async_trait; use futures_util::{SinkExt, StreamExt}; use prost::{Message as ProstMessage, bytes::Bytes}; -use regex::Regex; use serde::Deserialize; use tokio::sync::{Mutex, RwLock}; use tokio::task::JoinHandle; use tokio_util::sync::CancellationToken; 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::session::{ToolStatus, TurnBlock, TurnSnapshot, TurnStatus}; const FEISHU_API_BASE: &str = "https://open.feishu.cn/open-apis"; 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 { - fn strip_thinking_tags(content: &str) -> String { - use std::sync::LazyLock; - static THINK_RE: LazyLock = - LazyLock::new(|| Regex::new(r"(?s).*?").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. fn build_card_content(markdown: &str) -> String { serde_json::json!({ @@ -1849,7 +1841,7 @@ impl FeishuChannel { 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 split_at = search_region .rfind('\n') @@ -1886,7 +1878,7 @@ impl FeishuChannel { receive_id: &str, receive_id_type: &str, card_content: &str, - ) -> Result<(), ChannelError> { + ) -> Result { let token = self.get_tenant_access_token().await?; let resp = self @@ -1912,6 +1904,12 @@ impl FeishuChannel { struct SendResp { code: i32, msg: String, + data: Option, + } + + #[derive(Deserialize)] + struct SendData { + message_id: String, } 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(()) } } +#[async_trait] +trait FeishuTurnApi: Send { + async fn create_card(&mut self, markdown: &str) -> Result; + 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, +} + +#[async_trait] +impl FeishuTurnApi for FeishuTurnBackend { + async fn create_card(&mut self, markdown: &str) -> Result { + 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, + message_id: Option, + cleaned_up: bool, +} + +impl FeishuTurnSink { + fn new(api: Box) -> 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] impl Channel for FeishuChannel { fn name(&self) -> &str { @@ -2043,11 +2275,37 @@ impl Channel for FeishuChannel { self.running.try_read().map(|r| *r).unwrap_or(false) } - async fn send(&self, msg: OutboundMessage) -> Result<(), ChannelError> { - let msg = OutboundMessage { - content: Self::strip_thinking_tags(&msg.content), - ..msg + fn live_policy(&self) -> LivePolicy { + if self.config.live_updates { + LivePolicy::Snapshot { + 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, 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_") { &msg.chat_id } else { @@ -2252,6 +2510,53 @@ impl Channel for FeishuChannel { #[cfg(test)] mod tests { use super::*; + use crate::agent::TurnEvent; + use crate::delivery::{PresentationPolicy, project_snapshot}; + + #[derive(Default)] + struct MockTurnState { + created: Vec, + updated: Vec<(String, String)>, + cleanups: usize, + fail_updates: usize, + } + + struct MockTurnApi { + state: Arc>, + } + + #[async_trait] + impl FeishuTurnApi for MockTurnApi { + async fn create_card(&mut self, markdown: &str) -> Result { + 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>) -> FeishuTurnSink { + FeishuTurnSink::new(Box::new(MockTurnApi { state })) + } fn test_channel() -> FeishuChannel { FeishuChannel::new( @@ -2263,12 +2568,169 @@ mod tests { agent: String::new(), media_dir: String::new(), reaction_emoji: "THUMBSUP".to_string(), + live_updates: false, + live_update_interval_ms: 500, }, Path::new("/tmp"), ) .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] async fn stop_aborts_connection_task_that_ignores_cancellation() { let channel = test_channel(); diff --git a/src/channels/mod.rs b/src/channels/mod.rs index 9331bad..2eadad9 100644 --- a/src/channels/mod.rs +++ b/src/channels/mod.rs @@ -4,7 +4,7 @@ pub mod feishu; pub mod manager; pub mod slash_command; -pub use base::{Channel, ChannelError}; +pub use base::{Channel, ChannelError, LivePolicy, TurnSink, TurnTarget}; pub use cli_chat::CliChatChannel; pub use feishu::FeishuChannel; pub use manager::ChannelManager; diff --git a/src/client/mod.rs b/src/client/mod.rs index c976c2f..50d3dcf 100644 --- a/src/client/mod.rs +++ b/src/client/mod.rs @@ -253,6 +253,28 @@ async fn run_app( async fn handle_ws_message(app: &mut App, outbound: WsOutbound) { 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 { id, content, diff --git a/src/client/tui/app.rs b/src/client/tui/app.rs index 613791e..642b4da 100644 --- a/src/client/tui/app.rs +++ b/src/client/tui/app.rs @@ -1,6 +1,7 @@ use crate::protocol::{ HistoryMessage, MessageAttachment, SessionSummary, SlashCommandInfo, UploadDescriptor, }; +use crate::session::{TurnSnapshot, TurnStatus}; use std::collections::VecDeque; use tokio_tungstenite::tungstenite::Message; @@ -19,6 +20,8 @@ pub struct ChatMessage { pub id: String, pub role: MessageRole, pub content: String, + pub reasoning_content: Option, + pub completion_status: crate::bus::CompletionStatus, pub attachments: Vec, } @@ -63,6 +66,7 @@ pub struct App { pub selected_session: usize, pub show_archived: bool, pub messages: VecDeque, + pub active_turn: Option, pub input: String, /// UTF-8 byte offset. It is always maintained at a character boundary. pub input_cursor_pos: usize, @@ -95,6 +99,7 @@ impl App { selected_session: 0, show_archived: false, messages: VecDeque::new(), + active_turn: None, input: String::new(), input_cursor_pos: 0, focus: Focus::Input, @@ -132,6 +137,8 @@ impl App { id, role, content, + reasoning_content: None, + completion_status: crate::bus::CompletionStatus::Completed, attachments, }); while self.messages.len() > MAX_MESSAGES { @@ -144,6 +151,10 @@ impl App { if self.current_session_id.as_deref() != Some(session_id) { 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 .into_iter() .filter_map(|message| { @@ -157,6 +168,8 @@ impl App { id: message.id, role, content: message.content, + reasoning_content: message.reasoning_content, + completion_status: message.completion_status, attachments: message.attachments, }) }) @@ -166,6 +179,9 @@ impl App { } self.chat_scroll_from_bottom = 0; self.status_message = None; + if calibrates_terminal { + self.active_turn = None; + } } pub fn set_sessions(&mut self, sessions: Vec) { @@ -185,6 +201,7 @@ impl App { if self.current_session_id != session_id { self.current_session_id = session_id; self.messages.clear(); + self.active_turn = None; self.pending_uploads.clear(); 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 { self.current_session_id .as_ref() @@ -383,6 +415,21 @@ fn next_boundary(value: &str, offset: usize) -> Option { #[cfg(test)] mod tests { 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] fn unicode_cursor_edits_only_at_character_boundaries() { @@ -404,4 +451,57 @@ mod tests { app.set_history("old", Vec::new()); 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) + ); + } } diff --git a/src/client/tui/components/chat_history.rs b/src/client/tui/components/chat_history.rs index 38dbadf..3e48f1a 100644 --- a/src/client/tui/components/chat_history.rs +++ b/src/client/tui/components/chat_history.rs @@ -1,4 +1,5 @@ use crate::client::tui::app::{App, MessageRole}; +use crate::session::{ToolStatus, TurnBlock, TurnPhase, TurnStatus}; use ratatui::{ Frame, layout::Rect, @@ -23,6 +24,15 @@ pub fn render(frame: &mut Frame, area: Rect, app: &App) { label, 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() { let wrapped = textwrap::wrap(source_line, content_width); 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 { lines.push(Line::from(Span::styled( format!(" [附件 {}] {}", attachment.index + 1, attachment.name), @@ -43,7 +59,74 @@ pub fn render(frame: &mut Frame, area: Rect, app: &App) { } 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( "● 正在思考…", Style::default().fg(Color::Cyan), @@ -66,3 +149,16 @@ pub fn render(frame: &mut Frame, area: Rect, app: &App) { area, ); } + +fn push_wrapped(lines: &mut Vec>, 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))) + })); + } + } +} diff --git a/src/client/tui/ui.rs b/src/client/tui/ui.rs index 64e337b..da6a2a7 100644 --- a/src/client/tui/ui.rs +++ b/src/client/tui/ui.rs @@ -189,4 +189,38 @@ mod tests { ); 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(); + } } diff --git a/src/config/mod.rs b/src/config/mod.rs index 8ba72b2..b19c188 100644 --- a/src/config/mod.rs +++ b/src/config/mod.rs @@ -78,6 +78,11 @@ pub struct FeishuChannelConfig { /// Emoji type for message reactions (e.g. "THUMBSUP", "OK", "EYES"). #[serde(default = "default_reaction_emoji")] 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 { @@ -95,6 +100,10 @@ fn default_reaction_emoji() -> String { "Typing".to_string() } +fn default_feishu_live_update_interval_ms() -> u64 { + 500 +} + #[derive(Debug, Clone, Deserialize, Serialize)] pub struct ProviderConfig { #[serde(rename = "type")] diff --git a/src/delivery/coordinator.rs b/src/delivery/coordinator.rs new file mode 100644 index 0000000..789f123 --- /dev/null +++ b/src/delivery/coordinator.rs @@ -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>>>>, +} + +impl ConversationWriteLocks { + pub fn for_target(&self, channel: &str, chat_id: &str) -> Arc> { + 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>, + ) -> 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, + target: TurnTarget, + presentation: PresentationPolicy, + snapshots: watch::Receiver>, + ) -> 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, + target: TurnTarget, + presentation: PresentationPolicy, + snapshots: watch::Receiver>, + ) -> Result>, 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>, + mut sink: Box, + ) -> Result>, 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>, + mut sink: Box, + ) -> 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>, + 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>, + 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>, + terminal: TokioMutex>, + update_started: Notify, + release_update: Notify, + block_first_update: bool, + fail_updates: AtomicUsize, + fail_finish: AtomicUsize, + } + + struct RecordingSink(Arc); + + struct SinkChannel { + state: Arc, + 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) -> 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, 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) -> Box { + 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); + } +} diff --git a/src/delivery/mod.rs b/src/delivery/mod.rs new file mode 100644 index 0000000..10f594e --- /dev/null +++ b/src/delivery/mod.rs @@ -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; diff --git a/src/delivery/policy.rs b/src/delivery/policy.rs new file mode 100644 index 0000000..14b082b --- /dev/null +++ b/src/delivery/policy.rs @@ -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 { .. })); + } +} diff --git a/src/delivery/service.rs b/src/delivery/service.rs new file mode 100644 index 0000000..5c3303c --- /dev/null +++ b/src/delivery/service.rs @@ -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>, + ) -> 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(()) + } +} diff --git a/src/gateway/mod.rs b/src/gateway/mod.rs index d60b777..5d21d98 100644 --- a/src/gateway/mod.rs +++ b/src/gateway/mod.rs @@ -12,11 +12,12 @@ use crate::bus::{ControlMessage, MessageBus, OutboundDispatcher}; use crate::channels::base::ChannelError; use crate::channels::{ChannelManager, CliChatChannel}; use crate::config::{Config, ensure_workspace_dir, expand_path}; +use crate::delivery::{ConversationWriteLocks, DeliveryCoordinator, TurnDeliveryService}; use crate::logging; use crate::mcp; use crate::memory::MemoryManager; use crate::scheduler::Scheduler; -use crate::session::SessionManager; +use crate::session::{SessionManager, SessionManagerServices}; use crate::task_supervisor::TaskSupervisor; pub struct GatewayState { @@ -27,6 +28,7 @@ pub struct GatewayState { pub channel_manager: ChannelManager, pub storage: Arc, pub task_supervisor: TaskSupervisor, + pub delivery_coordinator: DeliveryCoordinator, pub connection_shutdown: tokio_util::sync::CancellationToken, pub auth: auth::AuthManager, pub uploads: uploads::UploadRegistry, @@ -44,6 +46,7 @@ impl GatewayState { config_path: std::path::PathBuf, ) -> Result> { let task_supervisor = TaskSupervisor::new(); + let delivery_coordinator = DeliveryCoordinator::new(ConversationWriteLocks::default()); let connection_shutdown = tokio_util::sync::CancellationToken::new(); let auth = auth::AuthManager::load( config.gateway.require_pairing, @@ -109,6 +112,20 @@ impl GatewayState { // Create MessageBus first (shared by SessionManager and ChannelManager) 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 { Some(config.browser.clone()) } else { @@ -119,22 +136,17 @@ impl GatewayState { let session_manager = SessionManager::new( provider_config.clone(), storage.clone(), - bus.clone(), - memory_manager, + SessionManagerServices::new( + bus.clone(), + memory_manager, + task_supervisor.clone(), + turn_delivery, + ), browser_config, config.gateway.max_concurrent_background_tasks, - task_supervisor.clone(), )?; 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 let available_channels = channel_manager.list_channel_names().await; let valid_channels = available_channels.clone(); @@ -209,6 +221,7 @@ impl GatewayState { channel_manager, storage, task_supervisor, + delivery_coordinator, connection_shutdown, auth, uploads, @@ -282,6 +295,7 @@ impl GatewayState { &inbound.chat_id, &inbound.content, inbound.media, + inbound.forwarded_metadata.clone(), ).await { Ok(crate::session::session::HandleResult::AgentResponse(content)) => { let outbound = crate::bus::OutboundMessage { @@ -342,6 +356,7 @@ impl GatewayState { bus_for_outbound, self.channel_manager.clone(), self.task_supervisor.clone(), + self.delivery_coordinator.write_locks(), ); self.task_supervisor diff --git a/src/gateway/ws.rs b/src/gateway/ws.rs index 5831e72..064e38d 100644 --- a/src/gateway/ws.rs +++ b/src/gateway/ws.rs @@ -56,10 +56,12 @@ async fn handle_socket( let _ = sender .send(WsOutbound::SessionEstablished { session_id: session_id.clone(), - capabilities: if state.uploads.enabled() { - vec!["file_transfer_v1".to_string()] - } else { - Vec::new() + capabilities: { + let mut capabilities = vec!["turn_snapshots_v1".to_string()]; + if state.uploads.enabled() { + capabilities.push("file_transfer_v1".to_string()); + } + capabilities }, }) .await; diff --git a/src/lib.rs b/src/lib.rs index 3aea9f1..9412e40 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -3,6 +3,7 @@ pub mod bus; pub mod channels; pub mod client; pub mod config; +pub mod delivery; pub mod gateway; pub mod logging; pub mod mcp; diff --git a/src/protocol.rs b/src/protocol.rs index 9a43174..f1221cd 100644 --- a/src/protocol.rs +++ b/src/protocol.rs @@ -62,6 +62,10 @@ pub struct HistoryMessage { pub seq: i64, pub role: String, pub content: String, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub reasoning_content: Option, + #[serde(default)] + pub completion_status: crate::bus::CompletionStatus, pub created_at: i64, #[serde(default, skip_serializing_if = "Option::is_none")] pub tool_call_id: Option, @@ -140,6 +144,10 @@ pub enum WsInbound { #[derive(Debug, Clone, Serialize, Deserialize)] #[serde(tag = "type")] pub enum WsOutbound { + #[serde(rename = "turn_updated")] + TurnUpdated { + snapshot: crate::session::TurnSnapshot, + }, #[serde(rename = "assistant_response")] AssistantResponse { id: String, @@ -222,3 +230,51 @@ pub fn serialize_inbound(msg: &WsInbound) -> Result { pub fn serialize_outbound(msg: &WsOutbound) -> Result { 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 + ); + } +} diff --git a/src/providers/anthropic.rs b/src/providers/anthropic.rs index 975d68d..8659156 100644 --- a/src/providers/anthropic.rs +++ b/src/providers/anthropic.rs @@ -1,12 +1,19 @@ use async_trait::async_trait; +use futures_util::stream; use reqwest::Client; -use serde::{Deserialize, Serialize}; -use std::collections::HashMap; +use serde::Serialize; +use serde_json::Value; +use std::collections::{BTreeMap, HashMap, VecDeque}; use std::time::Duration; +use thiserror::Error; +use super::stream::SseFramer; use super::traits::Usage; -use super::{ChatCompletionRequest, ChatCompletionResponse, LLMProvider, Message, Tool, ToolCall}; -use crate::bus::message::ContentBlock; +use super::{ + ChatCompletionRequest, DynProviderError, FinishReason, LLMProvider, Message, ProviderChunk, + ProviderStream, Tool, +}; +use crate::bus::{ProviderReasoningState, message::ContentBlock}; use crate::storage::Storage; use std::sync::Arc; @@ -127,6 +134,7 @@ struct AnthropicRequest { messages: Vec, max_tokens: u32, temperature: Option, + stream: bool, #[serde(skip_serializing_if = "Option::is_none")] tools: Option>, #[serde(flatten)] @@ -154,6 +162,8 @@ fn convert_messages(messages: &[Message]) -> Vec { "tool_use_id": tool_call_id, "content": convert_content_blocks(&message.content, false), })] + } else if let Some(native) = native_anthropic_content(message) { + native } else { let mut blocks = convert_content_blocks(&message.content, message.role == "system"); if let Some(tool_calls) = message @@ -177,6 +187,26 @@ fn convert_messages(messages: &[Message]) -> Vec { .collect() } +fn native_anthropic_content(message: &Message) -> Option> { + 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)] struct AnthropicTool { name: String, @@ -186,56 +216,298 @@ struct AnthropicTool { cache_control: Option, } -#[derive(Deserialize)] -struct AnthropicResponse { - id: Option, - model: Option, - #[serde(default)] - content: Vec, - #[serde(default)] - usage: Option, +#[derive(Debug, Error)] +enum AnthropicStreamError { + #[error("invalid UTF-8 in Anthropic SSE event: {0}")] + Utf8(#[from] std::string::FromUtf8Error), + #[error("invalid Anthropic SSE payload: {0}")] + Json(#[from] serde_json::Error), + #[error("Anthropic stream error: {0}")] + Api(String), + #[error("Anthropic stream ended without message_stop")] + MissingFinish, } -#[derive(Deserialize)] -#[serde(tag = "type", rename_all = "snake_case")] -enum AnthropicContent { - Text { - #[serde(alias = "content")] - text: String, - }, - Thinking { - #[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(Default)] +struct AnthropicSseDecoder { + framer: SseFramer, + blocks: BTreeMap, + tool_json: HashMap, + usage: Usage, + finish_reason: Option, + done_emitted: bool, } -#[derive(Deserialize)] -struct AnthropicUsage { - #[serde(default)] - input_tokens: u32, - #[serde(default)] - output_tokens: u32, - #[serde(default)] - cache_read_input_tokens: Option, - #[serde(default)] - cache_creation_input_tokens: Option, +impl AnthropicSseDecoder { + fn push(&mut self, bytes: &[u8]) -> Result, AnthropicStreamError> { + let frames = self.framer.push(bytes)?; + self.decode_frames(frames) + } + + fn finish(&mut self) -> Result, AnthropicStreamError> { + let frames = self.framer.finish()?; + let chunks = self.decode_frames(frames)?; + if !self.done_emitted { + return Err(AnthropicStreamError::MissingFinish); + } + Ok(chunks) + } + + fn decode_frames( + &mut self, + frames: Vec, + ) -> Result, 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) { + 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) { + 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) { + 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) { + 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) { + if self.done_emitted { + return; + } + let content = self.blocks.values().cloned().collect::>(); + 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 { + payload + .get("index") + .and_then(Value::as_u64) + .and_then(|value| usize::try_from(value).ok()) +} + +fn append_block_string(blocks: &mut BTreeMap, 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 { + 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, + reached_eof: bool, +} + +async fn next_anthropic_chunk( + mut state: AnthropicHttpStream, +) -> Result, 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] impl LLMProvider for AnthropicProvider { - async fn chat( + async fn stream( &self, request: ChatCompletionRequest, - ) -> Result> { + ) -> Result { let start = std::time::Instant::now(); let url = format!("{}/v1/messages", self.base_url); 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), max_tokens, temperature: request.temperature.or(self.temperature), + stream: true, tools, extra: self.model_extra.clone(), }; @@ -272,8 +545,16 @@ impl LLMProvider for AnthropicProvider { req_builder = req_builder.header(key.as_str(), value.as_str()); } - let req_body_str = serde_json::to_string_pretty(&body).unwrap_or_default(); - tracing::debug!(req_body = %req_body_str, "LLM request"); + let request_summary = super::stream::diagnostic_request_summary( + &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 is_timeout = e.is_timeout(); @@ -289,10 +570,8 @@ impl LLMProvider for AnthropicProvider { })?; let status = resp.status(); - let body_text = resp.text().await?; - tracing::debug!(status = %status, resp_body = %body_text, "LLM response"); - if !status.is_success() { + let body_text = resp.text().await?; let error_msg = serde_json::from_str::(&body_text) .ok() .and_then(|v| { @@ -315,7 +594,7 @@ impl LLMProvider for AnthropicProvider { .append_llm_call( &self.name, &self.model_id, - &req_body_str, + &request_summary, Some(&body_text), Some(&error_msg), 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()); } - - let anthropic_resp: AnthropicResponse = match serde_json::from_str(&body_text) { - Ok(response) => response, - Err(e) => { - let err_msg = format!("decode error: {} | body: {}", e, &body_text); - if let Some(ref storage) = self.storage { - 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(&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), + tracing::debug!(status = %status, "Anthropic streaming response started"); + Ok(Box::pin(stream::try_unfold( + AnthropicHttpStream { + response: resp, + decoder: AnthropicSseDecoder::default(), + pending: VecDeque::new(), + reached_eof: false, }, - }; - - 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) + next_anthropic_chunk, + ))) } fn ptype(&self) -> &str { @@ -446,6 +631,7 @@ impl LLMProvider for AnthropicProvider { #[cfg(test)] mod tests { use super::*; + use crate::providers::ProviderResponseAccumulator; use serde_json::json; #[test] @@ -487,6 +673,7 @@ mod tests { ContentBlock::image_url("data:image/png;base64,AAAA"), ], reasoning_content: None, + provider_state: None, tool_call_id: Some("call_1".to_string()), name: Some("file_read".to_string()), 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"]["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::(); + 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) + )); + } } diff --git a/src/providers/mod.rs b/src/providers/mod.rs index 57f70c7..beecd14 100644 --- a/src/providers/mod.rs +++ b/src/providers/mod.rs @@ -1,11 +1,18 @@ pub mod anthropic; pub mod openai; +pub mod stream; pub mod traits; pub use self::anthropic::AnthropicProvider; pub use self::openai::OpenAIProvider; 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::{ ChatCompletionRequest, ChatCompletionResponse, LLMProvider, Message, Tool, ToolCall, ToolFunction, Usage, diff --git a/src/providers/openai.rs b/src/providers/openai.rs index ceab97a..9efef17 100644 --- a/src/providers/openai.rs +++ b/src/providers/openai.rs @@ -1,12 +1,16 @@ use async_trait::async_trait; +use futures_util::stream; use reqwest::Client; -use serde::Deserialize; use serde_json::{Value, json}; -use std::collections::HashMap; +use std::collections::{HashMap, VecDeque}; use std::time::Duration; +use thiserror::Error; -use super::traits::Usage; -use super::{ChatCompletionRequest, ChatCompletionResponse, LLMProvider, Message, ToolCall}; +use super::stream::SseFramer; +use super::{ + ChatCompletionRequest, DynProviderError, FinishReason, LLMProvider, Message, ProviderChunk, + ProviderStream, Usage, +}; use crate::bus::message::ContentBlock; use crate::storage::Storage; use std::sync::Arc; @@ -208,82 +212,348 @@ impl OpenAIProvider { 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)] -struct OpenAIResponse { - id: String, - model: String, - choices: Vec, - #[serde(default)] - usage: OpenAIUsage, +#[derive(Debug, Error)] +enum OpenAIStreamError { + #[error("invalid UTF-8 in SSE event: {0}")] + Utf8(#[from] std::string::FromUtf8Error), + #[error("invalid OpenAI-compatible SSE payload: {0}")] + Json(#[from] serde_json::Error), + #[error("OpenAI-compatible stream ended without a finish marker")] + MissingFinish, } -#[derive(Deserialize)] -struct OpenAIChoice { - message: OpenAIMessage, +#[derive(Clone, Copy, PartialEq, Eq)] +enum InlineMode { + Text, + Reasoning, } -fn null_or_missing_tool_calls<'de, D>(deserializer: D) -> Result, D::Error> -where - D: serde::Deserializer<'de>, -{ - Ok(Option::>::deserialize(deserializer)?.unwrap_or_default()) +struct InlineReasoningParser { + mode: InlineMode, + pending: String, } -#[derive(Deserialize)] -struct OpenAIMessage { - #[serde(default)] - content: Option, - #[serde(default)] - reasoning_content: Option, - #[serde(default, deserialize_with = "null_or_missing_tool_calls")] - tool_calls: Vec, +impl Default for InlineReasoningParser { + fn default() -> Self { + Self { + mode: InlineMode::Text, + pending: String::new(), + } + } } -#[derive(Deserialize)] -struct OpenAIToolCall { - id: String, - #[serde(rename = "function")] - function: OAIFunction, +impl InlineReasoningParser { + const TAGS: [(&'static str, InlineMode); 4] = [ + ("", InlineMode::Reasoning), + ("", InlineMode::Reasoning), + ("", InlineMode::Text), + ("", InlineMode::Text), + ]; + + fn push(&mut self, delta: &str) -> Vec { + self.pending.push_str(delta); + self.drain(false) + } + + fn finish(&mut self) -> Vec { + self.drain(true) + } + + fn drain(&mut self, finish: bool) -> Vec { + 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) { + if text.is_empty() { + return; + } + chunks.push(match self.mode { + InlineMode::Text => ProviderChunk::Text(text), + InlineMode::Reasoning => ProviderChunk::Reasoning(text), + }); + } } -#[derive(Deserialize)] -struct OAIFunction { - name: String, - arguments: String, +fn longest_tag_prefix_suffix(value: &str, tags: &[(&str, InlineMode)]) -> usize { + let mut best = 0; + for boundary in value + .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)] -struct OpenAIUsage { - #[serde(default)] - prompt_tokens: u32, - #[serde(default)] - completion_tokens: u32, - #[serde(default)] - total_tokens: u32, - #[serde(default)] - cached_tokens: Option, - #[serde(default)] - prompt_tokens_details: Option, +#[derive(Default)] +struct PartialStreamTool { + id: Option, + name: Option, + started: bool, } -#[derive(Deserialize, Default)] -struct OpenAIPromptTokensDetails { - #[serde(default)] - cached_tokens: Option, +#[derive(Default)] +struct OpenAISseDecoder { + framer: SseFramer, + inline_reasoning: InlineReasoningParser, + tools: HashMap, + metadata_emitted: bool, + done_emitted: bool, +} + +impl OpenAISseDecoder { + fn push(&mut self, bytes: &[u8]) -> Result, OpenAIStreamError> { + let frames = self.framer.push(bytes)?; + self.decode_frames(frames) + } + + fn finish(&mut self) -> Result, 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, + ) -> Result, 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) { + 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) { + let mut indexes = self.tools.keys().copied().collect::>(); + 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, + reached_eof: bool, +} + +async fn next_openai_chunk( + mut state: OpenAIHttpStream, +) -> Result, 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] impl LLMProvider for OpenAIProvider { - async fn chat( + async fn stream( &self, request: ChatCompletionRequest, - ) -> Result> { + ) -> Result { let start = std::time::Instant::now(); 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) #[cfg(debug_assertions)] @@ -322,8 +592,11 @@ impl LLMProvider for OpenAIProvider { req_builder = req_builder.header(key.as_str(), value.as_str()); } - let req_body_str = serde_json::to_string_pretty(&body).unwrap_or_default(); - tracing::debug!(req_body = %req_body_str, "LLM request"); + let request_summary = super::stream::diagnostic_request_summary( + &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 is_timeout = e.is_timeout(); @@ -339,10 +612,8 @@ impl LLMProvider for OpenAIProvider { })?; let status = resp.status(); - let text = resp.text().await?; - tracing::debug!(status = %status, resp_body = %text, "LLM response"); - if !status.is_success() { + let text = resp.text().await?; let error = format!("API error {}: {}", status, text); tracing::error!( provider = %self.name, @@ -357,7 +628,7 @@ impl LLMProvider for OpenAIProvider { .append_llm_call( &self.name, &self.model_id, - &req_body_str, + &request_summary, Some(&text), Some(&error), start.elapsed().as_millis() as u64, @@ -369,93 +640,13 @@ impl LLMProvider for OpenAIProvider { return Err(error.into()); } - let openai_resp: OpenAIResponse = match serde_json::from_str(&text) { - Ok(response) => response, - Err(e) => { - let err_msg = format!("decode error: {} | body: {}", e, &text); - if let Some(ref storage) = self.storage { - 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()); - } + let state = OpenAIHttpStream { + response: resp, + decoder: OpenAISseDecoder::default(), + pending: VecDeque::new(), + reached_eof: false, }; - - 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 = 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) + Ok(Box::pin(stream::try_unfold(state, next_openai_chunk))) } fn ptype(&self) -> &str { @@ -474,7 +665,7 @@ impl LLMProvider for OpenAIProvider { #[cfg(test)] mod tests { use super::*; - use crate::providers::Message; + use crate::providers::{Message, ToolCall}; #[test] fn test_build_request_body_includes_assistant_tool_calls() { @@ -494,6 +685,7 @@ mod tests { role: "assistant".to_string(), content: vec![ContentBlock::text("calling tool")], reasoning_content: None, + provider_state: None, tool_call_id: None, name: None, tool_calls: Some(vec![ToolCall { @@ -519,6 +711,10 @@ mod tests { tool_calls[0]["function"]["arguments"], "{\"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] @@ -532,6 +728,7 @@ mod tests { ContentBlock::image_url("data:image/png;base64,AAAA"), ], reasoning_content: None, + provider_state: None, tool_call_id: Some("call_2".to_string()), name: Some("file_read".to_string()), tool_calls: None, @@ -552,76 +749,113 @@ mod tests { assert_eq!(converted[1]["content"], "second image"); } - #[test] - fn test_decode_response_accepts_null_tool_calls() { - let text = r#"{ - "id": "d21abaa6552741949e2aba76bde59359", - "choices": [{ - "finish_reason": "stop", - "index": 0, - "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.") + #[tokio::test] + async fn sse_decoder_handles_byte_boundaries_reasoning_content_and_usage() { + let input = concat!( + "data: {\"id\":\"r1\",\"model\":\"m1\",\"choices\":[{\"delta\":{\"reasoning_content\":\"why\",\"content\":\"你好\"},\"finish_reason\":null}]}\r\n\r\n", + "data: {\"choices\":[{\"delta\":{},\"finish_reason\":\"stop\"}]}\n\n", + "data: {\"choices\":[],\"usage\":{\"prompt_tokens\":10,\"completion_tokens\":5,\"total_tokens\":15,\"prompt_tokens_details\":{\"cached_tokens\":3}}}\n\n", + "data: [DONE]\n\n" ); - assert!(message.tool_calls.is_empty()); - assert_eq!(response.usage.total_tokens, 11806); + let mut decoder = OpenAISseDecoder::default(); + 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] - fn test_decode_response_exposes_cached_tokens() { - let text = r#"{ - "id": "d21abaa6552741949e2aba76bde59359", - "choices": [{ - "finish_reason": "stop", - "index": 0, - "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) + fn inline_reasoning_tags_may_span_sse_chunks() { + let input = concat!( + "data: {\"choices\":[{\"delta\":{\"content\":\"secretanswer\"},\"finish_reason\":\"stop\"}]}\n\n", + "data: [DONE]\n\n" ); + 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::(), + "secret" + ); + assert_eq!( + chunks + .iter() + .filter_map(|chunk| match chunk { + ProviderChunk::Text(value) => Some(value.as_str()), + _ => None, + }) + .collect::(), + "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) + )); } } diff --git a/src/providers/stream.rs b/src/providers/stream.rs new file mode 100644 index 0000000..c01da82 --- /dev/null +++ b/src/providers/stream.rs @@ -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; +pub type ProviderStreamItem = Result; +pub type ProviderStream = Pin + 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, +} + +impl SseFramer { + pub(crate) fn push(&mut self, bytes: &[u8]) -> Result, std::string::FromUtf8Error> { + self.buffer.extend_from_slice(bytes); + self.drain_frames(false) + } + + pub(crate) fn finish(&mut self) -> Result, std::string::FromUtf8Error> { + self.drain_frames(true) + } + + fn drain_frames(&mut self, finish: bool) -> Result, 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::>(); + 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) -> Result, 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::>() + .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, + name: Option, + }, + ToolCallArguments { + index: usize, + delta: String, + }, + ProviderState(ProviderReasoningState), + Usage(Usage), + Done(FinishReason), +} + +#[derive(Default)] +struct PartialToolCall { + id: Option, + name: Option, + arguments: String, +} + +#[derive(Default)] +pub struct ProviderResponseAccumulator { + id: String, + model: String, + content: String, + reasoning_content: String, + provider_state: Option, + usage: Usage, + tool_calls: BTreeMap, +} + +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 { + 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::(&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}) + ); + } +} diff --git a/src/providers/traits.rs b/src/providers/traits.rs index 84ce6ab..5258ce1 100644 --- a/src/providers/traits.rs +++ b/src/providers/traits.rs @@ -2,12 +2,17 @@ use crate::bus::message::ContentBlock; use async_trait::async_trait; use serde::{Deserialize, Serialize}; +use super::stream::{DynProviderError, ProviderStream, collect_provider_stream}; + #[derive(Debug, Clone, Serialize, Deserialize)] pub struct Message { pub role: String, pub content: Vec, #[serde(skip_serializing_if = "Option::is_none")] pub reasoning_content: Option, + /// Opaque state replayed only by the provider that produced it. + #[serde(skip)] + pub provider_state: Option, #[serde(skip_serializing_if = "Option::is_none")] pub tool_call_id: Option, #[serde(skip_serializing_if = "Option::is_none")] @@ -22,6 +27,7 @@ impl Message { role: "user".to_string(), content: vec![ContentBlock::text(content)], reasoning_content: None, + provider_state: None, tool_call_id: None, name: None, tool_calls: None, @@ -33,6 +39,7 @@ impl Message { role: "assistant".to_string(), content: vec![ContentBlock::text(content)], reasoning_content: None, + provider_state: None, tool_call_id: None, name: None, tool_calls: None, @@ -44,6 +51,7 @@ impl Message { role: "system".to_string(), content: vec![ContentBlock::text(content)], reasoning_content: None, + provider_state: None, tool_call_id: None, name: None, tool_calls: None, @@ -59,6 +67,7 @@ impl Message { role: "tool".to_string(), content: vec![ContentBlock::text(content)], reasoning_content: None, + provider_state: None, tool_call_id: Some(tool_call_id.into()), name: Some(tool_name.into()), tool_calls: None, @@ -101,11 +110,12 @@ pub struct ChatCompletionResponse { pub model: String, pub content: String, pub reasoning_content: Option, + pub provider_state: Option, pub tool_calls: Vec, pub usage: Usage, } -#[derive(Debug, Clone, Serialize, Deserialize)] +#[derive(Debug, Clone, Default, PartialEq, Eq, Serialize, Deserialize)] pub struct Usage { pub prompt_tokens: u32, pub completion_tokens: u32, @@ -120,10 +130,17 @@ pub struct Usage { #[async_trait] pub trait LLMProvider: Send + Sync { + async fn stream( + &self, + request: ChatCompletionRequest, + ) -> Result; + async fn chat( &self, request: ChatCompletionRequest, - ) -> Result>; + ) -> Result { + collect_provider_stream(self.stream(request).await?).await + } fn ptype(&self) -> &str; diff --git a/src/session/mod.rs b/src/session/mod.rs index 0445afa..7523334 100644 --- a/src/session/mod.rs +++ b/src/session/mod.rs @@ -7,9 +7,14 @@ mod persistence; #[allow(clippy::module_inception)] pub mod session; pub mod session_id; +pub mod turn; pub use commands::SessionCommand; pub use error::SessionError; 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 turn::{ + BlockId, ToolStatus, TurnBlock, TurnController, TurnId, TurnPhase, TurnSnapshot, TurnState, + TurnStatus, +}; diff --git a/src/session/persistence.rs b/src/session/persistence.rs index 303925a..4de586f 100644 --- a/src/session/persistence.rs +++ b/src/session/persistence.rs @@ -1,10 +1,12 @@ use std::sync::Arc; +use std::{fmt::Display, future::Future}; use tokio::sync::Mutex; use super::session::{MessagePersistSnapshot, Session}; use crate::bus::ChatMessage; use crate::storage::StorageError; +use crate::{providers::Usage, session::TurnController}; async fn persist_added_messages( snapshots: Vec>, @@ -63,3 +65,74 @@ pub(super) async fn append_persisted_messages( } 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( + controller: &TurnController, + usage: Option, + persistence: F, +) -> Result +where + F: Future>, + 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") + ); + } +} diff --git a/src/session/session.rs b/src/session/session.rs index 5ab3d5b..11bac93 100644 --- a/src/session/session.rs +++ b/src/session/session.rs @@ -3,8 +3,11 @@ use std::sync::Arc; use tokio::sync::{Mutex, mpsc, oneshot}; -use super::persistence::append_persisted_messages; -use crate::bus::{ChatMessage, MediaItem, MediaRef, MessageSource, OutboundMessage, SourceKind}; +use super::persistence::{append_persisted_messages, finalize_turn_after_persistence}; +use super::turn::{TurnBlock, TurnController, TurnSnapshot}; +use crate::bus::{ + ChatMessage, CompletionStatus, MediaItem, MediaRef, MessageSource, OutboundMessage, SourceKind, +}; use crate::mcp::get_mcp_status; use crate::storage::{Storage, StorageError}; use std::sync::Arc as StdArc; @@ -22,6 +25,15 @@ fn outbound_session_metadata(session_id: &str) -> HashMap { HashMap::from([("_session_id".to_string(), session_id.to_string())]) } +fn outbound_turn_metadata( + session_id: &str, + forwarded: &HashMap, +) -> HashMap { + let mut metadata = forwarded.clone(); + metadata.insert("_session_id".to_string(), session_id.to_string()); + metadata +} + tokio::task_local! { pub(super) static CURRENT_SOURCE_SESSION: Option; } @@ -37,10 +49,11 @@ pub enum HandleResult { } use crate::agent::context_compressor::ContextCompressionConfig; 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::config::BrowserConfig; use crate::config::LLMProviderConfig; +use crate::delivery::TurnDeliveryService; /// Check if an LLM error message indicates a context window overflow. 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("input is too long") } + +fn partial_assistant_message( + snapshot: &TurnSnapshot, + completion_status: CompletionStatus, +) -> Option { + 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>, + 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::providers::{LLMProvider, create_provider}; use crate::session::events::DialogInfo; @@ -93,6 +239,7 @@ pub struct Session { agent_tx: Option>, /// Cancel signal for the currently executing agent task current_cancel: Option>, + active_turn_emitter: Option, /// Monotonic counter to detect stale workers worker_generation: u64, /// Monotonic counter for in-memory session mutations. @@ -108,12 +255,18 @@ pub struct Session { pub(super) persistence_lock: Arc>, } +struct ActiveTurnEmitter { + turn_id: String, + emitter: TurnEmitter, +} + /// A task to be processed by the per-session agent worker struct AgentTask { channel: String, chat_id: String, content: String, media: Vec, + forwarded_metadata: HashMap, } #[derive(Clone)] @@ -123,6 +276,7 @@ struct AgentWorkerDeps { work_manager: Arc, skills_loader: Arc, task_supervisor: crate::task_supervisor::TaskSupervisor, + turn_delivery: TurnDeliveryService, } impl Session { @@ -178,6 +332,7 @@ impl Session { memory_manager, agent_tx: None, current_cancel: None, + active_turn_emitter: None, worker_generation: 0, state_version: 0, persistence_lock: Arc::new(Mutex::new(())), @@ -258,6 +413,12 @@ impl Session { role: m.role, content: m.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 .map(|refs| serde_json::from_str(&refs).unwrap_or_default()) @@ -293,6 +454,12 @@ impl Session { role: m.role, content: m.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 .map(|refs| serde_json::from_str(&refs).unwrap_or_default()) @@ -354,6 +521,7 @@ impl Session { memory_manager, agent_tx: None, current_cancel: None, + active_turn_emitter: None, worker_generation: 0, state_version: 0, persistence_lock: Arc::new(Mutex::new(())), @@ -386,6 +554,13 @@ impl Session { role: message.role.clone(), content: message.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() { None } else { @@ -897,6 +1072,31 @@ pub struct SessionManager { work_manager: Arc, sub_agent_manager: Arc, task_supervisor: crate::task_supervisor::TaskSupervisor, + turn_delivery: TurnDeliveryService, +} + +/// Gateway-owned runtime services shared by all Session workers. +pub struct SessionManagerServices { + bus: Arc, + memory_manager: Arc, + task_supervisor: crate::task_supervisor::TaskSupervisor, + turn_delivery: TurnDeliveryService, +} + +impl SessionManagerServices { + pub fn new( + bus: Arc, + memory_manager: Arc, + task_supervisor: crate::task_supervisor::TaskSupervisor, + turn_delivery: TurnDeliveryService, + ) -> Self { + Self { + bus, + memory_manager, + task_supervisor, + turn_delivery, + } + } } struct SessionManagerInner { @@ -1010,18 +1210,23 @@ impl SessionManager { work_manager: self.work_manager.clone(), skills_loader: self.skills_loader.clone(), task_supervisor: self.task_supervisor.clone(), + turn_delivery: self.turn_delivery.clone(), } } pub fn new( provider_config: LLMProviderConfig, storage: Arc, - bus: Arc, - memory_manager: Arc, + services: SessionManagerServices, browser_config: Option, max_concurrent_background_tasks: usize, - task_supervisor: crate::task_supervisor::TaskSupervisor, ) -> Result { + let SessionManagerServices { + bus, + memory_manager, + task_supervisor, + turn_delivery, + } = services; let mut skills_loader = SkillsLoader::new(); skills_loader.load_skills(); skills_loader.set_workspace_skills_dir(provider_config.workspace_dir.clone()); @@ -1108,6 +1313,7 @@ impl SessionManager { work_manager, sub_agent_manager, task_supervisor, + turn_delivery, }) } @@ -1454,6 +1660,9 @@ impl SessionManager { if guard.current_cancel.take().is_some() { 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() { msgs.push("消息队列已清空。".to_string()); } @@ -2001,6 +2210,7 @@ impl SessionManager { chat_id: &str, content: &str, media: Vec, + forwarded_metadata: HashMap, ) -> Result { let unified_id = self.resolve_dialog_id(channel, chat_id).await?; tracing::debug!(unified_id = %unified_id, "handle_message resolved unified_id"); @@ -2038,6 +2248,7 @@ impl SessionManager { chat_id: chat_id.to_string(), content: content.to_string(), media, + forwarded_metadata, }; let session_clone = session.clone(); let unified_str = unified_id.to_string(); @@ -2167,6 +2378,7 @@ fn spawn_agent_worker( work_manager, skills_loader, task_supervisor, + turn_delivery, } = deps; let worker_supervisor = task_supervisor.clone(); 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 { let task_chan = task.channel.clone(); let task_cid = task.chat_id.clone(); + let task_metadata = task.forwarded_metadata.clone(); let notification_session_id = unified_str.clone(); 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(), reply_to: None, media: vec![], - metadata: outbound_session_metadata(&unified_str), + metadata: outbound_turn_metadata(&unified_str, &task_metadata), delivery: None, }; let _ = bus.publish_outbound(err_outbound).await; @@ -2258,7 +2471,7 @@ fn spawn_agent_worker( .to_string(), reply_to: None, media: vec![], - metadata: outbound_session_metadata(&unified_str), + metadata: outbound_turn_metadata(&unified_str, &task_metadata), delivery: None, }; 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); } + 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 let session2 = session.clone(); let bus2 = bus.clone(); let chan2 = task_chan.clone(); let cid2 = task_cid.clone(); let unified_str2 = unified_str.clone(); + let task_metadata2 = task_metadata.clone(); + let turn_lifecycle = &turn_controller; let process_future = async move { let response_session_id = unified_str2.clone(); let process_result = crate::agent::sub_agent::DELEGATE_CONTEXT.scope( @@ -2386,7 +2651,7 @@ fn spawn_agent_worker( channel: chan2.clone(), chat_id: cid2.clone(), }, - agent.process(history_out.clone()), + agent.process_streaming(history_out.clone(), agent_turn.clone()), ).await; let result = match process_result { Ok(r) => r, @@ -2416,6 +2681,12 @@ fn spawn_agent_worker( Ok(r) => r, Err(e) => { 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 { channel: chan2, chat_id: cid2, @@ -2423,12 +2694,15 @@ fn spawn_agent_worker( .to_string(), reply_to: None, media: vec![], - metadata: outbound_session_metadata( + metadata: outbound_turn_metadata( &response_session_id, + &task_metadata2, ), delivery: None, }; - let _ = bus2.publish_outbound(err_outbound).await; + if !live_delivery_started { + let _ = bus2.publish_outbound(err_outbound).await; + } return; } }; @@ -2440,6 +2714,10 @@ fn spawn_agent_worker( session_id = %guard.id, "Session changed while retry-compressing after context overflow" ); + turn_lifecycle.cancel(Some( + "session changed during context overflow recovery" + .to_string(), + )); return; } guard.compressor.set_context_window(new_window); @@ -2471,55 +2749,100 @@ fn spawn_agent_worker( retry }; - match agent.process(retry_history).await { + match agent + .process_streaming(retry_history, agent_turn.clone()) + .await + { Ok(r) => r, Err(e) => { tracing::error!( error = %e, "Agent retry after overflow failed" ); + fail_turn_with_partial( + turn_lifecycle, + &session2, + e.to_string(), + ) + .await; let err_outbound = OutboundMessage { channel: chan2, chat_id: cid2, content: format!("Processing error: {}", e), reply_to: None, media: vec![], - metadata: outbound_session_metadata(&response_session_id), + metadata: outbound_turn_metadata( + &response_session_id, + &task_metadata2, + ), delivery: None, }; - let _ = bus2.publish_outbound(err_outbound).await; + if !live_delivery_started { + let _ = bus2.publish_outbound(err_outbound).await; + } return; } } } Err(e) => { tracing::error!(error = %e, "Agent processing error"); + fail_turn_with_partial( + turn_lifecycle, + &session2, + e.to_string(), + ) + .await; let err_outbound = OutboundMessage { channel: chan2, chat_id: cid2, content: format!("Processing error: {}", e), reply_to: None, media: vec![], - metadata: outbound_session_metadata(&response_session_id), + metadata: outbound_turn_metadata( + &response_session_id, + &task_metadata2, + ), delivery: None, }; - let _ = bus2.publish_outbound(err_outbound).await; + if !live_delivery_started { + let _ = bus2.publish_outbound(err_outbound).await; + } return; } }; let response_content = result.final_response.content; let total_tokens = result.total_tokens; - let response = - if let Err(e) = append_persisted_messages(&session2, result.emitted_messages).await { - tracing::error!(error = %e, "Failed to atomically persist agent turn"); - None - } else { + let usage = result.usage; + { + let guard = session2.lock().await; + if guard.worker_generation != worker_gen + || 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 sent_count = guard.messages.len(); guard.compressor.set_last_api_info(sent_count, total_tokens); Some(response_content) - }; + } + Err(e) => { + tracing::error!(error = %e, "Failed to atomically persist agent turn"); + None + } + }; let Some(response) = response else { let err_outbound = OutboundMessage { @@ -2529,10 +2852,15 @@ fn spawn_agent_worker( .to_string(), reply_to: None, media: vec![], - metadata: outbound_session_metadata(&response_session_id), + metadata: outbound_turn_metadata( + &response_session_id, + &task_metadata2, + ), delivery: None, }; - let _ = bus2.publish_outbound(err_outbound).await; + if !live_delivery_started { + let _ = bus2.publish_outbound(err_outbound).await; + } return; }; @@ -2540,27 +2868,60 @@ fn spawn_agent_worker( tracing::warn!("failed to generate title: {}", e); } - let outbound = OutboundMessage { - channel: chan2, - chat_id: cid2, - content: response, - reply_to: None, - media: vec![], - metadata: outbound_session_metadata(&response_session_id), - delivery: None, - }; - let _ = bus2.publish_outbound(outbound).await; + if !live_delivery_started { + let outbound = OutboundMessage { + channel: chan2, + chat_id: cid2, + content: response, + reply_to: None, + media: vec![], + metadata: outbound_turn_metadata( + &response_session_id, + &task_metadata2, + ), + delivery: None, + }; + let _ = bus2.publish_outbound(outbound).await; + } }; tokio::select! { () = process_future => {} _ = cancel_rx => { // 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 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 { guard.current_cancel = None; } diff --git a/src/session/turn.rs b/src/session/turn.rs new file mode 100644 index 0000000..a102394 --- /dev/null +++ b/src/session/turn.rs @@ -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, + }, +} + +#[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, + pub usage: Option, + pub error: Option, +} + +pub type TurnSnapshot = TurnState; + +struct TurnControllerInner { + state: TurnState, + snapshots: watch::Sender>, + 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, + error: Option, + ) -> 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>, +} + +impl TurnController { + pub fn start( + session_id: impl Into, + message_id: impl Into, + ) -> (Self, TurnEmitter, watch::Receiver>) { + 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> = 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 { + 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) -> bool { + lock_unpoisoned(&self.inner).transition_terminal(TurnStatus::Completed, usage, None) + } + + pub fn cancel(&self, reason: Option) -> bool { + lock_unpoisoned(&self.inner).transition_terminal(TurnStatus::Cancelled, None, reason) + } + + pub fn fail(&self, error: impl Into) -> bool { + lock_unpoisoned(&self.inner).transition_terminal( + TurnStatus::Failed, + None, + Some(error.into()), + ) + } +} + +fn lock_unpoisoned(mutex: &Mutex) -> 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>, + ) { + 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()); + } +} diff --git a/src/storage/message.rs b/src/storage/message.rs index 6ae0c29..11b3123 100644 --- a/src/storage/message.rs +++ b/src/storage/message.rs @@ -1,5 +1,7 @@ use serde::{Deserialize, Serialize}; +use crate::bus::CompletionStatus; + #[derive(Debug, Clone, Serialize, Deserialize)] pub struct MessageMeta { pub id: String, @@ -8,6 +10,10 @@ pub struct MessageMeta { pub role: String, pub content: String, pub reasoning_content: Option, + pub provider_state: Option, + pub turn_id: Option, + pub iteration: Option, + pub completion_status: CompletionStatus, pub media_refs: Option, pub tool_call_id: Option, pub tool_name: Option, diff --git a/src/storage/mod.rs b/src/storage/mod.rs index 1d3612d..c0331bb 100644 --- a/src/storage/mod.rs +++ b/src/storage/mod.rs @@ -9,15 +9,21 @@ pub use background_task::BackgroundTask; pub use error::StorageError; 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 std::path::Path; use tokio::time::{Duration, sleep}; -const SCHEMA_VERSION: i64 = 3; +const SCHEMA_VERSION: i64 = 4; 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) - VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) + INSERT INTO messages ( + 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>( @@ -31,6 +37,10 @@ fn insert_message_query<'a>( .bind(&msg.role) .bind(&msg.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.tool_call_id) .bind(&msg.tool_name) @@ -39,6 +49,28 @@ fn insert_message_query<'a>( .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(crate) pool: Pool, } @@ -111,6 +143,10 @@ impl Storage { tool_calls TEXT, source TEXT, reasoning_content TEXT, + provider_state TEXT, + turn_id TEXT, + iteration INTEGER, + completion_status TEXT NOT NULL DEFAULT 'completed', created_at INTEGER NOT NULL, FOREIGN KEY (session_id) REFERENCES sessions(id) ON DELETE CASCADE ) @@ -351,6 +387,14 @@ impl Storage { for (table, column, definition) in [ ("messages", "source", "source 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", @@ -829,7 +873,9 @@ impl Storage { ) -> Result, StorageError> { let rows = sqlx::query( 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 WHERE session_id = ? AND seq >= ? ORDER BY seq ASC @@ -840,23 +886,7 @@ impl Storage { .fetch_all(self.pool()) .await?; - Ok(rows - .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(rows.into_iter().map(message_meta_from_row).collect()) } pub async fn get_message( @@ -866,8 +896,9 @@ impl Storage { ) -> Result, StorageError> { let row = sqlx::query( 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 WHERE session_id = ? AND id = ? "#, @@ -877,20 +908,7 @@ impl Storage { .fetch_optional(self.pool()) .await?; - Ok(row.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"), - })) + Ok(row.map(message_meta_from_row)) } pub async fn get_max_message_seq(&self, session_id: &str) -> Result { @@ -912,11 +930,13 @@ impl Storage { let limit = limit.clamp(1, 2_000); let rows = sqlx::query( 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 ( - 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 WHERE session_id = ? ORDER BY seq DESC @@ -930,23 +950,7 @@ impl Storage { .fetch_all(self.pool()) .await?; - Ok(rows - .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(rows.into_iter().map(message_meta_from_row).collect()) } pub async fn load_messages_after_timestamp( @@ -956,7 +960,9 @@ impl Storage { ) -> Result, StorageError> { let rows = sqlx::query( 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 WHERE session_id = ? AND created_at > ? ORDER BY seq ASC @@ -967,23 +973,7 @@ impl Storage { .fetch_all(self.pool()) .await?; - Ok(rows - .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(rows.into_iter().map(message_meta_from_row).collect()) } pub async fn query_sessions_range( @@ -1040,7 +1030,9 @@ impl Storage { ) -> Result, StorageError> { let rows = sqlx::query( 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 WHERE session_id = ? ORDER BY seq DESC @@ -1052,23 +1044,7 @@ impl Storage { .fetch_all(self.pool()) .await?; - let mut messages: Vec<_> = rows - .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(); + let mut messages: Vec<_> = rows.into_iter().map(message_meta_from_row).collect(); messages.reverse(); Ok(messages) } @@ -1095,7 +1071,9 @@ impl Storage { ); let select_sql = format!( 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 WHERE session_id = ?{} ORDER BY seq ASC @@ -1127,23 +1105,7 @@ impl Storage { .fetch_all(self.pool()) .await?; - let messages: Vec<_> = rows - .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(); + let messages: Vec<_> = rows.into_iter().map(message_meta_from_row).collect(); Ok((messages, total)) } @@ -1616,7 +1578,17 @@ mod tests { let storage = Storage::new(&db_path).await.unwrap(); for (table, expected) in [ - ("messages", vec!["source", "reasoning_content"]), + ( + "messages", + vec![ + "source", + "reasoning_content", + "provider_state", + "turn_id", + "iteration", + "completion_status", + ], + ), ( "sessions", 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] async fn test_upsert_and_get_session() { let (storage, _dir) = create_test_storage().await; @@ -1784,6 +1831,10 @@ mod tests { role: "user".to_string(), content: "你好".to_string(), reasoning_content: None, + provider_state: None, + turn_id: None, + iteration: None, + completion_status: crate::bus::CompletionStatus::Completed, media_refs: None, tool_call_id: None, tool_name: None, @@ -1821,6 +1872,67 @@ mod tests { 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] async fn test_persist_message_batch_is_atomic() { let (storage, _dir) = create_test_storage().await; @@ -1848,6 +1960,10 @@ mod tests { role: "assistant".to_string(), content: "must roll back".to_string(), reasoning_content: None, + provider_state: None, + turn_id: None, + iteration: None, + completion_status: crate::bus::CompletionStatus::Completed, media_refs: None, tool_call_id: None, tool_name: None, diff --git a/src/tools/chat_manager.rs b/src/tools/chat_manager.rs index 06497d9..39d6e8e 100644 --- a/src/tools/chat_manager.rs +++ b/src/tools/chat_manager.rs @@ -358,6 +358,10 @@ mod tests { }, content: format!("消息内容 {}", i), reasoning_content: None, + provider_state: None, + turn_id: None, + iteration: None, + completion_status: crate::bus::CompletionStatus::Completed, media_refs: None, tool_call_id: None, tool_name: None, @@ -420,6 +424,10 @@ mod tests { }, content: format!("消息内容 {}", i), reasoning_content: None, + provider_state: None, + turn_id: None, + iteration: None, + completion_status: crate::bus::CompletionStatus::Completed, media_refs: None, tool_call_id: None, tool_name: None, @@ -476,6 +484,10 @@ mod tests { role: "user".to_string(), content: format!("消息内容 {}", i), reasoning_content: None, + provider_state: None, + turn_id: None, + iteration: None, + completion_status: crate::bus::CompletionStatus::Completed, media_refs: None, tool_call_id: None, tool_name: None, diff --git a/tests/test_request_format.rs b/tests/test_request_format.rs index a6a5be7..0deff88 100644 --- a/tests/test_request_format.rs +++ b/tests/test_request_format.rs @@ -135,6 +135,8 @@ fn test_bounded_session_history_protocol() { seq: 1, role: "user".to_string(), content: "你好".to_string(), + reasoning_content: None, + completion_status: picobot::bus::CompletionStatus::Completed, created_at: 123, tool_call_id: None, tool_name: None, @@ -182,6 +184,8 @@ fn test_session_history_preserves_tool_call_metadata() { seq: 2, role: "assistant".to_string(), content: String::new(), + reasoning_content: Some("checking".to_string()), + completion_status: picobot::bus::CompletionStatus::Completed, created_at: 124, tool_call_id: None, tool_name: None, @@ -197,6 +201,7 @@ fn test_session_history_preserves_tool_call_metadata() { let json = serde_json::to_string(&outbound).unwrap(); assert!(json.contains(r#""tool_calls""#)); assert!(json.contains(r#""read_file""#)); + assert!(json.contains(r#""reasoning_content":"checking""#)); let decoded: WsOutbound = serde_json::from_str(&json).unwrap(); match decoded { WsOutbound::SessionHistory { messages, .. } => { diff --git a/webui/src/lib/TurnView.svelte b/webui/src/lib/TurnView.svelte new file mode 100644 index 0000000..6573573 --- /dev/null +++ b/webui/src/lib/TurnView.svelte @@ -0,0 +1,45 @@ + + +
+
P
+
+ {#each turn.blocks as block (block.id)} + {#if block.type === "reasoning"} +
+ 思考过程 +
+
+ {:else if block.type === "assistant"} +
+ {:else if block.type === "tool"} +
+ ⌘{block.name}{toolStatus(block.status)} +
+ {#if block.arguments !== null}
{JSON.stringify(block.arguments, null, 2)}
{/if} + {#if block.preview}
{/if} +
+
+ {/if} + {/each} +
+ {#if turn.status === "running"}{/if}{statusLabel(turn.status)} +
+ {#if turn.error}
{turn.error}
{/if} +
+
diff --git a/webui/src/pages/ChatPage.svelte b/webui/src/pages/ChatPage.svelte index c3ee14b..6f49474 100644 --- a/webui/src/pages/ChatPage.svelte +++ b/webui/src/pages/ChatPage.svelte @@ -4,6 +4,7 @@ import { clientId, formatTime, randomId } from "../lib/api.js"; import Markdown from "../lib/Markdown.svelte"; import ToolCallCard from "../lib/ToolCallCard.svelte"; + import TurnView from "../lib/TurnView.svelte"; let { notify } = $props(); let socket = $state(null); @@ -17,6 +18,7 @@ let selectedCommand = $state(0); let commandMenuDismissed = $state(false); let thinking = $state(false); + let activeTurn = $state(null); let pendingUploads = $state([]); let fileInput; let plansBySession = $state({}); @@ -75,7 +77,7 @@ function handleFrame(frame) { switch (frame.type) { - case "session_established": currentId = frame.session_id; break; + case "session_established": currentId = frame.session_id; activeTurn = null; break; case "session_list": sessions = frame.sessions || []; if (frame.current_session_id) currentId = frame.current_session_id; @@ -84,16 +86,19 @@ case "session_created": currentId = frame.session_id; messages = []; + activeTurn = null; send({ type: "list_sessions", include_archived: false }); break; case "session_loaded": currentId = frame.session_id; + activeTurn = null; send({ type: "get_session_history", session_id: currentId, limit: 1000 }); send({ type: "get_session_plan", session_id: currentId }); break; case "session_history": if (frame.session_id === currentId) { messages = frame.messages || []; + if (activeTurn?.status !== "running" && messages.some((message) => message.id === activeTurn?.message_id)) activeTurn = null; scrollToBottom(); } break; @@ -133,6 +138,19 @@ } send({ type: "list_sessions", include_archived: false }); 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": if (!frame.session_id || frame.session_id === currentId) appendMessage("assistant", frame.content); break; @@ -156,6 +174,8 @@ currentId = id; clearPendingUploads(); messages = []; + activeTurn = null; + thinking = false; todoOpen = Boolean(unseenPlanSessions[id]); unseenPlanSessions[id] = false; send({ type: "load_session", session_id: id }); @@ -375,7 +395,7 @@
- {#if messages.length === 0} + {#if messages.length === 0 && !activeTurn}

今天想做些什么？

消息与 CLI 客户端使用同一套会话、记忆和工具能力。

{/if} {#each messages as message (message.id)} @@ -383,7 +403,16 @@
{message.role === "user" ? "你" : "P"}
+ {#if message.reasoning_content} +
+ 思考过程 +
+
+ {/if} {#if message.content}
{/if} + {#if message.completion_status && message.completion_status !== "completed"} + {message.completion_status === "cancelled" ? "已停止" : "回复中断"} + {/if} {#if message.attachments?.length}
{#each message.attachments as attachment (`${message.id}:${attachment.index}`)} @@ -408,7 +437,8 @@
{/if} {/each} - {#if thinking}
P
正在思考…
{/if} + {#if activeTurn} + {:else if thinking}
P
正在思考…
{/if}
{ event.preventDefault(); submit(); }} ondragover={(event) => event.preventDefault()} ondrop={dropFiles}> {#if commandSuggestions.length} diff --git a/webui/src/styles.css b/webui/src/styles.css index 63265cd..7de62da 100644 --- a/webui/src/styles.css +++ b/webui/src/styles.css @@ -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 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); } +.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; } @keyframes pulse { 50% { opacity: .25; transform: scale(.8); } } .markdown-body { min-width: 0; }