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

Compare commits

base: xiaoski:5079458f55b68f4de65e2f5f4b9137e2aa0be67f
Branches Tags
xiaoski:main
xiaoski:scheduled-run-design
xiaoski:feat/webui-p1
xiaoski:agent/tool-media-input
...
compare: xiaoski:27d126cf7310d2ac9bc7be30675d6dd4fb8a78d4
Branches Tags
xiaoski:main
xiaoski:scheduled-run-design
xiaoski:feat/webui-p1
xiaoski:agent/tool-media-input

15 Commits
5079458f55 ... 27d126cf73

Author SHA1 Message Date
xiaoxixi
27d126cf73 tui客户端增强。 2026-07-14 14:31:16 +08:00
xiaoxixi
901b622b40 更新文档 2026-07-14 13:02:06 +08:00
xiaoxixi
18f1e47f77 fix(feishu): bound shutdown during connection 2026-07-14 12:43:02 +08:00
xiaoxixi
3f1350c33b fix(channels): retry only transient delivery errors 2026-07-14 11:57:47 +08:00
xiaoxixi
f42e9d44cc refactor(session): isolate messaging and persistence 2026-07-14 11:56:19 +08:00
xiaoxixi
24d3e26b43 fix(runtime): eliminate unowned persistence tasks 2026-07-14 11:50:55 +08:00
xiaoxixi
59ecb27c06 删除陈旧文档 2026-07-14 11:49:42 +08:00
xiaoxixi
954bfd1d75 fix(messaging): make persistence and delivery explicit 2026-07-14 11:49:02 +08:00
xiaoxixi
c2d4fc5f09 refactor: remove dead code and unused state 2026-07-14 11:41:07 +08:00
xiaoxixi
560ace50c1 refactor(storage): remove legacy persistence paths 2026-07-14 11:33:48 +08:00
xiaoxixi
f691c1aad6 fix(agent): supervise background tasks atomically 2026-07-14 11:30:32 +08:00
xiaoxixi
a9c297764e fix(channels): unify lifecycle ownership 2026-07-14 11:24:03 +08:00
xiaoxixi
63d20d1eb8 refactor: eliminate build warnings and legacy evaluator 
Resolve strict Clippy findings across all targets, preserve public API compatibility with scoped lint exceptions, and fix sourced messages retaining media references. Replace meval and its future-incompatible nom dependency with a bounded internal expression parser and regression tests.
 2026-07-14 11:00:39 +08:00
xiaoxixi
3d580828b5 fix(runtime): harden persistence scheduling and shutdown 
Add versioned SQLite migrations, enforced runtime pragmas, and transactional message constraints. Claim scheduled jobs with durable leases, execute them concurrently within bounds, and commit completion atomically. Supervise background tasks and clean up WebSocket clients and writer tasks during ordered shutdown.
 2026-07-14 10:37:58 +08:00
xiaoxixi
b06bc4f025 fix(runtime): harden session persistence and outbound dispatch 2026-07-14 10:08:55 +08:00
74 changed files with 4938 additions and 2906 deletions
Show Stats Expand all files Collapse all files

71
AGENTS.md
View File

@ -1,5 +1,7 @@
# PicoBot
This file is the operational contract for coding agents working in this repository. Read [docs/ARCHITECTURE.md](docs/ARCHITECTURE.md) before making architectural or lifecycle changes. Code and tests are the final source of truth; update the document when an architectural invariant changes.
## Build & Run
- `cargo build` — build the binary
@ -8,20 +10,23 @@
## Config
- Config load order: `~/.picobot/config.json` then fallback to `./config.json` (`src/config/mod.rs:237-267`)
- Config load order: `~/.picobot/config.json` then fallback to `./config.json` (`Config::load_default` in `src/config/mod.rs`)
- `.env` (cwd) is loaded with a custom parser, not via dotenv crate; env var placeholders `<VAR_NAME>` in config JSON are substituted
- Config example: `resources/templates/config.example.json` (released to `~/.picobot/` on first run)
- CLI TUI identity is stored in `~/.picobot/tui_client_id`; it is a non-secret stable chat scope used to restore dialogs across reconnects
## Tests
- `cargo test --lib` — run unit tests (runs all `#[test]` in `src/`)
- `cargo test --test test_integration -- --ignored` — run integration tests (also `test_tool_calling`, `test_request_format`)
- **All** integration tests require `tests/test.env` with real API keys; copy from `tests/test.env.example` and fill in keys
- Integration tests are `#[ignore]` by default; use `-- --ignored` to run them
- `cargo clippy --all-targets --all-features -- -D warnings` — required for Rust changes
- `cargo test --test test_scheduler` and `cargo test --test test_request_format` — offline integration/protocol tests
- `cargo test --test test_integration -- --ignored` and `cargo test --test test_tool_calling -- --ignored` — model API integration tests
- API-backed tests require `tests/test.env` with real API keys; copy from `tests/test.env.example` and fill in keys
- API-backed tests are `#[ignore]` by default; use `-- --ignored` to run them
## Reference
- `reference/` — third-party reference implementations (nanobot, Mini-Agent, zeroclaw); not part of this project; do not modify
- `reference/` — third-party reference implementations; not part of this project; do not modify
## Architecture
@ -33,9 +38,12 @@
### Core Data Flow
```
Channel → MessageBus → SessionManager → AgentLoop → (tools) → SessionManager → MessageBus → OutboundDispatcher → Channel
↑
ControlChannel ──→ SessionManager (dialog ops: create/switch/archive/delete)
Channel → MessageBus.inbound → Gateway processor → SessionManager → per-session worker → AgentLoop
↑ │
└── Channel ← OutboundDispatcher ← per-conversation lane ← MessageBus.outbound
WebSocket/Channel → MessageBus.control → Gateway processor → SessionManager (dialog operations)
Scheduler → SessionManager.handle_cron_message → AgentLoop → send_message tool
```
### Modules
@ -45,8 +53,8 @@ Channel → MessageBus → SessionManager → AgentLoop → (tools) → SessionM
| `gateway` | Server lifecycle, HTTP/WS endpoints, owns `GatewayState` | `GatewayState`, `run()` |
| `client` | TUI rendering, WebSocket client for CLI chat | `App`, `run()` |
| `channels` | External integrations (Feishu, CLI chat) | `ChannelManager`, `Channel` trait |
| `bus` | Async message queue (inbound/outbound/control channels) | `MessageBus`, `InboundMessage`, `OutboundMessage`, `ControlMessage` |
| `session` | Conversation session lifecycle, dialog operations | `SessionManager`, `Session` |
| `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()` |
| `tools` | Agent tools (bash, file ops, http, web, get_skill) | `ToolRegistry`, `Tool` trait |
@ -54,25 +62,54 @@ Channel → MessageBus → SessionManager → AgentLoop → (tools) → SessionM
| `storage` | SQLite persistence for sessions and messages | `Storage`, `SessionMeta`, `MessageMeta` |
| `scheduler` | Cron-based job scheduling, next-run computation | `Scheduler`, `Schedule`, `next_run_for_schedule()` |
| `observability` | Observer pattern for agent/tool telemetry events | `Observer` trait, `ObserverEvent`, `MultiObserver` |
| `protocol` | WebSocket protocol message types | `WsInbound`, `WsOutbound`, `SessionSummary` |
| `protocol` | WebSocket protocol message types | `WsInbound`, `WsOutbound`, `SessionSummary`, `HistoryMessage` |
| `config` | Config loading, env substitution, path resolution | `Config`, `LLMProviderConfig` |
| `logging` | Tracing initialization with file rotation | `init_logging()`, `init_logging_console_only()` |
| `task_supervisor` | Owns, cancels, and boundedly joins gateway background tasks | `TaskSupervisor` |
### Functional Boundaries
- **Channels** only send/receive messages via `MessageBus`; they know nothing about sessions or LLM
- **MessageBus** is a pure async queue; it routes nothing, just passes messages
- **SessionManager** owns session state and dialog operations; it does NOT call LLM directly
- SessionManager is responsible for injecting skills prompt into conversation history
- **AgentLoop** receives dialog events from `SessionManager`, calls LLM via `providers`, executes tools, returns text responses
- AgentLoop is stateless; all state is managed by Session/SessionManager
- **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
- **AgentLoop** is stateless across turns; it receives prepared history, calls LLM providers, executes tools, and returns one result
- **Providers** are pure HTTP clients; no bus/session/channel awareness
- **Tools** are executed by `AgentLoop`; they receive raw arguments and return string results
### 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
- 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
- Long-lived gateway tasks must be owned by `TaskSupervisor`; connection-local tasks must be explicitly joined or aborted by their owner
- Connection, retry sleep, queue wait, and shutdown join paths must observe cancellation and have hard time bounds
- Retry only errors explicitly classified as transient; permanent failures must surface immediately
- Never log secrets, authorization headers, or full connection URLs containing temporary credentials
- The workspace cwd is not a filesystem sandbox: default file tools accept absolute paths and Bash inherits process permissions; add explicit canonical-path/process isolation when a hard boundary is required
### Key Constraints
- Gateway **changes working directory** to workspace on startup (`src/gateway/mod.rs:31`)
- Gateway **changes working directory** to workspace in `GatewayState::new` (`src/gateway/mod.rs`)
- Session/message persistence uses SQLite via `sqlx`; DB stored in workspace as `picobot.db` by default
- `ChannelManager` owns the `MessageBus` and all channel instances
- `OutboundDispatcher` routes outbound messages to the correct channel via `ChannelManager`
- Config `.env` loading uses `unsafe { env::set_var(...) }` — don't refactor to safer patterns without understanding side effects
## Change Workflow
1. Inspect the relevant implementation, tests, and `docs/ARCHITECTURE.md`; do not rely on names or old reports alone.
2. Search `reference/` only for comparison. Never edit it or copy behavior without checking PicoBot's boundaries.
3. Preserve unrelated user changes in a dirty worktree. Use `rg` for search and `apply_patch` for edits.
4. Add regression tests for bugs, especially cancellation, timeout, queue saturation, stale state, persistence failure, and retry classification.
5. For Rust changes run targeted tests, `cargo test --lib`, Clippy with warnings denied, and `cargo build`. Integration tests require real credentials.
6. For documentation-only changes verify links, commands, paths, and `git diff --check`.
7. Update README, this file, and the architecture document together when public behavior or an architectural invariant changes.
## Documentation Roles
- `README.md` — user-facing overview, setup, capabilities, and navigation
- `docs/ARCHITECTURE.md` — maintainer-facing runtime design, invariants, lifecycle, and extension guidance
- `AGENTS.md` — concise operational rules for repository agents
- `resources/skills/about-picobot/references/` — runtime knowledge shipped to PicoBot; update it only when the assistant's built-in product knowledge must change

7
Cargo.toml
View File

@ -5,7 +5,6 @@ edition = "2024"
[dependencies]
reqwest = { version = "0.13.3", default-features = false, features = ["json", "rustls", "multipart"] }
dotenv = "0.15"
serde = { version = "1.0", features = ["derive"] }
regex = "1.12"
serde_json = "1.0"
@ -30,13 +29,12 @@ base64 = "0.22"
tempfile = "3"
cron = "0.16"
chrono-tz = "0.10"
meval = "0.2"
ratatui = "0.30"
crossterm = { version = "0.29", features = ["event-stream"] }
termimad = "0.34"
textwrap = "0.16"
unicode-width = "0.2"
chrono = "0.4"
hostname = "0.4"
sqlx = { version = "0.8", features = ["sqlite", "macros", "chrono", "runtime-tokio"] }
jieba-rs = "0.9"
which = "8"
@ -53,6 +51,9 @@ tar = "0.4"
fantoccini = { version = "0.22", default-features = false, features = ["rustls-tls"] }
portable-pty = "0.9"
[dev-dependencies]
dotenv = "0.15"
[build-dependencies]
zstd = "0.13"
tar = "0.4"

45
README.md
View File

@ -4,7 +4,7 @@ PicoBot 是一个用 Rust 编写的个人 AI 助手运行时。它在本地启
它更像一个可扩展的“个人助手操作系统”：渠道负责收发消息，SessionManager 负责会话和上下文，AgentLoop 负责模型与工具循环，Storage 负责可靠落盘。
![PicoBot runtime architecture](docs/assets/runtime-architecture.svg)
完整的组件边界、并发不变量、启动/关停顺序和扩展指南见 [架构文档](docs/ARCHITECTURE.md)。
## 适合做什么
@ -91,11 +91,28 @@ cargo run -- chat
CLI 默认连接 `ws://127.0.0.1:19876/ws`。如需指定地址，可使用 `--gateway-url`。
TUI 会把一个随机客户端标识保存到 `~/.picobot/tui_client_id`，因此关闭并重新打开客户端后会恢复同一组 dialog 和最近使用的会话。界面支持历史回放、会话列表与归档筛选、命令补全、Unicode/中文编辑、括号粘贴和多行输入。
常用快捷键：
| 快捷键 | 操作 |
|--------|------|
| `F1` / `Ctrl+H` | 打开帮助 |
| `Tab` / `Ctrl+S` | 切换焦点 / 聚焦会话列表 |
| `Ctrl+N` | 新建会话 |
| `Ctrl+R` / `Ctrl+A` / `Ctrl+D` | 重命名 / 归档 / 删除所选会话 |
| `Ctrl+L` / `Ctrl+O` | 清空历史 / 显示归档会话 |
| `Enter` / `Shift+Enter` | 发送 / 换行 |
| `PageUp` / `PageDown` | 滚动对话历史 |
| 连按两次 `Ctrl+C` | 退出客户端 |
## 运行时数据流
用户消息进入 PicoBot 后，会被转换为统一的 inbound message，经由 MessageBus 交给 SessionManager。SessionManager 选择当前 dialog、组装上下文、调用 AgentLoop；AgentLoop 调用模型和工具，最终响应通过 outbound bus 回到原渠道。
![Message flow](docs/assets/message-flow.svg)
详细时序和失败语义见 [架构文档：消息与控制数据流](docs/ARCHITECTURE.md#4-消息与控制数据流)。
同一 session 的普通消息由专属有界队列串行处理，不同 session 可以并发；出站消息按 `(channel, chat_id)` 分 lane 保序，慢渠道不会阻塞其他目标。Gateway 的长生命周期任务统一由 `TaskSupervisor` 取消和限时回收。
核心边界：
@ -111,6 +128,7 @@ CLI 默认连接 `ws://127.0.0.1:19876/ws`。如需指定地址，可使用 `--g
| `scheduler` | 轮询 Cron 任务并把任务 prompt 送入目标会话 |
| `skills` | 加载 Skill，并把 Skill 指南注入系统提示 |
| `mcp` | 连接 MCP Server，将远端工具包装成普通 Tool |
| `task_supervisor` | 统一管理 Gateway 后台任务的取消和有界关停 |
## 核心能力
@ -156,7 +174,7 @@ PicoBot 有两类记忆：
| Knowledge | 偏好、事实、项目规则、长期可复用信息 | 长期保留，手动删除 |
| Timeline | 长对话压缩后的历史摘要 | 默认保留 90 天 |
每轮处理用户消息时，MemoryManager 会按用户输入召回最多 `memory.recall_limit` 条 Knowledge，并注入系统提示。上下文压缩产生的摘要会保存为 Timeline，后续可通过 `timeline_recall` 工具检索。
每轮处理用户消息时，MemoryManager 会按用户输入召回 Knowledge，并作为运行时上下文附加到本轮用户消息。当前召回上限固定为 5；`memory.recall_limit` 已支持解析但尚未接入 worker。上下文压缩产生的摘要会保存为 Timeline，后续可通过 `timeline_recall` 工具检索。
### 工具
@ -214,7 +232,7 @@ Skill 是包含 `SKILL.md` 的目录。加载优先级从高到低：
| `gateway.max_concurrent_background_tasks` | `10` |
| `gateway.scheduler.enabled` | `true` |
| `client.gateway_url` | `ws://127.0.0.1:19876/ws` |
| `memory.recall_limit` | `5` |
| `memory.recall_limit` | `5`（当前运行时固定为 5） |
| `memory.timeline_retention_days` | `90` |
| `mcp.tool_timeout_secs` | `180` |
| `browser.enabled` | `false` |
@ -239,13 +257,14 @@ Inbound 消息类型：
| `create_session` | 可选 `title` |
| `list_sessions` | `include_archived` |
| `load_session` | `session_id` |
| `get_session_history` | `session_id`，可选 `limit`（服务端限制为 1–2000） |
| `rename_session` | 可选 `session_id`，`title` |
| `archive_session` | 可选 `session_id` |
| `delete_session` | 可选 `session_id` |
| `get_slash_commands` | 无 |
| `ping` | 无 |
Outbound 消息类型包括 `assistant_response`、`error`、`session_established`、`session_created`、`session_list`、`session_loaded`、`session_renamed`、`session_archived`、`session_deleted`、`history_cleared`、`slash_commands_list`、`pong`、`command_executed` 和 `system_notification`。
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。
## 测试
@ -253,14 +272,17 @@ Outbound 消息类型包括 `assistant_response`、`error`、`session_establishe
# 单元测试
cargo test --lib
# 集成测试需要 tests/test.env 中有真实 API key
# 离线集成/协议测试
cargo test --test test_scheduler
cargo test --test test_request_format
# 模型 API 集成测试需要 tests/test.env 中有真实 API key
cp tests/test.env.example tests/test.env
cargo test --test test_integration -- --ignored
cargo test --test test_tool_calling -- --ignored
cargo test --test test_request_format -- --ignored
```
集成测试默认 `#[ignore]`，因为它们会真实调用模型 API。
会真实调用模型 API 的测试标记为 `#[ignore]`；离线集成测试默认执行。
## 项目结构
@ -281,11 +303,12 @@ src/
skills/ Skill 加载和内置 Skill 安装
storage/ SQLite schema 和 CRUD
tools/ Agent 工具实现
task_supervisor.rs Gateway 后台任务的生命周期管理
resources/
skills/ 构建时嵌入的内置 Skills
templates/ 首次运行释放的配置和用户模板
tests/ 单元测试和 ignored 集成测试
docs/ 分析报告、文档插图和补充资料
docs/ 面向维护者和 Agent 的架构与开发文档
```
## 关键依赖
@ -304,9 +327,9 @@ docs/ 分析报告、文档插图和补充资料
## 进一步阅读
- [架构机制](resources/skills/about-picobot/references/architecture.md)
- [维护者架构文档](docs/ARCHITECTURE.md)
- [内置 Skill：架构机制](resources/skills/about-picobot/references/architecture.md)
- [配置说明](resources/skills/about-picobot/references/config.md)
- [命令说明](resources/skills/about-picobot/references/commands.md)
- [工具说明](resources/skills/about-picobot/references/tools.md)
- [数据库结构](resources/skills/about-picobot/references/db-schema.md)
- [代码质量分析](docs/CODE_QUALITY_ANALYSIS.md)

287
docs/ARCHITECTURE.md Normal file
View File

@ -0,0 +1,287 @@
# PicoBot 架构
本文档描述 PicoBot 当前实现的运行时边界、数据流、并发模型和演进约束。它面向维护者和后续参与改进的 Agent，是代码架构的主入口；行为细节仍以代码和测试为最终依据。
## 1. 设计目标
PicoBot 是一个单进程、异步、可扩展的个人 AI 助手运行时。核心目标是：
- 用统一消息模型接入不同聊天渠道。
- 隔离渠道、会话、模型、工具和持久化职责。
- 同一会话内保持消息顺序，不同会话之间允许并发。
- 让外部 I/O、后台任务和程序关停都有明确边界。
- 通过 SQLite 保存会话、消息、记忆、定时任务及后台任务状态。
当前不是分布式系统。除调度任务使用数据库租约防止重复领取外，运行时会话状态由单个 Gateway 进程持有。
## 2. 运行模式与进程边界
PicoBot 只有一个二进制，提供两种模式：
| 模式 | 入口 | 职责 |
|------|------|------|
| Gateway | `cargo run -- gateway` | 组装服务、监听 HTTP/WebSocket、运行渠道、会话、调度器和后台任务 |
| CLI client | `cargo run -- chat` | 运行 Ratatui UI，通过 WebSocket 使用 Gateway，不持有业务状态 |
CLI TUI 在 `~/.picobot/tui_client_id` 保存非敏感客户端标识，并通过 WebSocket 查询参数 `client_id` 传给 Gateway。`cli_chat` 以该标识作为稳定 chat scope；重连时恢复内存中的当前 dialog，Gateway 重启后则恢复该 scope 最近活跃的未归档 dialog。无效或缺失的标识会退化为连接级随机 scope。
Gateway 启动时会切换进程工作目录到 `workspace_dir`。因此所有相对路径都应按 workspace 解释，不能假设仍位于源码仓库。
## 3. 组件关系
```mermaid
flowchart LR
External[CLI / Feishu] --> Channels[channels]
Channels -->|InboundMessage| Bus[MessageBus]
Bus --> Processor[Gateway message processor]
Processor --> Sessions[SessionManager]
Sessions --> Agent[AgentLoop]
Agent --> Providers[LLM providers]
Agent --> Tools[ToolRegistry / MCP]
Sessions <--> Storage[(SQLite)]
Sessions -->|OutboundMessage| Bus
Scheduler[Scheduler] --> Sessions
Bus --> Dispatcher[OutboundDispatcher]
Dispatcher --> Channels
Supervisor[TaskSupervisor] -. lifecycle .-> Processor
Supervisor -. lifecycle .-> Dispatcher
Supervisor -. lifecycle .-> Scheduler
Supervisor -. lifecycle .-> Sessions
```
### 模块职责
| 模块 | 拥有的职责 | 不应承担的职责 |
|------|------------|----------------|
| `gateway` | 依赖装配、HTTP/WS 入口、启动和关停顺序 | 业务规则、渠道协议细节 |
| `channels` | 外部协议适配、权限检查、媒体收发 | 会话选择、LLM 调用 |
| `bus` | 三条有界异步队列与出站投递协调 | 会话路由、业务状态 |
| `session` | dialog 路由、会话状态、串行工作队列、上下文和持久化协调 | 外部渠道协议 |
| `agent` | 单次无状态模型/工具循环、上下文压缩、子 Agent | 持有 dialog 生命周期 |
| `providers` | 把统一请求映射到模型 API | Session、Bus 或 Channel 感知 |
| `tools` / `mcp` | 工具定义、注册和执行适配 | 隐式修改会话路由 |
| `storage` | SQLite schema、迁移、原子 CRUD | 运行时调度策略 |
| `memory` | Knowledge/Timeline 的存取与召回 | 直接驱动消息发送 |
| `scheduler` | 领取到期任务、限并发执行、原子记录结果 | 复用聊天会话历史 |
| `task_supervisor` | 后台任务注册、取消、限时回收 | 业务级重试和结果语义 |
## 4. 消息与控制数据流
`MessageBus` 包含三条容量相同的 Tokio MPSC 队列：
- `inbound`：Channel → Gateway message processor。
- `outbound`：Session/Tool → `OutboundDispatcher`。
- `control`：WebSocket/Channel → Gateway message processor，用于 dialog 操作。
### 普通消息
```mermaid
sequenceDiagram
participant C as Channel
participant B as MessageBus
participant G as Message processor
participant S as SessionManager
participant W as Per-session worker
participant A as AgentLoop
participant D as OutboundDispatcher
C->>B: publish InboundMessage
B->>G: consume inbound
G->>S: handle_message
S->>W: try_send AgentTask
S-->>G: AgentProcessing
W->>A: process(history)
A-->>W: final response
W->>B: publish OutboundMessage
B->>D: consume outbound
D->>C: Channel::send
```
关键语义：
- Gateway 的主消息处理循环不等待模型完成；普通消息进入对应 session worker 后立即返回 `AgentProcessing`。
- 每个 session 有一条容量为 32 的队列，同一 session 串行处理，不同 session 的 worker 可并发执行。
- 队列满时明确拒绝新消息，不允许无界积压。
- Slash command 不进入 Agent 队列，由 `SessionManager` 直接执行，因此 `/stop` 等控制操作不会排在长模型调用之后。
### 出站投递
`OutboundDispatcher` 按 `(channel, chat_id)` 建立独立 lane：
- 同一目标的消息保持顺序。
- 慢目标不会阻塞其他目标。
- 每条 lane 容量为 64，空闲 300 秒后退出。
- 单次发送超时 30 秒；最多尝试 3 次，前两次失败后分别等待 1/2 秒。只有 `ConnectionError` 和 `SendError` 会重试。
- `deliver_outbound` 可等待渠道真实投递结果，等待上限 120 秒；普通 `publish_outbound` 只保证成功入队。
不要把“已进入 Bus”误认为“外部渠道已收到”。需要确认语义时必须使用 `deliver_outbound`。
### Control 消息
WebSocket dialog 操作通过 `ControlMessage` 携带一次性回复通道。Gateway 在统一 message processor 中调用 `SessionManager`，再将 `SessionEvent` 回传给发起者。Bus 只承载消息，不解释操作。
TUI 的历史回放同样走 control 队列：`get_session_history` 先校验 session 属于当前客户端 scope，再由 SessionManager 从 Storage 读取最近消息。单次查询限制为 1–2000 条，TUI 默认请求最近 1000 条；迟到的历史响应只有在目标仍是当前 dialog 时才允许更新界面。
Agent worker 发出的异步回复和通知在 OutboundMessage metadata 中标记来源 session，`cli_chat` 将其映射为 WebSocket `session_id`。TUI 切换 dialog 后不渲染其他 session 的迟到结果；结果仍按原 session 持久化，切回时通过历史回放显示。
## 5. 会话模型与并发不变量
Session ID 格式为：
```text
<channel>:<chat_id>:<dialog_id>
```
`SessionManagerInner` 保存：
- `sessions`：完整 Session ID 到内存 Session 的映射。
- `current_sessions`：`channel:chat_id` 到当前 dialog 的映射。
必须维护以下不变量：
1. 同一 Session 的 Agent 工作由一个 generation 对应的 worker 串行执行。
2. `/stop` 或 worker 替换会递增 `worker_generation`；旧 worker 不得再提交结果。
3. 慢操作（模型、记忆召回、压缩、SQLite I/O）不能长期持有 Session mutex。
4. 慢操作开始前记录 `state_version`，提交前重新验证，防止旧快照覆盖 `/clear`、`/delete` 等并发修改。
5. 持久化写入由 `persistence_lock` 串行化；多条相关记录应使用 Storage 的原子接口。
6. 内存先变更但持久化失败时，必须回滚精确匹配的消息后缀，不能删除无关的新状态。
SessionManager 负责组装会话上下文：系统提示、Skills、召回的 Knowledge、压缩后的 Timeline 和当前消息历史。`AgentLoop` 接收完整输入执行一次模型/工具循环，本身不拥有会话状态。
## 6. 持久化
`Storage` 使用 SQLx + SQLite，默认数据库为 `{workspace_dir}/picobot.db`。连接启用：
- WAL journal mode。
- foreign keys。
- 5 秒 busy timeout。
- schema version 迁移。
持久化范围包括 sessions、messages、memories、scheduled jobs、job runs 和 background tasks。修改 schema 时应：
1. 更新集中式 schema/迁移逻辑。
2. 保留已有数据库的升级路径。
3. 为新库初始化和旧库迁移分别增加测试。
4. 对“状态更新 + 执行记录”等复合写入使用事务。
### 安全边界
- API Key 和渠道凭据只来自配置占位符、`.env` 或进程环境，不得写入仓库。
- 日志不得输出 token、secret、Authorization header，或包含临时凭据的完整 URL；应记录脱敏后的 host/path 和必要诊断字段。
- Gateway 把 cwd 切到 workspace，因此相对文件路径和 Shell 默认从 workspace 开始；这不是硬沙箱。当前内置文件工具接受绝对路径，Bash 也可访问进程权限允许的位置。若某场景需要硬边界，必须显式配置/实现 allowed directory 和进程隔离。
- `http_request` 和 `web_fetch` 的私网/回环地址校验属于 SSRF 防线，重构网络层时不能绕过。
- 外部内容、Tool 输出和 MCP 响应均是不可信输入；解析错误应返回结构化失败，不能 panic。
## 7. 后台任务与生命周期
`TaskSupervisor` 是 Gateway 内部后台任务的统一所有者。message processor、outbound dispatcher、scheduler、session workers、outbound lanes、通知消费者和子 Agent 后台任务都应通过它注册。
两种注册方式：
- `spawn`：收到全局取消后直接丢弃任务 future，适合无需异步清理的任务。
- `spawn_graceful`：任务自己观察 cancellation token 并清理，Supervisor 在总宽限期结束后再强制 abort。
新增长生命周期任务时必须满足：
- 有明确 owner，禁止无法回收的裸 `tokio::spawn`。
- 能响应取消；外部连接、重试 sleep 和阻塞式等待也要纳入取消分支。
- 等待任务退出必须有硬超时，超时后 abort 并回收 JoinHandle。
- 任务 panic、超时和永久错误要可观测，且不能阻止其他组件清理。
WebSocket 每个客户端的 writer task 是连接局部任务，由连接 handler 自己限时回收；它不跨越连接生命周期。
## 8. 启动与关停顺序
### 启动
1. 加载配置和 `.env`，解析 workspace。
2. 创建并切换到 workspace。
3. 初始化 SQLite、MemoryManager、MessageBus 和 SessionManager。
4. 注册内置工具、渠道、MCP 工具和 Cron 工具。
5. 启动所有 Channel。
6. 通过 TaskSupervisor 启动 message processor、dispatcher 和 scheduler。
7. 绑定 Axum listener，开始接收请求。
### 关停
1. `Ctrl-C` 触发 Axum graceful shutdown，并取消所有 WebSocket 连接。
2. `ChannelManager::stop_all` 先停止外部消息入口并注销渠道。
3. 取消 TaskSupervisor，停止接受新后台任务。
4. 在共享的 10 秒总宽限期内等待任务退出，之后 abort 剩余任务。
渠道自己的 `stop()` 也必须有界。以飞书为例，端点请求、WebSocket 建连、重试等待和已连接循环共享 CancellationToken，另有 5 秒强制回收兜底。
## 9. 扩展指南
### 新增 Channel
1. 实现 `Channel` trait，仅处理外部协议和统一消息转换。
2. 在 `ChannelManager::init` 注册，并通过同一个 MessageBus 收发。
3. 将可重试错误表示为 `ConnectionError`/`SendError`，永久错误使用其他类型。
4. 为 start/stop 幂等性、取消建连、投递失败和媒体边界增加测试。
5. 不要从 Channel 直接调用 SessionManager 或 Provider。
### 新增 Tool
1. 实现 `Tool`，在集中注册点加入 `ToolRegistry`。
2. 参数 schema 必须明确；返回值保持可供模型消费的字符串协议。
3. 明确工具需要“workspace 默认目录”还是“不可逃逸的硬边界”；后者必须显式校验 canonical path，不能只依赖 cwd。
4. 网络工具必须保留 SSRF/私网地址校验。
5. 长操作应有超时；后台执行应交给 SubAgentManager/TaskSupervisor。
### 新增 Provider
1. 实现 `LLMProvider`，保持其为纯 HTTP/API 适配器。
2. 统一映射文本、媒体、tool calls、usage 和错误。
3. 不在 Provider 中访问 Session、Bus 或 Channel。
4. 为请求序列化和响应兼容性编写离线单元测试。
### 修改 Session
先画出锁、慢 I/O、版本检查和持久化顺序。任何跨 `await` 的 Session 锁都需要特别审查；任何由旧快照产生的结果都必须在提交前验证 generation/state version。
## 10. 验证策略
按改动范围选择最小但充分的验证：
| 改动 | 至少执行 |
|------|----------|
| 文档 | 检查链接、命令和源码路径；`git diff --check` |
| Rust 实现 | 相关定向测试、`cargo test --lib`、`cargo clippy --all-targets --all-features -- -D warnings` |
| 构建/依赖 | 上述检查加 `cargo build` |
| Provider/真实渠道 | 离线测试；有凭据时再运行 ignored integration tests |
| SQLite schema | 新库测试、迁移测试、原子性测试 |
| 生命周期/并发 | 成功、失败、超时、取消、队列满和重复 start/stop 测试 |
模型 API 集成测试需要 `tests/test.env` 中的真实 API Key，并默认 ignored；`test_scheduler` 和 `test_request_format` 是可直接运行的离线测试。不应在无凭据时假装已经验证真实 Provider。
## 11. 演进决策清单
提交架构性修改前，逐项确认：
- 是否仍遵守 Channel → Bus → Session → Agent 的依赖方向？
- 是否引入第二个 owner、重复状态或绕过统一注册点？
- 队列、并发、重试和等待是否全部有界？
- 取消信号能否覆盖建连、sleep、I/O 和清理阶段？
- 是否在持锁时执行了网络、模型或数据库慢操作？
- 内存与 SQLite 失败时能否保持一致？
- 错误是否区分瞬态和永久语义？
- 是否有测试覆盖正常路径与最危险的失败路径？
- README、AGENTS.md 和本文档是否需要同步更新？
## 12. 代码导航
| 主题 | 入口 |
|------|------|
| Gateway 装配和关停 | `src/gateway/mod.rs` |
| WebSocket 生命周期 | `src/gateway/ws.rs` |
| Bus 和消息类型 | `src/bus/mod.rs`, `src/bus/message.rs` |
| 出站并发与重试 | `src/bus/dispatcher.rs` |
| Channel 接口与注册 | `src/channels/base.rs`, `src/channels/manager.rs` |
| Session 核心 | `src/session/session.rs` |
| Session 命令/事件 | `src/session/commands.rs`, `src/session/events.rs` |
| Agent loop | `src/agent/agent_loop.rs` |
| 后台任务监督 | `src/task_supervisor.rs` |
| SQLite 初始化和迁移 | `src/storage/mod.rs` |
| Scheduler | `src/scheduler/mod.rs` |
| 配置加载 | `src/config/mod.rs` |

369
docs/CODE_QUALITY_ANALYSIS.md
View File

@ -1,369 +0,0 @@
# PicoBot 代码质量分析报告
审查日期：2026-06-15
## 结论摘要
PicoBot 的总体架构方向是清晰的：Gateway 负责装配，Channel 只做收发，MessageBus 解耦输入输出，SessionManager 管理会话，AgentLoop 保持无状态并执行工具，Storage 统一持久化。这条主线是成立的，也已经具备较完整的 AI 助手运行时能力。
当前主要质量风险集中在三类：
1. 会话/CLI 路由语义不一致，导致多客户端隔离、加载会话、当前会话追踪不可靠。
2. 若干公开控制接口是空实现或弱实现，协议层暴露的能力和后端实际行为不匹配。
3. 工具和后台任务的资源边界偏弱，文件、shell、HTTP、长期任务在异常情况下容易突破预期的安全或稳定性边界。
如果只安排一轮修复，优先处理会话路由和控制接口。这些问题会直接影响用户看到的行为；工具安全和大模块拆分可以作为第二阶段。
## 修复状态
- 已修复：CLI 会话路由现在按每个 WebSocket client 的稳定 `chat_id` 隔离，普通输入、创建、列表、加载和 outbound 投递不再混用完整 `session_id` 与 `chat_id`。
- 已修复：Dialog 控制接口已补齐当前会话查询、列表 current 标记、归档、清空历史和 `/delete` 删除当前会话后新建的行为；`include_archived` 现在由 Storage 查询生效。
- 已修复：Session 主处理路径不再在持有 session mutex 时执行 memory recall、上下文压缩、标题 LLM 生成、消息持久化、`/stop` sub-agent 取消或清历史存储操作；慢操作改为锁外执行并用 `state_version`/`worker_generation` 防止陈旧结果覆盖当前会话。
- 已修复：Bash 超时清理、文件读取大文件限制、HTTP DNS 私网校验、Bus 关闭退出、Cron `from` 语义和 PTY 工具接入等中等级问题已完成清扫。
- 待处理：工具文件边界仍是后续质量风险。
## 主要发现
### 已修复：CLI 会话路由会破坏会话连续性和多客户端隔离
位置：
- `src/channels/cli_chat.rs:113-126`
- `src/channels/cli_chat.rs:160-164`
- `src/channels/cli_chat.rs:225-249`
- `src/channels/cli_chat.rs:479-494`
- `src/session/session.rs:1305-1310`
问题：
`Client.current_session_id` 存的是完整 session id，但 CLI channel 在多个地方把它当作 `chat_id` 使用。普通用户输入如果没有显式传 `chat_id`，会在 `src/channels/cli_chat.rs:119` 生成新的短 ID，而不是复用当前 client 的 chat scope。`CreateSession` 又把当前完整 session id 当成新会话的 chat_id。`LoadSession` 解析了传入 session id，但随后调用 `GetCurrentDialog`，而后端 `get_current_dialog()` 固定返回 `None`。
同时，`send()` 会把所有 `OutboundMessage` 广播给所有 CLI WebSocket client，没有按 `msg.chat_id` 或 client 当前会话过滤。这意味着一个客户端的回复可能出现在另一个客户端里。
影响：
- CLI 多轮对话可能落入不同 chat scope。
- 创建/列出/加载会话得到的结果可能不符合 UI 预期。
- 多个 CLI 客户端同时连接时存在串话。
建议：
- 将 client 状态拆成 `chat_id` 和 `current_session_id`，不要混用。
- 注册 client 时生成稳定 `chat_id`，后续 `UserInput` 默认复用它。
- `send()` 按 `OutboundMessage.chat_id` 精确投递；必要时维护 `chat_id -> clients` 映射。
- `LoadSession` 应直接切换到指定 session，或通过 `SwitchDialog` 使用其中的 `dialog_id`。
- 为 CLI WebSocket 增加多客户端路由测试。
### 已修复：Dialog 控制接口与协议承诺不一致
位置：
- `src/session/session.rs:996-997`
- `src/session/session.rs:1305-1310`
- `src/session/session.rs:1329-1349`
- `src/session/session.rs:1378-1384`
- `src/channels/cli_chat.rs:128-158`
问题：
后端暴露了 create/list/load/rename/archive/delete/clear 等 dialog 操作，但部分行为是空实现或语义错位：
- `/delete` 只创建新 session，并没有删除当前 session。
- `get_current_dialog()` 固定返回 `Ok(None)`。
- `list_dialogs()` 忽略 `include_archived`，且总是返回 `current_dialog_id = None`。
- `archive_dialog()` 是空操作。
- `clear_dialog_history()` 直接返回不可用，但 WebSocket 协议仍暴露 `clear_history`。
影响：
用户通过 slash command 和 WebSocket 调用同一类能力时，会得到不一致结果。前端难以基于协议实现可靠状态同步。
建议：
- 明确“archive/clear 是否支持”。不支持就从协议和命令列表移除；支持就实现到底。
- `/delete` 应调用 `delete_dialog(current_session_id)`，再创建一个新的 current session。
- `get_current_dialog()` 应读取 `current_sessions[channel:chat_id]` 并解析为 `UnifiedSessionId`。
- `list_dialogs()` 返回真实 current dialog，并补上 archived 模型或移除 archived 参数。
### 高优先级：工具文件边界不符合“工作目录内工具”的架构约束
位置：
- `src/tools/mod.rs:56-62`
- `src/tools/path_utils.rs:3-23`
- `src/tools/bash.rs:146-185`
问题：
文件工具默认通过 `FileReadTool::new()`、`FileWriteTool::new()` 等注册，没有传入 workspace allowlist。`resolve_path()` 对绝对路径直接放行；即使传入 allowlist，也只是做 `Path::starts_with()` 的词法判断，没有 canonicalize，不能防御 `..`、符号链接等路径逃逸。
`bash` 默认工作目录是 `"."`，Gateway 启动时切到 workspace，这对相对路径有效，但 shell 命令仍然可以访问绝对路径。当前 denylist 只挡少数危险模式，不构成权限边界。
影响：
Agent 工具实际可以读写 workspace 外文件，和文档/架构里的“工作目录内操作”不一致。对于个人助手这可能是有意设计，但如果未来接入外部渠道、多用户或 MCP，风险会放大。
建议：
- 工具注册时传入 `workspace_dir`，默认所有文件工具限制在 workspace。
- `resolve_path()` 使用 `std::fs::canonicalize` 或 `path_absolutize` 风格逻辑，并处理目标文件不存在时的父目录 canonicalize。
- 写工具禁止跟随危险符号链接，或至少在文档中明确该能力是全文件系统权限。
- shell 工具如果保留，应在配置中显式开关，并区分本地可信模式和渠道暴露模式。
### 已修复：Session 锁内执行过多异步操作
位置：
- `src/session/session.rs:1001-1018`
- `src/session/session.rs:1604-1711`
问题：
`/compact` 在持有 session mutex 时执行压缩和持久化。agent worker 的 Phase 1 也在持有 session mutex 时执行用户消息落库、memory recall、上下文压缩、session meta 持久化和 agent 创建。其中 `compress_if_needed()` 可能触发 LLM 摘要，属于慢操作。
影响：
- 同一 session 的 slash command、stop、消息排队、状态查询会被慢操作阻塞。
- 当压缩或存储出现抖动时，用户感觉像“卡死”。
- 后续如果在这些慢操作里间接需要 session 状态，容易形成锁顺序问题。
已采取修复：
- 为 `Session` 增加 `state_version`，慢操作提交前检查会话是否已被 `/stop`、清历史或其它内存变更替换。
- `/compact` 改为锁内取 history 快照，锁外压缩，锁内提交压缩结果，锁外持久化 meta。
- agent worker Phase 1 改为锁内只创建用户消息、agent、cancel handle 和 history 快照；memory recall 与 context compression 都在锁外执行。
- context overflow retry 的二次压缩移到锁外。
- 标题生成改为锁内取 prompt/provider 快照，锁外调用 LLM，锁内应用标题，锁外持久化。
- `add_message` 拆出内存更新和持久化快照，主消息路径在释放 session 锁后写入 SQLite。
- `/stop` 和清历史不再持有 session 锁等待 sub-agent 取消或 Storage 操作。
### 已修复：Bash 超时不会显式终止子进程
位置：
- `src/tools/bash.rs:150-174`
- `src/tools/bash.rs:180-207`
问题：
`timeout()` 包裹的是 `run_command()` future。超时后 future 被取消，但代码没有持有 child 句柄并显式 `kill()` / `wait()`。对于已经启动的长运行命令或子进程树，可能留下后台进程。
影响：
长任务、服务进程或卡住的 shell 命令会泄漏进程和资源，后续工具调用的行为也会变得不可预测。
已采取修复：
- Bash 一次性命令改用 `wait_with_output()`，避免 stdout/stderr 顺序读取造成 pipe 阻塞。
- 子进程启用 `kill_on_drop(true)`，超时后丢弃等待 future 时会清理 child。
- 新增大 stderr 输出测试，覆盖不会因为 stderr pipe 填满而卡住。
- 持久/交互式进程通过已接入的 PTY 工具承载。
### 已修复：文件读取对大二进制文件没有输出上限
位置：
- `src/tools/file_read.rs:121-131`
- `src/tools/file_read.rs:214-229`
问题：
`file_read` 先 `std::fs::read()` 读取整个文件。文本路径有 `MAX_CHARS` 截断，但二进制路径会完整 base64 编码后返回，没有大小限制。
影响：
读取大文件会造成内存膨胀、响应膨胀、上下文污染，甚至拖垮进程。
已采取修复：
- `file_read` 在读取前检查 metadata size，超过安全阈值直接拒绝。
- 二进制 inline base64 增加单独大小上限，超限只返回错误和文件信息。
- 含 NUL 字节内容按二进制处理，避免全 0 文件被 UTF-8 路径误判为文本。
- 增加大文件和大二进制文件测试。
### 已修复：HTTP 私网防护只检查字面 host，未做 DNS 解析校验
位置：
- `src/tools/http_request.rs:31-59`
问题：
`http_request` 阻止 localhost、私网 IP 字面量和 `.local`，但普通域名不会解析后检查最终 IP。DNS rebinding 或内网域名解析到私网地址时，当前校验拦不住。
影响：
如果该工具暴露给非完全可信输入，存在 SSRF 风险。
已采取修复：
- `http_request` 和 `web_fetch` 在发送请求前通过 DNS 解析 host，并拒绝解析到 loopback、private、link-local、multicast、unspecified 的地址。
- IPv6 unique-local 和 link-local 地址也纳入私网判定。
- 禁用 reqwest 自动重定向，避免跳转到未校验的内网地址。
- 增加端口解析和 IPv6 私网判断测试。
### 已修复：后台任务和主循环缺少监督与优雅关闭
位置：
- `src/bus/mod.rs:51-99`
- `src/gateway/mod.rs:187-244`
- `src/gateway/mod.rs:247-266`
问题：
Gateway 中多个长期任务通过 `tokio::spawn` 启动后没有保存 JoinHandle，也没有统一 cancellation token。MessageBus 的 `consume_*()` 在 channel 关闭时使用 `expect()` panic。
影响：
- 某个后台 loop 异常退出后，Gateway 不一定能发现。
- 关闭流程只能 stop channel，无法系统性停止 scheduler、dispatcher、agent workers、notification publishers。
- bus channel 关闭时更像崩溃，而不是可恢复状态。
已采取修复：
- `MessageBus::consume_inbound/consume_outbound/consume_control` 不再在 channel 关闭时 `expect()` panic，改为返回 `Option<T>`。
- Gateway message processor 在 inbound/control bus 关闭时记录 warning 并退出 loop。
- OutboundDispatcher 在 outbound bus 关闭时记录 warning 并退出 loop。
- 这不是完整 runtime supervisor，但已消除 bus 关闭导致的 panic 崩溃路径，为后续集中 JoinHandle 管理留出接口。
### 已修复：Cron 计算函数没有按入参 `from` 计算 cron 下一次时间
位置：
- `src/scheduler/mod.rs:18-40`
问题：
`next_run_for_schedule(schedule, from)` 的注释说基于 `from` 计算，但 cron 分支创建了 `from_dt` 后没有传给 `cron_schedule`，实际使用的是 `upcoming(Utc)` 或 `upcoming(tz)` 的当前时间。
影响：
单元测试或补偿调度传入历史/未来时间时，结果不符合函数契约。线上 reschedule 当前使用 now，影响较小，但函数语义是错的。
已采取修复：
- cron 分支改用 `cron_schedule.after(&from_dt).next()`。
- timezone 分支用 `from_dt.with_timezone(&tz)` 作为计算起点。
- 增加 UTC 和 Asia/Shanghai 固定时间输入测试。
### 已修复：存在未接入或半接入代码，增加维护噪音
位置：
- `src/tools/pty.rs`
- `src/tools/mod.rs:1-20`
- `src/tools/mod.rs:49-88`
问题：
仓库里有完整 `pty.rs`，但 `tools/mod.rs` 没有声明 `pub mod pty`，`create_default_tools()` 也没有注册 PTY 工具。类似情况会让文档、计划和实现状态难以判断。
影响：
维护者会误以为功能已上线。未来改动容易遗漏测试和注册路径。
已采取修复：
- `src/tools/pty.rs` 已接入 `tools/mod.rs`，导出 `PtyManager`/`PtyTool`。
- `create_default_tools()` 默认注册共享 `PtyManager` 的 `PtyTool`。
- 修复 PTY 原本因未编译暴露不出的借用问题。
## 架构评价
### 做得好的地方
- 模块分层方向清楚：Channel、Bus、Session、Agent、Provider、Tool、Storage 边界基本可理解。
- AgentLoop 设计为无状态，历史由 SessionManager 管理，这一点利于恢复、压缩和测试。
- Provider 抽象简单直接，OpenAI-compatible 与 Anthropic 的差异被限制在 provider 层。
- Storage 集中初始化 schema，便于部署单二进制应用。
- Skill、memory、MCP、delegate 这几条扩展线已经形成统一的 ToolRegistry 接入点。
### 主要架构债务
- SessionManager 承担过多职责：会话生命周期、命令解析、memory recall、压缩、agent worker、任务取消、send_message 目标解析都在一个 2000 行文件内。
- Channel 和 Session 对 chat_id/session_id/dialog_id 的边界没有类型保护，导致 CLI 层混用字符串。
- Tool 权限模型不够显式：工具是否能访问全文件系统、是否能联网、是否能修改状态主要靠工具自身约定。
- 后台任务生命周期分散：gateway loop、agent worker、notification publisher、scheduler、sub-agent task 各自 spawn，缺少统一管理。
## 模块级分析
### gateway
`GatewayState::new()` 是清晰的装配中心：配置、workspace、storage、memory、bus、session manager、channels、MCP、scheduler 都在这里接线。问题是启动后任务监督不足，且 scheduler 默认 `unwrap_or_default()` 会在省略 `gateway.scheduler` 时启用调度器，这和“省略配置是否代表开启”需要产品层确认。
### channels
Feishu channel 功能较厚，单文件接近 2000 行，建议后续按 API client、message parsing、media handling、outbound rendering 拆分。CLI channel 目前是质量风险最高的 channel，核心问题是会话身份混用和广播投递。
### bus
MessageBus 简洁，但当前消费者 API 通过 mutex 包住 receiver 并 `expect()`，更像“单消费者内部队列”。这没问题，但应该把“只能有一个 consumer”写进类型/文档，并把关闭作为正常状态处理。
### session
这是系统核心，也是债务最集中的模块。建议把 `session.rs` 拆成：
- `manager.rs`：SessionManager 状态和 dialog 生命周期
- `worker.rs`：per-session agent worker 和 cancellation
- `commands.rs`：slash command 执行
- `outbound.rs`：OutboundMessenger 实现
- `restore.rs`：storage 恢复与 tool call chain repair
拆分之前，先补行为测试，尤其是 CLI/WS session lifecycle。
### agent
AgentLoop 的职责相对聚焦：请求模型、执行工具、回填 tool result、循环直到 final response。需要关注的是工具并发的语义：`read_only()` 目前是工具自己声明，副作用工具不能错标。LoopDetector 有帮助，但属于 runtime guard，不应替代工具层的资源限制。
### providers
Provider 层整体可维护。OpenAI/Anthropic 的请求构造逻辑可以继续保留在 provider 内。建议补充请求脱敏策略：当前 debug log 和 `llm_calls` 会持久化完整 request/response，可能包含用户隐私、API 返回内容和文件内容。
### tools
工具体系覆盖面很强，但需要明确权限模型。建议新增统一的 `ToolExecutionContext`，包含 workspace、channel、session_id、权限策略、网络策略、输出预算。现在很多策略散落在各工具构造函数里，默认值容易失控。
### storage
Storage schema 初始化实用，但迁移方式是“CREATE IF NOT EXISTS + ALTER IGNORE”，适合早期迭代，不适合长期演进。建议引入 schema version 表或 sqlx migrations，至少把每次迁移记录下来。
### skills
Skill 加载优先级清晰，内置 skill 打包也实用。需要注意 `SkillsLoader` 使用同步文件系统扫描和 `std::sync::Mutex`，在请求路径频繁 `reload_if_changed()` 时可能造成阻塞。短期可以接受，长期建议缓存刷新放到后台 watcher。
## 建议修复路线
### P0：先修会话正确性
1. 修正 CLI `chat_id/current_session_id` 数据模型。
2. 修正 CLI 出站按 client/chat_id 投递。
3. 实现 `get_current_dialog()`、`list_dialogs()` current 返回。
4. 修正 `/delete`、`clear_history`、`archive` 的真实行为或从协议移除。
5. 增加 WebSocket session lifecycle 测试。
### P1：收紧工具和资源边界
1. 文件工具默认限制 workspace，路径 canonicalize。
2. bash 超时杀进程，必要时引入进程组。
3. file_read 增加文件大小上限和二进制输出上限。
4. HTTP/web 工具增加 DNS 解析后的私网校验和重定向校验。
5. 明确高危工具的配置开关。
### P2：降低架构复杂度
1. 拆分 `session.rs`、`feishu.rs`、`storage/mod.rs`、`browser.rs`。
2. 引入任务 supervisor 和统一 shutdown token。
3. 引入正式数据库迁移。
4. 增加工具注册快照测试，避免死代码和文档漂移。
## 建议测试补充
- CLI 多客户端并发：两个 WebSocket client 同时发消息，互不串话。
- CLI 不传 chat_id 的连续对话：所有消息应进入同一 session。
- Load/switch/list/delete/clear 的完整 WebSocket 流程。
- `/delete` 后旧 session 软删除、新 session 成为 current。
- 文件路径逃逸：`../`、绝对路径、符号链接、workspace 前缀欺骗。
- bash timeout 后检查子进程不存在。
- cron `next_run_for_schedule()` 使用固定 `from` 的 deterministic 测试。
- HTTP 工具对 DNS 解析到 `127.0.0.1` / `10.0.0.0/8` 的域名拒绝测试。

496
docs/INTERACTIVE_CHANNEL_DESIGN.md
View File

@ -1,496 +0,0 @@
# PicoBot 跨渠道交互式消息规划
规划日期：2026-06-16
## 背景
飞书交互式卡片可以让用户直接在消息卡片上点击按钮、提交选择或触发回调。这个能力很适合用于工具调用审批、快捷回复、任务确认、表单收集等 agent 交互。
参考项目调研结果：
- `reference/zeroclaw` 已实现飞书/Lark 工具审批卡片：发送 Card JSON 2.0 按钮卡片，收到 `card.action.trigger` 后解析 `approval_id` 和 `decision`，唤醒等待中的 approval future，并 PATCH 原卡片为已处理状态。
- `reference/nanobot` 主要使用飞书 CardKit 做 agent 输出展示和流式更新，适合参考消息渲染体验，但没有完整的按钮回调驱动 agent 流程。
- `reference/openlark` 是 SDK/API 封装，支持发送 interactive card 和 CardKit API，不包含完整 agent channel 编排。
PicoBot 当前飞书渠道已经会把普通 markdown 回复发送成 interactive card，但还缺少“用户在卡片上操作 -> 统一交互事件 -> Session/Agent/Tool 流程继续”的抽象。
## 目标
1. 支持飞书交互式卡片按钮回调。
2. 设计成跨渠道能力，后续 Slack、Telegram、Discord、CLI chat 等渠道可以复用同一套交互语义。
3. 支持渠道降级：不支持按钮的渠道也能用纯文本命令完成同样操作。
4. 保持 PicoBot 现有边界：Channel 只做收发和渠道适配，SessionManager 管会话，AgentLoop 执行 LLM 和工具。
5. 为工具调用审批、快捷回复和未来表单交互预留扩展点。
非目标：
- 本阶段不立即实现完整功能。
- 不把飞书卡片细节泄漏到 AgentLoop 或工具层。
- 不要求所有渠道同时支持原生交互组件。
## 核心原则
交互语义和渠道渲染分离。
Agent、工具或 Session 层只表达“我要一个 approval/quick reply/form interaction”。具体是飞书卡片按钮、Slack Block Kit、Telegram inline keyboard，还是 CLI 里显示编号选项，由 Channel 根据能力渲染。
回调也要统一。
飞书的 `card.action.trigger`、Telegram 的 `callback_query`、Slack 的 interaction payload 都应归一化成 PicoBot 内部的 `InteractionEvent`，再交给统一的处理器。
## 数据模型
建议新增一个 `interaction` 模块，定义渠道无关的数据结构。
```rust
#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
pub enum InteractionKind {
QuickReply,
Approval,
FormSubmit,
Command,
}
#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
pub enum InteractionStyle {
Default,
Primary,
Danger,
}
#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
pub struct InteractionAction {
pub id: String,
pub label: String,
pub value: String,
pub style: InteractionStyle,
}
#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
pub struct InteractionPayload {
pub interaction_id: String,
pub kind: InteractionKind,
pub title: Option<String>,
pub body: String,
pub actions: Vec<InteractionAction>,
pub expires_at: Option<i64>,
}
#[derive(Debug, Clone)]
pub struct InteractionEvent {
pub channel: String,
pub chat_id: String,
pub sender_id: String,
pub interaction_id: String,
pub action_id: String,
pub action_value: String,
pub timestamp: i64,
pub metadata: std::collections::HashMap<String, String>,
}
```
`InteractionPayload` 用于 outbound 渲染，`InteractionEvent` 用于 inbound 回调。
## OutboundMessage 扩展
短期兼容方案：
- 继续使用 `OutboundMessage.metadata` 携带交互描述。
- 例如：
- `interaction.kind = "approval"`
- `interaction.id = "<uuid>"`
- `interaction.actions = "<json>"`
长期推荐方案：
```rust
pub struct OutboundMessage {
pub channel: String,
pub chat_id: String,
pub content: String,
pub reply_to: Option<String>,
pub media: Vec<MediaItem>,
pub metadata: HashMap<String, String>,
pub interaction: Option<InteractionPayload>,
}
```
推荐长期方案。它能避免把结构化交互塞进字符串 metadata，也让每个 Channel 的 `send()` 更清晰。
## Channel 能力声明
给 `Channel` 增加可选能力声明：
```rust
#[derive(Debug, Clone, Default)]
pub struct ChannelCapabilities {
pub interactive_buttons: bool,
pub forms: bool,
pub message_update: bool,
pub markdown_cards: bool,
}
pub trait Channel {
fn capabilities(&self) -> ChannelCapabilities {
ChannelCapabilities::default()
}
}
```
渠道能力示例：
| 渠道 | 原生按钮 | 表单 | 更新原消息 | 降级策略 |
|------|----------|------|------------|----------|
| Feishu | 是，interactive card | 可后续支持 | 是，PATCH message/card | 文本命令 |
| Slack | 是，Block Kit | 是 | 是 | 文本命令 |
| Telegram | 是，inline keyboard | 有限 | 可编辑消息 | 文本命令 |
| Discord | 是，components | 有限 | 可编辑消息 | 文本命令 |
| CLI chat | 否 | 否 | 局部可模拟 | 编号/命令输入 |
| Webhook/Email/SMS | 否 | 否 | 通常否 | 纯文本命令或链接 |
## 渲染策略
每个 channel 实现一个渠道内的渲染函数：
```rust
async fn send_interaction(
&self,
chat_id: &str,
payload: &InteractionPayload,
) -> Result<(), ChannelError>;
```
也可以先不改 trait，在 `send()` 内部判断 `msg.interaction`。
飞书渲染：
- 使用 Card JSON 2.0。
- `schema = "2.0"`。
- body 用 markdown 展示 `payload.body`。
- actions 渲染为 button。
- 每个按钮的 callback value 写入：
```json
{
"interaction_id": "...",
"action_id": "...",
"action_value": "approve"
}
```
需要兼容飞书 Card 2.0 回调路径：
- `/action/value`
- `/action/behaviors/0/value`
CLI 降级渲染：
```text
需要确认：
Tool: bash
Args: cargo test --lib
可选操作：
1. Approve
2. Deny
3. Always approve
回复：
/_interaction <interaction_id> approve
/_interaction <interaction_id> deny
/_interaction <interaction_id> always
```
纯文本渠道都可以复用这个 fallback renderer。
## Inbound 回调归一化
飞书 WebSocket 当前在 `src/channels/feishu.rs` 里处理 `im.message.receive_v1`。需要新增对 `card.action.trigger` 的识别：
1. ACK 仍要尽快发送，飞书要求 3 秒内响应。
2. 如果 event type 是 `card.action.trigger`，不要走普通消息解析。
3. 从 event payload 中解析 `interaction_id`、`action_id`、`action_value`。
4. 构造 `InteractionEvent` 发布给统一处理器。
5. 对未知、过期或重复 interaction 返回成功但记录日志，不应导致渠道重连或报错。
如果短期不新增 interaction bus，可以把回调转成特殊 `InboundMessage`：
```text
content = "/_interaction <interaction_id> <action_value>"
metadata["event.kind"] = "interaction"
metadata["interaction.id"] = "<interaction_id>"
metadata["interaction.action_id"] = "<action_id>"
metadata["interaction.action_value"] = "<action_value>"
```
但必须由 SessionManager 或 InteractionManager 先拦截，不能把 `/_interaction` 当普通用户文本直接送进 LLM。
长期推荐新增 bus 通道：
```rust
pub enum InboundEvent {
Message(InboundMessage),
Interaction(InteractionEvent),
}
```
或者在 `MessageBus` 上增加 `interaction_tx`。
## InteractionManager
建议新增 `InteractionManager`，集中管理 pending 交互状态，而不是让每个 Channel 各自维护。
职责：
- 生成 `interaction_id`。
- 保存 pending interaction。
- 处理超时和过期。
- 接收 `InteractionEvent` 并解析成业务结果。
- 对重复点击、未知 interaction、过期 interaction 做幂等处理。
- 必要时通知 channel 更新原消息。
内部状态示例：
```rust
pub struct PendingInteraction {
pub id: String,
pub kind: InteractionKind,
pub channel: String,
pub chat_id: String,
pub sender_id: Option<String>,
pub session_id: Option<String>,
pub created_at: i64,
pub expires_at: Option<i64>,
pub status: InteractionStatus,
pub responder: InteractionResponder,
pub message_ref: Option<InteractionMessageRef>,
}
pub struct InteractionMessageRef {
pub channel: String,
pub chat_id: String,
pub message_id: String,
pub metadata: HashMap<String, String>,
}
```
`InteractionResponder` 可以先支持 oneshot：
```rust
pub enum InteractionResponder {
Approval(tokio::sync::oneshot::Sender<ApprovalDecision>),
InboundMessage,
}
```
后续如果需要持久化长期交互，oneshot 不够，需要落库。
## 工具审批流程
工具审批是第一批最适合落地的交互类型。
推荐流程：
1. AgentLoop 准备执行需要审批的工具。
2. 调用 `InteractionManager::request_approval(...)`。
3. InteractionManager 创建 `InteractionPayload`，通过 outbound 发送到原 channel/chat。
4. AgentLoop 等待 oneshot，带 timeout。
5. 用户在飞书卡片上点击 Approve/Deny/Always。
6. FeishuChannel 收到 `card.action.trigger`，发布 `InteractionEvent`。
7. InteractionManager resolve pending approval。
8. AgentLoop 收到结果，继续执行或拒绝工具。
9. 如果 channel 支持更新消息，InteractionManager 或 Channel 把原卡片更新成 resolved 状态。
审批 action 建议：
```rust
approve -> ApprovalDecision::Approve
deny -> ApprovalDecision::Deny
always -> ApprovalDecision::AlwaysApprove
```
`DenyWithEdit` 可后续支持，适合 ACP/Web/CLI 这类能输入文本的渠道。
## 快捷回复流程
快捷回复不是阻塞工具执行，而是把用户点击转成新的用户输入。
示例：
```json
{
"kind": "QuickReply",
"body": "你想继续哪个操作？",
"actions": [
{ "label": "继续分析", "value": "继续分析" },
{ "label": "生成报告", "value": "生成报告" }
]
}
```
用户点击后：
- `InteractionEvent.action_value` 转成一条普通 `InboundMessage.content`。
- `sender_id` 和 `chat_id` 保留原用户和会话。
- metadata 标记来源为 interaction，供审计或 UI 使用。
## 消息更新
支持原消息更新的渠道应在交互完成后更新 UI，避免重复点击。
飞书：
- 发送卡片后保存 `data.message_id`。
- resolve 后 PATCH `/im/v1/messages/{message_id}`。
- 卡片 schema 发送和更新都使用 Card JSON 2.0，参考项目指出跨版本 PATCH 可能返回成功但客户端不重渲染。
不支持更新的渠道：
- 发送一条新消息提示“已批准/已拒绝”。
- 或仅在后台幂等拒绝重复点击。
## 安全和权限
交互回调必须校验：
- `interaction_id` 是否存在。
- 是否已过期。
- 是否已处理。
- 点击用户是否允许处理该 interaction。
- 当前 channel/chat 是否匹配。
对于工具审批，默认建议只有触发该 agent turn 的用户或允许列表用户可以审批。群聊里要特别注意 `sender_id`，不能只看 `chat_id`。
日志中避免记录原始飞书回调敏感字段：
- callback token
- operator open_id/union_id/user_id/tenant_key
- open_chat_id/open_message_id
可以记录脱敏后的 payload shape，用于排查飞书回调字段变化。
## 持久化策略
第一阶段可以只做内存 pending map：
- 适合短时工具审批。
- 进程重启后旧按钮点击会变成 unknown/expired。
- 实现简单。
后续如果要支持长期任务或跨重启交互，需要持久化：
- `interactions` 表保存 id、kind、channel、chat_id、sender_id、status、payload、created_at、expires_at。
- `interaction_actions` 可选，或直接 JSON 存在 payload 中。
- resolve 时事务更新 status，防止重复点击竞态。
## 与现有架构的关系
现有数据流：
```text
Channel -> MessageBus -> SessionManager -> AgentLoop -> tools -> SessionManager -> MessageBus -> OutboundDispatcher -> Channel
```
加入 interaction 后建议：
```text
Outbound:
AgentLoop/Tool approval -> InteractionManager -> MessageBus outbound -> OutboundDispatcher -> Channel renderer
Inbound:
Channel callback -> InteractionEvent -> InteractionManager -> pending waiter / synthetic InboundMessage
```
Channel 仍然只做渠道协议适配：
- 飞书负责 Card JSON 和 `card.action.trigger`。
- Slack 负责 Block Kit 和 signing secret。
- Telegram 负责 callback query。
- CLI 负责文本命令 fallback。
InteractionManager 负责语义：
- 这是 approval 还是 quick reply。
- 是否过期。
- 是否有权限。
- 应该唤醒哪个等待者。
SessionManager/AgentLoop 不需要知道飞书卡片格式。
## 分阶段实施计划
### 阶段 1：模型和 fallback
- 新增 `src/interaction/` 模块。
- 定义 `InteractionPayload`、`InteractionAction`、`InteractionEvent`。
- 增加 fallback text renderer。
- 为 `OutboundMessage` 增加 `interaction: Option<InteractionPayload>`，或短期使用 metadata。
- 增加单元测试覆盖序列化和 fallback 文本。
### 阶段 2：Feishu card action
- Feishu outbound 支持把 `InteractionPayload` 渲染为 Card JSON 2.0。
- Feishu inbound 在 WebSocket frame 中识别 `card.action.trigger`。
- 解析 `/action/value` 和 `/action/behaviors/0/value`。
- 发布统一 `InteractionEvent`。
- 保存 `message_id`，支持完成后 PATCH resolved card。
- 增加 fixtures 测试真实/模拟的 `card.action.trigger` payload。
### 阶段 3：InteractionManager 和审批
- 新增内存版 `InteractionManager`。
- 支持 request/resolve/timeout。
- 接入工具执行前审批点。
- 支持 Approve/Deny/AlwaysApprove。
- 未知、过期、重复点击保持幂等。
- 增加 agent/tool 审批单元测试。
### 阶段 4：其他渠道兼容
- CLI chat 支持 `/_interaction <id> <value>` fallback。
- 其他不支持原生按钮的渠道使用纯文本 fallback。
- 后续按需实现 Slack/Telegram/Discord 原生按钮。
### 阶段 5：持久化和高级交互
- 需要时落库 pending interaction。
- 支持 quick reply 生成 synthetic inbound message。
- 支持表单提交。
- 支持 `DenyWithEdit`。
- 支持长期任务交互和重启恢复。
## 测试计划
单元测试：
- `InteractionPayload` 序列化。
- fallback text renderer 输出。
- Feishu Card JSON 包含正确 callback value。
- Feishu 回调同时支持 `/action/value` 和 `/action/behaviors/0/value`。
- unknown/expired interaction 不报错。
- 重复点击只 resolve 一次。
集成测试：
- 模拟 Feishu `card.action.trigger`，验证 pending approval 被唤醒。
- 模拟超时，验证默认 deny。
- CLI fallback 输入 `/_interaction`，验证能 resolve。
- 不支持按钮的 channel 能收到可读 fallback 文本。
手工验证：
- 飞书群聊点击 Approve/Deny/Always。
- 私聊点击。
- 点击后原卡片更新为 resolved。
- 重复点击不会重复执行工具。
- 非触发用户点击时按权限策略处理。
## 开放问题
1. 工具审批应该由 AgentLoop 直接调用 InteractionManager，还是通过 SessionManager 代理？
2. 群聊中是否只允许原始触发者审批，还是允许配置中的所有 allowed user 审批？
3. `AlwaysApprove` 的作用域是本次会话、本 dialog、本 chat，还是全局工具策略？
4. 是否需要第一阶段就修改 `OutboundMessage` 结构，还是先用 metadata 降低改动面？
5. 飞书 CardKit 流式输出是否要与交互卡片统一，还是继续保持普通回复卡片和交互卡片两套路径？

101
docs/assets/message-flow.svg
View File

@ -1,101 +0,0 @@
<svg xmlns="http://www.w3.org/2000/svg" width="1120" height="460" viewBox="0 0 1120 460" role="img" aria-labelledby="title desc">
<title id="title">PicoBot message flow</title>
<desc id="desc">Message flow from channel input through bus, session manager, agent loop, tools, provider, storage, outbound dispatcher, and back to the channel.</desc>
<defs>
<style>
.bg { fill: #fbfbf8; }
.lane { fill: #ffffff; stroke: #d6d3d1; stroke-width: 2; rx: 18; }
.step { fill: #f4f4f5; stroke: #71717a; stroke-width: 1.6; rx: 12; }
.in { fill: #e0f2fe; stroke: #0284c7; }
.state { fill: #ecfdf5; stroke: #059669; }
.agent { fill: #fff7ed; stroke: #ea580c; }
.out { fill: #f5f3ff; stroke: #7c3aed; }
.text { font-family: ui-sans-serif, -apple-system, BlinkMacSystemFont, "Segoe UI", sans-serif; fill: #18181b; }
.title { font-size: 30px; font-weight: 800; }
.label { font-size: 17px; font-weight: 700; }
.small { font-size: 14px; fill: #52525b; }
.arrow { stroke: #3f3f46; stroke-width: 2.4; fill: none; marker-end: url(#arrow); }
.soft { stroke: #71717a; stroke-width: 2; fill: none; stroke-dasharray: 7 6; marker-end: url(#arrow-soft); }
.num { fill: #18181b; font-size: 13px; font-weight: 800; }
.badge { fill: #ffffff; stroke: #a1a1aa; }
</style>
<marker id="arrow" markerWidth="12" markerHeight="12" refX="10" refY="6" orient="auto">
<path d="M2,2 L10,6 L2,10 Z" fill="#3f3f46" />
</marker>
<marker id="arrow-soft" markerWidth="12" markerHeight="12" refX="10" refY="6" orient="auto">
<path d="M2,2 L10,6 L2,10 Z" fill="#71717a" />
</marker>
</defs>
<rect class="bg" width="1120" height="460" />
<text class="text title" x="48" y="56">Message Flow</text>
<text class="text small" x="48" y="82">A user message becomes a session-scoped agent run, then returns to the original channel.</text>
<rect class="lane" x="38" y="122" width="1044" height="236" />
<g transform="translate(70 182)">
<rect class="step in" width="130" height="78" />
<circle class="badge" cx="18" cy="18" r="13" />
<text class="text num" x="14" y="23">1</text>
<text class="text label" x="32" y="37">Channel</text>
<text class="text small" x="22" y="60">CLI / Feishu</text>
</g>
<g transform="translate(230 182)">
<rect class="step" width="130" height="78" />
<circle class="badge" cx="18" cy="18" r="13" />
<text class="text num" x="14" y="23">2</text>
<text class="text label" x="30" y="37">MessageBus</text>
<text class="text small" x="29" y="60">inbound queue</text>
</g>
<g transform="translate(390 182)">
<rect class="step state" width="150" height="78" />
<circle class="badge" cx="18" cy="18" r="13" />
<text class="text num" x="14" y="23">3</text>
<text class="text label" x="34" y="37">SessionManager</text>
<text class="text small" x="31" y="60">dialog + context</text>
</g>
<g transform="translate(580 182)">
<rect class="step agent" width="130" height="78" />
<circle class="badge" cx="18" cy="18" r="13" />
<text class="text num" x="14" y="23">4</text>
<text class="text label" x="34" y="37">AgentLoop</text>
<text class="text small" x="23" y="60">LLM/tool loop</text>
</g>
<g transform="translate(750 182)">
<rect class="step agent" width="130" height="78" />
<circle class="badge" cx="18" cy="18" r="13" />
<text class="text num" x="14" y="23">5</text>
<text class="text label" x="48" y="37">Tools</text>
<text class="text small" x="29" y="60">side effects</text>
</g>
<g transform="translate(920 182)">
<rect class="step out" width="130" height="78" />
<circle class="badge" cx="18" cy="18" r="13" />
<text class="text num" x="14" y="23">6</text>
<text class="text label" x="30" y="37">Response</text>
<text class="text small" x="28" y="60">outbound bus</text>
</g>
<path class="arrow" d="M200 221 H230" />
<path class="arrow" d="M360 221 H390" />
<path class="arrow" d="M540 221 H580" />
<path class="arrow" d="M710 221 H750" />
<path class="arrow" d="M880 221 H920" />
<path class="soft" d="M455 182 C460 120 620 116 652 181" />
<text class="text small" x="500" y="126">recall memory + build prompt</text>
<path class="soft" d="M645 260 C640 322 482 326 458 260" />
<text class="text small" x="496" y="334">persist messages and metadata</text>
<path class="soft" d="M815 182 C820 122 918 122 958 181" />
<text class="text small" x="828" y="126">provider calls</text>
<path class="arrow" d="M985 260 C972 390 139 392 135 260" />
<text class="text small" x="420" y="408">OutboundDispatcher routes the reply back to the same channel/chat scope.</text>
</svg>
Side by Side

Before

Width:  |  Height:  |  Size: 4.6 KiB

88
docs/assets/runtime-architecture.svg
View File

@ -1,88 +0,0 @@
<svg xmlns="http://www.w3.org/2000/svg" width="1120" height="620" viewBox="0 0 1120 620" role="img" aria-labelledby="title desc">
<title id="title">PicoBot runtime architecture</title>
<desc id="desc">High level architecture diagram showing channels, gateway, message bus, session manager, agent loop, tools, providers, storage, scheduler, skills, and MCP.</desc>
<defs>
<style>
.bg { fill: #f8fafc; }
.panel { fill: #ffffff; stroke: #cbd5e1; stroke-width: 2; rx: 18; }
.box { fill: #f1f5f9; stroke: #64748b; stroke-width: 1.5; rx: 12; }
.accent { fill: #e0f2fe; stroke: #0284c7; }
.warm { fill: #fff7ed; stroke: #ea580c; }
.green { fill: #ecfdf5; stroke: #059669; }
.violet { fill: #f5f3ff; stroke: #7c3aed; }
.text { font-family: ui-sans-serif, -apple-system, BlinkMacSystemFont, "Segoe UI", sans-serif; fill: #0f172a; }
.small { font-size: 16px; }
.label { font-size: 18px; font-weight: 700; }
.title { font-size: 30px; font-weight: 800; }
.note { font-size: 14px; fill: #475569; }
.arrow { stroke: #334155; stroke-width: 2.4; fill: none; marker-end: url(#arrow); }
.soft { stroke: #64748b; stroke-width: 2; fill: none; stroke-dasharray: 7 6; marker-end: url(#arrow-soft); }
</style>
<marker id="arrow" markerWidth="12" markerHeight="12" refX="10" refY="6" orient="auto">
<path d="M2,2 L10,6 L2,10 Z" fill="#334155" />
</marker>
<marker id="arrow-soft" markerWidth="12" markerHeight="12" refX="10" refY="6" orient="auto">
<path d="M2,2 L10,6 L2,10 Z" fill="#64748b" />
</marker>
</defs>
<rect class="bg" width="1120" height="620" />
<text class="text title" x="48" y="56">PicoBot Runtime Architecture</text>
<text class="text note" x="48" y="82">Channels stay thin, SessionManager owns conversation state, AgentLoop remains stateless.</text>
<rect class="panel" x="40" y="118" width="220" height="360" />
<text class="text label" x="70" y="154">Channels</text>
<rect class="box accent" x="70" y="185" width="160" height="58" />
<text class="text small" x="106" y="220">CLI TUI</text>
<rect class="box accent" x="70" y="263" width="160" height="58" />
<text class="text small" x="104" y="298">Feishu/Lark</text>
<rect class="box accent" x="70" y="341" width="160" height="58" />
<text class="text small" x="109" y="376">WebSocket</text>
<text class="text note" x="70" y="437">Only receive and send</text>
<rect class="panel" x="330" y="118" width="450" height="360" />
<text class="text label" x="360" y="154">Gateway Core</text>
<rect class="box" x="370" y="184" width="160" height="64" />
<text class="text small" x="414" y="222">MessageBus</text>
<rect class="box green" x="580" y="184" width="160" height="64" />
<text class="text small" x="606" y="222">SessionManager</text>
<rect class="box warm" x="580" y="288" width="160" height="64" />
<text class="text small" x="621" y="326">AgentLoop</text>
<rect class="box" x="370" y="288" width="160" height="64" />
<text class="text small" x="402" y="326">Outbound</text>
<text class="text small" x="399" y="346">Dispatcher</text>
<rect class="box violet" x="475" y="390" width="200" height="54" />
<text class="text small" x="514" y="423">Control Channel</text>
<rect class="panel" x="850" y="118" width="230" height="360" />
<text class="text label" x="880" y="154">Capabilities</text>
<rect class="box warm" x="880" y="185" width="170" height="50" />
<text class="text small" x="929" y="216">Tools</text>
<rect class="box warm" x="880" y="249" width="170" height="50" />
<text class="text small" x="916" y="280">Providers</text>
<rect class="box green" x="880" y="313" width="170" height="50" />
<text class="text small" x="926" y="344">SQLite</text>
<rect class="box violet" x="880" y="377" width="170" height="50" />
<text class="text small" x="925" y="408">Skills</text>
<text class="text note" x="880" y="457">MCP tools join ToolRegistry</text>
<rect class="box green" x="300" y="520" width="165" height="54" />
<text class="text small" x="340" y="553">Scheduler</text>
<rect class="box violet" x="505" y="520" width="165" height="54" />
<text class="text small" x="552" y="553">Memory</text>
<rect class="box accent" x="710" y="520" width="165" height="54" />
<text class="text small" x="759" y="553">MCP</text>
<path class="arrow" d="M260 216 H370" />
<path class="arrow" d="M530 216 H580" />
<path class="arrow" d="M660 248 V288" />
<path class="arrow" d="M740 320 H880" />
<path class="arrow" d="M880 274 H750" />
<path class="arrow" d="M580 320 H530" />
<path class="arrow" d="M370 320 H260" />
<path class="soft" d="M575 390 V352" />
<path class="soft" d="M660 248 C750 250 800 330 880 338" />
<path class="soft" d="M382 520 C420 470 520 465 610 352" />
<path class="soft" d="M588 520 C600 470 620 420 650 352" />
<path class="soft" d="M792 520 C825 470 860 430 880 402" />
</svg>
Side by Side

Before

Width:  |  Height:  |  Size: 4.9 KiB

5
resources/skills/about-picobot/SKILL.md
View File

@ -14,10 +14,11 @@ PicoBot 是一个基于 Rust 的个人 AI 助手运行时，包含本地 Gateway
| 文件 | 内容 |
|------|------|
| `references/config.md` | 配置字段详解：providers、models、agents、gateway、client、channels、memory、mcp、browser |
| `references/db-schema.md` | 数据库表结构：sessions、messages、memories、scheduled_jobs、llm_calls、background_tasks |
| `references/architecture.md` | 核心架构：数据流、会话系统、上下文压缩、记忆系统、Skill 优先级、MCP、子 Agent |
| `references/db-schema.md` | 数据库表结构与运行约束：sessions、messages、memories、scheduled_jobs、job_runs、llm_calls、background_tasks |
| `references/architecture.md` | 核心架构：消息并发、会话系统、持久化、生命周期、上下文压缩、记忆、MCP、子 Agent |
| `references/faq.md` | 常见问题：模型切换、渠道添加、Skill 安装、历史查询、定时任务、MCP 等 |
| `references/commands.md` | 常用命令：编译、启动网关、启动客户端、运行测试 |
| `references/tools.md` | 内置工具名称、参数和重要使用约束 |
| `assets/config.example.json` | config.json 完整示例 |
Skill 根目录路径见上方 **Skill Root Directory**。

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

@ -3,9 +3,12 @@
## 核心数据流
```
Channel → MessageBus → SessionManager → AgentLoop → (tools) → SessionManager → MessageBus → OutboundDispatcher → Channel
↑
ControlChannel → SessionManager (dialog 操作: 创建/切换/归档/删除)
Channel → MessageBus.inbound → Gateway processor → SessionManager → per-session worker → AgentLoop
↑ │
└── Channel ← OutboundDispatcher ← per-conversation lane ← MessageBus.outbound
WebSocket/Channel → MessageBus.control → Gateway processor → SessionManager (dialog 操作)
Scheduler → SessionManager.handle_cron_message → AgentLoop → send_message
```
## 模块职责
@ -15,8 +18,8 @@ Channel → MessageBus → SessionManager → AgentLoop → (tools) → SessionM
| `gateway` | HTTP/WebSocket 服务器，持有 GatewayState |
| `client` | TUI 聊天客户端 |
| `channels` | 外部集成（飞书、CLI），仅收发消息 |
| `bus` | 异步消息队列，纯队列不路由 |
| `session` | 会话生命周期管理、dialog 操作 |
| `bus` | 有界 inbound/outbound/control 队列；出站 dispatcher 与分目标 lane |
| `session` | 会话生命周期、dialog 操作、每 session 串行队列、上下文与持久化协调 |
| `agent` | LLM 调用循环、工具执行、上下文压缩、媒体处理、子 Agent |
| `providers` | LLM API 客户端（OpenAI 兼容、Anthropic） |
| `tools` | Agent 工具（bash、文件操作、搜索、HTTP、web、browser、memory、delegate 等） |
@ -28,13 +31,14 @@ Channel → MessageBus → SessionManager → AgentLoop → (tools) → SessionM
| `config` | 配置加载、环境变量替换、路径解析 |
| `memory` | 长期记忆存储与检索 |
| `mcp` | MCP（Model Context Protocol）工具集成 |
| `task_supervisor` | Gateway 后台任务注册、取消、限时等待和强制回收 |
## 功能边界
- Channels 仅收发消息，不感知 session 或 LLM
- MessageBus 是纯异步队列，不路由
- SessionManager 拥有 session 状态，不直接调 LLM；负责注入 skills prompt
- AgentLoop 无状态，接收 dialog 事件调用 LLM、执行工具
- MessageBus 本体持有三条有界队列；出站路由、顺序和重试由 `OutboundDispatcher` 负责
- SessionManager 拥有 session 状态、dialog 路由、上下文构建和每 session worker，并通过 worker 创建 AgentLoop
- AgentLoop 跨轮无状态，接收已准备的 history 调用 LLM、执行工具并返回一次结果
- Providers 是纯 HTTP 客户端，无 bus/session/channel 感知
- Tools 接收原始参数，返回字符串结果
- MCP 工具在 Gateway 初始化时连接服务器、发现工具，并包装成普通 Tool 注册到 ToolRegistry
@ -48,6 +52,11 @@ Channel → MessageBus → SessionManager → AgentLoop → (tools) → SessionM
- OutboundDispatcher 通过 ChannelManager 路由出站消息
- Config `.env` 加载使用 `unsafe { env::set_var(...) }`
- `browser` 工具只有在 `browser.enabled=true` 时注册，依赖 Chrome/Chromium 与 WebDriver
- 同一 session 的普通消息串行处理，不同 session 可并发；session 队列容量为 32，满时明确拒绝
- 出站消息按 `(channel, chat_id)` 分 lane 保序；lane 容量为 64，慢目标不阻塞其他目标
- 长生命周期后台任务由 TaskSupervisor 管理；连接局部任务由其 owner 限时 join 或 abort
- 外部建连、重试等待和关停 join 必须可取消且有硬超时
- 不得记录 API Key、Authorization header 或包含临时凭据的完整连接 URL
## 上下文压缩
@ -101,7 +110,7 @@ create → 存入 Storage → 载入 memory → 设为当前 dialog
| `list_dialogs` | 列出 `channel:chat_id` 下最近 10 个 session |
| `rename` | 更新标题，内存 + Storage 同步 |
| `delete` | 软删除（设 deleted_at），从内存移除 |
| `archive` | 当前为空操作 |
| `archive` | 设置 archived_at，从内存和当前 dialog 追踪中移除；可通过 include_archived 查询 |
### SessionManager 数据结构
@ -112,13 +121,17 @@ create → 存入 Storage → 载入 memory → 设为当前 dialog
消息到达时 `resolve_dialog_id()` 按顺序确定接收 session：当前 session → Storage 最近活跃 session → 新建。
### 消息处理三阶段
### 消息处理与并发
**阶段 1（持锁）**：斜杠命令检测 → 用户消息入库 → 提取记忆上下文 → 构建系统提示（skills + memory_context）→ 上下文压缩 → 创建 AgentLoop
普通消息先 `try_send` 到该 session 的有界 worker 队列，Gateway 主 processor 随即返回 `AgentProcessing`。Slash command 直接执行，不进入此队列，因此 `/stop` 不会排在长模型调用后。
**阶段 2（无锁）**：`agent.process(history)` → LLM 调用 + 工具执行。上下文溢出时自动重新压缩重试
Worker 的处理原则：
**阶段 3（持锁）**：持久化 agent 响应消息 → 自动生成标题（消息数 ≥ 5 且标题为"新对话"时）
1. 短暂持 Session 锁抓取快照并记录 `worker_generation`/`state_version`。
2. 释放锁后执行消息持久化、记忆召回、上下文压缩、LLM 和工具等慢操作。
3. 提交由旧快照产生的结果前重新验证 generation/version，防止 `/stop`、`/clear` 或 `/delete` 后写回陈旧状态。
4. Session 持久化由独立 `persistence_lock` 串行化；批量消息使用原子写入，失败时精确回滚内存后缀。
5. 上下文溢出时按 Provider 返回的真实限制重新压缩并重试。
### 会话恢复
@ -136,7 +149,7 @@ create → 存入 Storage → 载入 memory → 设为当前 dialog
| 类别 | 用途 | 生命周期 | 检索方式 |
|------|------|----------|----------|
| **Knowledge** | 事实、偏好、模式、洞察 | 长期保留，手动删除 | 每轮注入系统提示，关键词匹配 |
| **Timeline** | 历史会话摘要 | 自动清理（默认 90 天） | `timeline_recall` 工具按需检索 |
| **Timeline** | 历史会话摘要 | 配置预期保留 90 天；当前尚无自动清理循环 | `timeline_recall` 工具按需检索 |
### MemoryEntry 字段
@ -163,7 +176,7 @@ create → 存入 Storage → 载入 memory → 设为当前 dialog
→ MemoryManager::recall(content, 5, Knowledge)
返回最多 5 条匹配的知识记忆（按 importance DESC）
→ 格式化为 "- key: content"
→ 注入系统提示的 "记忆上下文" 部分
→ 作为运行时上下文附加到本轮 user message
→ LLM 可见，辅助回答
```
@ -191,10 +204,12 @@ LLM 对话上下文接近 token 限制 (默认 128K × 70%) 时自动触发压
| 时机 | 操作 |
|------|------|
| 每次消息处理 | `memory_manager.recall()` 提取 Knowledge 上下文 |
| 系统提示构建 | `MemorySection` 渲染记忆指南 + 匹配的记忆 |
| 系统提示构建 | `MemorySection` 渲染记忆工具指南；匹配的 Knowledge 附加到本轮 user message |
| 有压缩历史时 | `HistorySection` 提示 LLM 使用 `timeline_recall` |
| 压缩完成后 | 摘要自动存储为 Timeline 记忆 |
| 空闲时 | 可配置自动 consolidation（`idle_consolidation_minutes`） |
| 会话恢复 | 加载最近 Timeline 和压缩边界后的原始消息 |
`memory.recall_limit`、`idle_consolidation_minutes`、`timeline_retention_days` 和 `max_failures_before_degrade` 当前会被配置解析；其中每轮 Knowledge 召回在 worker 中仍固定为 5，其余自动维护策略尚未接入运行循环。不要把“配置可解析”误认为“行为已生效”。
---
@ -223,6 +238,23 @@ Gateway 初始化时读取 `config.mcp.servers`：
默认工具集是只读工具：`file_read`、`file_search`、`content_search`、`web_fetch`、`http_request`、`calculator`。调用时可通过 `allowed_tools` 显式放开其他工具。后台任务会写入 `background_tasks` 表，默认 24 小时后清理。
后台子 Agent 通过 `TaskSupervisor::spawn_graceful` 注册，受 `gateway.max_concurrent_background_tasks` 限制；Gateway 关停时先收到取消信号，再在总宽限期内清理。
---
## 出站投递与关停
`OutboundDispatcher` 对每个 `(channel, chat_id)` 创建独立 lane：同一目标保持顺序，单次发送超时 30 秒，最多尝试 3 次。只有 `ConnectionError` 和 `SendError` 被视为瞬态错误；永久错误不重试。`deliver_outbound` 等待真实渠道投递结果，`publish_outbound` 只表示成功入队。
Gateway 关停顺序：
1. Ctrl-C 停止 Axum 接入并取消 WebSocket 连接。
2. `ChannelManager::stop_all` 停止外部渠道并注销它们。
3. 取消 TaskSupervisor，并在共享的 10 秒宽限期内等待后台任务。
4. 超时后 abort 剩余任务并回收 JoinHandle。
渠道 `stop()` 也必须有界。飞书端点请求、WebSocket 建连、重试 sleep 和连接循环共享 CancellationToken，另有 5 秒强制终止兜底。
---
## 当前斜杠命令

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

@ -13,6 +13,16 @@ cargo run -- chat
# 运行单元测试
cargo test --lib
# 运行集成测试 (需配置 tests/test.env)
# 运行离线集成/协议测试
cargo test --test test_scheduler
cargo test --test test_request_format
# 运行代码检查（Rust 修改必跑）
cargo clippy --all-targets --all-features -- -D warnings
# 运行模型 API 集成测试（需配置 tests/test.env）
cargo test --test test_integration -- --ignored
cargo test --test test_tool_calling -- --ignored
```
`test_scheduler` 和 `test_request_format` 不需要 API Key，也没有标记 `#[ignore]`。只有会真实调用 Provider 的测试需要从 `tests/test.env.example` 创建 `tests/test.env` 后使用 `-- --ignored`。

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

@ -55,9 +55,9 @@
|------|------|------|------|
| `host` | string | 127.0.0.1 | 监听地址 |
| `port` | int | 19876 | 监听端口 |
| `session_ttl_hours` | int | - | 会话过期小时数 |
| `session_ttl_hours` | int | - | 兼容/预留字段；当前没有会话 TTL 清理循环 |
| `session_db_path` | string | - | SQLite 数据库路径，默认在 workspace 下 |
| `cleanup_interval_minutes` | int | - | 清理间隔 |
| `cleanup_interval_minutes` | int | - | 兼容/预留字段；当前没有按此间隔运行的 session 清理任务 |
| `max_concurrent_background_tasks` | int | 10 | delegate 后台子任务最大并发数 |
| `scheduler` | object | - | 调度器配置 |
@ -67,18 +67,21 @@
|------|------|------|------|
| `enabled` | bool | true | 是否启动调度器并注册 cron 工具 |
| `poll_interval_secs` | int | 60 | 检查到期任务的轮询间隔 |
| `max_concurrent` | int | 1 | 最大并发任务数，当前实现预留 |
| `max_concurrent` | int | 1 | 每批到期任务的最大并发数，运行时限制在 1–256 |
| `execution_timeout_secs` | int | 900 | 单个定时任务 Agent 执行的硬超时；租约会额外增加 30 秒 |
## memory 字段
| 字段 | 类型 | 默认 | 说明 |
|------|------|------|------|
| `consolidation_provider` | string | - | 记忆归并 LLM 提供商 |
| `consolidation_model` | string | - | 记忆归并 LLM 模型 |
| `recall_limit` | int | 5 | 每轮注入的知识记忆条数 |
| `idle_consolidation_minutes` | int | 10 | 空闲后触发归并的分钟数 |
| `timeline_retention_days` | int | 90 | 时间线记忆保留天数 |
| `max_failures_before_degrade` | int | 3 | 归并失败次数阈值 |
| `consolidation_provider` | string | 主 Agent provider | 当前记录在 MemoryManager 中供后续归并使用；压缩摘要仍使用 Session provider |
| `consolidation_model` | string | 主 Agent model | 当前记录在 MemoryManager 中供后续归并使用；压缩摘要仍使用 Session model |
| `recall_limit` | int | 5 | 预期的每轮知识召回上限；当前 worker 固定使用 5 |
| `idle_consolidation_minutes` | int | 10 | 预留的空闲归并阈值；当前无对应循环 |
| `timeline_retention_days` | int | 90 | 预留的 Timeline 保留期；当前无自动清理循环 |
| `max_failures_before_degrade` | int | 3 | 预留的归并失败阈值；当前无失败降级循环 |
注意：这些字段都会被解析，但当前 worker 的 Knowledge 召回数量仍固定为 5；idle consolidation、Timeline 自动清理和失败降级循环尚未接入。配置存在不等于对应后台行为已经生效。
## channels.feishu 字段
@ -89,7 +92,7 @@
| `app_secret` | string | - | 飞书应用密钥 |
| `allow_from` | []string | ["*"] | 允许交互的用户列表 |
| `agent` | string | - | 使用的 agent 名称 |
| `media_dir` | string | ~/.picobot/media/feishu | 媒体存储目录 |
| `media_dir` | string | ~/.picobot/media/feishu | 配置默认值；Gateway 注册渠道时会覆盖为 `{workspace}/media/feishu` |
| `reaction_emoji` | string | "Typing" | 回复意向表达的表情 |
## mcp 字段

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

@ -2,6 +2,8 @@
数据库为 SQLite，默认位于 workspace 下的 `picobot.db`。
连接启用 WAL、`synchronous=NORMAL`、foreign keys、5 秒 busy timeout，连接池最多 8 个连接。当前 `PRAGMA user_version=1`；启动时会在事务内补齐旧库字段和索引，遇到比程序更新的 schema version 会拒绝启动。
## sessions 表
会话表，一个 session 对应一个 (channel, chat_id, dialog_id) 组合。
@ -13,13 +15,16 @@
| `chat_id` | TEXT | 聊天/群组标识 |
| `dialog_id` | TEXT | 对话标识 |
| `title` | TEXT | 会话标题（默认 "新对话"） |
| `created_at` | INTEGER | 创建时间（unix 秒） |
| `last_active_at` | INTEGER | 最后活跃时间 |
| `created_at` | INTEGER | 创建时间（Unix 毫秒） |
| `last_active_at` | INTEGER | 最后活跃时间（Unix 毫秒） |
| `message_count` | INTEGER | 消息计数 |
| `routing_info` | TEXT | 路由信息 |
| `archived_at` | INTEGER | 归档时间（Unix 毫秒），NULL 表示未归档 |
| `deleted_at` | INTEGER | 软删除时间戳 |
| `last_consolidated_at` | INTEGER | 上次记忆归并时间 |
| `last_compressed_message_at` | INTEGER | 上次上下文压缩消息序号 |
| `last_compressed_message_at` | INTEGER | 上次上下文压缩边界时间戳 |
`(channel, chat_id, dialog_id)` 唯一。普通列表排除 `deleted_at`；是否包含归档记录由查询参数决定。
## messages 表
@ -35,9 +40,11 @@
| `tool_name` | TEXT | 工具名称 |
| `tool_calls` | TEXT | 工具调用参数 JSON |
| `source` | TEXT | 消息来源（跨会话消息时标记来源 session_id） |
| `created_at` | INTEGER | 创建时间（unix 秒） |
| `created_at` | INTEGER | 创建时间（Unix 毫秒） |
| `reasoning_content` | TEXT | provider 返回的推理内容（如有） |
`(session_id, seq)` 有唯一索引，防止并发写入重复序号。删除 session 会通过外键级联删除 messages。
## background_tasks 表
delegate 后台子任务表。`session_id` 不使用数据库外键，因为 session 使用软删除，关联关系由应用层维护。
@ -82,17 +89,24 @@ delegate 后台子任务表。`session_id` 不使用数据库外键，因为 ses
|------|------|------|
| `id` | TEXT PK | 任务 UUID |
| `name` | TEXT | 任务名称 |
| `schedule` | TEXT | 调度规则 JSON（once/every/cron） |
| `schedule` | TEXT | 调度规则 JSON（at/every/cron） |
| `prompt` | TEXT | 任务提示词 |
| `channel` | TEXT | 执行渠道 |
| `chat_id` | TEXT | 目标对话 |
| `model` | TEXT | 使用的模型（可选） |
| `model` | TEXT | 可选模型标记；当前会存储/展示，但 Scheduler 执行仍使用默认 Agent 模型 |
| `enabled` | INTEGER | 是否启用 (1/0) |
| `delete_after_run` | INTEGER | 执行后自动删除 (1/0) |
| `next_run_at` | INTEGER | 下次执行时间 |
| `last_run_at` | INTEGER | 上次执行时间 |
| `last_status` | TEXT | 上次执行状态 |
| `last_error` | TEXT | 上次错误信息 |
| `locked_at` | INTEGER | 本次领取时间 |
| `lock_owner` | TEXT | 领取任务的 Scheduler owner UUID |
| `lease_until` | INTEGER | 租约到期时间；进程崩溃后允许其他实例重新领取 |
| `created_at` | INTEGER | 创建时间（Unix 毫秒） |
| `updated_at` | INTEGER | 更新时间（Unix 毫秒） |
Scheduler 使用原子 `UPDATE ... RETURNING` 领取到期任务。任务结果、下次运行时间和租约释放在同一事务中提交，并校验 owner，防止过期 worker 覆盖已恢复的任务。
## job_runs 表
@ -121,3 +135,5 @@ delegate 后台子任务表。`session_id` 不使用数据库外键，因为 ses
| `response_body` | TEXT | 响应体 JSON |
| `error` | TEXT | 错误信息 |
| `duration_ms` | INTEGER | 耗时（毫秒） |
`request_body`/`response_body` 可能包含用户内容，排障和导出数据库时应按敏感数据处理。

14
resources/skills/about-picobot/references/faq.md
View File

@ -20,6 +20,8 @@
`~/.picobot/skills/about-picobot/` 下，SKILL.md 为索引，references/ 下为各详细文档，assets/ 下为 config 示例。如被删除，重启程序自动重新安装。
内置 Skill 只在目标目录不存在时释放，不会覆盖已安装目录。升级 PicoBot 后如需获取新版内置文档，应先备份自己的修改，再删除旧的 `~/.picobot/skills/about-picobot/` 并重启。也可把定制版放在 `{workspace}/skills/about-picobot/`，它的优先级更高。
## Q: 数据库文件在哪里？
默认 `{workspace}/picobot.db`，workspace 默认 `~/.picobot/workspace/`。
@ -30,7 +32,7 @@
## Q: 如何创建定时任务？
使用 `cron` 工具，支持一次性 (`once`)、周期性 (`every`) 和 cron 表达式调度。
使用 `cron_add` 创建，`cron_list` 查看，`cron_update`、`cron_enable`、`cron_disable`、`cron_remove` 管理。Schedule 类型为 `at`（Unix 毫秒时间戳）、`every`（毫秒间隔）或 6 段 `cron` 表达式，可指定 IANA 时区。
## Q: 上下文压缩是什么意思？
@ -43,3 +45,13 @@
## Q: 如何查看 LLM 调用日志？
LLM 调用记录存储在 `llm_calls` 表中。可通过 SQLite 客户端直接查询，或在代码中通过 storage 模块访问。
该表可能包含完整用户消息、工具参数和模型响应，按敏感数据处理；不要直接上传或粘贴到公开 issue。
## Q: 为什么修改了某些 memory 配置却没有看到行为变化？
当前 `recall_limit`、`idle_consolidation_minutes`、`timeline_retention_days` 和 `max_failures_before_degrade` 都能被配置解析，但每轮 Knowledge 召回仍固定为 5，自动 idle consolidation、Timeline 清理和失败降级循环尚未接入。以当前代码行为为准。
## Q: Gateway 为什么无法立即退出？
正常关停会先停止渠道，再取消受 TaskSupervisor 管理的后台任务，并最多等待 10 秒。飞书自身另有 5 秒连接任务兜底。如果持续超过这些上限，应检查是否新增了未受监督的 `tokio::spawn`、不可取消的外部 I/O，或没有 timeout 的 JoinHandle 等待。

55
resources/skills/about-picobot/references/tools.md
View File

@ -12,16 +12,8 @@
| `content` | 是 | 消息文本内容 |
| `files` | 否 | 文件路径列表 |
| `origin` | 否 | 消息来源标识，不填则自动使用当前 session_id |
| `file_types` | 否 | 指定文件发送类型，`{"路径": "audio"|"file"}`。未指定则自动判断 |
### file_types 说明
控制文件以何种消息类型发送，主要用于飞书渠道：
- `"audio"`：作为语音消息发送（仅 opus 格式支持）
- `"file"`：作为文件附件发送
飞书渠道限制：上传类型和消息类型必须一致。opus 文件以 `"audio"` 发送，其他音频（mp3、wav 等）只能以 `"file"` 发送。
`files` 支持绝对路径和 workspace 相对路径，媒体类型由文件扩展名/MIME 自动判断。目前 schema 不支持手工指定 `file_types`。
### 示例
@ -29,11 +21,12 @@
{
"target_chat_id": "feishu:oc_abc123",
"content": "这是生成的音乐文件",
"files": ["/workspace/music.mp3"],
"file_types": {"/workspace/music.mp3": "file"}
"files": ["/workspace/music.mp3"]
}
```
发送流程会先把消息写入目标 Session，再调用 `MessageBus::deliver_outbound` 等待真实渠道投递结果（上限 120 秒）。投递失败会作为工具失败返回，不应把“已入队”报告成“已送达”。
---
## chat_manager — 会话管理
@ -53,20 +46,28 @@
---
## cron — 定时任务管理
## Cron 定时任务工具
管理 cron 定时任务。
Cron 不是一个带 `action` 的统一工具，而是六个独立工具；仅在 `gateway.scheduler.enabled=true` 时注册。
### 参数
| 工具 | 主要参数 | 说明 |
|------|----------|------|
| `cron_add` | `schedule`, `prompt`, `channel`, `chat_id`; 可选 `name`, `model` | 创建任务 |
| `cron_list` | 可选 `status=all|enabled|disabled` | 列出任务 |
| `cron_update` | `job_id`; 可选 `prompt`, `schedule`, `channel`, `chat_id`, `model` | 更新指定字段 |
| `cron_remove` | `job_id` | 永久删除任务和关联 job runs |
| `cron_enable` | `job_id` | 启用并重新计算下次运行时间 |
| `cron_disable` | `job_id` | 禁用但保留任务 |
| 参数 | 必填 | 说明 |
|------|------|------|
| `action` | 是 | 操作: `add`, `list`, `update`, `remove`, `enable`, `disable` |
| `name` | add必须 | 任务名称 |
| `schedule` | add需要 | 调度规则: `once`(时间戳), `every`(间隔秒), `cron`(表达式) |
| `prompt` | add必须 | 任务提示词 |
| `channel` | add必须 | 执行渠道 |
| `chat_id` | add必须 | 目标对话 |
`schedule` 支持：
```json
{"type":"at","at":1750000000000}
{"type":"every","every_ms":3600000}
{"type":"cron","expr":"0 0 9 * * *","tz":"Asia/Shanghai"}
```
时间戳和间隔单位为毫秒；Cron 表达式为 6 段（秒、分、时、日、月、周）。定时 Agent 不复用聊天历史，`prompt` 必须包含完整上下文，并由 Agent 使用 `send_message` 投递结果。`model` 当前会持久化和展示，但执行仍使用默认 Agent Provider/Model，不能依赖它实现模型覆盖。
---
@ -168,15 +169,19 @@
## file_read / file_write / file_edit / file_search / content_search — 文件操作和搜索
工作目录内的文件读写编辑、文件名搜索和内容搜索。详细的参数定义见各工具的 parameters_schema。
文件读写编辑、文件名搜索和内容搜索。相对路径从 workspace cwd 解析；默认注册的文件工具也接受绝对路径，因此 workspace 不是硬沙箱。详细参数以各工具的 `parameters_schema` 为准。
## bash — 执行命令
在本地环境执行 bash 命令，有超时限制和安全检查（阻止 rm -rf /、fork bomb 等）。
默认从 workspace cwd 执行 bash 命令。参数为 `command` 和可选 `timeout`；默认 60 秒、最大 600 秒，输出最多约 50,000 字符。命令继承 Gateway 进程权限，可以访问 workspace 外路径；危险模式拦截不是完整沙箱。
## pty — 持久终端会话
用于交互式程序和需要保持状态的长运行命令。`action` 支持 `spawn`、`write`、`read`、`kill`、`list`；`write/read/kill` 需要 `session_id`。Gateway 进程退出时 PTY manager 会清理子进程。
## http_request / web_fetch — HTTP 和 Web 工具
发送 HTTP 请求和获取网页内容，有 URL 安全校验（阻止内网/本地访问）。
`http_request` 支持 GET/POST/PUT/DELETE/PATCH、headers 和字符串 body；`web_fetch` 提取 HTML/JSON 的可读文本。两者校验 URL 与 DNS 解析结果，阻止回环、私网、link-local 和本地域名，并禁用自动重定向，以降低 SSRF 风险。
## calculator — 计算器

16
src/agent/agent_loop.rs
View File

@ -362,16 +362,16 @@ impl AgentLoop {
let end = messages.len().saturating_sub(keep_recent);
let start = 1; // protect system message at [0] if present
let mut modified = 0;
for i in start..end {
if messages[i].role != "tool" {
for message in messages.iter_mut().take(end).skip(start) {
if message.role != "tool" {
continue;
}
if messages[i].content.len() <= max_chars {
if message.content.len() <= max_chars {
continue;
}
let tool_name = messages[i].tool_name.as_deref().unwrap_or("unknown");
let chars = messages[i].content.len();
messages[i].content = format!(
let tool_name = message.tool_name.as_deref().unwrap_or("unknown");
let chars = message.content.len();
message.content = format!(
"[Tool output ({}) — {} chars, omitted from context]",
tool_name, chars
);
@ -810,14 +810,14 @@ mod tests {
fn test_should_execute_in_parallel_single_tool() {
// Would need a proper setup with AgentLoop to test fully
// For now, just verify the logic: single tool should return false
let calls = vec![ToolCall {
let calls = [ToolCall {
id: "1".to_string(),
name: "test".to_string(),
arguments: serde_json::json!({}),
}];
// If there's only 1 tool, should return false regardless
assert_eq!(calls.len() <= 1, true);
assert!(calls.len() <= 1);
}
#[test]

60
src/agent/context_compressor.rs
View File

@ -403,18 +403,17 @@ impl ContextCompressor {
// Strip tool_calls from any assistant in the head whose results
// were dropped (previously in the middle section).
for msg in &mut truncated[..self.config.protect_first_n] {
if msg.role == "assistant" {
if let Some(ref tcs) = msg.tool_calls
&& !tcs.is_empty()
{
let names: Vec<&str> = tcs.iter().map(|tc| tc.name.as_str()).collect();
msg.content = format!(
"{}\n\n[Tool calls ({}) — results dropped during truncation]",
msg.content,
names.join(", ")
);
msg.tool_calls = None;
}
if msg.role == "assistant"
&& let Some(ref tcs) = msg.tool_calls
&& !tcs.is_empty()
{
let names: Vec<&str> = tcs.iter().map(|tc| tc.name.as_str()).collect();
msg.content = format!(
"{}\n\n[Tool calls ({}) — results dropped during truncation]",
msg.content,
names.join(", ")
);
msg.tool_calls = None;
}
}
@ -537,22 +536,19 @@ impl ContextCompressor {
summary
);
let key = format!("ctx_compressed_{}", uuid::Uuid::new_v4());
let mm = self.memory.clone();
let sid = self.session_id.clone();
tokio::spawn(async move {
if let Err(e) = mm
.store(
&key,
&timeline_content,
crate::memory::MemoryCategory::Timeline,
sid.as_deref(),
Some(0.3),
)
.await
{
tracing::warn!(error = %e, "Failed to store compressed context as timeline");
}
});
if let Err(e) = self
.memory
.store(
&key,
&timeline_content,
crate::memory::MemoryCategory::Timeline,
self.session_id.as_deref(),
Some(0.3),
)
.await
{
tracing::warn!(error = %e, "Failed to store compressed context as timeline");
}
// Add summary as a special user message
new_messages.push(ChatMessage::user(format!(
@ -564,9 +560,7 @@ impl ContextCompressor {
// Add last user and everything after (protected)
let last_user_idx = user_indices[user_indices.len() - 1];
for i in last_user_idx..history.len() {
new_messages.push(history[i].clone());
}
new_messages.extend_from_slice(&history[last_user_idx..]);
// Remove orphan tool results whose declaring tool_calls were compressed away
Self::repair_tool_pairs(&mut new_messages);
@ -786,7 +780,7 @@ mod tests {
let mut messages = vec![
ChatMessage::user("Hello"),
ChatMessage::tool("call1", "bash", &"x".repeat(200)),
ChatMessage::tool("call1", "bash", "x".repeat(200)),
];
let modified = compressor.fast_trim_tool_results(&mut messages, 2);
@ -820,7 +814,7 @@ mod tests {
let messages = vec![
ChatMessage::user("Hi"),
ChatMessage::tool("call1", "bash", &"x".repeat(3000)),
ChatMessage::tool("call1", "bash", "x".repeat(3000)),
];
let result = compressor

6
src/agent/media_handler.rs
View File

@ -67,6 +67,12 @@ pub struct MediaHandlerRegistry {
handlers: HashMap<String, Box<dyn MediaHandler>>,
}
impl Default for MediaHandlerRegistry {
fn default() -> Self {
Self::new()
}
}
impl MediaHandlerRegistry {
pub fn new() -> Self {
Self {

220
src/agent/sub_agent.rs
View File

@ -3,6 +3,7 @@ use std::sync::Arc;
use std::time::Instant;
use dashmap::DashMap;
use tokio::sync::Semaphore;
use tokio_util::sync::CancellationToken;
use uuid::Uuid;
@ -117,9 +118,11 @@ pub struct SubAgentManager {
full_tools: Arc<ToolRegistry>,
storage: Option<Arc<crate::storage::Storage>>,
active_tasks: Arc<DashMap<String, CancellationToken>>,
background_permits: Arc<Semaphore>,
notify_tx: tokio::sync::mpsc::UnboundedSender<TaskNotification>,
max_concurrent_background_tasks: usize,
skills_loader: Option<Arc<SkillsLoader>>,
task_supervisor: crate::task_supervisor::TaskSupervisor,
}
impl SubAgentManager {
@ -130,15 +133,18 @@ impl SubAgentManager {
notify_tx: tokio::sync::mpsc::UnboundedSender<TaskNotification>,
max_concurrent_background_tasks: usize,
skills_loader: Option<Arc<SkillsLoader>>,
task_supervisor: crate::task_supervisor::TaskSupervisor,
) -> Self {
Self {
provider_config,
full_tools,
storage,
active_tasks: Arc::new(DashMap::new()),
background_permits: Arc::new(Semaphore::new(max_concurrent_background_tasks)),
notify_tx,
max_concurrent_background_tasks,
skills_loader,
task_supervisor,
}
}
@ -158,12 +164,10 @@ impl SubAgentManager {
fn get_skills_prompt(&self, tools: &ToolRegistry) -> Option<String> {
let has_get_skill = tools.iter().iter().any(|(name, _)| name == "get_skill");
if has_get_skill {
if let Some(ref loader) = self.skills_loader {
let prompt = loader.build_skills_prompt();
if !prompt.is_empty() {
return Some(prompt);
}
if has_get_skill && let Some(ref loader) = self.skills_loader {
let prompt = loader.build_skills_prompt();
if !prompt.is_empty() {
return Some(prompt);
}
}
None
@ -300,7 +304,7 @@ impl SubAgentManager {
.collect();
let results = futures_util::future::join_all(futures).await;
Ok(results.into_iter().collect::<Result<Vec<_>, _>>()?)
results.into_iter().collect::<Result<Vec<_>, _>>()
}
pub async fn run_background(
@ -308,11 +312,11 @@ impl SubAgentManager {
config: SubAgentConfig,
ctx: DelegateContext,
) -> Result<String, SubAgentError> {
if self.active_tasks.len() >= self.max_concurrent_background_tasks {
return Err(SubAgentError::TooManyTasks(
self.max_concurrent_background_tasks,
));
}
let permit = self
.background_permits
.clone()
.try_acquire_owned()
.map_err(|_| SubAgentError::TooManyTasks(self.max_concurrent_background_tasks))?;
let task_id = generate_task_id();
let cancel_token = CancellationToken::new();
@ -370,6 +374,7 @@ impl SubAgentManager {
let storage = self.storage.clone();
let notify_tx = self.notify_tx.clone();
let active_tasks = Arc::clone(&self.active_tasks);
let shutdown = self.task_supervisor.cancellation_token();
let tid = task_id.clone();
let sess_id = ctx.session_id.clone();
@ -377,30 +382,34 @@ impl SubAgentManager {
let cid = ctx.chat_id.clone();
let prompt = config.prompt.clone();
tokio::spawn(async move {
let spawned = self.task_supervisor.spawn_graceful(
format!("sub-agent:{task_id}"),
async move {
let _permit = permit;
let started_at = chrono::Utc::now().timestamp_millis();
// Update DB: running
if let Some(ref s) = storage {
let _ = s
.update_background_task_status(
&tid,
"running",
None,
None,
Some(started_at),
None,
)
.update_background_task_status(&tid, crate::storage::background_task::BackgroundTaskUpdate {
status: "running",
result: None,
error: None,
started_at: Some(started_at),
finished_at: None,
tool_calls_count: None,
iterations: None,
})
.await;
}
let mut provider = create_provider(provider_config.clone()).ok();
if let Some(ref mut p) = provider {
if let Some(ref s) = storage {
p.set_storage(s.clone());
}
if let Some(ref mut p) = provider
&& let Some(ref s) = storage
{
p.set_storage(s.clone());
}
let provider_result: Option<Arc<dyn LLMProvider>> = provider.map(|p| Arc::from(p));
let provider_result: Option<Arc<dyn LLMProvider>> = provider.map(Arc::from);
let result = match provider_result {
Some(provider) => {
@ -425,14 +434,20 @@ impl SubAgentManager {
agent.process(history),
) => {
match r {
Ok(Ok(agent_result)) => SubAgentResult {
task_id: tid.clone(),
content: agent_result.final_response.content,
content_truncated: false,
status: TaskStatus::Completed,
tool_calls_count: 0,
iterations: 0,
duration_ms: 0,
Ok(Ok(agent_result)) => {
let tool_calls_count = agent_result.emitted_messages
.iter().filter(|m| m.tool_calls.is_some()).count();
let iterations = agent_result.emitted_messages
.iter().filter(|m| m.role == "assistant" && m.tool_calls.is_some()).count();
SubAgentResult {
task_id: tid.clone(),
content: agent_result.final_response.content,
content_truncated: false,
status: TaskStatus::Completed,
tool_calls_count,
iterations,
duration_ms: 0,
}
},
Ok(Err(e)) => SubAgentResult {
task_id: tid.clone(),
@ -463,6 +478,15 @@ impl SubAgentManager {
iterations: 0,
duration_ms: 0,
},
_ = shutdown.cancelled() => SubAgentResult {
task_id: tid.clone(),
content: String::new(),
content_truncated: false,
status: TaskStatus::Cancelled,
tool_calls_count: 0,
iterations: 0,
duration_ms: 0,
},
}
}
None => SubAgentResult {
@ -488,14 +512,15 @@ impl SubAgentManager {
if let Some(ref s) = storage {
let _ = s
.update_background_task_status(
&tid,
&status_str,
Some(&result.content),
error_val.as_deref(),
Some(started_at),
Some(finished_at),
)
.update_background_task_status(&tid, crate::storage::background_task::BackgroundTaskUpdate {
status: &status_str,
result: Some(&result.content),
error: error_val.as_deref(),
started_at: Some(started_at),
finished_at: Some(finished_at),
tool_calls_count: Some(result.tool_calls_count as i64),
iterations: Some(result.iterations as i64),
})
.await;
}
@ -511,6 +536,29 @@ impl SubAgentManager {
active_tasks.remove(&tid);
});
if !spawned {
self.active_tasks.remove(&task_id);
if let Some(ref storage) = self.storage {
let _ = storage
.update_background_task_status(
&task_id,
crate::storage::background_task::BackgroundTaskUpdate {
status: "cancelled",
result: None,
error: Some("gateway shutdown"),
started_at: None,
finished_at: Some(chrono::Utc::now().timestamp_millis()),
tool_calls_count: None,
iterations: None,
},
)
.await;
}
return Err(SubAgentError::Other(
"gateway is shutting down and cannot accept background tasks".to_string(),
));
}
Ok(task_id)
}
@ -520,11 +568,15 @@ impl SubAgentManager {
if let Some(ref s) = self.storage {
s.update_background_task_status(
task_id,
"cancelled",
None,
None,
None,
Some(chrono::Utc::now().timestamp_millis()),
crate::storage::background_task::BackgroundTaskUpdate {
status: "cancelled",
result: None,
error: None,
started_at: None,
finished_at: Some(chrono::Utc::now().timestamp_millis()),
tool_calls_count: None,
iterations: None,
},
)
.await
.map_err(|e| SubAgentError::Storage(e.to_string()))?;
@ -566,12 +618,12 @@ impl SubAgentManager {
pub async fn cancel_by_session(&self, session_id: &str) {
// Cancel all running tasks for a session by checking DB
if let Some(ref s) = self.storage {
if let Ok(tasks) = s.list_background_tasks(session_id).await {
for task in &tasks {
if task.status == "pending" || task.status == "running" {
let _ = self.cancel_task(&task.id).await;
}
if let Some(ref s) = self.storage
&& let Ok(tasks) = s.list_background_tasks(session_id).await
{
for task in &tasks {
if task.status == "pending" || task.status == "running" {
let _ = self.cancel_task(&task.id).await;
}
}
}
@ -621,3 +673,65 @@ fn summarize_for_notification(content: &str, _duration_ms: u64) -> String {
format!("{}...", &content[..truncate_at])
}
}
#[cfg(test)]
mod tests {
use super::*;
use std::collections::HashMap;
fn manager(max_tasks: usize) -> SubAgentManager {
let (notify_tx, _notify_rx) = tokio::sync::mpsc::unbounded_channel();
SubAgentManager::new(
LLMProviderConfig {
provider_type: "openai".into(),
name: "test".into(),
base_url: "http://localhost".into(),
api_key: "test".into(),
extra_headers: HashMap::new(),
model_id: "test".into(),
temperature: None,
max_tokens: None,
model_extra: HashMap::new(),
max_tool_iterations: 1,
token_limit: 4096,
workspace_dir: std::env::temp_dir(),
input_types: vec!["text".into()],
},
Arc::new(ToolRegistry::new()),
None,
notify_tx,
max_tasks,
None,
crate::task_supervisor::TaskSupervisor::new(),
)
}
#[tokio::test]
async fn background_limit_is_enforced_by_atomic_permit() {
let manager = manager(1);
let _permit = manager
.background_permits
.clone()
.try_acquire_owned()
.unwrap();
let error = manager
.run_background(
SubAgentConfig {
prompt: "test".into(),
mode: ExecutionMode::Background,
allowed_tools: None,
max_iterations: None,
timeout_secs: Some(1),
},
DelegateContext {
session_id: "cli:test:dialog".into(),
channel: "cli".into(),
chat_id: "test".into(),
},
)
.await
.unwrap_err();
assert!(matches!(error, SubAgentError::TooManyTasks(1)));
}
}

342
src/bus/dispatcher.rs
View File

@ -1,27 +1,44 @@
use std::collections::HashMap;
use std::sync::Arc;
use std::time::Duration;
use tokio::sync::mpsc;
use crate::bus::{MessageBus, OutboundMessage};
use crate::channels::ChannelManager;
use crate::channels::base::{Channel, ChannelError};
use crate::task_supervisor::TaskSupervisor;
/// OutboundDispatcher consumes outbound messages from the MessageBus
/// and dispatches them to the appropriate Channel
const LANE_CAPACITY: usize = 64;
const LANE_IDLE_TIMEOUT: Duration = Duration::from_secs(300);
const SEND_TIMEOUT: Duration = Duration::from_secs(30);
/// Dispatches outbound messages through independent per-conversation lanes.
/// Messages to the same channel/chat remain ordered, while a slow destination
/// cannot block delivery to unrelated conversations.
pub struct OutboundDispatcher {
bus: Arc<MessageBus>,
channel_manager: ChannelManager,
task_supervisor: TaskSupervisor,
}
impl OutboundDispatcher {
pub fn new(bus: Arc<MessageBus>, channel_manager: ChannelManager) -> Self {
pub fn new(
bus: Arc<MessageBus>,
channel_manager: ChannelManager,
task_supervisor: TaskSupervisor,
) -> Self {
Self {
bus,
channel_manager,
task_supervisor,
}
}
/// Run the dispatcher loop - consumes from bus and dispatches to channels
pub async fn run(&self) {
tracing::info!("OutboundDispatcher started");
tracing::info!(lane_capacity = LANE_CAPACITY, "OutboundDispatcher started");
let mut lanes: HashMap<String, mpsc::Sender<OutboundMessage>> = HashMap::new();
let mut messages_seen = 0_u64;
loop {
let Some(msg) = self.bus.consume_outbound().await else {
@ -29,46 +46,309 @@ impl OutboundDispatcher {
break;
};
let channel_name = msg.channel.clone();
let channel = self.channel_manager.get_channel(&channel_name).await;
messages_seen = messages_seen.wrapping_add(1);
if messages_seen.is_multiple_of(128) {
lanes.retain(|_, sender| !sender.is_closed());
}
match channel {
Some(ch) => {
if let Err(e) = self.send_with_retry(&*ch, msg).await {
tracing::error!(channel = %channel_name, error = %e, "Failed to send message after retries");
}
let lane_key = format!("{}\0{}", msg.channel, msg.chat_id);
let mut sender = lanes.get(&lane_key).cloned();
if sender.as_ref().is_none_or(mpsc::Sender::is_closed) {
let Some(channel) = self.channel_manager.get_channel(&msg.channel).await else {
tracing::warn!(channel = %msg.channel, "No channel found for message");
msg.complete_delivery(Err(format!("channel not found: {}", msg.channel)));
continue;
};
let (new_sender, receiver) = mpsc::channel(LANE_CAPACITY);
if !self.spawn_lane(channel, receiver, msg.channel.clone(), msg.chat_id.clone()) {
msg.complete_delivery(Err("dispatcher is shutting down".to_string()));
continue;
}
None => {
tracing::warn!(channel = %channel_name, "No channel found for message");
lanes.insert(lane_key.clone(), new_sender.clone());
sender = Some(new_sender);
}
let Some(sender) = sender else {
tracing::error!("Outbound lane creation did not produce a sender");
continue;
};
match sender.try_send(msg) {
Ok(()) => {}
Err(mpsc::error::TrySendError::Full(msg)) => {
tracing::error!(
channel = %msg.channel,
chat_id = %msg.chat_id,
capacity = LANE_CAPACITY,
"Outbound lane full; rejecting message instead of blocking other destinations"
);
msg.complete_delivery(Err("outbound lane is full".to_string()));
}
Err(mpsc::error::TrySendError::Closed(msg)) => {
// The lane may have expired between the closed check and
// enqueue. Recreate it once and preserve this message.
lanes.remove(&lane_key);
let Some(channel) = self.channel_manager.get_channel(&msg.channel).await else {
tracing::warn!(channel = %msg.channel, "No channel found for message");
msg.complete_delivery(Err(format!("channel not found: {}", msg.channel)));
continue;
};
let (new_sender, receiver) = mpsc::channel(LANE_CAPACITY);
if !self.spawn_lane(channel, receiver, msg.channel.clone(), msg.chat_id.clone())
{
msg.complete_delivery(Err("dispatcher is shutting down".to_string()));
continue;
}
match new_sender.try_send(msg) {
Ok(()) => {
lanes.insert(lane_key, new_sender);
}
Err(error) => {
error.into_inner().complete_delivery(Err(
"outbound lane could not be restarted during shutdown".to_string(),
));
}
}
}
}
}
}
/// Send a message with exponential retry
async fn send_with_retry(
fn spawn_lane(
&self,
channel: Arc<dyn Channel + Send + Sync>,
mut receiver: mpsc::Receiver<OutboundMessage>,
channel_name: String,
chat_id: String,
) -> bool {
self.task_supervisor.spawn(
format!("outbound-lane:{channel_name}:{chat_id}"),
async move {
loop {
let msg = match tokio::time::timeout(LANE_IDLE_TIMEOUT, receiver.recv()).await {
Ok(Some(msg)) => msg,
Ok(None) | Err(_) => break,
};
let result = Self::send_with_retry(&*channel, &msg).await;
if let Err(error) = &result {
tracing::error!(
channel = %channel_name,
chat_id = %chat_id,
error = %error,
"Failed to send message after retries"
);
}
msg.complete_delivery(result.map_err(|error| error.to_string()));
}
},
)
}
async fn send_with_retry(
channel: &dyn Channel,
msg: OutboundMessage,
msg: &OutboundMessage,
) -> Result<(), ChannelError> {
const DELAYS: &[u64] = &[1, 2, 4];
for (i, &delay) in DELAYS.iter().enumerate() {
match channel.send(msg.clone()).await {
Ok(()) => return Ok(()),
Err(e) if i < DELAYS.len() - 1 => {
tracing::warn!(
attempt = i + 1,
delay = delay,
error = %e,
"Send failed, retrying"
);
tokio::time::sleep(tokio::time::Duration::from_secs(delay)).await;
for (attempt, &delay) in DELAYS.iter().enumerate() {
let result = tokio::time::timeout(SEND_TIMEOUT, channel.send(msg.clone())).await;
match result {
Ok(Ok(())) => return Ok(()),
Ok(Err(error)) if attempt < DELAYS.len() - 1 && error.is_transient() => {
tracing::warn!(attempt = attempt + 1, delay, error = %error, "Send failed, retrying");
}
Ok(Err(error)) => return Err(error),
Err(_) if attempt < DELAYS.len() - 1 => {
tracing::warn!(attempt = attempt + 1, delay, "Send timed out, retrying");
}
Err(_) => {
return Err(ChannelError::Other(format!(
"send timed out after {} seconds",
SEND_TIMEOUT.as_secs()
)));
}
Err(e) => return Err(e),
}
tokio::time::sleep(Duration::from_secs(delay)).await;
}
// All retries exhausted - should not reach here as last iteration returns
Ok(())
unreachable!()
}
}
#[cfg(test)]
mod tests {
use super::*;
use async_trait::async_trait;
use std::sync::atomic::{AtomicUsize, Ordering};
use tokio::sync::{Mutex, Notify};
struct RecordingChannel {
sent: Mutex<Vec<String>>,
notify: Notify,
}
struct PermanentFailureChannel {
attempts: AtomicUsize,
}
#[async_trait]
impl Channel for PermanentFailureChannel {
fn name(&self) -> &str {
"permanent-failure"
}
fn is_running(&self) -> bool {
true
}
async fn start(&self, _bus: Arc<MessageBus>) -> Result<(), ChannelError> {
Ok(())
}
async fn stop(&self) -> Result<(), ChannelError> {
Ok(())
}
async fn send(&self, _msg: OutboundMessage) -> Result<(), ChannelError> {
self.attempts.fetch_add(1, Ordering::SeqCst);
Err(ChannelError::Other("invalid destination".to_string()))
}
}
#[async_trait]
impl Channel for RecordingChannel {
fn name(&self) -> &str {
"recording"
}
fn is_running(&self) -> bool {
true
}
async fn start(&self, _bus: Arc<MessageBus>) -> Result<(), ChannelError> {
Ok(())
}
async fn stop(&self) -> Result<(), ChannelError> {
Ok(())
}
async fn send(&self, msg: OutboundMessage) -> Result<(), ChannelError> {
if msg.chat_id == "slow" {
tokio::time::sleep(Duration::from_millis(50)).await;
}
self.sent.lock().await.push(msg.content);
self.notify.notify_waiters();
Ok(())
}
}
fn outbound(chat_id: &str, content: &str) -> OutboundMessage {
OutboundMessage {
channel: "recording".to_string(),
chat_id: chat_id.to_string(),
content: content.to_string(),
reply_to: None,
media: vec![],
metadata: HashMap::new(),
delivery: None,
}
}
#[tokio::test]
async fn slow_conversation_does_not_block_other_conversations() {
let bus = MessageBus::new(8);
let manager = ChannelManager::with_bus(
Arc::new(crate::channels::CliChatChannel::new()),
bus.clone(),
);
let channel = Arc::new(RecordingChannel {
sent: Mutex::new(Vec::new()),
notify: Notify::new(),
});
manager.register_channel("recording", channel.clone()).await;
let supervisor = TaskSupervisor::new();
let dispatcher = OutboundDispatcher::new(bus.clone(), manager, supervisor.clone());
let task = tokio::spawn(async move { dispatcher.run().await });
bus.publish_outbound(outbound("slow", "slow-1"))
.await
.unwrap();
bus.publish_outbound(outbound("slow", "slow-2"))
.await
.unwrap();
bus.publish_outbound(outbound("fast", "fast-1"))
.await
.unwrap();
tokio::time::timeout(Duration::from_secs(1), async {
loop {
if channel.sent.lock().await.len() == 3 {
break;
}
tokio::time::sleep(Duration::from_millis(5)).await;
}
})
.await
.unwrap();
let sent = channel.sent.lock().await.clone();
assert_eq!(sent[0], "fast-1");
assert_eq!(&sent[1..], &["slow-1", "slow-2"]);
task.abort();
supervisor.shutdown(Duration::from_secs(1)).await;
}
#[tokio::test]
async fn confirmed_delivery_reports_missing_channel() {
let bus = MessageBus::new(8);
let manager = ChannelManager::with_bus(
Arc::new(crate::channels::CliChatChannel::new()),
bus.clone(),
);
let supervisor = TaskSupervisor::new();
let dispatcher = OutboundDispatcher::new(bus.clone(), manager, supervisor.clone());
let task = tokio::spawn(async move { dispatcher.run().await });
let mut message = outbound("missing", "not delivered");
message.channel = "missing".to_string();
let error = bus.deliver_outbound(message).await.unwrap_err();
assert!(matches!(error, crate::bus::BusError::DeliveryFailed(_)));
task.abort();
supervisor.shutdown(Duration::from_secs(1)).await;
}
#[tokio::test]
async fn confirmed_delivery_waits_for_channel_send() {
let bus = MessageBus::new(8);
let manager = ChannelManager::with_bus(
Arc::new(crate::channels::CliChatChannel::new()),
bus.clone(),
);
let channel = Arc::new(RecordingChannel {
sent: Mutex::new(Vec::new()),
notify: Notify::new(),
});
manager.register_channel("recording", channel.clone()).await;
let supervisor = TaskSupervisor::new();
let dispatcher = OutboundDispatcher::new(bus.clone(), manager, supervisor.clone());
let task = tokio::spawn(async move { dispatcher.run().await });
bus.deliver_outbound(outbound("confirmed", "delivered"))
.await
.unwrap();
assert_eq!(channel.sent.lock().await.as_slice(), &["delivered"]);
task.abort();
supervisor.shutdown(Duration::from_secs(1)).await;
}
#[tokio::test]
async fn permanent_send_failure_is_not_retried() {
let channel = PermanentFailureChannel {
attempts: AtomicUsize::new(0),
};
let error = OutboundDispatcher::send_with_retry(&channel, &outbound("invalid", "message"))
.await
.unwrap_err();
assert!(matches!(error, ChannelError::Other(_)));
assert_eq!(channel.attempts.load(Ordering::SeqCst), 1);
}
}

13
src/bus/message.rs
View File

@ -264,12 +264,6 @@ pub struct InboundMessage {
pub forwarded_metadata: HashMap<String, String>,
}
impl InboundMessage {
pub fn session_key(&self) -> String {
format!("{}:{}", self.channel, self.chat_id)
}
}
// ============================================================================
// OutboundMessage - Message from Agent to Channel (bot response)
// ============================================================================
@ -282,11 +276,14 @@ pub struct OutboundMessage {
pub reply_to: Option<String>,
pub media: Vec<MediaItem>,
pub metadata: HashMap<String, String>,
pub(crate) delivery: Option<tokio::sync::watch::Sender<Option<Result<(), String>>>>,
}
impl OutboundMessage {
pub fn is_stream_delta(&self) -> bool {
self.metadata.get("_stream_delta").is_some()
pub(crate) fn complete_delivery(&self, result: Result<(), String>) {
if let Some(delivery) = &self.delivery {
delivery.send_replace(Some(result));
}
}
}

23
src/bus/mod.rs
View File

@ -68,6 +68,25 @@ impl MessageBus {
.map_err(|_| BusError::Closed)
}
/// Publish an outbound message and wait for the dispatcher to report the
/// actual channel delivery result.
pub async fn deliver_outbound(&self, mut msg: OutboundMessage) -> Result<(), BusError> {
let (delivery_tx, mut delivery_rx) = tokio::sync::watch::channel(None);
msg.delivery = Some(delivery_tx);
self.publish_outbound(msg).await?;
tokio::time::timeout(std::time::Duration::from_secs(120), async {
loop {
delivery_rx.changed().await.map_err(|_| BusError::Closed)?;
if let Some(result) = delivery_rx.borrow().clone() {
return result.map_err(BusError::DeliveryFailed);
}
}
})
.await
.map_err(|_| BusError::DeliveryTimedOut)?
}
/// Consume an outbound message (Dispatcher -> Bus)
pub async fn consume_outbound(&self) -> Option<OutboundMessage> {
self.outbound_rx.lock().await.recv().await
@ -95,12 +114,16 @@ impl MessageBus {
#[derive(Debug)]
pub enum BusError {
Closed,
DeliveryFailed(String),
DeliveryTimedOut,
}
impl std::fmt::Display for BusError {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
match self {
BusError::Closed => write!(f, "Bus channel closed"),
BusError::DeliveryFailed(error) => write!(f, "Outbound delivery failed: {error}"),
BusError::DeliveryTimedOut => write!(f, "Outbound delivery confirmation timed out"),
}
}
}

6
src/channels/base.rs
View File

@ -26,6 +26,12 @@ impl std::fmt::Display for ChannelError {
impl std::error::Error for ChannelError {}
impl ChannelError {
pub fn is_transient(&self) -> bool {
matches!(self, Self::ConnectionError(_) | Self::SendError(_))
}
}
impl From<BusError> for ChannelError {
fn from(e: BusError) -> Self {
ChannelError::BusError(e.to_string())

294
src/channels/cli_chat.rs
View File

@ -1,9 +1,10 @@
use async_trait::async_trait;
use std::collections::HashMap;
use std::sync::Arc;
use tokio::sync::{Mutex, mpsc};
use crate::bus::{ControlMessage, InboundMessage, MessageBus, OutboundMessage};
use crate::protocol::{SlashCommandInfo, WsInbound, WsOutbound, parse_inbound};
use crate::protocol::{HistoryMessage, SlashCommandInfo, WsInbound, WsOutbound, parse_inbound};
use crate::session::{SessionCommand, SessionEvent, UnifiedSessionId};
use super::base::{Channel, ChannelError};
@ -18,13 +19,19 @@ pub(crate) struct Client {
current_session_id: Mutex<Option<String>>,
}
impl Client {
pub(crate) fn chat_id(&self) -> &str {
&self.chat_id
}
}
// ============================================================================
// CliChatChannel - Channel implementation for CLI chat
// ============================================================================
pub struct CliChatChannel {
bus: std::sync::Mutex<Option<Arc<MessageBus>>>,
clients: Mutex<Vec<Arc<Client>>>,
clients: Mutex<HashMap<String, Arc<Client>>>,
}
impl Default for CliChatChannel {
@ -37,7 +44,7 @@ impl CliChatChannel {
pub fn new() -> Self {
Self {
bus: std::sync::Mutex::new(None),
clients: Mutex::new(Vec::new()),
clients: Mutex::new(HashMap::new()),
}
}
@ -45,25 +52,30 @@ impl CliChatChannel {
pub(crate) async fn register_client(
&self,
sender: mpsc::Sender<WsOutbound>,
requested_chat_id: Option<String>,
) -> (String, Arc<Client>) {
// Each WebSocket connection gets a stable chat scope. All user input and
// dialog controls for this client stay inside that scope unless the
// protocol explicitly carries a full session id.
let chat_id = crate::util::short_id();
let chat_id = requested_chat_id.unwrap_or_else(crate::util::short_id);
let client = Arc::new(Client {
sender,
chat_id: chat_id.clone(),
current_session_id: Mutex::new(None),
});
self.clients.lock().await.push(client.clone());
self.clients
.lock()
.await
.insert(chat_id.clone(), client.clone());
// Create initial session via control message
let session_id = match self.create_session_via_control(&chat_id, None).await {
Ok((id, _title)) => id,
// Resume the current/most-recent dialog for a stable TUI identity. Only
// create a dialog when this client has never connected before.
let session_id = match self.resume_session_via_control(&chat_id).await {
Ok(id) => id,
Err(e) => {
tracing::error!(error = %e, "Failed to create initial session");
UnifiedSessionId::new("cli_chat", &chat_id, &crate::util::short_id()).to_string()
tracing::error!(error = %e, "Failed to resume initial session");
UnifiedSessionId::new("cli_chat", &chat_id, crate::util::short_id()).to_string()
}
};
@ -76,6 +88,16 @@ impl CliChatChannel {
(session_id, client)
}
pub(crate) async fn unregister_client(&self, client: &Arc<Client>) {
let mut clients = self.clients.lock().await;
if clients
.get(client.chat_id())
.is_some_and(|registered| Arc::ptr_eq(registered, client))
{
clients.remove(client.chat_id());
}
}
/// Handle an inbound message from a client
pub(crate) async fn handle_inbound(&self, client: Arc<Client>, raw_msg: &str) {
match parse_inbound(raw_msg) {
@ -143,10 +165,13 @@ impl CliChatChannel {
} => {
let (reply_tx, mut reply_rx) = mpsc::channel(1);
let session_id = if let Some(session_id) = session_id {
UnifiedSessionId::parse(&session_id).ok_or_else(|| {
ChannelError::Other("Invalid session ID format".to_string())
})?
Self::parse_client_session(&client, &session_id)?
} else if let Some(chat_id) = chat_id {
if chat_id != client.chat_id {
return Err(ChannelError::Other(
"Chat does not belong to this client".to_string(),
));
}
let (current_tx, mut current_rx) = mpsc::channel(1);
bus.publish_control(ControlMessage {
op: SessionCommand::GetCurrentDialog {
@ -177,9 +202,7 @@ impl CliChatChannel {
let target = current_session_guard
.clone()
.ok_or_else(|| ChannelError::Other("No active session".to_string()))?;
UnifiedSessionId::parse(&target).ok_or_else(|| {
ChannelError::Other("Invalid session ID format".to_string())
})?
Self::parse_client_session(&client, &target)?
};
let target = session_id.to_string();
bus.publish_control(ControlMessage {
@ -324,14 +347,56 @@ impl CliChatChannel {
}
}
}
WsInbound::GetSessionHistory { session_id, limit } => {
let unified_id = Self::parse_client_session(&client, &session_id)?;
let (reply_tx, mut reply_rx) = mpsc::channel(1);
bus.publish_control(ControlMessage {
op: SessionCommand::GetDialogHistory {
session_id: unified_id,
limit: limit.unwrap_or(1_000).clamp(1, 2_000),
},
reply_tx,
})
.await?;
match reply_rx.recv().await {
Some(Ok(SessionEvent::DialogHistory {
session_id,
messages,
})) => {
let messages = messages
.into_iter()
.filter(|message| !message.content.is_empty())
.map(|message| HistoryMessage {
id: message.id,
seq: message.seq,
role: message.role,
content: message.content,
created_at: message.created_at,
})
.collect();
let _ = client
.sender
.send(WsOutbound::SessionHistory {
session_id: session_id.to_string(),
messages,
})
.await;
}
Some(Ok(_)) => {}
Some(Err(e)) => return Err(e),
None => {
return Err(ChannelError::Other("Control channel closed".to_string()));
}
}
}
WsInbound::RenameSession { session_id, title } => {
let target = session_id
.or(current_session_guard.clone())
.ok_or_else(|| ChannelError::Other("No active session".to_string()))?;
let (reply_tx, mut reply_rx) = mpsc::channel(1);
let unified_id = UnifiedSessionId::parse(&target)
.ok_or_else(|| ChannelError::Other("Invalid session ID format".to_string()))?;
let unified_id = Self::parse_client_session(&client, &target)?;
bus.publish_control(ControlMessage {
op: SessionCommand::RenameDialog {
session_id: unified_id,
@ -369,8 +434,7 @@ impl CliChatChannel {
let was_current = current_session_guard.as_deref() == Some(&target);
let (reply_tx, mut reply_rx) = mpsc::channel(1);
let unified_id = UnifiedSessionId::parse(&target)
.ok_or_else(|| ChannelError::Other("Invalid session ID format".to_string()))?;
let unified_id = Self::parse_client_session(&client, &target)?;
bus.publish_control(ControlMessage {
op: SessionCommand::ArchiveDialog {
session_id: unified_id,
@ -418,8 +482,7 @@ impl CliChatChannel {
.ok_or_else(|| ChannelError::Other("No active session".to_string()))?;
let (reply_tx, mut reply_rx) = mpsc::channel(1);
let unified_id = UnifiedSessionId::parse(&target)
.ok_or_else(|| ChannelError::Other("Invalid session ID format".to_string()))?;
let unified_id = Self::parse_client_session(&client, &target)?;
bus.publish_control(ControlMessage {
op: SessionCommand::DeleteDialog {
session_id: unified_id,
@ -555,6 +618,78 @@ impl CliChatChannel {
None => Err(ChannelError::Other("Control channel closed".to_string())),
}
}
async fn resume_session_via_control(&self, chat_id: &str) -> Result<String, ChannelError> {
let bus = {
let guard = self.bus.lock().unwrap();
guard
.clone()
.ok_or_else(|| ChannelError::Other("Channel not started".to_string()))?
};
let (reply_tx, mut reply_rx) = mpsc::channel(1);
bus.publish_control(ControlMessage {
op: SessionCommand::GetCurrentDialog {
channel: "cli_chat".to_string(),
chat_id: chat_id.to_string(),
},
reply_tx,
})
.await?;
if let Some(Ok(SessionEvent::CurrentDialog {
session_id: Some(session_id),
})) = reply_rx.recv().await
{
return Ok(session_id.to_string());
}
let (reply_tx, mut reply_rx) = mpsc::channel(1);
bus.publish_control(ControlMessage {
op: SessionCommand::ListDialogs {
channel: "cli_chat".to_string(),
chat_id: chat_id.to_string(),
include_archived: false,
},
reply_tx,
})
.await?;
if let Some(Ok(SessionEvent::DialogList { dialogs, .. })) = reply_rx.recv().await
&& let Some(dialog) = dialogs.first()
{
let session_id = dialog.session_id.clone();
let (reply_tx, mut reply_rx) = mpsc::channel(1);
bus.publish_control(ControlMessage {
op: SessionCommand::SwitchDialog {
channel: session_id.channel.clone(),
chat_id: session_id.chat_id.clone(),
dialog_id: session_id.dialog_id.clone(),
},
reply_tx,
})
.await?;
if let Some(Ok(SessionEvent::DialogSwitched { session_id })) = reply_rx.recv().await {
return Ok(session_id.to_string());
}
}
self.create_session_via_control(chat_id, None)
.await
.map(|(session_id, _)| session_id)
}
fn parse_client_session(
client: &Client,
session_id: &str,
) -> Result<UnifiedSessionId, ChannelError> {
let unified_id = UnifiedSessionId::parse(session_id)
.ok_or_else(|| ChannelError::Other("Invalid session ID format".to_string()))?;
if unified_id.channel != "cli_chat" || unified_id.chat_id != client.chat_id {
return Err(ChannelError::Other(
"Session does not belong to this client".to_string(),
));
}
Ok(unified_id)
}
}
#[async_trait]
@ -574,29 +709,112 @@ impl Channel for CliChatChannel {
async fn stop(&self) -> Result<(), ChannelError> {
*self.bus.lock().unwrap() = None;
self.clients.lock().await.clear();
Ok(())
}
async fn send(&self, msg: OutboundMessage) -> Result<(), ChannelError> {
let clients = self.clients.lock().await.clone();
for client in clients {
if client.chat_id != msg.chat_id {
continue;
let client = self.clients.lock().await.get(&msg.chat_id).cloned();
let Some(client) = client else {
tracing::debug!(chat_id = %msg.chat_id, "No active CLI client for outbound message");
return Ok(());
};
let message_type = msg.metadata.get("_type").map(String::as_str);
let session_id = msg.metadata.get("_session_id").cloned();
let outbound = if message_type == Some("notification") {
WsOutbound::SystemNotification {
content: msg.content,
session_id,
}
let outbound = if msg.metadata.get("_type").map(|v| v.as_str()) == Some("notification")
} else if message_type == Some("command") {
WsOutbound::CommandExecuted {
message: msg.content,
}
} else {
WsOutbound::AssistantResponse {
id: crate::util::short_id(),
content: msg.content,
role: "assistant".to_string(),
session_id,
}
};
if client.sender.send(outbound).await.is_err() {
let mut clients = self.clients.lock().await;
if clients
.get(&msg.chat_id)
.is_some_and(|registered| Arc::ptr_eq(registered, &client))
{
WsOutbound::SystemNotification {
content: msg.content.clone(),
}
} else {
WsOutbound::AssistantResponse {
id: crate::util::short_id(),
content: msg.content.clone(),
role: "assistant".to_string(),
}
};
let _ = client.sender.send(outbound).await;
clients.remove(&msg.chat_id);
}
}
Ok(())
}
}
#[cfg(test)]
mod tests {
use super::*;
#[tokio::test]
async fn failed_sender_is_pruned_from_client_registry() {
let channel = CliChatChannel::new();
let (sender, receiver) = mpsc::channel(1);
drop(receiver);
let client = Arc::new(Client {
sender,
chat_id: "dead-client".to_string(),
current_session_id: Mutex::new(None),
});
channel
.clients
.lock()
.await
.insert("dead-client".to_string(), client);
channel
.send(OutboundMessage {
channel: "cli_chat".to_string(),
chat_id: "dead-client".to_string(),
content: "message".to_string(),
reply_to: None,
media: Vec::new(),
metadata: Default::default(),
delivery: None,
})
.await
.unwrap();
assert!(channel.clients.lock().await.is_empty());
}
#[tokio::test]
async fn stale_connection_cannot_unregister_replacement() {
let channel = CliChatChannel::new();
let (old_sender, _old_receiver) = mpsc::channel(1);
let (new_sender, _new_receiver) = mpsc::channel(1);
let old = Arc::new(Client {
sender: old_sender,
chat_id: "stable-client".to_string(),
current_session_id: Mutex::new(None),
});
let replacement = Arc::new(Client {
sender: new_sender,
chat_id: "stable-client".to_string(),
current_session_id: Mutex::new(None),
});
channel
.clients
.lock()
.await
.insert("stable-client".to_string(), replacement.clone());
channel.unregister_client(&old).await;
let registered = channel.clients.lock().await;
assert!(
registered
.get("stable-client")
.is_some_and(|client| Arc::ptr_eq(client, &replacement))
);
}
}

286
src/channels/feishu.rs
View File

@ -8,7 +8,9 @@ use futures_util::{SinkExt, StreamExt};
use prost::{Message as ProstMessage, bytes::Bytes};
use regex::Regex;
use serde::Deserialize;
use tokio::sync::{RwLock, broadcast};
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};
@ -20,6 +22,11 @@ const FEISHU_WS_BASE: &str = "https://open.feishu.cn";
/// Heartbeat timeout for WS connection — must be larger than ping_interval (default 120 s).
/// If no binary frame (pong or event) is received within this window, reconnect.
const WS_HEARTBEAT_TIMEOUT: Duration = Duration::from_secs(300);
const CHANNEL_STOP_GRACE: Duration = if cfg!(test) {
Duration::from_millis(100)
} else {
Duration::from_secs(5)
};
/// Refresh tenant token this many seconds before the announced expiry.
const TOKEN_REFRESH_SKEW: Duration = Duration::from_secs(120);
/// Default tenant token TTL when `expire`/`expires_in` is absent.
@ -89,7 +96,6 @@ struct LarkEvent {
#[derive(Deserialize)]
struct LarkEventHeader {
event_type: String,
#[allow(dead_code)]
event_id: String,
}
@ -121,19 +127,13 @@ struct LarkSenderId {
}
#[derive(Deserialize)]
#[allow(dead_code)]
struct LarkMessage {
message_id: String,
chat_id: String,
chat_type: String,
message_type: String,
#[serde(default)]
content: String,
#[serde(default)]
mentions: Vec<serde_json::Value>,
#[serde(default)]
root_id: Option<String>,
#[serde(default)]
parent_id: Option<String>,
}
@ -151,7 +151,8 @@ pub struct FeishuChannel {
config: FeishuChannelConfig,
http_client: reqwest::Client,
running: Arc<RwLock<bool>>,
shutdown_tx: Arc<RwLock<Option<broadcast::Sender<()>>>>,
shutdown: Arc<RwLock<Option<CancellationToken>>>,
run_task: Arc<Mutex<Option<JoinHandle<()>>>>,
connected: Arc<RwLock<bool>>,
/// Cached tenant access token with proactive refresh.
tenant_token: Arc<RwLock<Option<CachedTenantToken>>>,
@ -184,7 +185,8 @@ impl FeishuChannel {
config,
http_client: reqwest::Client::new(),
running: Arc::new(RwLock::new(false)),
shutdown_tx: Arc::new(RwLock::new(None)),
shutdown: Arc::new(RwLock::new(None)),
run_task: Arc::new(Mutex::new(None)),
connected: Arc::new(RwLock::new(false)),
tenant_token: Arc::new(RwLock::new(None)),
seen_message_ids: Arc::new(RwLock::new(HashMap::new())),
@ -1170,18 +1172,22 @@ impl FeishuChannel {
async fn run_ws_loop(
&self,
bus: Arc<MessageBus>,
mut shutdown_rx: broadcast::Receiver<()>,
shutdown: CancellationToken,
) -> Result<(), ChannelError> {
let (wss_url, client_config) = self.get_ws_endpoint(&self.http_client).await?;
let (wss_url, client_config) = tokio::select! {
result = self.get_ws_endpoint(&self.http_client) => result?,
_ = shutdown.cancelled() => return Ok(()),
};
let service_id = Self::extract_service_id(&wss_url);
tracing::info!(url = %wss_url, "Connecting to Feishu WebSocket");
let (ws_stream, _) = tokio_tungstenite::connect_async(&wss_url)
.await
.map_err(|e| {
let (ws_stream, _) = tokio::select! {
result = tokio_tungstenite::connect_async(&wss_url) => result.map_err(|e| {
ChannelError::ConnectionError(format!("WebSocket connection failed: {}", e))
})?;
})?,
_ = shutdown.cancelled() => return Ok(()),
};
*self.connected.write().await = true;
tracing::info!("Feishu WebSocket connected");
@ -1200,14 +1206,14 @@ impl FeishuChannel {
}],
payload: None,
};
write
.send(tokio_tungstenite::tungstenite::Message::Binary(
tokio::select! {
result = write.send(tokio_tungstenite::tungstenite::Message::Binary(
ping_frame.encode_to_vec().into(),
))
.await
.map_err(|e| {
)) => result.map_err(|e| {
ChannelError::ConnectionError(format!("Failed to send initial ping: {}", e))
})?;
})?,
_ = shutdown.cancelled() => return Ok(()),
};
let ping_interval = client_config.ping_interval.unwrap_or(120).max(10);
let mut ping_interval_tok =
@ -1256,29 +1262,24 @@ impl FeishuChannel {
forwarded_metadata.insert("feishu.parent_id".to_string(), pid.clone());
}
// Publish to bus asynchronously
let channel = self.clone();
let bus = bus.clone();
tokio::spawn(async move {
#[cfg(debug_assertions)]
tracing::debug!(open_id = %parsed.open_id, chat_id = %parsed.chat_id, content_len = %parsed.content.len(), media_count = %parsed.media.len(), "Publishing message to bus");
let msg = crate::bus::InboundMessage {
channel: "feishu".to_string(),
sender_id: parsed.open_id.clone(),
chat_id: parsed.chat_id.clone(),
content: parsed.content.clone(),
timestamp: crate::bus::message::current_timestamp(),
media: parsed.media.clone(),
metadata: std::collections::HashMap::new(),
forwarded_metadata,
};
if let Err(e) = self.handle_and_publish(&bus, &msg).await {
tracing::error!(error = %e, open_id = %parsed.open_id, chat_id = %parsed.chat_id, "Failed to publish Feishu message to bus");
} else {
#[cfg(debug_assertions)]
tracing::debug!(open_id = %parsed.open_id, chat_id = %parsed.chat_id, content_len = %parsed.content.len(), media_count = %parsed.media.len(), "Publishing message to bus");
let msg = crate::bus::InboundMessage {
channel: "feishu".to_string(),
sender_id: parsed.open_id.clone(),
chat_id: parsed.chat_id.clone(),
content: parsed.content.clone(),
timestamp: crate::bus::message::current_timestamp(),
media: parsed.media.clone(),
metadata: std::collections::HashMap::new(),
forwarded_metadata,
};
if let Err(e) = channel.handle_and_publish(&bus, &msg).await {
tracing::error!(error = %e, open_id = %parsed.open_id, chat_id = %parsed.chat_id, "Failed to publish Feishu message to bus");
} else {
#[cfg(debug_assertions)]
tracing::debug!(open_id = %parsed.open_id, chat_id = %parsed.chat_id, "Message published to bus successfully");
}
});
tracing::debug!(open_id = %parsed.open_id, chat_id = %parsed.chat_id, "Message published to bus successfully");
}
}
Ok(None) => {}
Err(e) => {
@ -1345,7 +1346,7 @@ impl FeishuChannel {
let mut seen = self.seen_message_ids.write().await;
seen.retain(|_, ts| now.duration_since(*ts) < DEDUP_CACHE_TTL);
}
_ = shutdown_rx.recv() => {
_ = shutdown.cancelled() => {
tracing::info!("Feishu channel shutdown signal received");
break;
}
@ -1357,52 +1358,6 @@ impl FeishuChannel {
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn collect_post_image_keys_finds_nested_images() {
let content = serde_json::json!({
"zh_cn": {
"title": "",
"content": [[
{"tag": "img", "image_key": "img_v3_001"},
{"tag": "text", "text": "这是哪里？"},
{"tag": "img", "image_key": "img_v3_002"},
{"tag": "img", "image_key": "img_v3_001"}
]]
}
})
.to_string();
assert_eq!(
collect_post_image_keys(&content),
vec!["img_v3_001".to_string(), "img_v3_002".to_string()]
);
}
#[test]
fn parse_post_content_preserves_image_positions() {
let content = serde_json::json!({
"zh_cn": {
"title": "",
"content": [[
{"tag": "text", "text": "这是一张图："},
{"tag": "img", "image_key": "img_v3_001"},
{"tag": "text", "text": "看完继续说"}
]]
}
})
.to_string();
assert_eq!(
parse_post_content(&content),
"这是一张图：[image]看完继续说"
);
}
}
fn parse_post_content(content: &str) -> String {
/// Extract text from a single post element (text, link, at-mention).
fn extract_element(el: &serde_json::Value, out: &mut Vec<String>) {
@ -1732,13 +1687,8 @@ fn collect_list_items(items: &[serde_json::Value], lines: &mut Vec<String>, dept
collect_list_items(children, lines, depth + 1);
}
} else if let Some(children_arr) = item.as_array().and_then(|arr| {
arr.iter().find_map(|child| {
if child.as_object().and_then(|o| o.get("children")).is_some() {
Some(child)
} else {
None
}
})
arr.iter()
.find(|child| child.as_object().and_then(|o| o.get("children")).is_some())
}) && let Some(children) = children_arr
.as_object()
.and_then(|o| o.get("children"))
@ -1819,13 +1769,12 @@ fn resolve_image_ext(content_type: &str) -> &str {
}
fn resolve_file_ext(content_json: &serde_json::Value) -> String {
if let Some(name) = content_json.get("file_name").and_then(|v| v.as_str()) {
if let Some(ext) = std::path::Path::new(name)
if let Some(name) = content_json.get("file_name").and_then(|v| v.as_str())
&& let Some(ext) = std::path::Path::new(name)
.extension()
.and_then(|e| e.to_str())
{
return ext.to_string();
}
{
return ext.to_string();
}
String::new()
}
@ -2005,14 +1954,22 @@ impl Channel for FeishuChannel {
));
}
let mut run_task = self.run_task.lock().await;
if run_task.as_ref().is_some_and(|task| !task.is_finished()) {
return Ok(());
}
if let Some(finished) = run_task.take() {
let _ = finished.await;
}
*self.running.write().await = true;
let (shutdown_tx, _) = broadcast::channel(1);
*self.shutdown_tx.write().await = Some(shutdown_tx.clone());
let shutdown = CancellationToken::new();
*self.shutdown.write().await = Some(shutdown.clone());
let channel = self.clone();
let bus = bus.clone();
tokio::spawn(async move {
*run_task = Some(tokio::spawn(async move {
let mut consecutive_failures = 0;
let max_failures = 3;
@ -2021,8 +1978,7 @@ impl Channel for FeishuChannel {
break;
}
let shutdown_rx = shutdown_tx.subscribe();
match channel.run_ws_loop(bus.clone(), shutdown_rx).await {
match channel.run_ws_loop(bus.clone(), shutdown.clone()).await {
Ok(_) => {
tracing::info!("Feishu WebSocket disconnected");
}
@ -2036,17 +1992,20 @@ impl Channel for FeishuChannel {
}
}
if !*channel.running.read().await {
if !*channel.running.read().await || shutdown.is_cancelled() {
break;
}
tracing::info!("Feishu channel retrying in 5s...");
tokio::time::sleep(tokio::time::Duration::from_secs(5)).await;
tokio::select! {
_ = tokio::time::sleep(tokio::time::Duration::from_secs(5)) => {}
_ = shutdown.cancelled() => break,
}
}
*channel.running.write().await = false;
tracing::info!("Feishu channel stopped");
});
}));
tracing::info!("Feishu channel started");
Ok(())
@ -2054,9 +2013,27 @@ impl Channel for FeishuChannel {
async fn stop(&self) -> Result<(), ChannelError> {
*self.running.write().await = false;
*self.connected.write().await = false;
if let Some(tx) = self.shutdown_tx.write().await.take() {
let _ = tx.send(());
if let Some(shutdown) = self.shutdown.write().await.take() {
shutdown.cancel();
}
let task = { self.run_task.lock().await.take() };
if let Some(mut task) = task {
match tokio::time::timeout(CHANNEL_STOP_GRACE, &mut task).await {
Ok(result) => result.map_err(|error| {
ChannelError::Other(format!("Feishu channel task failed to join: {error}"))
})?,
Err(_) => {
tracing::warn!(
grace_ms = CHANNEL_STOP_GRACE.as_millis(),
"Feishu channel did not stop in time; aborting connection task"
);
task.abort();
let _ = task.await;
}
}
}
Ok(())
@ -2271,3 +2248,86 @@ impl Channel for FeishuChannel {
Ok(())
}
}
#[cfg(test)]
mod tests {
use super::*;
fn test_channel() -> FeishuChannel {
FeishuChannel::new(
FeishuChannelConfig {
enabled: true,
app_id: "test-app".to_string(),
app_secret: "test-secret".to_string(),
allow_from: vec!["*".to_string()],
agent: String::new(),
media_dir: String::new(),
reaction_emoji: "THUMBSUP".to_string(),
},
Path::new("/tmp"),
)
.expect("test channel should be valid")
}
#[tokio::test]
async fn stop_aborts_connection_task_that_ignores_cancellation() {
let channel = test_channel();
let shutdown = CancellationToken::new();
*channel.running.write().await = true;
*channel.connected.write().await = true;
*channel.shutdown.write().await = Some(shutdown.clone());
*channel.run_task.lock().await = Some(tokio::spawn(std::future::pending()));
tokio::time::timeout(Duration::from_secs(1), channel.stop())
.await
.expect("stop must have a hard deadline")
.expect("stop should succeed after aborting the stuck task");
assert!(shutdown.is_cancelled());
assert!(!channel.is_running());
assert!(!*channel.connected.read().await);
assert!(channel.run_task.lock().await.is_none());
}
#[test]
fn collect_post_image_keys_finds_nested_images() {
let content = serde_json::json!({
"zh_cn": {
"title": "",
"content": [[
{"tag": "img", "image_key": "img_v3_001"},
{"tag": "text", "text": "这是哪里？"},
{"tag": "img", "image_key": "img_v3_002"},
{"tag": "img", "image_key": "img_v3_001"}
]]
}
})
.to_string();
assert_eq!(
collect_post_image_keys(&content),
vec!["img_v3_001".to_string(), "img_v3_002".to_string()]
);
}
#[test]
fn parse_post_content_preserves_image_positions() {
let content = serde_json::json!({
"zh_cn": {
"title": "",
"content": [[
{"tag": "text", "text": "这是一张图："},
{"tag": "img", "image_key": "img_v3_001"},
{"tag": "text", "text": "看完继续说"}
]]
}
})
.to_string();
assert_eq!(
parse_post_content(&content),
"这是一张图：[image]看完继续说"
);
}
}

157
src/channels/manager.rs
View File

@ -2,7 +2,7 @@ use std::collections::HashMap;
use std::sync::Arc;
use tokio::sync::RwLock;
use crate::bus::{MessageBus, OutboundMessage};
use crate::bus::MessageBus;
use crate::channels::base::{Channel, ChannelError};
use crate::channels::feishu::FeishuChannel;
use crate::config::Config;
@ -16,14 +16,6 @@ pub struct ChannelManager {
}
impl ChannelManager {
pub fn new(cli_chat_channel: Arc<crate::channels::CliChatChannel>) -> Self {
Self {
channels: Arc::new(RwLock::new(HashMap::new())),
cli_chat_channel,
bus: MessageBus::new(100),
}
}
pub fn with_bus(
cli_chat_channel: Arc<crate::channels::CliChatChannel>,
bus: Arc<MessageBus>,
@ -91,27 +83,53 @@ impl ChannelManager {
}
pub async fn start_all(&self) -> Result<(), ChannelError> {
let channels = self.channels.read().await;
let channels: Vec<_> = self
.channels
.read()
.await
.iter()
.map(|(name, channel)| (name.clone(), channel.clone()))
.collect();
let bus = self.bus.clone();
for (name, channel) in channels.iter() {
let mut failures = Vec::new();
for (name, channel) in channels {
tracing::info!(channel = %name, "Starting channel");
if let Err(e) = channel.start(bus.clone()).await {
tracing::error!(channel = %name, error = %e, "Failed to start channel");
failures.push(format!("{name}: {e}"));
}
}
Ok(())
if failures.is_empty() {
Ok(())
} else {
Err(ChannelError::Other(format!(
"failed to start channels: {}",
failures.join("; ")
)))
}
}
pub async fn stop_all(&self) -> Result<(), ChannelError> {
let mut channels = self.channels.write().await;
for (name, channel) in channels.iter() {
let channels: Vec<_> = {
let mut registered = self.channels.write().await;
registered.drain().collect()
};
let mut failures = Vec::new();
for (name, channel) in channels {
tracing::info!(channel = %name, "Stopping channel");
if let Err(e) = channel.stop().await {
tracing::error!(channel = %name, error = %e, "Error stopping channel");
failures.push(format!("{name}: {e}"));
}
}
channels.clear();
Ok(())
if failures.is_empty() {
Ok(())
} else {
Err(ChannelError::Other(format!(
"failed to stop channels: {}",
failures.join("; ")
)))
}
}
pub async fn get_channel(&self, name: &str) -> Option<Arc<dyn Channel + Send + Sync>> {
@ -122,17 +140,104 @@ impl ChannelManager {
pub async fn list_channel_names(&self) -> Vec<String> {
self.channels.read().await.keys().cloned().collect()
}
}
/// Dispatch an outbound message to the appropriate channel
pub async fn dispatch(&self, msg: OutboundMessage) -> Result<(), ChannelError> {
let channel_name = &msg.channel;
if let Some(channel) = self.get_channel(channel_name).await {
channel.send(msg).await
} else {
Err(ChannelError::Other(format!(
"Channel not found: {}",
channel_name
)))
#[cfg(test)]
mod tests {
use super::*;
use async_trait::async_trait;
use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering};
struct TestChannel {
name: &'static str,
fail_start: bool,
fail_stop: bool,
running: AtomicBool,
starts: AtomicUsize,
stops: AtomicUsize,
}
impl TestChannel {
fn new(name: &'static str, fail_start: bool, fail_stop: bool) -> Self {
Self {
name,
fail_start,
fail_stop,
running: AtomicBool::new(false),
starts: AtomicUsize::new(0),
stops: AtomicUsize::new(0),
}
}
}
#[async_trait]
impl Channel for TestChannel {
fn name(&self) -> &str {
self.name
}
fn is_running(&self) -> bool {
self.running.load(Ordering::SeqCst)
}
async fn start(&self, _bus: Arc<MessageBus>) -> Result<(), ChannelError> {
self.starts.fetch_add(1, Ordering::SeqCst);
if self.fail_start {
return Err(ChannelError::ConnectionError("unavailable".into()));
}
self.running.store(true, Ordering::SeqCst);
Ok(())
}
async fn stop(&self) -> Result<(), ChannelError> {
self.stops.fetch_add(1, Ordering::SeqCst);
self.running.store(false, Ordering::SeqCst);
if self.fail_stop {
return Err(ChannelError::Other("stop failed".into()));
}
Ok(())
}
async fn send(&self, _msg: crate::bus::OutboundMessage) -> Result<(), ChannelError> {
Ok(())
}
}
fn manager() -> ChannelManager {
ChannelManager::with_bus(
Arc::new(crate::channels::CliChatChannel::new()),
MessageBus::new(8),
)
}
#[tokio::test]
async fn start_all_reports_channel_failures_and_starts_the_rest() {
let manager = manager();
let healthy = Arc::new(TestChannel::new("healthy", false, false));
let broken = Arc::new(TestChannel::new("broken", true, false));
manager.register_channel("healthy", healthy.clone()).await;
manager.register_channel("broken", broken.clone()).await;
let error = manager.start_all().await.unwrap_err().to_string();
assert!(error.contains("broken"));
assert_eq!(healthy.starts.load(Ordering::SeqCst), 1);
assert_eq!(broken.starts.load(Ordering::SeqCst), 1);
}
#[tokio::test]
async fn stop_all_reports_failures_and_unregisters_every_channel() {
let manager = manager();
let healthy = Arc::new(TestChannel::new("healthy", false, false));
let broken = Arc::new(TestChannel::new("broken", false, true));
manager.register_channel("healthy", healthy.clone()).await;
manager.register_channel("broken", broken.clone()).await;
let error = manager.stop_all().await.unwrap_err().to_string();
assert!(error.contains("broken"));
assert_eq!(healthy.stops.load(Ordering::SeqCst), 1);
assert_eq!(broken.stops.load(Ordering::SeqCst), 1);
assert!(manager.list_channel_names().await.is_empty());
}
}

171
src/client/mod.rs
View File

@ -3,31 +3,37 @@ pub use crate::protocol::{WsInbound, WsOutbound, serialize_inbound, serialize_ou
mod tui;
use crate::client::tui::app::{App, MessageRole};
use crate::client::tui::event::handle_key_event;
use crate::client::tui::event::{
handle_key_event, handle_paste, request_history, request_session_list, send,
};
use crate::client::tui::ui::render_ui;
use crossterm::{
event::{self, Event},
event::{self, DisableBracketedPaste, EnableBracketedPaste, Event, KeyEventKind},
execute,
terminal::{EnterAlternateScreen, LeaveAlternateScreen, disable_raw_mode, enable_raw_mode},
};
use futures_util::{SinkExt, StreamExt};
use futures_util::StreamExt;
use ratatui::{Terminal, prelude::CrosstermBackend};
use std::io;
use std::{fs, path::PathBuf};
use tokio_tungstenite::{connect_async, tungstenite::Message};
pub async fn run(gateway_url: &str) -> Result<(), Box<dyn std::error::Error>> {
let (ws_stream, _) = connect_async(gateway_url).await?;
tracing::info!(url = %gateway_url, "Connected to gateway");
let client_id = load_or_create_client_id();
let separator = if gateway_url.contains('?') { '&' } else { '?' };
let connect_url = format!("{gateway_url}{separator}client_id={client_id}");
let (ws_stream, _) = connect_async(&connect_url).await?;
tracing::info!("Connected to gateway");
let (ws_sender, ws_receiver) = ws_stream.split();
let mut app = App::new(gateway_url.to_string());
let mut app = App::new();
app.ws_sender = Some(ws_sender);
app.ws_receiver = Some(ws_receiver);
enable_raw_mode()?;
let mut stdout = io::stdout();
execute!(stdout, EnterAlternateScreen)?;
execute!(stdout, EnterAlternateScreen, EnableBracketedPaste)?;
let backend = CrosstermBackend::new(stdout);
let mut terminal = Terminal::new(backend)?;
terminal.clear()?;
@ -35,33 +41,57 @@ pub async fn run(gateway_url: &str) -> Result<(), Box<dyn std::error::Error>> {
let result = run_app(&mut terminal, app).await;
// Cleanup terminal, ignore errors
let _ = execute!(terminal.backend_mut(), LeaveAlternateScreen);
let _ = execute!(
terminal.backend_mut(),
DisableBracketedPaste,
LeaveAlternateScreen
);
let _ = disable_raw_mode();
let _ = terminal.show_cursor();
result
}
fn load_or_create_client_id() -> String {
let generated = uuid::Uuid::new_v4().simple().to_string();
let Some(home) = dirs::home_dir() else {
return generated;
};
let dir = home.join(".picobot");
let path: PathBuf = dir.join("tui_client_id");
if let Ok(value) = fs::read_to_string(&path) {
let value = value.trim();
if !value.is_empty()
&& value.len() <= 64
&& value
.bytes()
.all(|byte| byte.is_ascii_alphanumeric() || byte == b'-' || byte == b'_')
{
return value.to_string();
}
}
if fs::create_dir_all(dir).is_ok() {
let _ = fs::write(path, &generated);
}
generated
}
async fn run_app(
terminal: &mut Terminal<CrosstermBackend<io::Stdout>>,
mut app: App,
) -> Result<(), Box<dyn std::error::Error>> {
let mut ws_receiver = app.ws_receiver.take().unwrap();
let mut event_reader = event::EventStream::new();
let mut ws_open = true;
// Request command list on startup
if let Some(sender) = &mut app.ws_sender {
let inbound = WsInbound::GetSlashCommands;
if let Ok(text) = serialize_inbound(&inbound) {
let _ = sender.send(Message::Text(text.into())).await;
}
}
send(&mut app, WsInbound::GetSlashCommands).await;
request_session_list(&mut app).await;
loop {
terminal.draw(|f| render_ui(f, &app))?;
tokio::select! {
msg = ws_receiver.next() => {
msg = ws_receiver.next(), if ws_open => {
match msg {
Some(Ok(Message::Text(text))) => {
if let Ok(outbound) = serde_json::from_str::<WsOutbound>(&text) {
@ -70,14 +100,30 @@ async fn run_app(
}
Some(Ok(Message::Close(_))) | None => {
tracing::info!("Gateway disconnected");
app.quit();
app.connected = false;
app.ws_sender = None;
ws_open = false;
app.status_message = Some("Gateway 连接已关闭；按两次 Ctrl+C 退出".to_string());
}
Some(Err(error)) => {
app.connected = false;
app.ws_sender = None;
ws_open = false;
app.status_message = Some(format!("Gateway 连接错误：{error}"));
}
_ => {}
}
}
event_result = event_reader.next() => {
if let Some(Ok(Event::Key(key))) = event_result {
handle_key_event(&mut app, key).await;
match event_result {
Some(Ok(Event::Key(key))) if key.kind != KeyEventKind::Release => {
handle_key_event(&mut app, key).await;
}
Some(Ok(Event::Paste(text))) => handle_paste(&mut app, &text),
Some(Err(error)) => {
app.status_message = Some(format!("终端输入错误：{error}"));
}
_ => {}
}
}
}
@ -92,17 +138,39 @@ async fn run_app(
async fn handle_ws_message(app: &mut App, outbound: WsOutbound) {
match outbound {
WsOutbound::AssistantResponse { content, .. } => {
app.add_message(MessageRole::Assistant, content);
WsOutbound::AssistantResponse {
content,
session_id,
..
} => {
app.pending_responses = app.pending_responses.saturating_sub(1);
app.status_message = None;
if session_id
.as_ref()
.is_none_or(|session_id| app.current_session_id.as_ref() == Some(session_id))
{
app.add_message(MessageRole::Assistant, content);
} else {
app.status_message = Some("另一个会话已完成响应".to_string());
}
request_session_list(app).await;
}
WsOutbound::Error { message, .. } => {
app.pending_responses = app.pending_responses.saturating_sub(1);
app.status_message = Some(message.clone());
app.add_message(MessageRole::System, format!("Error: {}", message));
}
WsOutbound::SessionEstablished { session_id } => {
app.set_current_session(Some(session_id));
app.connected = true;
app.set_current_session(Some(session_id.clone()));
request_history(app, session_id).await;
request_session_list(app).await;
}
WsOutbound::SessionCreated { session_id, .. } => {
app.set_current_session(Some(session_id));
app.set_current_session(Some(session_id.clone()));
app.status_message = None;
request_history(app, session_id).await;
request_session_list(app).await;
}
WsOutbound::SessionList {
sessions,
@ -110,31 +178,72 @@ async fn handle_ws_message(app: &mut App, outbound: WsOutbound) {
} => {
app.set_sessions(sessions);
if let Some(id) = current_session_id {
app.set_current_session(Some(id));
let changed = app.current_session_id.as_deref() != Some(&id);
app.set_current_session(Some(id.clone()));
if changed {
request_history(app, id).await;
}
}
}
WsOutbound::SessionLoaded { session_id, .. } => {
app.set_current_session(Some(session_id));
app.set_current_session(Some(session_id.clone()));
request_history(app, session_id).await;
request_session_list(app).await;
}
WsOutbound::SessionHistory {
session_id,
messages,
} => app.set_history(&session_id, messages),
WsOutbound::SessionRenamed { session_id, title } => {
if let Some(session) = app
.sessions
.iter_mut()
.find(|session| session.session_id == session_id)
{
session.title = title;
}
request_session_list(app).await;
}
WsOutbound::SessionArchived { session_id } => {
app.sessions
.retain(|session| session.session_id != session_id);
request_session_list(app).await;
}
WsOutbound::SessionRenamed { .. } => {}
WsOutbound::SessionArchived { .. } => {}
WsOutbound::SessionDeleted { session_id } => {
if app.current_session_id.as_ref() == Some(&session_id) {
app.set_current_session(None);
}
app.sessions
.retain(|session| session.session_id != session_id);
request_session_list(app).await;
}
WsOutbound::HistoryCleared { .. } => {
app.messages.clear();
WsOutbound::HistoryCleared { session_id } => {
if app.current_session_id.as_deref() == Some(&session_id) {
app.messages.clear();
}
app.status_message = None;
request_session_list(app).await;
}
WsOutbound::SlashCommandsList { commands } => {
app.set_commands(commands);
}
WsOutbound::Pong => {}
WsOutbound::CommandExecuted { message } => {
app.pending_responses = app.pending_responses.saturating_sub(1);
app.status_message = None;
app.add_message(MessageRole::System, message);
request_session_list(app).await;
}
WsOutbound::SystemNotification { content } => {
app.add_message(MessageRole::System, content);
WsOutbound::SystemNotification {
content,
session_id,
} => {
if session_id
.as_ref()
.is_none_or(|session_id| app.current_session_id.as_ref() == Some(session_id))
{
app.add_message(MessageRole::System, content);
}
}
}
}

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

@ -1,17 +1,11 @@
#![allow(dead_code)]
use crate::protocol::{SessionSummary, SlashCommandInfo};
use crate::protocol::{HistoryMessage, SessionSummary, SlashCommandInfo};
use std::collections::VecDeque;
use tokio_tungstenite::tungstenite::Message;
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum FocusArea {
TitleBar,
SessionList,
ChatHistory,
InputArea,
}
const MAX_MESSAGES: usize = 2_000;
const MAX_INPUT_BYTES: usize = 16 * 1024;
#[derive(Debug, Clone)]
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum MessageRole {
User,
Assistant,
@ -24,8 +18,26 @@ pub struct ChatMessage {
pub content: String,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum Focus {
Input,
Sessions,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum ConfirmAction {
Archive,
Delete,
ClearHistory,
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum Modal {
Rename { input: String, cursor: usize },
Confirm(ConfirmAction),
}
pub struct App {
pub gateway_url: String,
pub ws_sender: Option<
futures_util::stream::SplitSink<
tokio_tungstenite::WebSocketStream<
@ -41,48 +53,49 @@ pub struct App {
>,
>,
>,
pub current_session_id: Option<String>,
pub sessions: Vec<SessionSummary>,
pub selected_session: usize,
pub show_archived: bool,
pub messages: VecDeque<ChatMessage>,
pub focus: FocusArea,
pub input: String,
/// UTF-8 byte offset. It is always maintained at a character boundary.
pub input_cursor_pos: usize,
pub focus: Focus,
pub modal: Option<Modal>,
pub show_help: bool,
pub chat_scroll_offset: u16,
pub session_scroll_offset: u16,
pub chat_scroll_from_bottom: u16,
pub should_quit: bool,
// Quit confirmation state (double Ctrl+C to exit)
pub ctrl_c_count: u8,
pub pending_quit: bool,
// Command menu state
pub connected: bool,
pub pending_responses: usize,
pub status_message: Option<String>,
pub commands: Vec<SlashCommandInfo>,
pub show_command_menu: bool,
pub selected_command_idx: u16,
pub selected_command_idx: usize,
}
impl App {
pub fn new(gateway_url: String) -> Self {
pub fn new() -> Self {
Self {
gateway_url,
ws_sender: None,
ws_receiver: None,
current_session_id: None,
sessions: Vec::new(),
selected_session: 0,
show_archived: false,
messages: VecDeque::new(),
focus: FocusArea::InputArea,
input: String::new(),
input_cursor_pos: 0,
focus: Focus::Input,
modal: None,
show_help: false,
chat_scroll_offset: 0,
session_scroll_offset: 0,
chat_scroll_from_bottom: 0,
should_quit: false,
ctrl_c_count: 0,
pending_quit: false,
connected: true,
pending_responses: 0,
status_message: Some("正在加载会话…".to_string()),
commands: Vec::new(),
show_command_menu: false,
selected_command_idx: 0,
@ -91,67 +104,161 @@ impl App {
pub fn add_message(&mut self, role: MessageRole, content: String) {
self.messages.push_back(ChatMessage { role, content });
self.chat_scroll_offset = 0;
while self.messages.len() > MAX_MESSAGES {
self.messages.pop_front();
}
self.chat_scroll_from_bottom = 0;
}
pub fn set_history(&mut self, session_id: &str, messages: Vec<HistoryMessage>) {
if self.current_session_id.as_deref() != Some(session_id) {
return;
}
self.messages = messages
.into_iter()
.filter_map(|message| {
let role = match message.role.as_str() {
"user" => MessageRole::User,
"assistant" => MessageRole::Assistant,
"system" | "tool" => MessageRole::System,
_ => return None,
};
Some(ChatMessage {
role,
content: message.content,
})
})
.collect();
while self.messages.len() > MAX_MESSAGES {
self.messages.pop_front();
}
self.chat_scroll_from_bottom = 0;
self.status_message = None;
}
pub fn set_sessions(&mut self, sessions: Vec<SessionSummary>) {
self.sessions = sessions;
if let Some(current) = &self.current_session_id
&& let Some(index) = self
.sessions
.iter()
.position(|session| &session.session_id == current)
{
self.selected_session = index;
}
self.clamp_session_selection();
}
pub fn set_current_session(&mut self, session_id: Option<String>) {
self.current_session_id = session_id;
self.messages.clear();
if self.current_session_id != session_id {
self.current_session_id = session_id;
self.messages.clear();
self.chat_scroll_from_bottom = 0;
}
if let Some(current) = &self.current_session_id
&& let Some(index) = self
.sessions
.iter()
.position(|session| &session.session_id == current)
{
self.selected_session = index;
}
}
pub fn scroll_chat_up(&mut self) {
self.chat_scroll_offset = self.chat_scroll_offset.saturating_add(1);
pub fn current_title(&self) -> &str {
self.current_session_id
.as_ref()
.and_then(|id| {
self.sessions
.iter()
.find(|session| &session.session_id == id)
})
.map(|session| session.title.as_str())
.unwrap_or("新对话")
}
pub fn scroll_chat_down(&mut self) {
self.chat_scroll_offset = self.chat_scroll_offset.saturating_sub(1);
pub fn selected_session_id(&self) -> Option<String> {
self.sessions
.get(self.selected_session)
.map(|session| session.session_id.clone())
}
pub fn scroll_session_up(&mut self) {
self.session_scroll_offset = self.session_scroll_offset.saturating_add(1);
pub fn select_next_session(&mut self) {
if !self.sessions.is_empty() {
self.selected_session = (self.selected_session + 1).min(self.sessions.len() - 1);
}
}
pub fn scroll_session_down(&mut self) {
self.session_scroll_offset = self.session_scroll_offset.saturating_sub(1);
pub fn select_previous_session(&mut self) {
self.selected_session = self.selected_session.saturating_sub(1);
}
fn clamp_session_selection(&mut self) {
self.selected_session = self
.selected_session
.min(self.sessions.len().saturating_sub(1));
}
pub fn scroll_chat_up(&mut self, lines: u16) {
self.chat_scroll_from_bottom = self.chat_scroll_from_bottom.saturating_add(lines);
}
pub fn scroll_chat_down(&mut self, lines: u16) {
self.chat_scroll_from_bottom = self.chat_scroll_from_bottom.saturating_sub(lines);
}
pub fn input_insert_char(&mut self, c: char) {
self.input.insert(self.input_cursor_pos, c);
self.input_cursor_pos += 1;
if self.input.len() + c.len_utf8() <= MAX_INPUT_BYTES {
self.input.insert(self.input_cursor_pos, c);
self.input_cursor_pos += c.len_utf8();
}
}
pub fn input_insert_str(&mut self, text: &str) {
let remaining = MAX_INPUT_BYTES.saturating_sub(self.input.len());
let mut end = text.len().min(remaining);
while !text.is_char_boundary(end) {
end -= 1;
}
self.input.insert_str(self.input_cursor_pos, &text[..end]);
self.input_cursor_pos += end;
}
pub fn input_delete_char(&mut self) {
if self.input_cursor_pos > 0 {
self.input.remove(self.input_cursor_pos - 1);
self.input_cursor_pos -= 1;
if let Some(previous) = previous_boundary(&self.input, self.input_cursor_pos) {
self.input.drain(previous..self.input_cursor_pos);
self.input_cursor_pos = previous;
}
}
pub fn input_delete_forward(&mut self) {
if let Some(next) = next_boundary(&self.input, self.input_cursor_pos) {
self.input.drain(self.input_cursor_pos..next);
}
}
pub fn input_move_cursor_left(&mut self) {
self.input_cursor_pos = self.input_cursor_pos.saturating_sub(1);
if let Some(previous) = previous_boundary(&self.input, self.input_cursor_pos) {
self.input_cursor_pos = previous;
}
}
pub fn input_move_cursor_right(&mut self) {
if self.input_cursor_pos < self.input.len() {
self.input_cursor_pos += 1;
if let Some(next) = next_boundary(&self.input, self.input_cursor_pos) {
self.input_cursor_pos = next;
}
}
pub fn input_move_cursor_to_start(&mut self) {
self.input_cursor_pos = 0;
let line_start = self.input[..self.input_cursor_pos]
.rfind('\n')
.map_or(0, |index| index + 1);
self.input_cursor_pos = line_start;
}
pub fn input_move_cursor_to_end(&mut self) {
self.input_cursor_pos = self.input.len();
}
pub fn input_clear(&mut self) {
self.input.clear();
self.input_cursor_pos = 0;
let tail = &self.input[self.input_cursor_pos..];
self.input_cursor_pos += tail.find('\n').unwrap_or(tail.len());
}
pub fn take_input(&mut self) -> String {
@ -160,86 +267,108 @@ impl App {
input
}
pub fn toggle_help(&mut self) {
self.show_help = !self.show_help;
}
pub fn quit(&mut self) {
self.should_quit = true;
}
/// Handle Ctrl+C for quit confirmation (requires double press)
pub fn handle_ctrl_c_for_quit(&mut self) -> bool {
pub fn handle_ctrl_c_for_quit(&mut self) {
if self.pending_quit {
self.ctrl_c_count += 1;
if self.ctrl_c_count >= 2 {
self.should_quit = true;
return true;
}
false
self.should_quit = true;
} else {
self.pending_quit = true;
self.ctrl_c_count = 1;
false
self.status_message = Some("再次按 Ctrl+C 退出".to_string());
}
}
/// Cancel pending quit if user presses any other key
pub fn cancel_pending_quit(&mut self) {
self.pending_quit = false;
self.ctrl_c_count = 0;
if self.pending_quit {
self.pending_quit = false;
self.status_message = None;
}
}
// Command menu methods
pub fn set_commands(&mut self, commands: Vec<SlashCommandInfo>) {
self.commands = commands;
}
pub fn get_filtered_commands(&self) -> Vec<&SlashCommandInfo> {
let input_lower = self.input.to_lowercase();
let query = self
.input
.split_whitespace()
.next()
.unwrap_or("")
.to_lowercase();
self.commands
.iter()
.filter(|cmd| {
cmd.name.to_lowercase().contains(&input_lower)
|| cmd.description.to_lowercase().contains(&input_lower)
|| cmd
.filter(|command| {
command.name.to_lowercase().contains(&query)
|| command.description.to_lowercase().contains(&query)
|| command
.aliases
.iter()
.any(|a| a.to_lowercase().contains(&input_lower))
.any(|alias| alias.to_lowercase().starts_with(&query))
})
.collect()
}
pub fn select_next_command(&mut self) {
let filtered = self.get_filtered_commands();
if !filtered.is_empty() {
self.selected_command_idx = (self.selected_command_idx + 1) % filtered.len() as u16;
let len = self.get_filtered_commands().len();
if len > 0 {
self.selected_command_idx = (self.selected_command_idx + 1) % len;
}
}
pub fn select_prev_command(&mut self) {
let filtered = self.get_filtered_commands();
if !filtered.is_empty() {
self.selected_command_idx = if self.selected_command_idx == 0 {
filtered.len() as u16 - 1
} else {
self.selected_command_idx - 1
};
pub fn select_previous_command(&mut self) {
let len = self.get_filtered_commands().len();
if len > 0 {
self.selected_command_idx = (self.selected_command_idx + len - 1) % len;
}
}
pub fn get_selected_command(&self) -> Option<&SlashCommandInfo> {
let filtered = self.get_filtered_commands();
filtered.get(self.selected_command_idx as usize).copied()
}
pub fn insert_command(&mut self) {
if let Some(cmd) = self.get_selected_command() {
// Use the first alias as the command to insert
if let Some(alias) = cmd.aliases.first() {
self.input = alias.clone();
self.input_cursor_pos = self.input.len();
}
pub fn insert_selected_command(&mut self) {
let command = self
.get_filtered_commands()
.get(self.selected_command_idx)
.and_then(|command| command.aliases.first().cloned());
if let Some(command) = command {
self.input = format!("{command} ");
self.input_cursor_pos = self.input.len();
}
}
}
fn previous_boundary(value: &str, offset: usize) -> Option<usize> {
value[..offset]
.char_indices()
.next_back()
.map(|(index, _)| index)
}
fn next_boundary(value: &str, offset: usize) -> Option<usize> {
value[offset..]
.chars()
.next()
.map(|character| offset + character.len_utf8())
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn unicode_cursor_edits_only_at_character_boundaries() {
let mut app = App::new();
app.input_insert_str("你a🙂");
app.input_move_cursor_left();
app.input_delete_char();
assert_eq!(app.input, "你🙂");
assert!(app.input.is_char_boundary(app.input_cursor_pos));
app.input_delete_forward();
assert_eq!(app.input, "你");
}
#[test]
fn stale_history_does_not_replace_the_active_dialog() {
let mut app = App::new();
app.set_current_session(Some("new".to_string()));
app.add_message(MessageRole::User, "keep".to_string());
app.set_history("old", Vec::new());
assert_eq!(app.messages.len(), 1);
}
}

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

@ -3,35 +3,60 @@ use ratatui::{
Frame,
layout::Rect,
style::{Color, Modifier, Style},
text::Line,
widgets::{Block, Borders, List, ListItem},
text::{Line, Span},
widgets::{Block, Borders, Paragraph},
};
pub fn render(f: &mut Frame, area: Rect, app: &App) {
let items: Vec<ListItem> = app
.messages
.iter()
.map(|msg| {
let (prefix, color) = match msg.role {
MessageRole::User => ("[User] ", Color::Blue),
MessageRole::Assistant => ("[Assistant] ", Color::Green),
MessageRole::System => ("[System] ", Color::Red),
};
pub fn render(frame: &mut Frame, area: Rect, app: &App) {
let content_width = area.width.saturating_sub(4).max(1) as usize;
let mut lines = Vec::new();
if app.messages.is_empty() {
lines.push(Line::from("开始一段对话，或从左侧选择已有会话。"));
}
for message in &app.messages {
let (label, color) = match message.role {
MessageRole::User => ("你", Color::Blue),
MessageRole::Assistant => ("PicoBot", Color::Green),
MessageRole::System => ("系统", Color::Yellow),
};
lines.push(Line::from(Span::styled(
label,
Style::default().fg(color).add_modifier(Modifier::BOLD),
)));
for source_line in message.content.lines() {
let wrapped = textwrap::wrap(source_line, content_width);
if wrapped.is_empty() {
lines.push(Line::from(""));
} else {
lines.extend(
wrapped
.into_iter()
.map(|line| Line::from(line.into_owned())),
);
}
}
lines.push(Line::from(""));
}
if app.pending_responses > 0 {
lines.push(Line::from(Span::styled(
"● 正在思考…",
Style::default().fg(Color::Cyan),
)));
}
let content = vec![
Line::from(vec![ratatui::text::Span::styled(
prefix,
Style::default().fg(color).add_modifier(Modifier::BOLD),
)]),
Line::from(msg.content.as_str()),
Line::from(""),
];
ListItem::new(content)
})
.collect();
let list = List::new(items).block(Block::default().title("Conversation").borders(Borders::ALL));
f.render_widget(list, area);
let visible_height = area.height.saturating_sub(2);
let line_count = u16::try_from(lines.len()).unwrap_or(u16::MAX);
let max_scroll = line_count.saturating_sub(visible_height);
let scroll = max_scroll.saturating_sub(app.chat_scroll_from_bottom.min(max_scroll));
let title = if app.chat_scroll_from_bottom > 0 {
" 对话 · 已暂停自动滚动 "
} else {
" 对话 "
};
frame.render_widget(
Paragraph::new(lines)
.scroll((scroll, 0))
.block(Block::default().title(title).borders(Borders::ALL)),
area,
);
}

7
src/client/tui/components/command_menu.rs
View File

@ -4,7 +4,7 @@ use ratatui::{
layout::Rect,
style::{Color, Modifier, Style},
text::{Line, Span},
widgets::{Block, Borders, List, ListItem},
widgets::{Block, Borders, List, ListItem, ListState},
};
pub fn render(f: &mut Frame, area: Rect, app: &App) {
@ -18,7 +18,7 @@ pub fn render(f: &mut Frame, area: Rect, app: &App) {
.iter()
.enumerate()
.map(|(i, cmd)| {
let is_selected = i == app.selected_command_idx as usize;
let is_selected = i == app.selected_command_idx;
let style = if is_selected {
Style::default()
.fg(Color::White)
@ -48,5 +48,6 @@ pub fn render(f: &mut Frame, area: Rect, app: &App) {
)
.highlight_style(Style::default().add_modifier(Modifier::BOLD));
f.render_widget(list, area);
let mut state = ListState::default().with_selected(Some(app.selected_command_idx));
f.render_stateful_widget(list, area, &mut state);
}

75
src/client/tui/components/help_popup.rs
View File

@ -2,41 +2,52 @@ use ratatui::{
Frame,
layout::Rect,
style::{Color, Modifier, Style},
widgets::{Block, Borders, Clear, List, ListItem},
text::{Line, Span},
widgets::{Block, Borders, Clear, Paragraph, Wrap},
};
pub fn render(f: &mut Frame, area: Rect) {
f.render_widget(Clear, area);
let help_text = vec![
ListItem::new("Commands:"),
ListItem::new(" /new [title] - Archive current, start new"),
ListItem::new(" /sessions - List all conversations"),
ListItem::new(" /switch <id> - Switch to conversation"),
ListItem::new(" /rename <t> - Rename current conversation"),
ListItem::new(" /archive - Archive current conversation"),
ListItem::new(" /delete - Delete current conversation"),
ListItem::new(" /compact - Trigger context compression"),
ListItem::new(" /info - Show session information"),
ListItem::new(""),
ListItem::new("Keyboard:"),
ListItem::new(" Enter - Send message"),
ListItem::new(" Ctrl+C ×2 - Quit"),
ListItem::new(" ? - Show help"),
ListItem::new(" Arrow keys - Navigate"),
ListItem::new(" / - Show command menu"),
pub fn render(frame: &mut Frame, area: Rect) {
frame.render_widget(Clear, area);
let lines = vec![
Line::from(Span::styled(
"全局",
Style::default().add_modifier(Modifier::BOLD),
)),
Line::from(" F1 / Ctrl+H 帮助 Tab 切换输入/会话焦点"),
Line::from(" Ctrl+N 新会话 Ctrl+S 聚焦会话列表"),
Line::from(" Ctrl+R 重命名 Ctrl+O 显示/隐藏归档"),
Line::from(" Ctrl+A 归档 Ctrl+D 删除"),
Line::from(" Ctrl+L 清空历史 Ctrl+C 两次 退出"),
Line::from(""),
Line::from(Span::styled(
"输入",
Style::default().add_modifier(Modifier::BOLD),
)),
Line::from(" Enter 发送 Shift/Alt+Enter 换行"),
Line::from(" / 打开命令菜单 Tab 补全命令"),
Line::from(" PageUp/PageDown 滚动对话，Ctrl+↑/↓ 微调"),
Line::from(""),
Line::from(Span::styled(
"会话列表",
Style::default().add_modifier(Modifier::BOLD),
)),
Line::from(" ↑/↓ 或 j/k 选择 Enter 切换"),
Line::from(" n 新建 · r 重命名 · a 归档 · d 删除"),
Line::from(""),
Line::from(Span::styled(
"Esc / F1 关闭帮助",
Style::default().fg(Color::DarkGray),
)),
];
let list = List::new(help_text).block(
Block::default()
.title("Help")
.title_style(
Style::default()
.fg(Color::Cyan)
.add_modifier(Modifier::BOLD),
frame.render_widget(
Paragraph::new(lines)
.block(
Block::default()
.title(" 帮助 ")
.title_style(Style::default().fg(Color::Cyan))
.borders(Borders::ALL),
)
.borders(Borders::ALL),
.wrap(Wrap { trim: false }),
area,
);
f.render_widget(list, area);
}

84
src/client/tui/components/input_area.rs
View File

@ -1,21 +1,81 @@
use crate::client::tui::app::App;
use crate::client::tui::app::{App, Focus};
use ratatui::{
Frame,
layout::Rect,
style::{Color, Style},
widgets::{Block, Borders, Paragraph},
widgets::{Block, Borders, Paragraph, Wrap},
};
use unicode_width::UnicodeWidthChar;
pub fn render(f: &mut Frame, area: Rect, app: &App) {
let input = Paragraph::new(app.input.as_str())
.style(Style::default().fg(Color::White))
.block(Block::default().title("Input").borders(Borders::ALL));
pub fn render(frame: &mut Frame, area: Rect, app: &App) {
let active = app.focus == Focus::Input && app.modal.is_none() && !app.show_help;
let border_style = if active {
Style::default().fg(Color::Cyan)
} else {
Style::default()
};
let title = if app.connected {
" 输入 · Enter 发送 / Shift+Enter 换行 "
} else {
" 输入 · Gateway 已断开 "
};
let inner_width = area.width.saturating_sub(2).max(1);
let inner_height = area.height.saturating_sub(2).max(1);
let (cursor_row, cursor_col) = cursor_position(&app.input[..app.input_cursor_pos], inner_width);
let vertical_scroll = cursor_row.saturating_sub(inner_height.saturating_sub(1));
frame.render_widget(
Paragraph::new(app.input.as_str())
.scroll((vertical_scroll, 0))
.wrap(Wrap { trim: false })
.style(Style::default().fg(Color::White))
.block(
Block::default()
.title(title)
.borders(Borders::ALL)
.border_style(border_style),
),
area,
);
f.render_widget(input, area);
let cursor_x = area.x + 1 + app.input_cursor_pos as u16;
let cursor_y = area.y + 1;
if cursor_x < area.right() && cursor_y < area.bottom() {
f.set_cursor_position((cursor_x, cursor_y));
if active {
let x = area.x + 1 + cursor_col.min(inner_width.saturating_sub(1));
let y = area.y + 1 + cursor_row.saturating_sub(vertical_scroll);
if x < area.right() && y < area.bottom() {
frame.set_cursor_position((x, y));
}
}
}
fn cursor_position(value: &str, width: u16) -> (u16, u16) {
let mut row = 0_u16;
let mut column = 0_u16;
for character in value.chars() {
if character == '\n' {
row = row.saturating_add(1);
column = 0;
continue;
}
let character_width = character.width().unwrap_or(0) as u16;
if column > 0 && column.saturating_add(character_width) > width {
row = row.saturating_add(1);
column = 0;
}
column = column.saturating_add(character_width);
if column >= width {
row = row.saturating_add(column / width);
column %= width;
}
}
(row, column)
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn cursor_accounts_for_wide_characters_and_wrapping() {
assert_eq!(cursor_position("你a", 10), (0, 3));
assert_eq!(cursor_position("1234你", 5), (1, 2));
}
}

88
src/client/tui/components/session_list.rs
View File

@ -1,42 +1,60 @@
use crate::client::tui::app::App;
use crate::client::tui::app::{App, Focus};
use ratatui::{
Frame,
layout::Rect,
style::{Color, Modifier, Style},
widgets::{Block, Borders, List, ListItem},
text::{Line, Span},
widgets::{Block, Borders, List, ListItem, ListState},
};
pub fn render(f: &mut Frame, area: Rect, app: &App) {
let items: Vec<ListItem> = app
.sessions
.iter()
.map(|session| {
let is_current = app.current_session_id.as_ref() == Some(&session.session_id);
let archived = session.archived_at.is_some();
let mut content = if is_current {
format!("• {}", session.title)
} else {
format!(" {}", session.title)
};
if archived {
content.push_str(" [archived]");
}
let style = if is_current {
Style::default()
.fg(Color::Yellow)
.add_modifier(Modifier::BOLD)
} else {
Style::default().fg(Color::White)
};
ListItem::new(content).style(style)
})
.collect();
let list = List::new(items).block(Block::default().title("Sessions").borders(Borders::ALL));
f.render_widget(list, area);
pub fn render(frame: &mut Frame, area: Rect, app: &App) {
let items = app.sessions.iter().enumerate().map(|(index, session)| {
let selected = app.focus == Focus::Sessions && index == app.selected_session;
let current = app.current_session_id.as_ref() == Some(&session.session_id);
let marker = if current { "●" } else { " " };
let archived = if session.archived_at.is_some() {
" 归档"
} else {
""
};
let style = if selected {
Style::default().fg(Color::Black).bg(Color::Cyan)
} else if current {
Style::default()
.fg(Color::Yellow)
.add_modifier(Modifier::BOLD)
} else {
Style::default()
};
ListItem::new(vec![
Line::from(vec![
Span::raw(format!("{marker} ")),
Span::styled(session.title.clone(), style),
]),
Line::from(Span::styled(
format!(" {} 条{archived}", session.message_count),
Style::default().fg(Color::DarkGray),
)),
])
.style(style)
});
let mode = if app.show_archived {
"全部"
} else {
"活跃"
};
let border = if app.focus == Focus::Sessions {
Style::default().fg(Color::Cyan)
} else {
Style::default()
};
let list = List::new(items).block(
Block::default()
.title(format!(" 会话 · {mode} "))
.borders(Borders::ALL)
.border_style(border),
);
let selected = (!app.sessions.is_empty()).then_some(app.selected_session);
let mut state = ListState::default().with_selected(selected);
frame.render_stateful_widget(list, area, &mut state);
}

60
src/client/tui/components/title_bar.rs
View File

@ -3,44 +3,34 @@ use ratatui::{
Frame,
layout::Rect,
style::{Color, Modifier, Style},
text::{Line, Span},
widgets::{Block, Borders, Paragraph},
};
pub fn render(f: &mut Frame, area: Rect, app: &App) {
let (title, style) = if app.pending_quit {
let msg = if let Some(session_id) = &app.current_session_id {
format!(
"PicoBot | Session: {} | Press Ctrl+C again to quit",
session_id
)
} else {
"PicoBot | Press Ctrl+C again to quit".to_string()
};
(
msg,
Style::default()
.fg(Color::Yellow)
.add_modifier(Modifier::BOLD),
)
} else if let Some(session_id) = &app.current_session_id {
(
format!("PicoBot | Session: {}", session_id),
Style::default()
.fg(Color::Cyan)
.add_modifier(Modifier::BOLD),
)
pub fn render(frame: &mut Frame, area: Rect, app: &App) {
let connection = if app.connected {
Span::styled("● 已连接", Style::default().fg(Color::Green))
} else {
(
"PicoBot".to_string(),
Style::default()
.fg(Color::Cyan)
.add_modifier(Modifier::BOLD),
)
Span::styled("● 已断开", Style::default().fg(Color::Red))
};
let paragraph = Paragraph::new(title)
.style(style)
.block(Block::default().borders(Borders::ALL));
f.render_widget(paragraph, area);
let pending = if app.pending_responses > 0 {
format!(" · {} 个请求处理中", app.pending_responses)
} else {
String::new()
};
frame.render_widget(
Paragraph::new(Line::from(vec![
Span::styled(
" PicoBot ",
Style::default()
.fg(Color::Cyan)
.add_modifier(Modifier::BOLD),
),
Span::raw(format!("{} ", app.current_title())),
connection,
Span::styled(pending, Style::default().fg(Color::DarkGray)),
]))
.block(Block::default().borders(Borders::ALL)),
area,
);
}

397
src/client/tui/event.rs
View File

@ -1,134 +1,341 @@
use crate::client::tui::app::{App, MessageRole};
use crate::protocol::WsInbound;
use crate::protocol::serialize_inbound;
use crossterm::event::{KeyCode, KeyEvent};
use crate::client::tui::app::{App, ConfirmAction, Focus, MessageRole, Modal};
use crate::protocol::{WsInbound, serialize_inbound};
use crossterm::event::{KeyCode, KeyEvent, KeyModifiers};
use futures_util::SinkExt;
use tokio_tungstenite::tungstenite::Message;
pub async fn handle_key_event(app: &mut App, key: KeyEvent) {
if app.show_help {
match key.code {
KeyCode::Esc | KeyCode::Char('q') => {
app.toggle_help();
}
_ => {}
if matches!(key.code, KeyCode::Esc | KeyCode::F(1))
|| (key.code == KeyCode::Char('q') && key.modifiers.is_empty())
{
app.show_help = false;
}
return;
}
if app.modal.is_some() {
handle_modal_key(app, key).await;
return;
}
let ctrl = key.modifiers.contains(KeyModifiers::CONTROL);
if ctrl && key.code == KeyCode::Char('c') {
app.handle_ctrl_c_for_quit();
return;
}
app.cancel_pending_quit();
if matches!(key.code, KeyCode::F(1)) || (ctrl && key.code == KeyCode::Char('h')) {
app.show_help = true;
return;
}
if app.show_command_menu {
match key.code {
KeyCode::Esc => {
app.show_command_menu = false;
app.selected_command_idx = 0;
}
KeyCode::Up => {
app.select_prev_command();
}
KeyCode::Down => {
app.select_next_command();
}
KeyCode::Enter => {
app.insert_command();
app.show_command_menu = false;
app.selected_command_idx = 0;
}
KeyCode::Esc => close_command_menu(app),
KeyCode::Up => app.select_previous_command(),
KeyCode::Down => app.select_next_command(),
KeyCode::Tab => {
app.insert_command();
app.insert_selected_command();
close_command_menu(app);
}
_ => {
// Handle normal input and check if menu should stay open
handle_normal_input(app, key).await;
KeyCode::Enter if key.modifiers.is_empty() => {
app.insert_selected_command();
close_command_menu(app);
}
_ => handle_input_key(app, key).await,
}
return;
}
handle_normal_input(app, key).await;
}
async fn handle_normal_input(app: &mut App, key: KeyEvent) {
// Handle Ctrl+C for quit (double press to exit)
let is_ctrl_c = key.code == KeyCode::Char('c')
&& key
.modifiers
.contains(crossterm::event::KeyModifiers::CONTROL);
if is_ctrl_c {
if app.handle_ctrl_c_for_quit() {
return;
if ctrl {
match key.code {
KeyCode::Char('n') => {
send(app, WsInbound::CreateSession { title: None }).await;
app.status_message = Some("正在创建会话…".to_string());
}
KeyCode::Char('s') => app.focus = Focus::Sessions,
KeyCode::Char('r') => open_rename(app),
KeyCode::Char('a') => app.modal = Some(Modal::Confirm(ConfirmAction::Archive)),
KeyCode::Char('d') => app.modal = Some(Modal::Confirm(ConfirmAction::Delete)),
KeyCode::Char('l') => app.modal = Some(Modal::Confirm(ConfirmAction::ClearHistory)),
KeyCode::Char('o') => {
app.show_archived = !app.show_archived;
request_session_list(app).await;
}
KeyCode::Char('u') if app.focus == Focus::Input => {
app.input.clear();
app.input_cursor_pos = 0;
}
KeyCode::Up => app.scroll_chat_up(3),
KeyCode::Down => app.scroll_chat_down(3),
_ => {}
}
} else {
app.cancel_pending_quit();
return;
}
match key.code {
KeyCode::Char('?') => {
app.toggle_help();
KeyCode::Tab => {
app.focus = match app.focus {
Focus::Input => Focus::Sessions,
Focus::Sessions => Focus::Input,
};
}
KeyCode::Esc => app.focus = Focus::Input,
KeyCode::PageUp => app.scroll_chat_up(10),
KeyCode::PageDown => app.scroll_chat_down(10),
KeyCode::Home if app.focus == Focus::Sessions => app.selected_session = 0,
KeyCode::End if app.focus == Focus::Sessions => {
app.selected_session = app.sessions.len().saturating_sub(1);
}
_ if app.focus == Focus::Sessions => handle_session_key(app, key).await,
_ => handle_input_key(app, key).await,
}
}
pub fn handle_paste(app: &mut App, text: &str) {
if let Some(Modal::Rename { input, cursor }) = &mut app.modal {
let remaining = 256_usize.saturating_sub(input.len());
let mut end = text.len().min(remaining);
while !text.is_char_boundary(end) {
end -= 1;
}
input.insert_str(*cursor, &text[..end]);
*cursor += end;
} else if app.focus == Focus::Input {
app.input_insert_str(text);
update_command_menu(app);
}
}
async fn handle_session_key(app: &mut App, key: KeyEvent) {
match key.code {
KeyCode::Up | KeyCode::Char('k') => app.select_previous_session(),
KeyCode::Down | KeyCode::Char('j') => app.select_next_session(),
KeyCode::Enter => {
if let Some(session_id) = app.selected_session_id()
&& app.current_session_id.as_deref() != Some(&session_id)
{
app.status_message = Some("正在载入会话…".to_string());
send(app, WsInbound::LoadSession { session_id }).await;
}
}
KeyCode::Char('n') => {
send(app, WsInbound::CreateSession { title: None }).await;
}
KeyCode::Char('r') => open_rename(app),
KeyCode::Char('a') => app.modal = Some(Modal::Confirm(ConfirmAction::Archive)),
KeyCode::Char('d') => app.modal = Some(Modal::Confirm(ConfirmAction::Delete)),
_ => {}
}
}
async fn handle_input_key(app: &mut App, key: KeyEvent) {
match key.code {
KeyCode::Char(c) => {
app.input_insert_char(c);
// Show command menu when input starts with /
if !app.show_command_menu
&& (app.input == "/" || (app.input.len() > 1 && app.input.starts_with('/')))
{
app.show_command_menu = true;
app.selected_command_idx = 0;
} else if app.show_command_menu && !app.input.starts_with('/') {
app.show_command_menu = false;
}
update_command_menu(app);
}
KeyCode::Backspace => {
app.input_delete_char();
// Hide menu if input no longer starts with /
if app.show_command_menu && !app.input.starts_with('/') {
app.show_command_menu = false;
app.selected_command_idx = 0;
}
update_command_menu(app);
}
KeyCode::Left => {
app.input_move_cursor_left();
}
KeyCode::Right => {
app.input_move_cursor_right();
}
KeyCode::Home => {
app.input_move_cursor_to_start();
}
KeyCode::End => {
app.input_move_cursor_to_end();
}
KeyCode::Up => {
app.scroll_chat_up();
}
KeyCode::Down => {
app.scroll_chat_down();
KeyCode::Delete => app.input_delete_forward(),
KeyCode::Left => app.input_move_cursor_left(),
KeyCode::Right => app.input_move_cursor_right(),
KeyCode::Home => app.input_move_cursor_to_start(),
KeyCode::End => app.input_move_cursor_to_end(),
KeyCode::Up => app.scroll_chat_up(1),
KeyCode::Down => app.scroll_chat_down(1),
KeyCode::Enter
if key.modifiers.contains(KeyModifiers::SHIFT)
|| key.modifiers.contains(KeyModifiers::ALT) =>
{
app.input_insert_char('\n');
}
KeyCode::Enter => {
let input = app.take_input();
app.show_command_menu = false;
app.selected_command_idx = 0;
if !input.is_empty() {
process_input(app, input).await;
close_command_menu(app);
if !input.trim().is_empty() {
app.add_message(MessageRole::User, input.clone());
app.pending_responses = app.pending_responses.saturating_add(1);
app.status_message = Some("PicoBot 正在处理…".to_string());
let sent = send(
app,
WsInbound::UserInput {
content: input,
channel: None,
// Session routing is owned by the server. A full session
// id is not a chat id and must never be sent here.
chat_id: None,
sender_id: None,
},
)
.await;
if !sent {
app.pending_responses = app.pending_responses.saturating_sub(1);
}
}
}
_ => {}
}
}
async fn process_input(app: &mut App, input: String) {
app.add_message(MessageRole::User, input.clone());
if let Some(sender) = &mut app.ws_sender {
let inbound = WsInbound::UserInput {
content: input,
chat_id: app.current_session_id.clone(),
channel: None,
sender_id: None,
};
if let Ok(text) = serialize_inbound(&inbound) {
let _ = sender
.send(tokio_tungstenite::tungstenite::Message::Text(text.into()))
.await;
}
async fn handle_modal_key(app: &mut App, key: KeyEvent) {
match app.modal.take() {
Some(Modal::Confirm(action)) => match key.code {
KeyCode::Char('y') | KeyCode::Char('Y') | KeyCode::Enter => {
let target = target_session_id(app);
let message = match action {
ConfirmAction::Archive => target.map(|session_id| WsInbound::ArchiveSession {
session_id: Some(session_id),
}),
ConfirmAction::Delete => target.map(|session_id| WsInbound::DeleteSession {
session_id: Some(session_id),
}),
ConfirmAction::ClearHistory => {
target.map(|session_id| WsInbound::ClearHistory {
chat_id: None,
session_id: Some(session_id),
})
}
};
if let Some(message) = message {
send(app, message).await;
app.status_message = Some("正在更新会话…".to_string());
}
}
KeyCode::Esc | KeyCode::Char('n') | KeyCode::Char('N') => {}
_ => app.modal = Some(Modal::Confirm(action)),
},
Some(Modal::Rename {
mut input,
mut cursor,
}) => match key.code {
KeyCode::Esc => {}
KeyCode::Enter => {
let title = input.trim().to_string();
if !title.is_empty() {
send(
app,
WsInbound::RenameSession {
session_id: target_session_id(app),
title,
},
)
.await;
}
}
KeyCode::Char(character) => {
input.insert(cursor, character);
cursor += character.len_utf8();
app.modal = Some(Modal::Rename { input, cursor });
}
KeyCode::Backspace => {
if let Some((index, _)) = input[..cursor].char_indices().next_back() {
input.drain(index..cursor);
cursor = index;
}
app.modal = Some(Modal::Rename { input, cursor });
}
KeyCode::Delete => {
if let Some(character) = input[cursor..].chars().next() {
input.drain(cursor..cursor + character.len_utf8());
}
app.modal = Some(Modal::Rename { input, cursor });
}
KeyCode::Left => {
if let Some((index, _)) = input[..cursor].char_indices().next_back() {
cursor = index;
}
app.modal = Some(Modal::Rename { input, cursor });
}
KeyCode::Right => {
if let Some(character) = input[cursor..].chars().next() {
cursor += character.len_utf8();
}
app.modal = Some(Modal::Rename { input, cursor });
}
_ => app.modal = Some(Modal::Rename { input, cursor }),
},
None => {}
}
}
fn open_rename(app: &mut App) {
if target_session_id(app).is_some() {
let input = if app.focus == Focus::Sessions {
app.sessions
.get(app.selected_session)
.map(|session| session.title.clone())
.unwrap_or_default()
} else {
app.current_title().to_string()
};
let cursor = input.len();
app.modal = Some(Modal::Rename { input, cursor });
}
}
fn target_session_id(app: &App) -> Option<String> {
if app.focus == Focus::Sessions {
app.selected_session_id()
} else {
app.current_session_id.clone()
}
}
fn update_command_menu(app: &mut App) {
app.show_command_menu = app.input.starts_with('/') && !app.input.contains('\n');
app.selected_command_idx = app
.selected_command_idx
.min(app.get_filtered_commands().len().saturating_sub(1));
}
fn close_command_menu(app: &mut App) {
app.show_command_menu = false;
app.selected_command_idx = 0;
}
pub async fn request_session_list(app: &mut App) {
send(
app,
WsInbound::ListSessions {
include_archived: app.show_archived,
},
)
.await;
}
pub async fn request_history(app: &mut App, session_id: String) {
send(
app,
WsInbound::GetSessionHistory {
session_id,
limit: Some(1_000),
},
)
.await;
}
pub async fn send(app: &mut App, inbound: WsInbound) -> bool {
let serialized = match serialize_inbound(&inbound) {
Ok(serialized) => serialized,
Err(error) => {
app.status_message = Some(format!("请求编码失败：{error}"));
return false;
}
};
let Some(sender) = &mut app.ws_sender else {
app.connected = false;
app.status_message = Some("Gateway 已断开".to_string());
return false;
};
if let Err(error) = sender.send(Message::Text(serialized.into())).await {
app.connected = false;
app.status_message = Some(format!("发送失败：{error}"));
return false;
}
true
}

5
src/client/tui/markdown.rs
View File

@ -1,5 +0,0 @@
#![allow(dead_code)]
pub fn render_markdown(content: &str) -> String {
content.to_string()
}

1
src/client/tui/mod.rs
View File

@ -1,5 +1,4 @@
pub mod app;
pub mod components;
pub mod event;
pub mod markdown;
pub mod ui;

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

@ -1,73 +1,165 @@
use crate::client::tui::app::App;
use crate::client::tui::app::{App, ConfirmAction, Focus, Modal};
use crate::client::tui::components::*;
use ratatui::{
Frame,
layout::{Constraint, Direction, Layout, Rect},
style::{Color, Modifier, Style},
text::{Line, Span},
widgets::{Block, Borders, Clear, Paragraph, Wrap},
};
use unicode_width::UnicodeWidthStr;
pub fn render_ui(f: &mut Frame, app: &App) {
let size = f.area();
let chunks = Layout::default()
pub fn render_ui(frame: &mut Frame, app: &App) {
let area = frame.area();
let input_height = app.input.lines().count().clamp(1, 6) as u16 + 2;
let rows = Layout::default()
.direction(Direction::Vertical)
.constraints([
Constraint::Length(3),
Constraint::Min(0),
Constraint::Length(5),
Constraint::Min(3),
Constraint::Length(input_height),
Constraint::Length(1),
])
.split(size);
.split(area);
title_bar::render(f, chunks[0], app);
title_bar::render(frame, rows[0], app);
if area.width >= 80 {
let columns = Layout::default()
.direction(Direction::Horizontal)
.constraints([Constraint::Length(28), Constraint::Min(30)])
.split(rows[1]);
session_list::render(frame, columns[0], app);
chat_history::render(frame, columns[1], app);
} else {
chat_history::render(frame, rows[1], app);
}
input_area::render(frame, rows[2], app);
render_footer(frame, rows[3], app);
let middle_chunks = Layout::default()
.direction(Direction::Horizontal)
.constraints([Constraint::Percentage(25), Constraint::Percentage(75)])
.split(chunks[1]);
session_list::render(f, middle_chunks[0], app);
chat_history::render(f, middle_chunks[1], app);
input_area::render(f, chunks[2], app);
// Render command menu if needed - position above input area
if app.show_command_menu && !app.get_filtered_commands().is_empty() {
let menu_area = menu_above_input(chunks[2]);
command_menu::render(f, menu_area, app);
let height = (app.get_filtered_commands().len().min(6) + 2) as u16;
let menu_area = Rect::new(
rows[2].x.saturating_add(1),
rows[2].y.saturating_sub(height),
rows[2].width.saturating_sub(2),
height,
);
command_menu::render(frame, menu_area, app);
}
if app.show_help {
let help_area = centered_rect(60, 60, size);
help_popup::render(f, help_area);
help_popup::render(frame, centered_rect(72, 26, area));
}
if let Some(modal) = &app.modal {
render_modal(frame, centered_rect(64, 7, area), modal);
}
}
fn menu_above_input(input_area: Rect) -> Rect {
let max_commands = 6; // Show up to 6 commands
let menu_height = max_commands + 2; // +2 for borders
fn render_footer(frame: &mut Frame, area: Rect, app: &App) {
let focus = match app.focus {
Focus::Input => "输入",
Focus::Sessions => "会话",
};
let status =
app.status_message
.as_deref()
.unwrap_or(if app.connected { "就绪" } else { "已断开" });
let line = Line::from(vec![
Span::styled(
format!(" {focus} "),
Style::default().fg(Color::Black).bg(Color::Cyan),
),
Span::raw(format!(" {status}")),
Span::styled(
" F1 帮助 Tab 切换焦点 Ctrl+N 新会话 ",
Style::default().fg(Color::DarkGray),
),
]);
frame.render_widget(Paragraph::new(line), area);
}
Rect {
x: input_area.x + 1,
y: input_area.y.saturating_sub(menu_height),
width: input_area.width.saturating_sub(2),
height: menu_height,
fn render_modal(frame: &mut Frame, area: Rect, modal: &Modal) {
frame.render_widget(Clear, area);
let block = Block::default()
.borders(Borders::ALL)
.border_style(Style::default().fg(Color::Yellow));
match modal {
Modal::Rename { input, cursor } => {
frame.render_widget(
Paragraph::new(vec![
Line::from(Span::styled(
"重命名会话",
Style::default().add_modifier(Modifier::BOLD),
)),
Line::from(""),
Line::from(input.as_str()),
Line::from(Span::styled(
"Enter 保存 · Esc 取消",
Style::default().fg(Color::DarkGray),
)),
])
.block(block)
.wrap(Wrap { trim: false }),
area,
);
let cursor_x = area.x
+ 1
+ UnicodeWidthStr::width(&input[..*cursor])
.min(area.width.saturating_sub(3) as usize) as u16;
let cursor_y = area.y.saturating_add(3);
if cursor_x < area.right() && cursor_y < area.bottom() {
frame.set_cursor_position((cursor_x, cursor_y));
}
}
Modal::Confirm(action) => {
let prompt = match action {
ConfirmAction::Archive => "归档所选会话？",
ConfirmAction::Delete => "永久删除所选会话？此操作不可撤销。",
ConfirmAction::ClearHistory => "清空所选会话的全部历史？",
};
frame.render_widget(
Paragraph::new(vec![
Line::from(""),
Line::from(prompt),
Line::from(""),
Line::from(Span::styled(
"Y / Enter 确认 · N / Esc 取消",
Style::default().fg(Color::DarkGray),
)),
])
.block(block)
.wrap(Wrap { trim: false }),
area,
);
}
}
}
fn centered_rect(percent_x: u16, percent_y: u16, r: Rect) -> Rect {
let popup_layout = Layout::default()
.direction(Direction::Vertical)
.constraints([
Constraint::Percentage((100 - percent_y) / 2),
Constraint::Percentage(percent_y),
Constraint::Percentage((100 - percent_y) / 2),
])
.split(r);
Layout::default()
.direction(Direction::Horizontal)
.constraints([
Constraint::Percentage((100 - percent_x) / 2),
Constraint::Percentage(percent_x),
Constraint::Percentage((100 - percent_x) / 2),
])
.split(popup_layout[1])[1]
fn centered_rect(max_width: u16, max_height: u16, area: Rect) -> Rect {
let width = max_width.min(area.width.saturating_sub(2)).max(1);
let height = max_height.min(area.height.saturating_sub(2)).max(1);
Rect::new(
area.x + area.width.saturating_sub(width) / 2,
area.y + area.height.saturating_sub(height) / 2,
width,
height,
)
}
#[cfg(test)]
mod tests {
use super::*;
use ratatui::{Terminal, backend::TestBackend};
#[test]
fn narrow_terminal_renders_without_sidebar_or_panic() {
let backend = TestBackend::new(48, 14);
let mut terminal = Terminal::new(backend).unwrap();
let mut app = App::new();
app.input_insert_str("中文 input");
app.add_message(
crate::client::tui::app::MessageRole::Assistant,
"一条很长的响应，用于验证窄终端换行。".repeat(4),
);
terminal.draw(|frame| render_ui(frame, &app)).unwrap();
}
}

14
src/config/mod.rs
View File

@ -19,10 +19,10 @@ pub fn get_default_workspace_dir() -> PathBuf {
/// Expand ~ in path to user home directory
pub fn expand_path(path: &str) -> PathBuf {
if path.starts_with("~/") {
if let Some(path) = path.strip_prefix("~/") {
dirs::home_dir()
.unwrap_or_else(|| PathBuf::from("."))
.join(&path[2..])
.join(path)
} else {
PathBuf::from(path)
}
@ -180,9 +180,12 @@ pub struct SchedulerConfig {
/// Poll interval in seconds (how often to check for due jobs)
#[serde(default = "default_poll_interval_secs")]
pub poll_interval_secs: u64,
/// Maximum concurrent job executions (currently sequential, reserved for future)
/// Maximum concurrent job executions.
#[serde(default = "default_max_concurrent")]
pub max_concurrent: usize,
/// Hard timeout for one scheduled execution.
#[serde(default = "default_execution_timeout_secs")]
pub execution_timeout_secs: u64,
}
fn default_scheduler_enabled() -> bool {
@ -197,12 +200,17 @@ fn default_max_concurrent() -> usize {
1
}
fn default_execution_timeout_secs() -> u64 {
900
}
impl Default for SchedulerConfig {
fn default() -> Self {
Self {
enabled: true,
poll_interval_secs: 60,
max_concurrent: 1,
execution_timeout_secs: 900,
}
}
}

86
src/gateway/mod.rs
View File

@ -6,7 +6,7 @@ use std::sync::Arc;
use tokio::net::TcpListener;
use crate::bus::{ControlMessage, MessageBus, OutboundDispatcher};
use crate::channels::base::{Channel, ChannelError};
use crate::channels::base::ChannelError;
use crate::channels::{ChannelManager, CliChatChannel};
use crate::config::{Config, ensure_workspace_dir, expand_path};
use crate::logging;
@ -14,6 +14,7 @@ use crate::mcp;
use crate::memory::MemoryManager;
use crate::scheduler::Scheduler;
use crate::session::SessionManager;
use crate::task_supervisor::TaskSupervisor;
pub struct GatewayState {
pub config: Config,
@ -21,11 +22,15 @@ pub struct GatewayState {
pub session_manager: Arc<SessionManager>,
pub channel_manager: ChannelManager,
pub storage: Arc<crate::storage::Storage>,
pub task_supervisor: TaskSupervisor,
pub connection_shutdown: tokio_util::sync::CancellationToken,
}
impl GatewayState {
pub async fn new() -> Result<Self, Box<dyn std::error::Error>> {
let config = Config::load_default()?;
let task_supervisor = TaskSupervisor::new();
let connection_shutdown = tokio_util::sync::CancellationToken::new();
// Initialize workspace directory: expand path and ensure it exists
let workspace_path = expand_path(&config.workspace_dir);
@ -98,6 +103,7 @@ impl GatewayState {
memory_manager,
browser_config,
config.gateway.max_concurrent_background_tasks,
task_supervisor.clone(),
)?;
let session_manager = Arc::new(session_manager);
@ -171,6 +177,8 @@ impl GatewayState {
session_manager: session_manager.clone(),
channel_manager,
storage,
task_supervisor,
connection_shutdown,
})
}
@ -190,15 +198,9 @@ impl GatewayState {
let bus_for_outbound = bus.clone();
let session_manager = self.session_manager.clone();
// Start CLI Chat Channel (it's already registered in ChannelManager)
let cli_chat_channel = self.cli_chat_channel();
if let Err(e) = cli_chat_channel.start(bus.clone()).await {
tracing::error!(error = %e, "Failed to start CLI chat channel");
}
// Spawn unified message processor
// This handles both inbound AI messages and control messages in one loop
tokio::spawn(async move {
self.task_supervisor.spawn("message-processor", async move {
tracing::info!("Message processor started");
loop {
@ -224,19 +226,23 @@ impl GatewayState {
reply_to: None,
media: vec![],
metadata: inbound.forwarded_metadata,
delivery: None,
};
if let Err(e) = bus.publish_outbound(outbound).await {
tracing::error!(error = %e, "Failed to publish outbound");
}
}
Ok(crate::session::session::HandleResult::CommandOutput(content)) => {
let mut metadata = inbound.forwarded_metadata;
metadata.insert("_type".to_string(), "command".to_string());
let outbound = crate::bus::OutboundMessage {
channel: inbound.channel.clone(),
chat_id: inbound.chat_id.clone(),
content,
reply_to: None,
media: vec![],
metadata: inbound.forwarded_metadata,
metadata,
delivery: None,
};
if let Err(e) = bus.publish_outbound(outbound).await {
tracing::error!(error = %e, "Failed to publish outbound");
@ -267,12 +273,17 @@ impl GatewayState {
});
// Spawn outbound dispatcher
let dispatcher = OutboundDispatcher::new(bus_for_outbound, self.channel_manager.clone());
let dispatcher = OutboundDispatcher::new(
bus_for_outbound,
self.channel_manager.clone(),
self.task_supervisor.clone(),
);
tokio::spawn(async move {
tracing::info!("Outbound dispatcher started");
dispatcher.run().await;
});
self.task_supervisor
.spawn("outbound-dispatcher", async move {
tracing::info!("Outbound dispatcher started");
dispatcher.run().await;
});
// Spawn scheduler background task if enabled
let scheduler_config = self.config.gateway.scheduler.clone().unwrap_or_default();
@ -282,7 +293,7 @@ impl GatewayState {
self.session_manager.clone(),
scheduler_config,
));
tokio::spawn(async move {
self.task_supervisor.spawn("scheduler", async move {
sched.run().await;
});
tracing::info!("Scheduler background task spawned");
@ -330,6 +341,14 @@ impl GatewayState {
.await
.map(|session_id| SessionEvent::DialogSwitched { session_id })
.map_err(|e| ChannelError::Other(e.to_string())),
GetDialogHistory { session_id, limit } => session_manager
.get_dialog_history(&session_id, limit)
.await
.map(|messages| SessionEvent::DialogHistory {
session_id,
messages,
})
.map_err(|e| ChannelError::Other(e.to_string())),
RenameDialog { session_id, title } => session_manager
.rename_dialog(&session_id, &title)
.await
@ -412,24 +431,27 @@ pub async fn run(
let listener = TcpListener::bind(&addr).await?;
tracing::info!(address = %addr, "Gateway listening");
// Graceful shutdown using oneshot channel
let (shutdown_tx, shutdown_rx) = tokio::sync::oneshot::channel::<()>();
let channel_manager = state.channel_manager.clone();
// Spawn ctrl_c handler
tokio::spawn(async move {
tokio::signal::ctrl_c().await.ok();
tracing::info!("Shutdown signal received");
let _ = channel_manager.stop_all().await;
let _ = shutdown_tx.send(());
});
// Serve with graceful shutdown
axum::serve(listener, app)
.with_graceful_shutdown(async {
shutdown_rx.await.ok();
let connection_shutdown = state.connection_shutdown.clone();
let serve_result = axum::serve(listener, app)
.with_graceful_shutdown(async move {
if let Err(error) = tokio::signal::ctrl_c().await {
tracing::error!(error = %error, "Failed to listen for shutdown signal");
}
tracing::info!("Shutdown signal received");
connection_shutdown.cancel();
})
.await?;
.await;
// Stop external intake before waiting for internal work to finish.
if let Err(error) = state.channel_manager.stop_all().await {
tracing::error!(error = %error, "Failed to stop channels cleanly");
}
state.task_supervisor.cancel();
state
.task_supervisor
.shutdown(std::time::Duration::from_secs(10))
.await;
serve_result?;
Ok(())
}

93
src/gateway/ws.rs
View File

@ -1,20 +1,41 @@
use super::GatewayState;
use crate::protocol::WsOutbound;
use crate::protocol::serialize_outbound;
use axum::extract::State;
use axum::extract::ws::{Message as WsMessage, WebSocket, WebSocketUpgrade};
use axum::extract::{Query, State};
use axum::response::Response;
use futures_util::{SinkExt, StreamExt};
use serde::Deserialize;
use std::sync::Arc;
use tokio::sync::mpsc;
use tokio::time::{Duration, timeout};
pub async fn ws_handler(ws: WebSocketUpgrade, State(state): State<Arc<GatewayState>>) -> Response {
#[derive(Debug, Default, Deserialize)]
pub struct WsQuery {
client_id: Option<String>,
}
pub async fn ws_handler(
ws: WebSocketUpgrade,
Query(query): Query<WsQuery>,
State(state): State<Arc<GatewayState>>,
) -> Response {
ws.on_upgrade(|socket| async move {
handle_socket(socket, state).await;
handle_socket(socket, state, valid_client_id(query.client_id)).await;
})
}
async fn handle_socket(ws: WebSocket, state: Arc<GatewayState>) {
fn valid_client_id(client_id: Option<String>) -> Option<String> {
client_id.filter(|value| {
!value.is_empty()
&& value.len() <= 64
&& value
.bytes()
.all(|byte| byte.is_ascii_alphanumeric() || byte == b'-' || byte == b'_')
})
}
async fn handle_socket(ws: WebSocket, state: Arc<GatewayState>, client_id: Option<String>) {
// Create channel for sending outbound messages to this client
let (sender, mut receiver) = mpsc::channel::<WsOutbound>(100);
@ -22,8 +43,9 @@ async fn handle_socket(ws: WebSocket, state: Arc<GatewayState>) {
let cli_chat_channel = state.cli_chat_channel();
// Register client with CliChatChannel and get initial session id
let (session_id, client) = cli_chat_channel.register_client(sender.clone()).await;
let (session_id, client) = cli_chat_channel
.register_client(sender.clone(), client_id)
.await;
// Send session established message
let _ = sender
.send(WsOutbound::SessionEstablished {
@ -36,7 +58,7 @@ async fn handle_socket(ws: WebSocket, state: Arc<GatewayState>) {
let (mut ws_sender, mut ws_receiver) = ws.split();
// Task: forward from receiver to WebSocket
tokio::spawn(async move {
let mut writer_task = tokio::spawn(async move {
while let Some(msg) = receiver.recv().await {
if let Ok(text) = serialize_outbound(&msg)
&& ws_sender.send(WsMessage::Text(text.into())).await.is_err()
@ -47,18 +69,59 @@ async fn handle_socket(ws: WebSocket, state: Arc<GatewayState>) {
});
// Main loop: receive WebSocket messages and forward to CliChatChannel
while let Some(msg) = ws_receiver.next().await {
match msg {
Ok(WsMessage::Text(text)) => {
cli_chat_channel.handle_inbound(client.clone(), &text).await;
}
Ok(WsMessage::Close(_)) | Err(_) => {
tracing::debug!(session_id = %session_id, "WebSocket closed");
let cancellation = state.connection_shutdown.clone();
let mut writer_finished = false;
loop {
tokio::select! {
_ = cancellation.cancelled() => break,
result = &mut writer_task => {
writer_finished = true;
if let Err(error) = result {
tracing::warn!(session_id = %session_id, error = %error, "WebSocket writer task failed");
}
break;
}
_ => {}
msg = ws_receiver.next() => {
match msg {
Some(Ok(WsMessage::Text(text))) => {
cli_chat_channel.handle_inbound(client.clone(), &text).await;
}
Some(Ok(WsMessage::Close(_))) | Some(Err(_)) | None => {
tracing::debug!(session_id = %session_id, "WebSocket closed");
break;
}
_ => {}
}
}
}
}
cli_chat_channel.unregister_client(&client).await;
drop(client);
drop(sender);
if !writer_finished
&& timeout(Duration::from_secs(2), &mut writer_task)
.await
.is_err()
{
writer_task.abort();
let _ = writer_task.await;
}
tracing::info!(session_id = %session_id, "CLI session ended");
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn client_id_is_strictly_bounded() {
assert_eq!(
valid_client_id(Some("client_123-abc".to_string())).as_deref(),
Some("client_123-abc")
);
assert!(valid_client_id(Some("bad/query".to_string())).is_none());
assert!(valid_client_id(Some("x".repeat(65))).is_none());
assert!(valid_client_id(Some(String::new())).is_none());
}
}

1
src/lib.rs
View File

@ -14,5 +14,6 @@ pub mod scheduler;
pub mod session;
pub mod skills;
pub mod storage;
pub mod task_supervisor;
pub mod tools;
pub mod util;

17
src/logging.rs
View File

@ -59,20 +59,3 @@ pub fn init_logging() {
tracing::info!("Logging initialized. Log directory: {}", log_dir.display());
}
/// Initialize logging without file output (console only)
pub fn init_logging_console_only() {
let env_filter = EnvFilter::try_from_default_env().unwrap_or_else(|_| EnvFilter::new("info"));
let console_layer = fmt::layer()
.with_timer(LocalTime::rfc_3339())
.with_target(true)
.with_level(true);
tracing_subscriber::registry()
.with(env_filter)
.with(console_layer)
.init();
tracing::info!("Logging initialized (console only)");
}

4
src/mcp/mod.rs
View File

@ -47,8 +47,6 @@ fn update_mcp_status(servers: Vec<McpServerStatus>) {
/// A connected MCP server. Holds a clonable Peer handle for tool calls,
/// and keeps the underlying service alive via a background task.
pub struct McpConnection {
#[allow(dead_code)]
pub name: String,
peer: Peer<RoleClient>,
/// Keep the service alive. When dropped, the MCP connection is closed.
_service: Option<Box<dyn std::any::Any + Send + Sync>>,
@ -230,7 +228,6 @@ async fn connect_server(config: &McpServerConfig) -> anyhow::Result<McpConnectio
let peer = service.peer().clone();
Ok(McpConnection {
name: config.name.clone(),
peer,
_service: Some(Box::new(service)),
})
@ -268,7 +265,6 @@ async fn connect_server(config: &McpServerConfig) -> anyhow::Result<McpConnectio
let peer = service.peer().clone();
Ok(McpConnection {
name: config.name.clone(),
peer,
_service: Some(Box::new(service)),
})

8
src/memory/types.rs
View File

@ -18,7 +18,7 @@ impl MemoryCategory {
}
}
pub fn from_str(s: &str) -> Option<Self> {
pub fn parse(s: &str) -> Option<Self> {
match s {
"knowledge" => Some(Self::Knowledge),
"timeline" => Some(Self::Timeline),
@ -78,13 +78,13 @@ mod tests {
#[test]
fn test_memory_category_from_str() {
assert_eq!(
MemoryCategory::from_str("knowledge"),
MemoryCategory::parse("knowledge"),
Some(MemoryCategory::Knowledge)
);
assert_eq!(
MemoryCategory::from_str("timeline"),
MemoryCategory::parse("timeline"),
Some(MemoryCategory::Timeline)
);
assert_eq!(MemoryCategory::from_str("invalid"), None);
assert_eq!(MemoryCategory::parse("invalid"), None);
}
}

34
src/observability/mod.rs
View File

@ -70,16 +70,6 @@ impl ToolExecutionOutcome {
}
}
/// Create a successful outcome with duration.
pub fn success_with_duration(output: String, duration: Duration) -> Self {
Self {
output,
success: true,
error_reason: None,
duration,
}
}
/// Create a failed outcome with zero duration.
pub fn failure(output: String, error_reason: Option<String>) -> Self {
Self {
@ -89,20 +79,6 @@ impl ToolExecutionOutcome {
duration: Duration::ZERO,
}
}
/// Create a failed outcome with duration.
pub fn failure_with_duration(
output: String,
error_reason: Option<String>,
duration: Duration,
) -> Self {
Self {
output,
success: false,
error_reason,
duration,
}
}
}
/// MultiObserver broadcasts events to multiple observers.
@ -204,16 +180,6 @@ mod tests {
assert_eq!(outcome.duration, Duration::ZERO);
}
#[test]
fn test_tool_execution_outcome_success_with_duration() {
let outcome = ToolExecutionOutcome::success_with_duration(
"output content".to_string(),
Duration::from_millis(100),
);
assert!(outcome.success);
assert_eq!(outcome.duration, Duration::from_millis(100));
}
#[test]
fn test_tool_execution_outcome_failure() {
let outcome = ToolExecutionOutcome::failure(

28
src/protocol.rs
View File

@ -19,6 +19,15 @@ pub struct SlashCommandInfo {
pub aliases: Vec<String>,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct HistoryMessage {
pub id: String,
pub seq: i64,
pub role: String,
pub content: String,
pub created_at: i64,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(tag = "type")]
pub enum WsInbound {
@ -51,6 +60,12 @@ pub enum WsInbound {
},
#[serde(rename = "load_session")]
LoadSession { session_id: String },
#[serde(rename = "get_session_history")]
GetSessionHistory {
session_id: String,
#[serde(default, skip_serializing_if = "Option::is_none")]
limit: Option<u32>,
},
#[serde(rename = "rename_session")]
RenameSession {
#[serde(default, skip_serializing_if = "Option::is_none")]
@ -81,6 +96,8 @@ pub enum WsOutbound {
id: String,
content: String,
role: String,
#[serde(default, skip_serializing_if = "Option::is_none")]
session_id: Option<String>,
},
#[serde(rename = "error")]
Error { code: String, message: String },
@ -100,6 +117,11 @@ pub enum WsOutbound {
title: String,
message_count: i64,
},
#[serde(rename = "session_history")]
SessionHistory {
session_id: String,
messages: Vec<HistoryMessage>,
},
#[serde(rename = "session_renamed")]
SessionRenamed { session_id: String, title: String },
#[serde(rename = "session_archived")]
@ -115,7 +137,11 @@ pub enum WsOutbound {
#[serde(rename = "command_executed")]
CommandExecuted { message: String },
#[serde(rename = "system_notification")]
SystemNotification { content: String },
SystemNotification {
content: String,
#[serde(default, skip_serializing_if = "Option::is_none")]
session_id: Option<String>,
},
}
pub fn parse_inbound(raw: &str) -> Result<WsInbound, serde_json::Error> {

42
src/providers/anthropic.rs
View File

@ -87,6 +87,8 @@ pub struct AnthropicProvider {
}
impl AnthropicProvider {
// Keep this constructor aligned with OpenAIProvider and LLMProviderConfig.
#[allow(clippy::too_many_arguments)]
pub fn new(
name: String,
api_key: String,
@ -165,7 +167,6 @@ enum AnthropicContent {
},
Thinking {
#[serde(alias = "content")]
#[allow(dead_code)]
thinking: String,
},
#[serde(rename = "tool_use")]
@ -329,24 +330,29 @@ impl LLMProvider for AnthropicProvider {
return Err(format!("API error ({}): {}", status.as_u16(), error_msg).into());
}
let anthropic_resp: AnthropicResponse = serde_json::from_str(&body_text).map_err(|e| {
let err_msg = format!("decode error: {} | body: {}", e, &body_text);
if let Some(ref storage) = self.storage {
let name = self.name.clone();
let model = self.model_id.clone();
let req = req_body_str.clone();
let resp_body = body_text.clone();
let dur = start.elapsed().as_millis() as u64;
let err = err_msg.clone();
let s = storage.clone();
tokio::spawn(async move {
let _ = s
.append_llm_call(&name, &model, &req, Some(&resp_body), Some(&err), dur)
.await;
});
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());
}
err_msg
})?;
};
let mut content = String::new();
let mut reasoning = None;

49
src/providers/openai.rs
View File

@ -46,6 +46,9 @@ pub struct OpenAIProvider {
}
impl OpenAIProvider {
// Provider construction mirrors the independently configurable fields in
// LLMProviderConfig; grouping them again would only duplicate that API.
#[allow(clippy::too_many_arguments)]
pub fn new(
name: String,
api_key: String,
@ -112,10 +115,10 @@ impl OpenAIProvider {
"role": m.role,
"content": convert_content_blocks(&m.content)
});
if m.role == "assistant" {
if let Some(ref rc) = m.reasoning_content {
msg["reasoning_content"] = json!(rc);
}
if m.role == "assistant"
&& let Some(ref rc) = m.reasoning_content
{
msg["reasoning_content"] = json!(rc);
}
msg
}
@ -295,27 +298,29 @@ impl LLMProvider for OpenAIProvider {
return Err(error.into());
}
let openai_resp: OpenAIResponse = serde_json::from_str(&text).map_err(|e| {
let err_msg = format!("decode error: {} | body: {}", e, &text);
if let Some(ref storage) = self.storage {
let name = self.name.clone();
let model = self.model_id.clone();
let req = req_body_str.clone();
let resp = text.clone();
let dur = start.elapsed().as_millis() as u64;
let err = err_msg.clone();
let s = storage.clone();
tokio::spawn(async move {
if let Err(e) = s
.append_llm_call(&name, &model, &req, Some(&resp), Some(&err), dur)
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): {}", e);
tracing::warn!("failed to persist LLM call (decode error): {}", error);
}
});
}
return Err(err_msg.into());
}
err_msg
})?;
};
let first_choice = openai_resp
.choices
@ -358,7 +363,7 @@ impl LLMProvider for OpenAIProvider {
prompt_tokens: usage.prompt_tokens,
completion_tokens: usage.completion_tokens,
total_tokens: usage.total_tokens,
cached_tokens: cached_tokens,
cached_tokens,
cache_read_input_tokens: None,
cache_creation_input_tokens: None,
},

11
src/providers/traits.rs
View File

@ -28,17 +28,6 @@ impl Message {
}
}
pub fn user_with_blocks(content: Vec<ContentBlock>) -> Self {
Self {
role: "user".to_string(),
content,
reasoning_content: None,
tool_call_id: None,
name: None,
tool_calls: None,
}
}
pub fn assistant(content: impl Into<String>) -> Self {
Self {
role: "assistant".to_string(),

282
src/scheduler/mod.rs
View File

@ -2,6 +2,8 @@ pub mod types;
use std::sync::Arc;
use std::time::Instant;
use futures_util::stream::{self, StreamExt};
use tokio::time;
use crate::config::SchedulerConfig;
@ -55,6 +57,7 @@ pub struct Scheduler {
storage: Arc<Storage>,
session_manager: Arc<SessionManager>,
config: SchedulerConfig,
owner: String,
}
impl Scheduler {
@ -67,209 +70,150 @@ impl Scheduler {
storage,
session_manager,
config,
owner: uuid::Uuid::new_v4().to_string(),
}
}
/// Run the scheduler loop. This is a long-running async function meant to be
/// spawned as a tokio background task.
/// Claim due jobs with a durable lease, then execute the claimed batch with
/// bounded concurrency.
pub async fn run(self: Arc<Self>) {
let poll_duration = time::Duration::from_secs(self.config.poll_interval_secs);
let poll_duration = time::Duration::from_secs(self.config.poll_interval_secs.max(1));
let mut interval = time::interval(poll_duration);
interval.tick().await;
interval.set_missed_tick_behavior(time::MissedTickBehavior::Skip);
// Keep accidental configuration values from claiming an unbounded
// batch and overwhelming the runtime or SQLite parameter conversion.
let max_concurrent = self.config.max_concurrent.clamp(1, 256);
tracing::info!(
"Scheduler started (poll interval: {}s, max concurrent: {})",
self.config.poll_interval_secs,
self.config.max_concurrent,
poll_interval_secs = self.config.poll_interval_secs,
max_concurrent,
scheduler_owner = %self.owner,
"Scheduler started"
);
loop {
interval.tick().await;
let now = now_ms();
let due = match self
let lease_ms = self
.config
.execution_timeout_secs
.saturating_add(30)
.saturating_mul(1000)
.min(i64::MAX as u64) as i64;
let lease_until = now.saturating_add(lease_ms);
let jobs = match self
.storage
.due_scheduled_jobs(now, self.config.max_concurrent)
.claim_due_scheduled_jobs(now, lease_until, &self.owner, max_concurrent)
.await
{
Ok(jobs) => jobs,
Err(e) => {
tracing::error!("scheduler: failed to query due jobs: {}", e);
Err(error) => {
tracing::error!(error = %error, "scheduler: failed to claim due jobs");
continue;
}
};
if due.is_empty() {
if jobs.is_empty() {
continue;
}
tracing::info!(count = jobs.len(), "scheduler: claimed due jobs");
tracing::info!("scheduler: found {} due job(s)", due.len());
for job in &due {
let start = Instant::now();
let started_at = now_ms();
if let Err(e) = self
.storage
.touch_scheduled_job_last_run(&job.id, started_at)
.await
{
tracing::error!(job_id = %job.id, "scheduler: failed to touch last_run_at: {}", e);
continue;
}
tracing::info!(
job_id = %job.id,
job_name = %job.name,
"scheduler: executing cron job"
);
let result = self
.session_manager
.handle_cron_message(
&job.channel,
&job.chat_id,
&job.prompt,
&job.id,
&job.name,
)
.await;
let finished_at = now_ms();
let duration_ms = start.elapsed().as_millis() as i64;
match result {
Ok(HandleResult::AgentResponse(output)) => {
let output_truncated = if output.len() > 8000 {
format!(
"{}...[truncated]",
&output[..output.ceil_char_boundary(8000)]
)
} else {
output.clone()
};
let run = JobRun {
id: 0,
job_id: job.id.clone(),
started_at,
finished_at,
status: "ok".to_string(),
output: Some(output_truncated),
error: None,
duration_ms,
};
if let Err(e) = self.storage.record_scheduled_job_run(&run).await {
tracing::error!(job_id = %job.id, "scheduler: failed to record run: {}", e);
}
if let Err(e) = self
.storage
.set_scheduled_job_last_status(&job.id, "ok", None)
.await
{
tracing::error!(job_id = %job.id, "scheduler: failed to set last_status: {}", e);
}
tracing::info!(
job_id = %job.id,
duration_ms = %duration_ms,
"scheduler: job completed successfully"
);
}
Ok(HandleResult::CommandOutput(output)) => {
let run = JobRun {
id: 0,
job_id: job.id.clone(),
started_at,
finished_at,
status: "ok".to_string(),
output: Some(output),
error: None,
duration_ms,
};
let _ = self.storage.record_scheduled_job_run(&run).await;
}
Ok(HandleResult::AgentProcessing) => {
tracing::warn!(job_id = %job.id, "scheduler: unexpected AgentProcessing from cron — response sent via bus");
}
Err(e) => {
let error_str = e.to_string();
let run = JobRun {
id: 0,
job_id: job.id.clone(),
started_at,
finished_at,
status: "error".to_string(),
output: None,
error: Some(error_str.clone()),
duration_ms,
};
if let Err(e2) = self.storage.record_scheduled_job_run(&run).await {
tracing::error!(job_id = %job.id, "scheduler: failed to record error run: {}", e2);
}
if let Err(e2) = self
.storage
.set_scheduled_job_last_status(&job.id, "error", Some(&error_str))
.await
{
tracing::error!(job_id = %job.id, "scheduler: failed to set error status: {}", e2);
}
tracing::error!(
job_id = %job.id,
duration_ms = %duration_ms,
error = %error_str,
"scheduler: job failed"
);
}
}
if let Err(e) = self.reschedule_after_run(job).await {
tracing::error!(job_id = %job.id, "scheduler: failed to reschedule: {}", e);
}
}
stream::iter(jobs)
.for_each_concurrent(max_concurrent, |job| {
let scheduler = self.clone();
async move { scheduler.execute_claimed_job(job).await }
})
.await;
}
}
/// After a job runs, compute its next execution time or disable/delete it.
async fn reschedule_after_run(&self, job: &ScheduledJob) -> anyhow::Result<()> {
let now = now_ms();
async fn execute_claimed_job(self: Arc<Self>, job: ScheduledJob) {
let start = Instant::now();
let started_at = now_ms();
tracing::info!(job_id = %job.id, job_name = %job.name, "scheduler: executing claimed job");
match &job.schedule {
Schedule::At { .. } => {
if job.delete_after_run {
self.storage.remove_scheduled_job(&job.id).await?;
tracing::info!(job_id = %job.id, "scheduler: one-shot job deleted after run");
let execution = self.session_manager.handle_cron_message(
&job.channel,
&job.chat_id,
&job.prompt,
&job.id,
&job.name,
);
let result = time::timeout(
time::Duration::from_secs(self.config.execution_timeout_secs.max(1)),
execution,
)
.await;
let finished_at = now_ms();
let duration_ms = start.elapsed().as_millis() as i64;
let (status, output, error) = match result {
Ok(Ok(HandleResult::AgentResponse(output) | HandleResult::CommandOutput(output))) => {
let output = if output.len() > 8000 {
format!(
"{}...[truncated]",
&output[..output.ceil_char_boundary(8000)]
)
} else {
self.storage
.set_scheduled_job_enabled(&job.id, false)
.await?;
tracing::info!(job_id = %job.id, "scheduler: one-shot job disabled after run");
}
output
};
("ok".to_string(), Some(output), None)
}
Ok(Ok(HandleResult::AgentProcessing)) => (
"error".to_string(),
None,
Some("cron execution returned asynchronous processing".to_string()),
),
Ok(Err(error)) => ("error".to_string(), None, Some(error.to_string())),
Err(_) => (
"timeout".to_string(),
None,
Some(format!(
"execution exceeded {} seconds",
self.config.execution_timeout_secs.max(1)
)),
),
};
let (next_run_at, disable, delete) = match &job.schedule {
Schedule::At { .. } => (None, !job.delete_after_run, job.delete_after_run),
Schedule::Every { .. } | Schedule::Cron { .. } => {
if let Some(next) = next_run_for_schedule(&job.schedule, now) {
self.storage
.set_scheduled_job_next_run(&job.id, next)
.await?;
tracing::info!(job_id = %job.id, next_run_at = %next, "scheduler: job rescheduled");
} else {
tracing::error!(job_id = %job.id, "scheduler: could not compute next run -- disabling job");
self.storage
.set_scheduled_job_enabled(&job.id, false)
.await?;
match next_run_for_schedule(&job.schedule, finished_at) {
Some(next) => (Some(next), false, false),
None => (None, true, false),
}
}
};
let run = JobRun {
id: 0,
job_id: job.id.clone(),
started_at,
finished_at,
status,
output,
error,
duration_ms,
};
if let Err(error) = self
.storage
.complete_scheduled_job(&run, &self.owner, next_run_at, disable, delete)
.await
{
tracing::error!(job_id = %job.id, error = %error, "scheduler: failed to commit job completion");
let _ = self
.storage
.release_scheduled_job_lease(&job.id, &self.owner)
.await;
return;
}
Ok(())
tracing::info!(
job_id = %job.id,
status = %run.status,
duration_ms,
"scheduler: job completed"
);
}
}

5
src/session/commands.rs
View File

@ -21,6 +21,11 @@ pub enum SessionCommand {
chat_id: String,
dialog_id: String,
},
/// Load persisted messages for a dialog.
GetDialogHistory {
session_id: UnifiedSessionId,
limit: u32,
},
/// Get the current dialog for a chat
GetCurrentDialog { channel: String, chat_id: String },
/// Rename a dialog

5
src/session/events.rs
View File

@ -31,6 +31,11 @@ pub enum SessionEvent {
},
/// Dialog switched successfully
DialogSwitched { session_id: UnifiedSessionId },
/// Persisted dialog messages, ordered by sequence.
DialogHistory {
session_id: UnifiedSessionId,
messages: Vec<crate::storage::message::MessageMeta>,
},
/// Dialog renamed
DialogRenamed {
session_id: UnifiedSessionId,

113
src/session/messenger.rs Normal file
View File

@ -0,0 +1,113 @@
use std::collections::HashMap;
use crate::bus::{ChatMessage, MediaItem, MessageSource, OutboundMessage};
use crate::session::UnifiedSessionId;
use crate::tools::OutboundMessenger;
use super::persistence::append_persisted_messages;
use super::session::{CURRENT_SOURCE_SESSION, SessionManager};
#[async_trait::async_trait]
impl OutboundMessenger for SessionManager {
async fn send_message(
&self,
channel: &str,
chat_id: &str,
dialog_id: Option<&str>,
content: &str,
mut source: MessageSource,
media: Vec<MediaItem>,
) -> Result<(), String> {
if source.from_session.is_none() {
source.from_session = CURRENT_SOURCE_SESSION
.try_with(|value| value.clone())
.ok()
.flatten();
}
let (target_sid, session) = if let Some(dialog_id) = dialog_id {
let session_id = UnifiedSessionId::new(channel, chat_id, dialog_id);
let session = self
.get_or_activate_session(&session_id)
.await
.map_err(|error| error.to_string())?;
(session_id, session)
} else {
let session_id = self
.resolve_dialog_id(channel, chat_id)
.await
.map_err(|error| error.to_string())?;
let session = self
.get_or_create_session(&session_id)
.await
.map_err(|error| error.to_string())?;
(session_id, session)
};
let origin = source.from_session.as_deref().unwrap_or("unknown");
let origin_id = source.from_session.clone();
let same_session = source.from_session.as_deref() == Some(target_sid.to_string().as_str());
let marked_content = if content.trim().is_empty() && !media.is_empty() && same_session {
String::new()
} else {
format!("[message from {origin}] \n{content}")
};
let message = outbound_history_message(marked_content.clone(), source, &media);
append_persisted_messages(&session, vec![message])
.await
.map_err(|error| error.to_string())?;
if let Some(origin_id) = origin_id {
self.restore_origin_dialog(&origin_id, &target_sid).await;
}
self.bus
.deliver_outbound(OutboundMessage {
channel: channel.to_string(),
chat_id: chat_id.to_string(),
content: marked_content,
reply_to: None,
media,
metadata: HashMap::new(),
delivery: None,
})
.await
.map_err(|error| error.to_string())
}
}
fn outbound_history_message(
content: impl Into<String>,
source: MessageSource,
media: &[MediaItem],
) -> ChatMessage {
let mut message = ChatMessage::assistant_with_source(content, source);
message.media_refs = media.iter().map(MediaItem::to_media_ref).collect();
message
}
#[cfg(test)]
mod tests {
use super::*;
use crate::bus::SourceKind;
#[test]
fn outbound_history_preserves_delivered_media() {
let source = MessageSource {
kind: SourceKind::CrossChannel,
from_channel: Some("cli_chat".into()),
from_session: Some("cli_chat:source:dialog".into()),
from_user_id: None,
system_name: None,
task_id: None,
};
let media = vec![MediaItem::new("/tmp/report.pdf", "file")];
let message = outbound_history_message("report", source, &media);
assert_eq!(message.media_refs.len(), 1);
assert_eq!(message.media_refs[0].path, "/tmp/report.pdf");
assert_eq!(message.media_refs[0].media_type, "file");
}
}

4
src/session/mod.rs
View File

@ -1,6 +1,10 @@
pub mod commands;
pub mod error;
pub mod events;
mod messenger;
mod persistence;
// The public `session::session` path is retained for API compatibility.
#[allow(clippy::module_inception)]
pub mod session;
pub mod session_id;

65
src/session/persistence.rs Normal file
View File

@ -0,0 +1,65 @@
use std::sync::Arc;
use tokio::sync::Mutex;
use super::session::{MessagePersistSnapshot, Session};
use crate::bus::ChatMessage;
use crate::storage::StorageError;
async fn persist_added_messages(
snapshots: Vec<Option<MessagePersistSnapshot>>,
) -> Result<(), StorageError> {
let mut storage = None;
let mut session_id = None;
let mut messages = Vec::new();
let mut final_meta = None;
for snapshot in snapshots.into_iter().flatten() {
let (snapshot_storage, snapshot_session_id, message, meta) = snapshot;
if let Some(ref expected) = session_id
&& expected != &snapshot_session_id
{
return Err(StorageError::Serialization(
"attempted to persist messages from different sessions in one turn".to_string(),
));
}
storage = Some(snapshot_storage);
session_id = Some(snapshot_session_id);
messages.push(message);
final_meta = Some(meta);
}
let (Some(storage), Some(session_id), Some(final_meta)) = (storage, session_id, final_meta)
else {
return Ok(());
};
storage
.persist_message_batch_with_retry(&session_id, &messages, &final_meta)
.await
}
pub(super) async fn append_persisted_messages(
session: &Arc<Mutex<Session>>,
messages: Vec<ChatMessage>,
) -> Result<(), StorageError> {
if messages.is_empty() {
return Ok(());
}
let persistence_lock = { session.lock().await.persistence_lock.clone() };
let _persistence_guard = persistence_lock.lock().await;
let message_ids: Vec<_> = messages.iter().map(|message| message.id.clone()).collect();
let snapshots = {
let mut guard = session.lock().await;
messages
.into_iter()
.map(|message| guard.add_message_in_memory(message, true))
.collect()
};
if let Err(error) = persist_added_messages(snapshots).await {
session.lock().await.rollback_message_suffix(&message_ids);
return Err(error);
}
Ok(())
}

701
src/session/session.rs
View File

File diff suppressed because it is too large Load Diff

29
src/session/session_id.rs
View File

@ -5,13 +5,8 @@
/// Examples:
/// - CLI: `"cli_chat:sid_abc123:dialog_xyz"`
/// - Feishu: `"feishu:oc_123456:dialog_xyz"`
///
/// For simple cases where only one dialog exists per chat:
/// - `dialog_id` defaults to `"default"`
use serde::{Deserialize, Serialize};
pub const DEFAULT_DIALOG_ID: &str = "default";
#[derive(Debug, Clone, PartialEq, Eq, Hash, Serialize, Deserialize)]
pub struct UnifiedSessionId {
pub channel: String,
@ -33,15 +28,6 @@ impl UnifiedSessionId {
}
}
/// Create with default dialog_id ("default")
pub fn with_default_dialog(channel: impl Into<String>, chat_id: impl Into<String>) -> Self {
Self {
channel: channel.into(),
chat_id: chat_id.into(),
dialog_id: DEFAULT_DIALOG_ID.to_string(),
}
}
/// Parse from string format "channel:chat_id:dialog_id"
pub fn parse(s: &str) -> Option<Self> {
let parts: Vec<&str> = s.split(':').collect();
@ -55,11 +41,6 @@ impl UnifiedSessionId {
})
}
/// Convert to string format "channel:chat_id:dialog_id"
pub fn to_string(&self) -> String {
format!("{}:{}:{}", self.channel, self.chat_id, self.dialog_id)
}
/// Get the session key without dialog_id (channel:chat_id)
/// This is used to group all dialogs within a chat
pub fn chat_scope(&self) -> String {
@ -69,7 +50,7 @@ impl UnifiedSessionId {
impl std::fmt::Display for UnifiedSessionId {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
write!(f, "{}", self.to_string())
write!(f, "{}:{}:{}", self.channel, self.chat_id, self.dialog_id)
}
}
@ -87,14 +68,6 @@ mod tests {
assert_eq!(id.dialog_id, "dialog456");
}
#[test]
fn test_with_default_dialog() {
let id = UnifiedSessionId::with_default_dialog("feishu", "oc123");
assert_eq!(id.channel, "feishu");
assert_eq!(id.chat_id, "oc123");
assert_eq!(id.dialog_id, "default");
}
#[test]
fn test_parse() {
let id = UnifiedSessionId::parse("cli_chat:sid123:dialog456").unwrap();

10
src/storage/background_task.rs
View File

@ -17,3 +17,13 @@ pub struct BackgroundTask {
pub finished_at: Option<i64>,
pub created_at: i64,
}
pub(crate) struct BackgroundTaskUpdate<'a> {
pub status: &'a str,
pub result: Option<&'a str>,
pub error: Option<&'a str>,
pub started_at: Option<i64>,
pub finished_at: Option<i64>,
pub tool_calls_count: Option<i64>,
pub iterations: Option<i64>,
}

25
src/storage/error.rs
View File

@ -13,4 +13,29 @@ pub enum StorageError {
#[error("serialization error: {0}")]
Serialization(String),
#[error("schema migration error: {0}")]
Migration(String),
#[error("storage conflict: {0}")]
Conflict(String),
}
impl StorageError {
/// Only retry failures that can plausibly clear without changing the data.
pub fn is_transient(&self) -> bool {
let Self::Database(error) = self else {
return false;
};
match error {
sqlx::Error::PoolTimedOut => true,
sqlx::Error::Database(database) => {
matches!(database.code().as_deref(), Some("5" | "6" | "261" | "262")) || {
let message = database.message().to_ascii_lowercase();
message.contains("database is locked") || message.contains("database is busy")
}
}
_ => false,
}
}
}

2
src/storage/memory.rs
View File

@ -282,7 +282,7 @@ fn parse_memory_rows(rows: &[sqlx::sqlite::SqliteRow]) -> Result<Vec<MemoryEntry
id: row.try_get("id")?,
key: row.try_get("key")?,
content: row.try_get("content")?,
category: MemoryCategory::from_str(&row.try_get::<String, _>("category")?)
category: MemoryCategory::parse(&row.try_get::<String, _>("category")?)
.unwrap_or(MemoryCategory::Knowledge),
importance: row.try_get::<f64, _>("importance")?,
session_id: row.try_get::<Option<String>, _>("session_id")?,

620
src/storage/mod.rs
View File

@ -9,10 +9,36 @@ pub use background_task::BackgroundTask;
pub use error::StorageError;
pub use scheduler::{JobRun, ScheduledJob};
use sqlx::{Pool, Row, Sqlite, SqlitePool};
use sqlx::sqlite::{SqliteConnectOptions, SqliteJournalMode, SqlitePoolOptions, SqliteSynchronous};
use sqlx::{Pool, Row, Sqlite};
use std::path::Path;
use tokio::time::{Duration, sleep};
const SCHEMA_VERSION: i64 = 1;
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 (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
"#;
fn insert_message_query<'a>(
session_id: &'a str,
msg: &'a crate::storage::message::MessageMeta,
) -> sqlx::query::Query<'a, Sqlite, sqlx::sqlite::SqliteArguments<'a>> {
sqlx::query(INSERT_MESSAGE_SQL)
.bind(&msg.id)
.bind(session_id)
.bind(msg.seq)
.bind(&msg.role)
.bind(&msg.content)
.bind(&msg.reasoning_content)
.bind(&msg.media_refs)
.bind(&msg.tool_call_id)
.bind(&msg.tool_name)
.bind(&msg.tool_calls)
.bind(&msg.source)
.bind(msg.created_at)
}
pub struct Storage {
pub(crate) pool: Pool<Sqlite>,
}
@ -20,8 +46,17 @@ pub struct Storage {
impl Storage {
/// 打开或创建数据库
pub async fn new(db_path: &Path) -> Result<Self, StorageError> {
let database_url = format!("sqlite:{}?mode=rwc", db_path.display());
let pool = SqlitePool::connect(&database_url).await?;
let options = SqliteConnectOptions::new()
.filename(db_path)
.create_if_missing(true)
.journal_mode(SqliteJournalMode::Wal)
.synchronous(SqliteSynchronous::Normal)
.busy_timeout(Duration::from_secs(5))
.foreign_keys(true);
let pool = SqlitePoolOptions::new()
.max_connections(8)
.connect_with(options)
.await?;
let storage = Self { pool };
storage.init_schema().await?;
@ -75,6 +110,7 @@ impl Storage {
tool_name TEXT,
tool_calls TEXT,
source TEXT,
reasoning_content TEXT,
created_at INTEGER NOT NULL,
FOREIGN KEY (session_id) REFERENCES sessions(id) ON DELETE CASCADE
)
@ -92,18 +128,6 @@ impl Storage {
.execute(&self.pool)
.await?;
// Migration: add source column if upgrading from older schema
sqlx::query(r#"ALTER TABLE messages ADD COLUMN source TEXT"#)
.execute(&self.pool)
.await
.ok();
// Migration: add reasoning_content column if upgrading from older schema
sqlx::query(r#"ALTER TABLE messages ADD COLUMN reasoning_content TEXT"#)
.execute(&self.pool)
.await
.ok();
// Background tasks table — for async sub-agent tasks.
// Note: No FOREIGN KEY on session_id because sessions use soft delete (deleted_at IS NULL).
// Session and task association is maintained at the application level.
@ -163,6 +187,12 @@ impl Storage {
.execute(&self.pool)
.await?;
let memory_fts_exists: bool = sqlx::query_scalar(
"SELECT EXISTS(SELECT 1 FROM sqlite_master WHERE type = 'table' AND name = 'memory_fts')",
)
.fetch_one(&self.pool)
.await?;
// FTS5 virtual table for full-text search on memories
sqlx::query(
r#"
@ -212,40 +242,14 @@ impl Storage {
.execute(&self.pool)
.await?;
// Rebuild FTS5 index for any existing records
sqlx::query("INSERT INTO memory_fts(memory_fts) VALUES ('rebuild')")
.execute(&self.pool)
.await?;
// Migration: add last_consolidated_at column if not exists
sqlx::query(
r#"
ALTER TABLE sessions ADD COLUMN archived_at INTEGER
"#,
)
.execute(&self.pool)
.await
.ok();
// Migration: add last_consolidated_at column if not exists
sqlx::query(
r#"
ALTER TABLE sessions ADD COLUMN last_consolidated_at INTEGER
"#,
)
.execute(&self.pool)
.await
.ok();
// Migration: add last_compressed_message_at column if not exists
sqlx::query(
r#"
ALTER TABLE sessions ADD COLUMN last_compressed_message_at INTEGER
"#,
)
.execute(&self.pool)
.await
.ok();
// Only a newly-created index needs a backfill. Triggers keep an
// existing index current, so rebuilding it on every startup is wasted
// work proportional to the total memory corpus.
if !memory_fts_exists {
sqlx::query("INSERT INTO memory_fts(memory_fts) VALUES ('rebuild')")
.execute(&self.pool)
.await?;
}
sqlx::query(
r#"
@ -264,13 +268,88 @@ impl Storage {
.execute(&self.pool)
.await?;
if let Err(e) = Self::init_scheduler_schema(&self.pool).await {
tracing::warn!(
"Failed to init scheduler schema (tables may already exist): {}",
e
);
Self::init_scheduler_schema(&self.pool).await?;
self.migrate_schema().await?;
Ok(())
}
/// Apply ordered, atomic migrations. Existing installations predate
/// `user_version`, so each step also checks the actual table shape.
async fn migrate_schema(&self) -> Result<(), StorageError> {
let current: i64 = sqlx::query_scalar("PRAGMA user_version")
.fetch_one(&self.pool)
.await?;
if current > SCHEMA_VERSION {
return Err(StorageError::Migration(format!(
"database schema version {current} is newer than supported version {SCHEMA_VERSION}"
)));
}
if current == SCHEMA_VERSION {
return Ok(());
}
let mut tx = self.pool.begin().await?;
for (table, column, definition) in [
("messages", "source", "source TEXT"),
("messages", "reasoning_content", "reasoning_content TEXT"),
("sessions", "archived_at", "archived_at INTEGER"),
(
"sessions",
"last_consolidated_at",
"last_consolidated_at INTEGER",
),
(
"sessions",
"last_compressed_message_at",
"last_compressed_message_at INTEGER",
),
("scheduled_jobs", "locked_at", "locked_at INTEGER"),
("scheduled_jobs", "lock_owner", "lock_owner TEXT"),
("scheduled_jobs", "lease_until", "lease_until INTEGER"),
] {
let pragma = format!("PRAGMA table_info({table})");
let columns = sqlx::query(&pragma).fetch_all(&mut *tx).await?;
if !columns
.iter()
.any(|row| row.get::<String, _>("name") == column)
{
let alter = format!("ALTER TABLE {table} ADD COLUMN {definition}");
sqlx::query(&alter).execute(&mut *tx).await?;
}
}
let duplicate: Option<(String, i64, i64)> = sqlx::query_as(
r#"
SELECT session_id, seq, COUNT(*)
FROM messages
GROUP BY session_id, seq
HAVING COUNT(*) > 1
LIMIT 1
"#,
)
.fetch_optional(&mut *tx)
.await?;
if let Some((session_id, seq, count)) = duplicate {
return Err(StorageError::Migration(format!(
"cannot enforce unique message sequence: session {session_id} has {count} rows at seq {seq}"
)));
}
sqlx::query(
"CREATE UNIQUE INDEX IF NOT EXISTS idx_messages_session_seq_unique ON messages(session_id, seq)",
)
.execute(&mut *tx)
.await?;
sqlx::query(
"CREATE INDEX IF NOT EXISTS idx_jobs_claimable ON scheduled_jobs(enabled, next_run_at, lease_until)",
)
.execute(&mut *tx)
.await?;
sqlx::query(&format!("PRAGMA user_version = {SCHEMA_VERSION}"))
.execute(&mut *tx)
.await?;
tx.commit().await?;
Ok(())
}
@ -292,6 +371,9 @@ impl Storage {
last_run_at INTEGER,
last_status TEXT,
last_error TEXT,
locked_at INTEGER,
lock_owner TEXT,
lease_until INTEGER,
created_at INTEGER NOT NULL,
updated_at INTEGER NOT NULL
)
@ -583,41 +665,85 @@ impl Storage {
session_id: &str,
msg: &crate::storage::message::MessageMeta,
) -> Result<i64, StorageError> {
sqlx::query(
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 (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
"#,
)
.bind(&msg.id)
.bind(session_id)
.bind(msg.seq)
.bind(&msg.role)
.bind(&msg.content)
.bind(&msg.reasoning_content)
.bind(&msg.media_refs)
.bind(&msg.tool_call_id)
.bind(&msg.tool_name)
.bind(&msg.tool_calls)
.bind(&msg.source)
.bind(msg.created_at)
.execute(self.pool())
.await?;
insert_message_query(session_id, msg)
.execute(self.pool())
.await?;
Ok(msg.seq)
}
pub async fn append_messages(
/// Atomically persist all messages produced by one logical turn together
/// with the resulting session metadata. A turn is either fully visible
/// after restart or not visible at all.
pub async fn persist_message_batch(
&self,
session_id: &str,
msgs: &[crate::storage::message::MessageMeta],
) -> Result<Vec<i64>, StorageError> {
let mut seqs = Vec::with_capacity(msgs.len());
meta: &crate::storage::session::SessionMeta,
) -> Result<(), StorageError> {
let mut tx = self.pool.begin().await?;
for msg in msgs {
let seq = self.append_message(session_id, msg).await?;
seqs.push(seq);
insert_message_query(session_id, msg)
.execute(&mut *tx)
.await?;
}
Ok(seqs)
sqlx::query(
r#"
INSERT INTO sessions (id, channel, chat_id, dialog_id, title, created_at, last_active_at, message_count, routing_info, archived_at, deleted_at, last_consolidated_at, last_compressed_message_at)
VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
ON CONFLICT(id) DO UPDATE SET
title = excluded.title,
last_active_at = excluded.last_active_at,
message_count = excluded.message_count,
routing_info = excluded.routing_info,
archived_at = excluded.archived_at,
deleted_at = excluded.deleted_at,
last_consolidated_at = excluded.last_consolidated_at,
last_compressed_message_at = excluded.last_compressed_message_at
"#,
)
.bind(&meta.id)
.bind(&meta.channel)
.bind(&meta.chat_id)
.bind(&meta.dialog_id)
.bind(&meta.title)
.bind(meta.created_at)
.bind(meta.last_active_at)
.bind(meta.message_count)
.bind(&meta.routing_info)
.bind(meta.archived_at)
.bind(meta.deleted_at)
.bind(meta.last_consolidated_at)
.bind(meta.last_compressed_message_at)
.execute(&mut *tx)
.await?;
tx.commit().await?;
Ok(())
}
/// Persist a turn with bounded retry. Retrying the whole transaction keeps
/// message rows and metadata consistent on transient SQLite failures.
pub async fn persist_message_batch_with_retry(
&self,
session_id: &str,
msgs: &[crate::storage::message::MessageMeta],
meta: &crate::storage::session::SessionMeta,
) -> Result<(), StorageError> {
let delays = [100, 200, 300];
for (attempt, delay) in delays.iter().enumerate() {
match self.persist_message_batch(session_id, msgs, meta).await {
Ok(()) => return Ok(()),
Err(error) if attempt < delays.len() - 1 && error.is_transient() => {
tracing::warn!(attempt = attempt + 1, error = %error, "Turn persistence failed; retrying");
sleep(Duration::from_millis(*delay)).await;
}
Err(error) => return Err(error),
}
}
unreachable!()
}
pub async fn load_messages(
@ -667,6 +793,52 @@ impl Storage {
Ok(row.get::<i64, _>("max_seq"))
}
/// Load a bounded tail of one session while preserving chronological order.
pub async fn load_recent_session_messages(
&self,
session_id: &str,
limit: u32,
) -> Result<Vec<crate::storage::message::MessageMeta>, StorageError> {
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
FROM (
SELECT id, session_id, seq, role, content, reasoning_content, media_refs,
tool_call_id, tool_name, tool_calls, source, created_at
FROM messages
WHERE session_id = ?
ORDER BY seq DESC
LIMIT ?
)
ORDER BY seq ASC
"#,
)
.bind(session_id)
.bind(i64::from(limit))
.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())
}
pub async fn load_messages_after_timestamp(
&self,
session_id: &str,
@ -930,28 +1102,28 @@ impl Storage {
Ok(())
}
pub async fn update_background_task_status(
pub(crate) async fn update_background_task_status(
&self,
id: &str,
status: &str,
result: Option<&str>,
error: Option<&str>,
started_at: Option<i64>,
finished_at: Option<i64>,
update: crate::storage::background_task::BackgroundTaskUpdate<'_>,
) -> Result<(), StorageError> {
sqlx::query(
r#"
UPDATE background_tasks
SET status = ?, result = COALESCE(?, result), error = COALESCE(?, error),
started_at = COALESCE(?, started_at), finished_at = COALESCE(?, finished_at)
started_at = COALESCE(?, started_at), finished_at = COALESCE(?, finished_at),
tool_calls_count = COALESCE(?, tool_calls_count),
iterations = COALESCE(?, iterations)
WHERE id = ?
"#,
)
.bind(status)
.bind(result)
.bind(error)
.bind(started_at)
.bind(finished_at)
.bind(update.status)
.bind(update.result)
.bind(update.error)
.bind(update.started_at)
.bind(update.finished_at)
.bind(update.tool_calls_count)
.bind(update.iterations)
.bind(id)
.execute(self.pool())
.await?;
@ -1054,6 +1226,200 @@ mod tests {
(storage, dir)
}
#[tokio::test]
async fn sqlite_runtime_guards_are_enabled() {
let (storage, _dir) = create_test_storage().await;
let journal_mode: String = sqlx::query_scalar("PRAGMA journal_mode")
.fetch_one(storage.pool())
.await
.unwrap();
let foreign_keys: i64 = sqlx::query_scalar("PRAGMA foreign_keys")
.fetch_one(storage.pool())
.await
.unwrap();
let busy_timeout: i64 = sqlx::query_scalar("PRAGMA busy_timeout")
.fetch_one(storage.pool())
.await
.unwrap();
let schema_version: i64 = sqlx::query_scalar("PRAGMA user_version")
.fetch_one(storage.pool())
.await
.unwrap();
assert_eq!(journal_mode, "wal");
assert_eq!(foreign_keys, 1);
assert_eq!(busy_timeout, 5000);
assert_eq!(schema_version, SCHEMA_VERSION);
let orphan = sqlx::query(
r#"
INSERT INTO messages (id, session_id, seq, role, content, created_at)
VALUES ('orphan', 'missing', 1, 'user', 'no parent', 1)
"#,
)
.execute(storage.pool())
.await;
assert!(orphan.is_err());
}
#[tokio::test]
async fn reopening_database_does_not_rebuild_existing_fts_index() {
let dir = tempfile::tempdir().unwrap();
let db_path = dir.path().join("fts.db");
let storage = Storage::new(&db_path).await.unwrap();
sqlx::query(
"INSERT INTO memory_fts(rowid, key, content) VALUES (999999, 'startup_sentinel', 'startup_sentinel')",
)
.execute(storage.pool())
.await
.unwrap();
drop(storage);
let reopened = Storage::new(&db_path).await.unwrap();
let sentinel_count: i64 = sqlx::query_scalar(
"SELECT count(*) FROM memory_fts WHERE memory_fts MATCH 'startup_sentinel'",
)
.fetch_one(reopened.pool())
.await
.unwrap();
assert_eq!(sentinel_count, 1);
}
#[tokio::test]
async fn background_task_completion_persists_execution_metrics() {
let (storage, _dir) = create_test_storage().await;
let task = crate::storage::BackgroundTask {
id: "task-metrics".into(),
session_id: "cli:test:dialog".into(),
channel: "cli".into(),
chat_id: "test".into(),
prompt: "measure".into(),
allowed_tools: None,
status: "pending".into(),
result: None,
error: None,
tool_calls_count: 0,
iterations: 0,
started_at: None,
finished_at: None,
created_at: 1,
};
storage.create_background_task(&task).await.unwrap();
storage
.update_background_task_status(
&task.id,
crate::storage::background_task::BackgroundTaskUpdate {
status: "completed",
result: Some("done"),
error: None,
started_at: Some(2),
finished_at: Some(3),
tool_calls_count: Some(4),
iterations: Some(5),
},
)
.await
.unwrap();
let persisted = storage.get_background_task(&task.id).await.unwrap();
assert_eq!(persisted.tool_calls_count, 4);
assert_eq!(persisted.iterations, 5);
}
#[tokio::test]
async fn legacy_schema_is_migrated_without_rebuild() {
let dir = tempfile::tempdir().unwrap();
let db_path = dir.path().join("legacy.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, deleted_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, media_refs TEXT,
tool_call_id TEXT, tool_name TEXT, tool_calls TEXT,
created_at INTEGER NOT NULL
)
"#,
)
.execute(&pool)
.await
.unwrap();
sqlx::query(
r#"
CREATE TABLE scheduled_jobs (
id TEXT PRIMARY KEY, name TEXT NOT NULL, schedule TEXT NOT NULL,
prompt TEXT NOT NULL, channel TEXT NOT NULL, chat_id TEXT NOT NULL,
model TEXT, enabled INTEGER NOT NULL DEFAULT 1,
delete_after_run INTEGER NOT NULL DEFAULT 0, next_run_at INTEGER NOT NULL,
last_run_at INTEGER, last_status TEXT, last_error TEXT,
created_at INTEGER NOT NULL, updated_at INTEGER NOT NULL
)
"#,
)
.execute(&pool)
.await
.unwrap();
drop(pool);
let storage = Storage::new(&db_path).await.unwrap();
for (table, expected) in [
("messages", vec!["source", "reasoning_content"]),
(
"sessions",
vec![
"archived_at",
"last_consolidated_at",
"last_compressed_message_at",
],
),
(
"scheduled_jobs",
vec!["locked_at", "lock_owner", "lease_until"],
),
] {
let columns = sqlx::query(&format!("PRAGMA table_info({table})"))
.fetch_all(storage.pool())
.await
.unwrap();
for column in expected {
assert!(
columns
.iter()
.any(|row| row.get::<String, _>("name") == column),
"missing migrated column {table}.{column}"
);
}
}
let schema_version: i64 = sqlx::query_scalar("PRAGMA user_version")
.fetch_one(storage.pool())
.await
.unwrap();
assert_eq!(schema_version, SCHEMA_VERSION);
}
#[tokio::test]
async fn test_upsert_and_get_session() {
let (storage, _dir) = create_test_storage().await;
@ -1101,8 +1467,8 @@ mod tests {
chat_id: "sid123".to_string(),
dialog_id: format!("dialog{}", i),
title: format!("会话{}", i),
created_at: i as i64 * 1000,
last_active_at: i as i64 * 1000,
created_at: i * 1000,
last_active_at: i * 1000,
message_count: i,
routing_info: None,
archived_at: None,
@ -1195,6 +1561,72 @@ mod tests {
let loaded = storage.load_messages(&session_meta.id, 0).await.unwrap();
assert_eq!(loaded.len(), 1);
assert_eq!(loaded[0].content, "你好");
for seq in 2..=5 {
let mut message = msg.clone();
message.id = format!("msg{seq}");
message.seq = seq;
message.content = format!("message {seq}");
storage
.append_message(&session_meta.id, &message)
.await
.unwrap();
}
let recent = storage
.load_recent_session_messages(&session_meta.id, 2)
.await
.unwrap();
assert_eq!(recent.len(), 2);
assert_eq!(recent[0].seq, 4);
assert_eq!(recent[1].seq, 5);
}
#[tokio::test]
async fn test_persist_message_batch_is_atomic() {
let (storage, _dir) = create_test_storage().await;
let session_meta = crate::storage::session::SessionMeta {
id: "cli_chat:atomic:dialog1".to_string(),
channel: "cli_chat".to_string(),
chat_id: "atomic".to_string(),
dialog_id: "dialog1".to_string(),
title: "Atomic turn".to_string(),
created_at: 1000,
last_active_at: 2000,
message_count: 1,
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: "duplicate-id".to_string(),
session_id: session_meta.id.clone(),
seq: 1,
role: "assistant".to_string(),
content: "must roll back".to_string(),
reasoning_content: None,
media_refs: None,
tool_call_id: None,
tool_name: None,
tool_calls: None,
source: None,
created_at: 2000,
};
let result = storage
.persist_message_batch(&session_meta.id, &[message.clone(), message], &session_meta)
.await;
assert!(result.is_err());
assert!(
storage
.load_messages(&session_meta.id, 0)
.await
.unwrap()
.is_empty()
);
}
#[tokio::test]

402
src/storage/scheduler.rs
View File

@ -2,6 +2,7 @@ use serde::{Deserialize, Serialize};
use sqlx::Row;
use crate::scheduler::Schedule;
use crate::storage::StorageError;
/// A scheduled job stored in the database.
#[derive(Debug, Clone, Serialize, Deserialize)]
@ -39,8 +40,8 @@ pub struct JobRun {
impl crate::storage::Storage {
/// Insert a new scheduled job.
pub async fn add_scheduled_job(&self, job: &ScheduledJob) -> anyhow::Result<()> {
let schedule_json = serde_json::to_string(&job.schedule)?;
pub async fn add_scheduled_job(&self, job: &ScheduledJob) -> Result<(), StorageError> {
let schedule_json = serialize_schedule(&job.schedule)?;
sqlx::query(
r#"
INSERT INTO scheduled_jobs
@ -71,17 +72,17 @@ impl crate::storage::Storage {
}
/// Fetch a single scheduled job by ID.
pub async fn get_scheduled_job(&self, id: &str) -> anyhow::Result<ScheduledJob> {
pub async fn get_scheduled_job(&self, id: &str) -> Result<ScheduledJob, StorageError> {
let row = sqlx::query("SELECT * FROM scheduled_jobs WHERE id = ?")
.bind(id)
.fetch_optional(self.pool())
.await?
.ok_or_else(|| anyhow::anyhow!("job not found: {id}"))?;
.ok_or_else(|| StorageError::NotFound(format!("scheduled job {id}")))?;
row_to_job(&row)
}
/// List all scheduled jobs, ordered by next_run_at ascending.
pub async fn list_scheduled_jobs(&self) -> anyhow::Result<Vec<ScheduledJob>> {
pub async fn list_scheduled_jobs(&self) -> Result<Vec<ScheduledJob>, StorageError> {
let rows = sqlx::query("SELECT * FROM scheduled_jobs ORDER BY next_run_at ASC")
.fetch_all(self.pool())
.await?;
@ -89,7 +90,7 @@ impl crate::storage::Storage {
}
/// Delete a scheduled job (cascades to job_runs).
pub async fn remove_scheduled_job(&self, id: &str) -> anyhow::Result<()> {
pub async fn remove_scheduled_job(&self, id: &str) -> Result<(), StorageError> {
sqlx::query("DELETE FROM scheduled_jobs WHERE id = ?")
.bind(id)
.execute(self.pool())
@ -98,7 +99,11 @@ impl crate::storage::Storage {
}
/// Enable or disable a scheduled job.
pub async fn set_scheduled_job_enabled(&self, id: &str, enabled: bool) -> anyhow::Result<()> {
pub async fn set_scheduled_job_enabled(
&self,
id: &str,
enabled: bool,
) -> Result<(), StorageError> {
sqlx::query("UPDATE scheduled_jobs SET enabled = ?, updated_at = ? WHERE id = ?")
.bind(enabled as i32)
.bind(now_ms())
@ -117,7 +122,7 @@ impl crate::storage::Storage {
channel: Option<String>,
chat_id: Option<String>,
model: Option<String>,
) -> anyhow::Result<()> {
) -> Result<(), StorageError> {
let now = now_ms();
if let Some(p) = prompt {
@ -129,7 +134,7 @@ impl crate::storage::Storage {
.await?;
}
if let Some(s) = schedule {
let json = serde_json::to_string(&s)?;
let json = serialize_schedule(&s)?;
sqlx::query("UPDATE scheduled_jobs SET schedule = ?, updated_at = ? WHERE id = ?")
.bind(&json)
.bind(now)
@ -169,7 +174,7 @@ impl crate::storage::Storage {
&self,
id: &str,
next_run_at: i64,
) -> anyhow::Result<()> {
) -> Result<(), StorageError> {
let now = now_ms();
sqlx::query(
"UPDATE scheduled_jobs SET next_run_at = ?, last_run_at = ?, updated_at = ? WHERE id = ?",
@ -183,68 +188,132 @@ impl crate::storage::Storage {
Ok(())
}
/// Set last_run_at for a job (used when starting execution).
pub async fn touch_scheduled_job_last_run(&self, id: &str, at: i64) -> anyhow::Result<()> {
sqlx::query("UPDATE scheduled_jobs SET last_run_at = ?, updated_at = ? WHERE id = ?")
.bind(at)
.bind(at)
.bind(id)
.execute(self.pool())
.await?;
Ok(())
}
/// Set last_status and last_error after job completion.
pub async fn set_scheduled_job_last_status(
&self,
id: &str,
status: &str,
error: Option<&str>,
) -> anyhow::Result<()> {
let now = now_ms();
sqlx::query(
"UPDATE scheduled_jobs SET last_status = ?, last_error = ?, updated_at = ? WHERE id = ?",
)
.bind(status)
.bind(error)
.bind(now)
.bind(id)
.execute(self.pool())
.await?;
Ok(())
}
/// Fetch enabled jobs whose next_run_at <= now, up to `limit`.
pub async fn due_scheduled_jobs(
/// Atomically claim due jobs for one scheduler instance. A crashed worker's
/// claims become eligible again after `lease_until`.
pub async fn claim_due_scheduled_jobs(
&self,
now: i64,
lease_until: i64,
owner: &str,
limit: usize,
) -> anyhow::Result<Vec<ScheduledJob>> {
) -> Result<Vec<ScheduledJob>, StorageError> {
if limit == 0 {
return Ok(Vec::new());
}
let rows = sqlx::query(
"SELECT * FROM scheduled_jobs WHERE enabled = 1 AND next_run_at <= ? ORDER BY next_run_at ASC LIMIT ?",
r#"
UPDATE scheduled_jobs
SET locked_at = ?, lock_owner = ?, lease_until = ?, last_run_at = ?, updated_at = ?
WHERE id IN (
SELECT id FROM scheduled_jobs
WHERE enabled = 1
AND next_run_at <= ?
AND (lease_until IS NULL OR lease_until <= ?)
ORDER BY next_run_at ASC
LIMIT ?
)
AND (lease_until IS NULL OR lease_until <= ?)
RETURNING *
"#,
)
.bind(now)
.bind(owner)
.bind(lease_until)
.bind(now)
.bind(now)
.bind(now)
.bind(now)
.bind(limit as i64)
.bind(now)
.fetch_all(self.pool())
.await?;
rows.iter().map(row_to_job).collect()
}
/// Record a job execution run.
pub async fn record_scheduled_job_run(&self, run: &JobRun) -> anyhow::Result<()> {
/// Persist the run result, reschedule/disable the job, and release its
/// lease in one transaction. The owner check prevents a stale worker from
/// completing a claim that has already been recovered elsewhere.
pub async fn complete_scheduled_job(
&self,
run: &JobRun,
owner: &str,
next_run_at: Option<i64>,
disable: bool,
delete: bool,
) -> Result<(), StorageError> {
let mut tx = self.pool().begin().await?;
if delete {
let result = sqlx::query("DELETE FROM scheduled_jobs WHERE id = ? AND lock_owner = ?")
.bind(&run.job_id)
.bind(owner)
.execute(&mut *tx)
.await?;
if result.rows_affected() != 1 {
return Err(StorageError::Conflict(format!(
"scheduled job lease lost before delete: {}",
run.job_id
)));
}
} else {
sqlx::query(
r#"
INSERT INTO job_runs (job_id, started_at, finished_at, status, output, error, duration_ms)
VALUES (?, ?, ?, ?, ?, ?, ?)
"#,
)
.bind(&run.job_id)
.bind(run.started_at)
.bind(run.finished_at)
.bind(&run.status)
.bind(&run.output)
.bind(&run.error)
.bind(run.duration_ms)
.execute(&mut *tx)
.await?;
let result = sqlx::query(
r#"
UPDATE scheduled_jobs
SET next_run_at = COALESCE(?, next_run_at),
enabled = CASE WHEN ? THEN 0 ELSE enabled END,
last_status = ?, last_error = ?,
locked_at = NULL, lock_owner = NULL, lease_until = NULL,
updated_at = ?
WHERE id = ? AND lock_owner = ?
"#,
)
.bind(next_run_at)
.bind(disable)
.bind(&run.status)
.bind(&run.error)
.bind(run.finished_at)
.bind(&run.job_id)
.bind(owner)
.execute(&mut *tx)
.await?;
if result.rows_affected() != 1 {
return Err(StorageError::Conflict(format!(
"scheduled job lease lost before completion: {}",
run.job_id
)));
}
}
tx.commit().await?;
Ok(())
}
pub async fn release_scheduled_job_lease(
&self,
job_id: &str,
owner: &str,
) -> Result<(), StorageError> {
sqlx::query(
r#"
INSERT INTO job_runs (job_id, started_at, finished_at, status, output, error, duration_ms)
VALUES (?, ?, ?, ?, ?, ?, ?)
"#,
"UPDATE scheduled_jobs SET locked_at = NULL, lock_owner = NULL, lease_until = NULL WHERE id = ? AND lock_owner = ?",
)
.bind(&run.job_id)
.bind(run.started_at)
.bind(run.finished_at)
.bind(&run.status)
.bind(&run.output)
.bind(&run.error)
.bind(run.duration_ms)
.bind(job_id)
.bind(owner)
.execute(self.pool())
.await?;
Ok(())
@ -255,7 +324,7 @@ impl crate::storage::Storage {
&self,
job_id: &str,
limit: usize,
) -> anyhow::Result<Vec<JobRun>> {
) -> Result<Vec<JobRun>, StorageError> {
let rows = sqlx::query(
"SELECT * FROM job_runs WHERE job_id = ? ORDER BY finished_at DESC LIMIT ?",
)
@ -280,7 +349,7 @@ impl crate::storage::Storage {
}
/// Delete disabled jobs whose updated_at is before `before`.
pub async fn cleanup_disabled_scheduled_jobs(&self, before: i64) -> anyhow::Result<()> {
pub async fn cleanup_disabled_scheduled_jobs(&self, before: i64) -> Result<(), StorageError> {
sqlx::query("DELETE FROM scheduled_jobs WHERE enabled = 0 AND updated_at < ?")
.bind(before)
.execute(self.pool())
@ -296,9 +365,14 @@ fn now_ms() -> i64 {
.as_millis() as i64
}
fn row_to_job(row: &sqlx::sqlite::SqliteRow) -> anyhow::Result<ScheduledJob> {
fn serialize_schedule(schedule: &Schedule) -> Result<String, StorageError> {
serde_json::to_string(schedule).map_err(|error| StorageError::Serialization(error.to_string()))
}
fn row_to_job(row: &sqlx::sqlite::SqliteRow) -> Result<ScheduledJob, StorageError> {
let schedule_json: String = row.try_get("schedule")?;
let schedule: Schedule = serde_json::from_str(&schedule_json)?;
let schedule: Schedule = serde_json::from_str(&schedule_json)
.map_err(|error| StorageError::Serialization(error.to_string()))?;
Ok(ScheduledJob {
id: row.try_get("id")?,
name: row.try_get("name")?,
@ -464,114 +538,6 @@ mod tests {
assert!(!got.enabled);
}
#[tokio::test]
async fn test_due_jobs_only_returns_enabled_and_overdue() {
let storage = setup_storage().await;
let t = now();
let jobs = vec![
ScheduledJob {
id: "due".into(),
name: "due".into(),
schedule: Schedule::At { at: t },
prompt: "1".into(),
channel: "cli_chat".into(),
chat_id: "c".into(),
model: None,
enabled: true,
delete_after_run: false,
next_run_at: t - 1000,
last_run_at: None,
last_status: None,
last_error: None,
created_at: t,
updated_at: t,
},
ScheduledJob {
id: "future".into(),
name: "future".into(),
schedule: Schedule::At { at: t + 99999999 },
prompt: "2".into(),
channel: "cli_chat".into(),
chat_id: "c".into(),
model: None,
enabled: true,
delete_after_run: false,
next_run_at: t + 99999999,
last_run_at: None,
last_status: None,
last_error: None,
created_at: t,
updated_at: t,
},
ScheduledJob {
id: "disabled-due".into(),
name: "disabled due".into(),
schedule: Schedule::At { at: t },
prompt: "3".into(),
channel: "cli_chat".into(),
chat_id: "c".into(),
model: None,
enabled: false,
delete_after_run: false,
next_run_at: t - 1000,
last_run_at: None,
last_status: None,
last_error: None,
created_at: t,
updated_at: t,
},
];
for j in &jobs {
storage.add_scheduled_job(j).await.unwrap();
}
let due = storage.due_scheduled_jobs(t, 10).await.unwrap();
assert_eq!(due.len(), 1);
assert_eq!(due[0].id, "due");
}
#[tokio::test]
async fn test_record_run_and_list_runs() {
let storage = setup_storage().await;
let t = now();
let job = ScheduledJob {
id: "job-run".into(),
name: "run test".into(),
schedule: Schedule::Every { every_ms: 1000 },
prompt: "hi".into(),
channel: "cli_chat".into(),
chat_id: "c".into(),
model: None,
enabled: true,
delete_after_run: false,
next_run_at: t,
last_run_at: None,
last_status: None,
last_error: None,
created_at: t,
updated_at: t,
};
storage.add_scheduled_job(&job).await.unwrap();
let run = super::JobRun {
id: 0,
job_id: "job-run".into(),
started_at: t,
finished_at: t + 500,
status: "ok".into(),
output: Some("hello".into()),
error: None,
duration_ms: 500,
};
storage.record_scheduled_job_run(&run).await.unwrap();
let runs = storage
.list_scheduled_job_runs("job-run", 10)
.await
.unwrap();
assert_eq!(runs.len(), 1);
assert_eq!(runs[0].status, "ok");
assert_eq!(runs[0].output.as_deref(), Some("hello"));
}
#[tokio::test]
async fn test_update_job() {
let storage = setup_storage().await;
@ -608,4 +574,100 @@ mod tests {
let got = storage.get_scheduled_job("job-update").await.unwrap();
assert_eq!(got.prompt, "new prompt");
}
#[tokio::test]
async fn claim_is_exclusive_until_lease_expires() {
let storage = setup_storage().await;
let t = now();
let job = ScheduledJob {
id: "leased-job".into(),
name: "leased".into(),
schedule: Schedule::Every { every_ms: 1000 },
prompt: "run".into(),
channel: "cli_chat".into(),
chat_id: "c".into(),
model: None,
enabled: true,
delete_after_run: false,
next_run_at: t,
last_run_at: None,
last_status: None,
last_error: None,
created_at: t,
updated_at: t,
};
storage.add_scheduled_job(&job).await.unwrap();
let first = storage
.claim_due_scheduled_jobs(t, t + 100, "owner-1", 1)
.await
.unwrap();
let duplicate = storage
.claim_due_scheduled_jobs(t, t + 100, "owner-2", 1)
.await
.unwrap();
let recovered = storage
.claim_due_scheduled_jobs(t + 101, t + 201, "owner-2", 1)
.await
.unwrap();
assert_eq!(first.len(), 1);
assert!(duplicate.is_empty());
assert_eq!(recovered.len(), 1);
}
#[tokio::test]
async fn completion_is_atomic_and_releases_lease() {
let storage = setup_storage().await;
let t = now();
let job = ScheduledJob {
id: "complete-job".into(),
name: "complete".into(),
schedule: Schedule::Every { every_ms: 1000 },
prompt: "run".into(),
channel: "cli_chat".into(),
chat_id: "c".into(),
model: None,
enabled: true,
delete_after_run: false,
next_run_at: t,
last_run_at: None,
last_status: None,
last_error: None,
created_at: t,
updated_at: t,
};
storage.add_scheduled_job(&job).await.unwrap();
storage
.claim_due_scheduled_jobs(t, t + 1000, "owner", 1)
.await
.unwrap();
let run = super::JobRun {
id: 0,
job_id: job.id.clone(),
started_at: t,
finished_at: t + 10,
status: "ok".into(),
output: Some("done".into()),
error: None,
duration_ms: 10,
};
storage
.complete_scheduled_job(&run, "owner", Some(t + 2000), false, false)
.await
.unwrap();
let completed = storage.get_scheduled_job(&job.id).await.unwrap();
let runs = storage.list_scheduled_job_runs(&job.id, 10).await.unwrap();
let lease: (Option<String>, Option<i64>) =
sqlx::query_as("SELECT lock_owner, lease_until FROM scheduled_jobs WHERE id = ?")
.bind(&job.id)
.fetch_one(storage.pool())
.await
.unwrap();
assert_eq!(completed.next_run_at, t + 2000);
assert_eq!(completed.last_status.as_deref(), Some("ok"));
assert_eq!(runs.len(), 1);
assert_eq!(lease, (None, None));
}
}

202
src/task_supervisor.rs Normal file
View File

@ -0,0 +1,202 @@
use std::future::Future;
use std::panic::AssertUnwindSafe;
use std::sync::{Arc, Mutex};
use std::time::Duration;
use futures_util::FutureExt;
use tokio::task::JoinHandle;
use tokio::time::{Instant, timeout_at};
use tokio_util::sync::CancellationToken;
#[derive(Clone)]
pub struct TaskSupervisor {
inner: Arc<Inner>,
}
struct Inner {
cancellation: CancellationToken,
state: Mutex<State>,
}
impl Drop for Inner {
fn drop(&mut self) {
self.cancellation.cancel();
}
}
#[derive(Default)]
struct State {
stopping: bool,
tasks: Vec<ManagedTask>,
}
struct ManagedTask {
name: String,
handle: JoinHandle<()>,
}
impl Default for TaskSupervisor {
fn default() -> Self {
Self::new()
}
}
impl TaskSupervisor {
pub fn new() -> Self {
Self {
inner: Arc::new(Inner {
cancellation: CancellationToken::new(),
state: Mutex::new(State::default()),
}),
}
}
pub fn cancellation_token(&self) -> CancellationToken {
self.inner.cancellation.clone()
}
/// Register a task before shutdown begins. Cancellation drops the task
/// future, so task code should keep externally visible state transactional.
pub fn spawn<F>(&self, name: impl Into<String>, future: F) -> bool
where
F: Future<Output = ()> + Send + 'static,
{
let name = name.into();
let cancellation = self.inner.cancellation.clone();
let mut state = self.inner.state.lock().unwrap_or_else(|e| e.into_inner());
if state.stopping {
return false;
}
// Completed handles no longer need to occupy the registry. Panics are
// observed and logged inside the wrapper below.
state.tasks.retain(|task| !task.handle.is_finished());
let task_name = name.clone();
let handle = tokio::spawn(async move {
tracing::debug!(task = %task_name, "Background task started");
let outcome = tokio::select! {
_ = cancellation.cancelled() => None,
outcome = AssertUnwindSafe(future).catch_unwind() => Some(outcome),
};
match outcome {
Some(Ok(())) => tracing::debug!(task = %task_name, "Background task finished"),
Some(Err(_)) => tracing::error!(task = %task_name, "Background task panicked"),
None => tracing::debug!(task = %task_name, "Background task cancelled"),
}
});
state.tasks.push(ManagedTask { name, handle });
true
}
/// Register a task that performs its own cooperative cancellation and
/// cleanup. The supervisor broadcasts cancellation during shutdown, but
/// does not drop this future until the grace period expires.
pub fn spawn_graceful<F>(&self, name: impl Into<String>, future: F) -> bool
where
F: Future<Output = ()> + Send + 'static,
{
let name = name.into();
let mut state = self.inner.state.lock().unwrap_or_else(|e| e.into_inner());
if state.stopping {
return false;
}
state.tasks.retain(|task| !task.handle.is_finished());
let task_name = name.clone();
let handle = tokio::spawn(async move {
tracing::debug!(task = %task_name, "Graceful background task started");
match AssertUnwindSafe(future).catch_unwind().await {
Ok(()) => tracing::debug!(task = %task_name, "Graceful background task finished"),
Err(_) => tracing::error!(task = %task_name, "Graceful background task panicked"),
}
});
state.tasks.push(ManagedTask { name, handle });
true
}
pub fn cancel(&self) {
let mut state = self.inner.state.lock().unwrap_or_else(|e| e.into_inner());
state.stopping = true;
self.inner.cancellation.cancel();
}
/// Stop accepting tasks, broadcast cancellation, and wait up to `grace`.
/// Remaining tasks are aborted so shutdown has a deterministic upper bound.
pub async fn shutdown(&self, grace: Duration) {
let mut tasks = {
let mut state = self.inner.state.lock().unwrap_or_else(|e| e.into_inner());
state.stopping = true;
self.inner.cancellation.cancel();
std::mem::take(&mut state.tasks)
};
let deadline = Instant::now() + grace;
for index in 0..tasks.len() {
let result = timeout_at(deadline, &mut tasks[index].handle).await;
match result {
Ok(Ok(())) => {}
Ok(Err(error)) if error.is_cancelled() => {}
Ok(Err(error)) => {
tracing::error!(task = %tasks[index].name, error = %error, "Background task join failed");
}
Err(_) => {
for task in &tasks[index..] {
if !task.handle.is_finished() {
tracing::warn!(task = %task.name, "Aborting background task after shutdown grace period");
task.handle.abort();
}
}
for task in &mut tasks[index..] {
let _ = (&mut task.handle).await;
}
break;
}
}
}
}
}
#[cfg(test)]
mod tests {
use super::*;
use std::sync::atomic::{AtomicBool, Ordering};
#[tokio::test]
async fn shutdown_cancels_registered_task() {
let supervisor = TaskSupervisor::new();
let dropped = Arc::new(AtomicBool::new(false));
let marker = dropped.clone();
supervisor.spawn("pending", async move {
struct DropMarker(Arc<AtomicBool>);
impl Drop for DropMarker {
fn drop(&mut self) {
self.0.store(true, Ordering::SeqCst);
}
}
let _marker = DropMarker(marker);
std::future::pending::<()>().await;
});
tokio::task::yield_now().await;
supervisor.cancel();
assert!(!supervisor.spawn("late", async {}));
supervisor.shutdown(Duration::from_secs(1)).await;
assert!(dropped.load(Ordering::SeqCst));
}
#[tokio::test]
async fn graceful_task_observes_cancellation_before_shutdown_returns() {
let supervisor = TaskSupervisor::new();
let cancellation = supervisor.cancellation_token();
let cleaned_up = Arc::new(std::sync::atomic::AtomicBool::new(false));
let task_cleaned_up = cleaned_up.clone();
assert!(supervisor.spawn_graceful("graceful", async move {
cancellation.cancelled().await;
task_cleaned_up.store(true, std::sync::atomic::Ordering::SeqCst);
}));
supervisor.shutdown(Duration::from_secs(1)).await;
assert!(cleaned_up.load(std::sync::atomic::Ordering::SeqCst));
}
}

33
src/tools/browser.rs
View File

@ -39,11 +39,11 @@ struct BrowserState {
impl Drop for BrowserTool {
fn drop(&mut self) {
if let Ok(mut driver) = self.driver.lock() {
if let Some(ref mut child) = driver.take() {
tracing::debug!("Stopping chromedriver process");
let _ = child.start_kill();
}
if let Ok(mut driver) = self.driver.lock()
&& let Some(ref mut child) = driver.take()
{
tracing::debug!("Stopping chromedriver process");
let _ = child.start_kill();
}
}
}
@ -454,10 +454,7 @@ impl BrowserState {
} => {
let client = self.active_client()?;
let result: Value = client
.execute(
&snapshot_script(interactive_only, compact, depth.map(i64::from)),
vec![],
)
.execute(&snapshot_script(interactive_only, compact, depth), vec![])
.await?;
let output = serde_json::to_string_pretty(&result)?;
Ok(ToolResult {
@ -826,11 +823,11 @@ impl BrowserState {
if let Some(client) = self.client.take() {
let _ = client.close().await;
}
if let Ok(mut guard) = driver.lock() {
if let Some(ref mut child) = guard.take() {
tracing::debug!("Stopping chromedriver process");
let _ = child.start_kill();
}
if let Ok(mut guard) = driver.lock()
&& let Some(ref mut child) = guard.take()
{
tracing::debug!("Stopping chromedriver process");
let _ = child.start_kill();
}
}
@ -957,10 +954,10 @@ fn launch_chromedriver(
}
fn kill_driver_guard(driver: &std::sync::Mutex<Option<tokio::process::Child>>) {
if let Ok(mut guard) = driver.lock() {
if let Some(ref mut child) = guard.take() {
let _ = child.start_kill();
}
if let Ok(mut guard) = driver.lock()
&& let Some(ref mut child) = guard.take()
{
let _ = child.start_kill();
}
}

2
src/tools/calculator.rs
View File

@ -380,7 +380,7 @@ fn calc_evaluate(args: &serde_json::Value) -> Result<String, String> {
.and_then(|v| v.as_str())
.ok_or_else(|| "Missing required parameter: expression".to_string())?;
meval::eval_str(expression)
super::expression::evaluate(expression)
.map(format_num)
.map_err(|e| format!("Expression evaluation error: {e}"))
}

10
src/tools/chat_manager.rs
View File

@ -299,8 +299,8 @@ mod tests {
chat_id: format!("sid{}", i),
dialog_id: format!("dialog{}", i),
title: format!("会话{}", i),
created_at: now - i * 3600_000,
last_active_at: now - i * 3600_000,
created_at: now - i * 3_600_000,
last_active_at: now - i * 3_600_000,
message_count: i * 5,
routing_info: None,
archived_at: None,
@ -350,7 +350,7 @@ mod tests {
let msg = crate::storage::message::MessageMeta {
id: format!("msg{}", i),
session_id: session_id.to_string(),
seq: i as i64 + 1,
seq: i + 1,
role: if i == 0 {
"user".to_string()
} else {
@ -412,7 +412,7 @@ mod tests {
let msg = crate::storage::message::MessageMeta {
id: format!("msg{}", i),
session_id: session_id.to_string(),
seq: i as i64 + 1,
seq: i + 1,
role: if i % 2 == 0 {
"user".to_string()
} else {
@ -472,7 +472,7 @@ mod tests {
let msg = crate::storage::message::MessageMeta {
id: format!("msg{}", i),
session_id: session_id.to_string(),
seq: i as i64 + 1,
seq: i + 1,
role: "user".to_string(),
content: format!("消息内容 {}", i),
reasoning_content: None,

10
src/tools/delegate.rs
View File

@ -303,11 +303,11 @@ impl DelegateTool {
if let Some(ref error) = task.error {
output.push_str(&format!("\n错误: {}", error));
}
if let Some(started) = task.started_at {
if let Some(finished) = task.finished_at {
let duration = (finished - started) as f64 / 1000.0;
output.push_str(&format!("\n耗时: {:.1}s", duration));
}
if let Some(started) = task.started_at
&& let Some(finished) = task.finished_at
{
let duration = (finished - started) as f64 / 1000.0;
output.push_str(&format!("\n耗时: {:.1}s", duration));
}
Ok(ToolResult {
success: true,

277
src/tools/expression.rs Normal file
View File

@ -0,0 +1,277 @@
/// Evaluate a self-contained mathematical expression without executing code or
/// resolving external variables.
pub(super) fn evaluate(input: &str) -> Result<f64, String> {
const MAX_EXPRESSION_BYTES: usize = 4096;
if input.len() > MAX_EXPRESSION_BYTES {
return Err(format!(
"expression exceeds the {MAX_EXPRESSION_BYTES}-byte limit"
));
}
let mut parser = Parser {
input,
position: 0,
depth: 0,
};
let value = parser.parse_expression()?;
parser.skip_whitespace();
if parser.position != input.len() {
return Err(parser.error("unexpected trailing input"));
}
Ok(value)
}
struct Parser<'a> {
input: &'a str,
position: usize,
depth: usize,
}
impl Parser<'_> {
fn parse_expression(&mut self) -> Result<f64, String> {
let mut value = self.parse_term()?;
loop {
if self.consume(b'+') {
value += self.parse_term()?;
} else if self.consume(b'-') {
value -= self.parse_term()?;
} else {
return Ok(value);
}
}
}
fn parse_term(&mut self) -> Result<f64, String> {
let mut value = self.parse_unary()?;
loop {
if self.consume(b'*') {
value *= self.parse_unary()?;
} else if self.consume(b'/') {
value /= self.parse_unary()?;
} else if self.consume(b'%') {
value %= self.parse_unary()?;
} else {
return Ok(value);
}
}
}
fn parse_unary(&mut self) -> Result<f64, String> {
if self.consume(b'+') {
self.nested(Self::parse_unary)
} else if self.consume(b'-') {
Ok(-self.nested(Self::parse_unary)?)
} else {
self.parse_power()
}
}
fn parse_power(&mut self) -> Result<f64, String> {
let base = self.parse_primary()?;
if self.consume(b'^') {
Ok(base.powf(self.nested(Self::parse_unary)?))
} else {
Ok(base)
}
}
fn parse_primary(&mut self) -> Result<f64, String> {
self.skip_whitespace();
match self.peek() {
Some(b'(') => {
self.position += 1;
let value = self.nested(Self::parse_expression)?;
if !self.consume(b')') {
return Err(self.error("expected ')'"));
}
Ok(value)
}
Some(byte) if byte.is_ascii_digit() || byte == b'.' => self.parse_number(),
Some(byte) if byte.is_ascii_alphabetic() || byte == b'_' => self.parse_identifier(),
Some(_) => Err(self.error("expected a number, constant, function, or '('")),
None => Err(self.error("unexpected end of expression")),
}
}
fn parse_number(&mut self) -> Result<f64, String> {
self.skip_whitespace();
let start = self.position;
let mut digits = 0;
while self.peek().is_some_and(|byte| byte.is_ascii_digit()) {
self.position += 1;
digits += 1;
}
if self.peek() == Some(b'.') {
self.position += 1;
while self.peek().is_some_and(|byte| byte.is_ascii_digit()) {
self.position += 1;
digits += 1;
}
}
if digits == 0 {
return Err(self.error("invalid number"));
}
if matches!(self.peek(), Some(b'e' | b'E')) {
self.position += 1;
if matches!(self.peek(), Some(b'+' | b'-')) {
self.position += 1;
}
let exponent_start = self.position;
while self.peek().is_some_and(|byte| byte.is_ascii_digit()) {
self.position += 1;
}
if self.position == exponent_start {
return Err(self.error("invalid numeric exponent"));
}
}
self.input[start..self.position]
.parse::<f64>()
.map_err(|_| self.error("invalid number"))
}
fn parse_identifier(&mut self) -> Result<f64, String> {
self.skip_whitespace();
let start = self.position;
while self
.peek()
.is_some_and(|byte| byte.is_ascii_alphanumeric() || byte == b'_')
{
self.position += 1;
}
let name = self.input[start..self.position].to_ascii_lowercase();
self.skip_whitespace();
if self.peek() != Some(b'(') {
return match name.as_str() {
"pi" => Ok(std::f64::consts::PI),
"e" => Ok(std::f64::consts::E),
_ => Err(self.error(&format!("unknown constant or variable '{name}'"))),
};
}
self.position += 1;
let mut arguments = Vec::new();
self.skip_whitespace();
if self.peek() != Some(b')') {
loop {
arguments.push(self.nested(Self::parse_expression)?);
if self.consume(b',') {
continue;
}
break;
}
}
if !self.consume(b')') {
return Err(self.error("expected ')' after function arguments"));
}
apply_function(&name, &arguments).map_err(|message| self.error(&message))
}
fn consume(&mut self, expected: u8) -> bool {
self.skip_whitespace();
if self.peek() == Some(expected) {
self.position += 1;
true
} else {
false
}
}
fn skip_whitespace(&mut self) {
while self.peek().is_some_and(|byte| byte.is_ascii_whitespace()) {
self.position += 1;
}
}
fn peek(&self) -> Option<u8> {
self.input.as_bytes().get(self.position).copied()
}
fn nested<T>(&mut self, parse: fn(&mut Self) -> Result<T, String>) -> Result<T, String> {
const MAX_PARSE_DEPTH: usize = 128;
if self.depth >= MAX_PARSE_DEPTH {
return Err(self.error("expression nesting limit exceeded"));
}
self.depth += 1;
let result = parse(self);
self.depth -= 1;
result
}
fn error(&self, message: &str) -> String {
format!("{message} at byte {}", self.position)
}
}
fn apply_function(name: &str, arguments: &[f64]) -> Result<f64, String> {
let unary = |function: fn(f64) -> f64| match arguments {
[value] => Ok(function(*value)),
_ => Err(format!("function '{name}' expects one argument")),
};
match name {
"sqrt" => unary(f64::sqrt),
"abs" => unary(f64::abs),
"exp" => unary(f64::exp),
"ln" => unary(f64::ln),
"log2" => unary(f64::log2),
"log10" => unary(f64::log10),
"sin" => unary(f64::sin),
"cos" => unary(f64::cos),
"tan" => unary(f64::tan),
"asin" => unary(f64::asin),
"acos" => unary(f64::acos),
"atan" => unary(f64::atan),
"sinh" => unary(f64::sinh),
"cosh" => unary(f64::cosh),
"tanh" => unary(f64::tanh),
"asinh" => unary(f64::asinh),
"acosh" => unary(f64::acosh),
"atanh" => unary(f64::atanh),
"floor" => unary(f64::floor),
"ceil" => unary(f64::ceil),
"round" => unary(f64::round),
"signum" => unary(f64::signum),
"atan2" => match arguments {
[y, x] => Ok(y.atan2(*x)),
_ => Err("function 'atan2' expects two arguments".to_string()),
},
"min" => arguments
.iter()
.copied()
.reduce(f64::min)
.ok_or_else(|| "function 'min' expects at least one argument".to_string()),
"max" => arguments
.iter()
.copied()
.reduce(f64::max)
.ok_or_else(|| "function 'max' expects at least one argument".to_string()),
_ => Err(format!("unknown function '{name}'")),
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn respects_precedence_and_right_associative_power() {
assert_eq!(evaluate("15*3+5^(2+1)").unwrap(), 170.0);
assert_eq!(evaluate("2^3^2").unwrap(), 512.0);
assert_eq!(evaluate("-2^2").unwrap(), -4.0);
}
#[test]
fn supports_constants_functions_and_scientific_notation() {
assert_eq!(evaluate("sqrt(1.44e2)").unwrap(), 12.0);
assert_eq!(evaluate("max(1, 2, 3) + min(4, 5)").unwrap(), 7.0);
assert!((evaluate("sin(pi / 2)").unwrap() - 1.0).abs() < f64::EPSILON);
}
#[test]
fn rejects_unknown_names_and_trailing_input() {
assert!(evaluate("unknown").is_err());
assert!(evaluate("1 + 2 garbage").is_err());
assert!(evaluate("sqrt() ").is_err());
assert!(evaluate(&"(".repeat(129)).is_err());
assert!(evaluate(&"1+".repeat(3000)).is_err());
}
}

9
src/tools/mod.rs
View File

@ -5,6 +5,7 @@ pub mod chat_manager;
pub mod content_search;
pub mod cron;
pub mod delegate;
mod expression;
pub mod file_edit;
pub mod file_read;
pub mod file_search;
@ -77,10 +78,10 @@ pub fn create_default_tools(
registry.register(TimelineRecallTool::new(memory.clone()));
registry.register(MemoryForgetTool::new(memory.clone()));
if let Some(cfg) = browser_config {
if cfg.enabled {
registry.register(BrowserTool::new(cfg));
}
if let Some(cfg) = browser_config
&& cfg.enabled
{
registry.register(BrowserTool::new(cfg));
}
if let Some(mgr) = sub_agent_manager {

17
src/tools/pty.rs
View File

@ -96,11 +96,7 @@ enum SessionStatus {
}
struct PtySession {
#[allow(dead_code)]
id: String,
#[allow(dead_code)]
command: String,
#[allow(dead_code)]
started_at: Instant,
status: SessionStatus,
child: Arc<Mutex<Option<Box<dyn portable_pty::Child + Send + Sync>>>>,
@ -111,13 +107,11 @@ struct PtySession {
impl PtySession {
fn new(
id: String,
command: String,
child: Box<dyn portable_pty::Child + Send + Sync>,
writer: Box<dyn Write + Send>,
) -> Self {
Self {
id,
command,
started_at: Instant::now(),
status: SessionStatus::Running,
@ -146,6 +140,12 @@ pub struct PtyManager {
sessions: Mutex<HashMap<String, Arc<Mutex<PtySession>>>>,
}
impl Default for PtyManager {
fn default() -> Self {
Self::new()
}
}
impl PtyManager {
pub fn new() -> Self {
Self {
@ -199,7 +199,7 @@ impl PtyManager {
.map_err(|e| format!("Failed to open PTY: {}", e))?;
let mut cmd = portable_pty::CommandBuilder::new("bash");
cmd.args(&["-c", command]);
cmd.args(["-c", command]);
cmd.cwd(cwd);
let child = pty_pair
@ -219,8 +219,7 @@ impl PtyManager {
.try_clone_reader()
.map_err(|e| format!("Failed to clone reader: {}", e))?;
let session_id_clone = session_id.clone();
let session = PtySession::new(session_id_clone, command.to_string(), child, writer);
let session = PtySession::new(command.to_string(), child, writer);
let session = Arc::new(Mutex::new(session));
sessions.insert(session_id.clone(), session.clone());

2
src/tools/send_message.rs
View File

@ -169,7 +169,7 @@ fn parse_files_arg(args: &serde_json::Value) -> Vec<MediaItem> {
files
.iter()
.filter_map(|v| v.as_str())
.map(|path| path_to_media_item(path))
.map(path_to_media_item)
.collect()
}

35
tests/test_request_format.rs
View File

@ -1,4 +1,4 @@
use picobot::protocol::{SessionSummary, WsInbound, WsOutbound};
use picobot::protocol::{HistoryMessage, SessionSummary, WsInbound, WsOutbound};
use picobot::providers::{ChatCompletionRequest, Message};
/// Test that message with special characters is properly escaped
@ -18,7 +18,7 @@ fn test_message_special_characters() {
/// Test that multi-line system prompt is preserved
#[test]
fn test_multiline_system_prompt() {
let messages = vec![
let messages = [
Message::system(
"You are a helpful assistant.\n\nFollow these rules:\n1. Be kind\n2. Be accurate",
),
@ -116,3 +116,34 @@ fn test_clear_history_with_session_id_serialization() {
assert!(json.contains(r#""type":"clear_history""#));
assert!(json.contains(r#""session_id":"session-1""#));
}
#[test]
fn test_bounded_session_history_protocol() {
let inbound = WsInbound::GetSessionHistory {
session_id: "cli_chat:client:dialog".to_string(),
limit: Some(1000),
};
let json = serde_json::to_string(&inbound).unwrap();
assert!(json.contains(r#""type":"get_session_history""#));
assert!(json.contains(r#""limit":1000"#));
let outbound = WsOutbound::SessionHistory {
session_id: "cli_chat:client:dialog".to_string(),
messages: vec![HistoryMessage {
id: "m1".to_string(),
seq: 1,
role: "user".to_string(),
content: "你好".to_string(),
created_at: 123,
}],
};
let decoded: WsOutbound =
serde_json::from_str(&serde_json::to_string(&outbound).unwrap()).unwrap();
match decoded {
WsOutbound::SessionHistory { messages, .. } => {
assert_eq!(messages.len(), 1);
assert_eq!(messages[0].content, "你好");
}
other => panic!("unexpected decoded variant: {other:?}"),
}
}

5
tests/test_scheduler.rs
View File

@ -1,6 +1,5 @@
/// Integration tests for the scheduled tasks (cron) system.
/// Run with: cargo test --test test_scheduler
use serde_json::json;
//! Integration tests for the scheduled tasks (cron) system.
//! Run with: `cargo test --test test_scheduler`.
/// Verify that Schedule types (de)serialize correctly.
#[tokio::test]