Compare commits
No commits in common. "7de9a8a054f0d557b79bf196d65c21bb86d0c17f" and "c7ee6bb519071cbc56a72a895fe6e122840d6d00" have entirely different histories.
7de9a8a054
...
c7ee6bb519
98
.github/workflows/ci.yml
vendored
98
.github/workflows/ci.yml
vendored
@ -1,98 +0,0 @@
|
||||
name: CI
|
||||
|
||||
on:
|
||||
push:
|
||||
branches: [main, master]
|
||||
pull_request:
|
||||
branches: [main, master]
|
||||
|
||||
env:
|
||||
CARGO_TERM_COLOR: always
|
||||
# 跳过前端构建（build.rs 会触发 npm install + npm run build）
|
||||
# CI 中独立处理前端检查，避免 cargo check 时重复构建
|
||||
SKIP_FRONTEND_BUILD: "1"
|
||||
|
||||
jobs:
|
||||
rust-checks:
|
||||
name: Rust fmt + clippy + test
|
||||
runs-on: ${{ matrix.os }}
|
||||
strategy:
|
||||
fail-fast: false
|
||||
matrix:
|
||||
os: [ubuntu-latest, windows-latest]
|
||||
steps:
|
||||
- uses: actions/checkout@v4
|
||||
|
||||
- name: Install Rust toolchain
|
||||
uses: dtolnay/rust-toolchain@stable
|
||||
with:
|
||||
components: rustfmt, clippy
|
||||
|
||||
- name: Cache cargo registry
|
||||
uses: actions/cache@v4
|
||||
with:
|
||||
path: |
|
||||
~/.cargo/registry
|
||||
~/.cargo/git
|
||||
target
|
||||
key: ${{ runner.os }}-cargo-${{ hashFiles('**/Cargo.lock') }}
|
||||
|
||||
- name: Check formatting
|
||||
run: cargo fmt --all -- --check
|
||||
|
||||
# 不使用 -D warnings：现有代码有大量 warn 级别的存量问题
|
||||
# （unwrap/expect/clone 等），一上线会让 CI 一直红。
|
||||
# 仅 correctness 类别（clippy 默认 deny）会失败，先渐进收集。
|
||||
# 等存量清理后再考虑加 -- -D warnings
|
||||
- name: Clippy
|
||||
run: cargo clippy --all-targets --all-features
|
||||
|
||||
- name: Build
|
||||
# 构建 lib + 二进制入口（main.rs），确保 gateway 二进制也被编译验证
|
||||
run: cargo build
|
||||
|
||||
- name: Run tests
|
||||
# 运行 lib 单元测试 + tests/ 集成测试
|
||||
# test_integration.rs / test_tool_calling.rs 中的 #[ignore] 测试
|
||||
# 需要真实 API key，会被跳过；test_request_format.rs 的测试会实际执行
|
||||
run: cargo test
|
||||
|
||||
- name: Security audit
|
||||
run: cargo install cargo-audit --locked && cargo audit
|
||||
|
||||
frontend-checks:
|
||||
name: Frontend build + test
|
||||
runs-on: ubuntu-latest
|
||||
steps:
|
||||
- uses: actions/checkout@v4
|
||||
|
||||
- name: Setup Node
|
||||
uses: actions/setup-node@v4
|
||||
with:
|
||||
node-version: "20"
|
||||
cache: "npm"
|
||||
cache-dependency-path: web/package-lock.json
|
||||
|
||||
- name: Install dependencies
|
||||
working-directory: web
|
||||
run: npm ci
|
||||
|
||||
- name: Lint (eslint)
|
||||
working-directory: web
|
||||
run: npm run lint
|
||||
|
||||
- name: Format check (prettier)
|
||||
working-directory: web
|
||||
run: npm run format:check
|
||||
|
||||
- name: Type check
|
||||
working-directory: web
|
||||
run: npx tsc --noEmit
|
||||
|
||||
- name: Run tests
|
||||
working-directory: web
|
||||
run: npm run test
|
||||
|
||||
- name: Security audit
|
||||
working-directory: web
|
||||
run: npm audit --audit-level=high
|
||||
5
.gitignore
vendored
5
.gitignore
vendored
@ -11,11 +11,6 @@ web/.cache
|
||||
web/coverage
|
||||
web/*.local
|
||||
|
||||
# Secrets — never commit real credentials
|
||||
*.env
|
||||
!*.env.example
|
||||
tests/test.env
|
||||
|
||||
# Build output
|
||||
static
|
||||
|
||||
|
||||
116
ARCHITECTURE.md
116
ARCHITECTURE.md
@ -1,116 +0,0 @@
|
||||
# PicoBot 架构
|
||||
|
||||
> 本文档聚焦"为什么这样设计"和"数据如何流动"，不是代码导读。
|
||||
> 代码是唯一真相源，文档可能滞后；有冲突以代码为准。
|
||||
|
||||
## 一句话定位
|
||||
|
||||
PicoBot 是一个**多渠道接入的 Agent 网关**：外部消息（微信/飞书/Web/CLI）经统一总线进入，
|
||||
由会话管理器路由到对应 Agent 循环，Agent 调用 LLM + 工具完成任务，结果原路返回。
|
||||
|
||||
## 核心数据流
|
||||
|
||||
```
|
||||
┌─────────┐ ┌──────────┐ ┌──────────────┐ ┌────────────┐ ┌──────────┐
|
||||
│ Channel │──▶│ MessageBus│──▶│InboundProc. │──▶│ Session │──▶│AgentLoop │
|
||||
│ (微信/ │ │ (解耦) │ │(并发+限流) │ │ Manager │ │(LLM+工具)│
|
||||
│ 飞书/Web)│ │ │ │ │ │ │ │ │
|
||||
└─────────┘ └──────────┘ └──────────────┘ └────────────┘ └────┬─────┘
|
||||
▲ │
|
||||
│ ▼
|
||||
┌────┴────┐ ┌──────────┐ ┌────────────┐ ┌──────────┐
|
||||
│Outbound │◀──│ MessageBus│◀─── SessionMessageSender ◀─│ storage │◀─│ tools │
|
||||
│Dispatcher│ │ │ │ (SQLite) │ │(bash/file│
|
||||
└─────────┘ └──────────┘ └────────────┘ │ /memory) │
|
||||
└──────────┘
|
||||
```
|
||||
|
||||
**入站**：Channel → MessageBus → InboundProcessor（并发控制 + 限流）→ SessionManager → AgentLoop
|
||||
**出站**：AgentLoop → SessionMessageSender → MessageBus → OutboundDispatcher → Channel
|
||||
**持久化**：AgentLoop / Tools → SessionStore（SQLite + r2d2 连接池）
|
||||
|
||||
## 关键设计决策
|
||||
|
||||
### 1. MessageBus 解耦 Channel 与 Session
|
||||
|
||||
**问题**：多个 Channel（微信/飞书/Web）接入，每个 Channel 协议不同，但 Session 处理逻辑相同。
|
||||
**决策**：引入 MessageBus 作为中间件，Channel 只负责协议适配和收发，Session 不关心消息来自哪个 Channel。
|
||||
**代价**：多一层间接。换来的是新增 Channel（如钉钉）只需实现 Channel trait，不碰 Session 逻辑。
|
||||
|
||||
### 2. SessionPool 会话隔离
|
||||
|
||||
**问题**：多用户同时对话，会话状态不能串扰。
|
||||
**决策**：SessionManager 持有 SessionPool，按 (channel_name, chat_id) 路由到独立 Session。
|
||||
每个 Session 有自己的 AgentLoop、消息历史、工具上下文。
|
||||
**代价**：内存占用随活跃会话数增长。用 session_ttl_hours 过期回收。
|
||||
|
||||
### 3. AgentLoop 的工具循环
|
||||
|
||||
**问题**：LLM 需要多轮工具调用才能完成任务（如"读文件→分析→写文件"）。
|
||||
**决策**：AgentLoop 是一个有界循环（max_tool_iterations，默认 1000），每轮：
|
||||
1. 把消息历史 + 工具定义发给 LLM
|
||||
2. LLM 返回文本或 tool_call
|
||||
3. 如果是 tool_call，执行工具，把结果加入历史，回到 1
|
||||
4. 如果是文本，结束循环
|
||||
**代价**：单次对话可能很长。用 CancelManager 支持中途取消。
|
||||
|
||||
### 4. SQLite + r2d2 连接池
|
||||
|
||||
**问题**：需要持久化会话历史、话题、记忆、待办，且要支持并发读写的 sub-agent 场景。
|
||||
**决策**：SQLite（bundled）+ r2d2 连接池（max_size=8）+ busy_timeout(30s)。
|
||||
**代价**：SQLite 写并发有限。通过 busy_timeout + 事务隔离级别（之前的修复）缓解锁冲突。
|
||||
**不选 Postgres 的原因**：单机部署、零外部依赖、足够用。
|
||||
|
||||
### 5. 重启机制（watch channel）
|
||||
|
||||
**问题**：配置变更后需要重启 gateway，但不能要求用户手动杀进程。
|
||||
**决策**：main.rs 用 `while should_restart` 循环，gateway 通过 `watch::Sender<bool>` 通知是否需要重启。
|
||||
配置 API `/api/restart` 触发 graceful shutdown，main 收到 should_restart=true 后重新初始化。
|
||||
**代价**：重启期间短暂不可用。比热重载简单且可靠。
|
||||
|
||||
### 6. 嵌入式静态文件
|
||||
|
||||
**问题**：Web 前端需要随二进制分发，但不想要求用户额外下载。
|
||||
**决策**：build.rs 在 cargo build 时执行 npm run build，产物通过 rust-embed 编译进二进制。
|
||||
开发时设 `STATIC_DIR` 环境变量走磁盘文件，支持热更新。
|
||||
**代价**：二进制体积增大。换来的是单文件部署。
|
||||
|
||||
### 7. Safety Guard（命令安全护栏）
|
||||
|
||||
**问题**：Agent 可以调用 bash 工具执行任意命令，需要防止误操作（如 `format C:`、`rm -rf`）。
|
||||
**决策**：platform 模块按平台注入危险命令正则，执行前匹配拦截。
|
||||
**关键教训**：正则要精确（曾因 `\bformat\s+` 误拦 `dart format`），按平台分组避免跨平台误伤。
|
||||
|
||||
## 模块职责速查
|
||||
|
||||
| 模块 | 职责 | 关键文件 |
|
||||
|------|------|---------|
|
||||
| gateway | HTTP/WS 服务、路由、生命周期 | `gateway/mod.rs`, `gateway/runtime.rs` |
|
||||
| gateway/session | 会话管理、路由、池化 | `gateway/session.rs`, `session_pool.rs` |
|
||||
| gateway/processor | 入站消息处理、并发控制 | `gateway/processor.rs` |
|
||||
| agent | Agent 循环、上下文压缩 | `agent/agent_loop.rs` |
|
||||
| providers | LLM Provider 抽象（OpenAI/Anthropic） | `providers/openai.rs`, `anthropic.rs` |
|
||||
| tools | 工具实现与注册 | `tools/` (bash/file/memory/task/...) |
|
||||
| storage | SQLite 持久化 | `storage/mod.rs`, `migrations.rs` |
|
||||
| channels | 渠道适配（微信/飞书/CLI） | `channels/` |
|
||||
| bus | 消息总线（解耦 channel 与 session） | `bus/message.rs` |
|
||||
| command | 前端命令处理（话题/会话/记忆 CRUD） | `command/handlers/` |
|
||||
| mcp | Model Context Protocol 客户端 | `mcp/client.rs` |
|
||||
| scheduler | 定时任务调度 | `scheduler/mod.rs` |
|
||||
| skills | 技能加载与激活 | `skills/mod.rs` |
|
||||
| experts | 专家配置运行时 | `experts/mod.rs` |
|
||||
|
||||
## 测试策略
|
||||
|
||||
- **单元测试**：与代码同文件 `#[cfg(test)] mod tests`，覆盖纯函数和逻辑分支
|
||||
- **集成测试**：`tests/` 目录，覆盖跨模块请求格式
|
||||
- **测试密度**（test 行 / 总行）：tools 98.8%、gateway 96.4%、storage 91.8% 为高覆盖区；
|
||||
providers 31%（anthropic 曾为 0%，已补）、bus 18.7%、cli 4.5% 为薄弱区
|
||||
- **不强制覆盖率工具**：静态审计 + 高风险区定向补测，比全量 tarpaulin 更务实
|
||||
|
||||
## 工程化基线
|
||||
|
||||
- **格式化**：rustfmt（Rust）+ prettier（前端），CI 强制 `--check`
|
||||
- **静态检查**：clippy（Rust）+ eslint（前端），CI 强制
|
||||
- **CI**：GitHub Actions，双平台（ubuntu + windows）跑 fmt + clippy + test + eslint + tsc + vitest
|
||||
- **本地**：`make check` 与 CI 完全对齐
|
||||
4
Cargo.lock
generated
4
Cargo.lock
generated
@ -1635,7 +1635,7 @@ dependencies = [
|
||||
|
||||
[[package]]
|
||||
name = "picobot"
|
||||
version = "0.3.3"
|
||||
version = "0.2.0"
|
||||
dependencies = [
|
||||
"anyhow",
|
||||
"async-trait",
|
||||
@ -1656,7 +1656,6 @@ dependencies = [
|
||||
"libc",
|
||||
"meval",
|
||||
"mime_guess",
|
||||
"parking_lot",
|
||||
"prost",
|
||||
"r2d2",
|
||||
"r2d2_sqlite",
|
||||
@ -1670,7 +1669,6 @@ dependencies = [
|
||||
"serde",
|
||||
"serde_json",
|
||||
"serde_yaml",
|
||||
"subtle",
|
||||
"tempfile",
|
||||
"thiserror 2.0.18",
|
||||
"tokio",
|
||||
|
||||
39
Cargo.toml
39
Cargo.toml
@ -1,25 +1,8 @@
|
||||
[package]
|
||||
name = "picobot"
|
||||
version = "0.3.3"
|
||||
version = "0.2.0"
|
||||
edition = "2024"
|
||||
|
||||
[lints.rust]
|
||||
# 编译期硬错误：避免明显的内存安全/正确性隐患
|
||||
unsafe_op_in_unsafe_fn = "warn"
|
||||
rust_2018_idioms = "warn"
|
||||
|
||||
[lints.clippy]
|
||||
# 渐进式策略：
|
||||
# - 不直接声明 lint group（correctness/suspicious/complexity/perf），
|
||||
# 因为 lint group 在 [lints] 中需用 priority 语法，简单 = "warn" 会报错；
|
||||
# 且 clippy 默认已把 correctness 设为 deny，无需重复声明。
|
||||
# - 仅显式 warn 少量高价值且存量不大的具体规则，避免一上线淹没在噪音中。
|
||||
# 后续随着存量问题清理，可逐步把 unwrap_used/expect_used 升级为 warn。
|
||||
redundant_clone = "warn"
|
||||
dbg_macro = "warn"
|
||||
print_stderr = "warn"
|
||||
print_stdout = "warn"
|
||||
|
||||
[dependencies]
|
||||
reqwest = { version = "0.13.2", default-features = false, features = ["json", "rustls", "multipart", "stream"] }
|
||||
dotenv = "0.15"
|
||||
@ -55,12 +38,18 @@ rusqlite = { version = "0.39", features = ["bundled"] }
|
||||
r2d2 = "0.8"
|
||||
r2d2_sqlite = "0.34"
|
||||
rustls = { version = "0.23", features = ["ring"] }
|
||||
subtle = "2.6"
|
||||
parking_lot = "0.12"
|
||||
wechatbot = { path = "vendor/wechatbot" }
|
||||
encoding_rs = "0.8"
|
||||
libc = "0.2"
|
||||
gray_matter = { version = "0.2", default-features = false, features = ["yaml"] }
|
||||
|
||||
[target.'cfg(windows)'.dependencies]
|
||||
windows-sys = { version = "0.59", features = [
|
||||
"Win32_System_Threading",
|
||||
"Win32_System_Diagnostics_Debug",
|
||||
"Win32_Foundation",
|
||||
"Win32_System_Kernel",
|
||||
] }
|
||||
# MCP (Model Context Protocol) support
|
||||
rmcp = { version = "1.7", features = [
|
||||
"client",
|
||||
@ -70,13 +59,5 @@ rmcp = { version = "1.7", features = [
|
||||
] }
|
||||
schemars = "1.0"
|
||||
http = "1"
|
||||
tower-http = { version = "0.6", features = ["fs", "cors"] }
|
||||
tower-http = { version = "0.6", features = ["fs"] }
|
||||
rust-embed = "8"
|
||||
|
||||
[target.'cfg(windows)'.dependencies]
|
||||
windows-sys = { version = "0.59", features = [
|
||||
"Win32_System_Threading",
|
||||
"Win32_System_Diagnostics_Debug",
|
||||
"Win32_Foundation",
|
||||
"Win32_System_Kernel",
|
||||
] }
|
||||
|
||||
20
Makefile
20
Makefile
@ -1,6 +1,6 @@
|
||||
# PicoBot Web UI Makefile
|
||||
|
||||
.PHONY: dev dev-backend dev-frontend build clean install check fmt fix help
|
||||
.PHONY: dev dev-backend dev-frontend build clean install
|
||||
|
||||
# Default target
|
||||
all: build
|
||||
@ -47,23 +47,11 @@ clean:
|
||||
|
||||
# Check code formatting and linting
|
||||
check:
|
||||
@echo "Checking formatting..."
|
||||
cargo fmt --all -- --check
|
||||
@echo "Checking frontend (lint + format + build)..."
|
||||
cd web && npm run lint
|
||||
cd web && npm run format:check
|
||||
@echo "Checking frontend..."
|
||||
cd web && npm run build
|
||||
@echo "Checking Rust code..."
|
||||
cargo check
|
||||
cargo clippy --all-targets --all-features
|
||||
|
||||
# Format all Rust code in place
|
||||
fmt:
|
||||
cargo fmt --all
|
||||
|
||||
# Auto-fix clippy lints where possible
|
||||
fix:
|
||||
cargo clippy --fix --all-targets --allow-dirty --allow-no-vcs
|
||||
cargo clippy
|
||||
|
||||
# Help
|
||||
help:
|
||||
@ -78,6 +66,4 @@ help:
|
||||
@echo " make run - Run production build"
|
||||
@echo " make clean - Clean build artifacts"
|
||||
@echo " make check - Check code formatting and linting"
|
||||
@echo " make fmt - Format all Rust code in place"
|
||||
@echo " make fix - Auto-fix clippy lints where possible"
|
||||
@echo " make help - Show this help message"
|
||||
|
||||
@ -24,7 +24,7 @@
|
||||
}
|
||||
},
|
||||
"gateway": {
|
||||
"host": "127.0.0.1",
|
||||
"host": "0.0.0.0",
|
||||
"port": 19876,
|
||||
"agent_prompt_reinject_every": 100
|
||||
},
|
||||
|
||||
@ -2,273 +2,6 @@
|
||||
|
||||
本文件记录 Picobot 各版本的显著变更，遵循 [Keep a Changelog](https://keepachangelog.com/zh-CN/1.1.0/) 风格。
|
||||
|
||||
## [0.3.3] - 2026-08-07
|
||||
|
||||
较 [0.3.2] 的 12 个 commit 迭代，聚焦 **前端渲染性能飞跃**、**流式通信健壮性** 与 **代码质量治理** 三大方向。
|
||||
|
||||
### 新增功能
|
||||
|
||||
#### 执行状态对账端点
|
||||
- 新增 `GET /api/executions` HTTP 端点，返回后端权威的执行中话题集合。前端 WebSocket 重连时调用此端点对账，修正断连期间丢失的 `execution_completed` 信号导致的状态漂移。
|
||||
|
||||
### 性能优化
|
||||
|
||||
#### 消息列表虚拟化（长对话不再卡顿）
|
||||
- 引入 `@tanstack/react-virtual` 替换原 `messages.map` 全量渲染，只渲染视口内 + overscan=6 条消息，DOM 节点数恒定。
|
||||
- `measureElement` 动态测量变高度消息（一行文本 vs 50 行代码块），`estimateSize=120` 提供初始估计，测量后自动校正。
|
||||
- Firefox 特殊处理 `measureElement`（`getBoundingClientRect`）。
|
||||
- 适配现有功能：自动滚到底部改用 `virtualizer.scrollToIndex`；`viewKey` 滚动位置记忆保留；`highlightedMessageId` 先渲染目标项再 rAF 加 class；流式 delta 不计数（消息条数不变）。
|
||||
- 性能对比（500 条消息）：DOM 节点 10000+ → ~200；首次渲染数百 ms → <16ms；流式 diff 从全量列表降至仅可见项。
|
||||
|
||||
#### 流式 delta 批处理（避免每 token 全数组拷贝）
|
||||
- 用 `streamingRef` 累加 delta chunks + `requestAnimationFrame` 批量 flush，将每秒数十次 O(n) 数组拷贝合并为一帧一次状态更新。
|
||||
- `useMessages.ts` 新增 `streamingRef` 累加器 + rAF 调度，所有 delta 只 push 到 ref，由 `flushStreaming` 批量落盘。
|
||||
- `useChat.ts` 5 处 `setMessages([])` 改为 `clearMessages()`，切话题/会话/通道时重置流式 ref，避免脏状态残留。
|
||||
- 复杂度：O(n²·L) → O(n) 每帧。
|
||||
|
||||
### 改进
|
||||
|
||||
#### 压缩阈值默认值下调
|
||||
- 工程化压缩阈值 `threshold_ratio` 从 0.7 下调至 0.5（50%），LLM 压缩阈值 `llm_compaction_threshold_ratio` 从 0.5 下调至 0.3（30%）。
|
||||
- 降低触发门槛使压缩更积极介入，减少上下文溢出风险。同步更新相关注释、日志信息和前端 hint 文案。
|
||||
|
||||
### 修复
|
||||
|
||||
#### 输入框/发送按钮状态按话题隔离
|
||||
- 将单一全局 `isLoading` 布尔值重构为按 `topic_id` 跟踪的 `processingTopicIds` 集合，`isLoading` 派生自当前选中话题是否在集合中。
|
||||
- 导航响应不再清空处理状态，切换话题后切回原话题仍能正确禁用输入。
|
||||
- 重连时通过 `/api/executions` 端点对账后端权威执行状态，修正断连期间丢失的 `execution_completed` 信号。
|
||||
|
||||
#### WebSocket 断连时清理流式状态
|
||||
- 修复 WebSocket 在流式输出中途断连时 `streamingRef`（contentChunks/index）不重置，导致重连后新消息 delta 追加到旧消息错误位置。
|
||||
- 修复链路：`useWebSocket` 的 `onDisconnect` 回调 → `useChat` 的 `finishStreaming` → `useMessages` 的 `finishStreaming`（重置 ref + flush pending chunks）。
|
||||
|
||||
#### 会话消息泄露到定时任务视图
|
||||
- `handleSchedulerMessage`（Tier 1 路由）原先无条件吞掉所有可转换的 chat 消息，导致正在执行的实时消息被错误追加到调度器视图。
|
||||
- 修复：调度器视图只接收不带 `topic_id` 的历史加载消息；实时流式消息和带 `topic_id` 的实时消息 fall through 到主视图。
|
||||
|
||||
### 重构
|
||||
|
||||
#### 公共工具函数提取与 DRY
|
||||
- 提取 `src/utils.rs` 公共工具函数模块，消除跨模块重复逻辑。
|
||||
- 修复飞书 channel 正则表达式重复编译问题，改为 `Lazy<Regex>` 一次编译。
|
||||
- 网关路由逻辑 DRY 化，减少 `gateway/mod.rs` 重复代码。
|
||||
- `agent_loop` / `providers` / `config` / `storage` 等多模块统一引用公共函数，20 个文件受影响。
|
||||
- CI 新增安全审计步骤（`cargo audit`）。
|
||||
|
||||
### 内部改进
|
||||
- 存储模块应用 `rustfmt` 格式化（`migrations.rs` / `row_mapping.rs` / `tests.rs`）。
|
||||
- 降低 `list_todos` handler 与 gateway 转发日志级别为 debug，减少子代理多次更新待办时的 INFO 日志噪音。
|
||||
- 删除临时优化计划文件 `OPTIMIZATION_PLAN.md`。
|
||||
|
||||
### 测试
|
||||
- 新增 MessageList 虚拟化对抗性 smoke test（5 个用例）：空消息渲染、100 条消息不崩溃、highlight 视口外不崩溃、viewKey 切换不崩溃、流式 rerender 不崩溃。
|
||||
- 新增流式断连对抗性测试（2 个用例）：`stream_end` 幂等性验证、多周期断连-重连循环下 `finishStreaming` 幂等验证。
|
||||
- 新增断连后新 delta 不污染旧消息测试（用例 17）。
|
||||
- 新增定时任务视图消息隔离测试（用例 16）。
|
||||
- 新增批量 delta 累加与 `stream_end` 后 `assistant_response` 替换对抗性测试。
|
||||
|
||||
---
|
||||
|
||||
## [0.3.2] - 2026-08-06
|
||||
|
||||
较 [0.3.1] 的 7 个 commit 迭代，聚焦 **上下文压缩可配置化**、**网关安全加固**、**并发模型优化** 与 **错误防御性提升** 四大方向。
|
||||
|
||||
### 新增功能
|
||||
|
||||
#### 压缩算法关键参数暴露到设置页面
|
||||
- 将 4 个上下文压缩参数从硬编码提取到 `config.json` 顶层 `compaction` 节：
|
||||
- `threshold_ratio`（0.7）：工程化压缩触发阈值，占 `context_window` 的比例
|
||||
- `llm_compaction_threshold_ratio`（0.5）：LLM 压缩触发阈值
|
||||
- `truncate_max_tokens`（100）：工程化压缩时 tool 结果截断到的 token 数
|
||||
- `preserve_count`（5）：LLM 三段压缩保留最旧/最新的 unit 数
|
||||
- 前端新增"上下文压缩"设置标签页（Archive 图标），4 个 number input 分两组展示。
|
||||
- `ContextCompressor` 内聚所有压缩参数，新增 `with_compaction_config()` 构造函数和 `truncate_tool_results()` 方法。
|
||||
- `AgentLoop` 通过调用 `compressor.truncate_tool_results()` 实现零参数耦合，压缩参数变更仅影响 `ContextCompressor`。
|
||||
|
||||
#### HTTP API 认证与 CORS 防护
|
||||
- 新增 HTTP API 认证机制，未授权请求直接返回 401。
|
||||
- CORS 防护：明确指定允许的 origin，阻止跨域未授权访问。
|
||||
- 修复密钥泄露风险：敏感信息不再出现在错误响应中。
|
||||
|
||||
### 改进
|
||||
|
||||
#### 上下文压缩重构为两阶段+双阈值+三段保留
|
||||
- **两阶段压缩**：工具调用完成后触发压缩流程，先用 `prompt_tokens` 判断 70% 阈值。
|
||||
- **阶段 1（工程化压缩）**：截断非子代理 tool 结果到约 100 token，仅改内存数据。
|
||||
- **阶段 2（LLM 压缩）**：工程化压缩后若估算 token 仍超 50% 阈值，才调用 LLM 压缩。
|
||||
- **三段保留**：保留最旧 5 个 unit + 中间段 LLM 摘要 + 最新 5 个 unit，通过 `HistoryUnit` 原子单元保证 `ToolRound` 不被拆分。
|
||||
- LLM 压缩失败时降级为仅工程化压缩，不中断 agent loop。
|
||||
|
||||
#### 迁移 parking_lot 锁并优化阻塞 IO 与内存管理
|
||||
- 将 `std::sync` 锁替换为 `parking_lot` 实现，减少锁竞争开销。
|
||||
- 优化阻塞 IO 路径，避免热路径无谓阻塞。
|
||||
- 内存管理改进：减少不必要的克隆与分配。
|
||||
|
||||
### 修复
|
||||
|
||||
#### 子智能体任务状态卡死
|
||||
- 修复已完成子智能体任务卡在"子智能体正在执行..."状态，前端无法正确展示结果。
|
||||
|
||||
#### 飞书 channel 消息丢失
|
||||
- 修正飞书 reaction 语义，与飞书 API 规范对齐。
|
||||
- 增容 message bus 队列，避免高频消息场景下队列满导致最后一条消息丢失。
|
||||
|
||||
#### 防御性错误处理
|
||||
- 替换 `unreachable!()` 宏调用为防御性错误返回，避免运行时 panic。
|
||||
- agent/gateway 关键路径改用 `Result` 传播错误，提升系统鲁棒性。
|
||||
|
||||
---
|
||||
|
||||
## [0.3.1] - 2026-08-05
|
||||
|
||||
较 [0.3.0] 的 19 个 commit 迭代，聚焦 **历史压缩可追溯性**、**token 统计准确性**、**前端渲染稳定性**、**工具执行可靠性** 与 **网关吞吐优化** 五大方向。
|
||||
|
||||
### 新增功能
|
||||
|
||||
#### 历史压缩保留原消息（双视图分离）
|
||||
- 压缩流程从物理删除原消息改为标记位 `is_compacted` 分离两个视图：
|
||||
- **LLM 视角**（`load_messages_for_topic`）：只看 `is_compacted=0` 的压缩摘要 + 新消息，context 仍精简。
|
||||
- **UI 视角**（`load_messages_for_topic_full`）：看原消息（含 `is_compacted=1`）+ 新消息，排除压缩摘要，前端切回老话题仍能展示完整原始对话。
|
||||
- 新增 `compact_topic_history`：不删除原消息，仅打标记 + 插入摘要，同事务内删除旧摘要避免累积。修复 token 统计因压缩丢失 `usage` 的问题。
|
||||
- `messages` 表新增 `is_compacted INTEGER NOT NULL DEFAULT 0` 列，migration 幂等。
|
||||
|
||||
#### Topic 维度 token 统计
|
||||
- 新增 `batch_topic_token_stats`，按 `topic_id` 聚合 token 消耗，替代原按 `session_id` 聚合导致的同 session 多 topic 共享总和问题。
|
||||
- 前端侧边栏每个 topic 显示独立 token 消耗与 context 占用百分比（Coins 图标 + 数值，绿/黄/红三色指示）。
|
||||
- 主代理执行结束时触发 `list_topics` 刷新（500ms 防抖），实时更新侧边栏 token 统计；子代理每轮执行也支持刷新。
|
||||
|
||||
#### 设置页重启按钮
|
||||
- 设置页 Footer 区新增重启按钮，样式为 `bg-tertiary` + 边框，与主保存按钮区分。
|
||||
- 确认对话框区分触发来源：保存后触发显示"配置已保存 / 稍后"，手动触发显示"重启服务 / 取消"并警告未保存改动。
|
||||
- 状态管理改用 `restartDialogMode: 'saved' | 'manual' | null`，按钮在重启期间显示 loading 并禁用。
|
||||
|
||||
#### 设置页懒加载标签页与子代理管理
|
||||
- 拆分设置页为懒加载标签页，减少首屏渲染负担。
|
||||
- 前端支持子代理创建与删除，配合 `PUT /api/subagents/update` 端点。
|
||||
|
||||
#### LLM 主调用可配置重试机制（0.3.0 收尾补充）
|
||||
- `ProviderConfig` 增加 `max_retries` 字段（serde default=3，向后兼容）。
|
||||
- `agent_loop` 流式 + summary 两个调用点实现重试循环：指数退避 1s/2s/4s，仅对 429/502/503/504/timeout/connection reset 重试；流式仅在未 emit delta 时重试，避免重复输出。
|
||||
|
||||
### 修复
|
||||
|
||||
#### 前端渲染稳定性
|
||||
- `MessageBubble` 使用 `React.memo` 避免流式期间不必要重渲染。
|
||||
- `useWebSocket` 回调 ref-化，避免父组件重渲染触发 WebSocket 无谓重连。
|
||||
- 修复触发滚动到底部按钮时新消息计数错误。
|
||||
- 修复话题重命名预填字段与写入字段不一致。
|
||||
- 修复流式期间消息重渲染、WebSocket 无谓重连与新消息计数错误。
|
||||
|
||||
#### 子智能体结果展示
|
||||
- 失败/超时子智能体返回结构化结果，前端可点击查看详情，而非静默丢失。
|
||||
|
||||
#### Token 统计准确性
|
||||
- 修复 token 统计按 `session_id` 聚合导致同 session 多 topic 共享总和：改为按 `topic_id` 聚合，子代理消息通过 `session_id NOT LIKE 'sub:%'` 排除。
|
||||
- 修复主代理执行结束后 token 统计不刷新：`useMessages` 在 `execution_completed` 事件时触发 `bumpTopicRefreshTrigger()`。
|
||||
|
||||
#### 工具执行可靠性
|
||||
- 修复 shell 工具超时未生效：`adb start-server` 守护进程继承 stdout 管道导致 EOF 阻塞，`adb devices` 误判 socket 等待为 stdin 等待绕过 deadline。新增 `STREAM_DRAIN_MS` 常量、改进 `kill_and_reap`、移除 `should_return_pending` 中的错误检查，确保 timeout 参数完全覆盖默认值。
|
||||
|
||||
#### 安全加固
|
||||
- `web_fetch` 加固 SSRF 防护，禁用 HTTP 重定向，与 `http_request` 实现对齐。
|
||||
|
||||
#### MCP 兼容性
|
||||
- 限制 MCP server 名称字符并兜底 `tool_name` 清洗，避免非法字符导致 OpenAI API 400。
|
||||
|
||||
### 性能优化
|
||||
|
||||
#### 网关吞吐
|
||||
- `send_with_retry` 对 `ChannelFull` 错误短路（不重试、不退避），加速队列消费，避免 dispatcher 在满队列时 7s 退避阻塞。
|
||||
- 新增 `ChannelError::ChannelFull` 变体，`CliChannel` 在 `try_send` 返回 `Full` 时映射为 `ChannelFull`。
|
||||
|
||||
#### 存储查询
|
||||
- `get_topic_message_count` 改用 `SELECT COUNT(*)` 在数据库侧计数，避免将所有消息（含 `content`、`tool_calls_json` 等大字段）加载到内存。
|
||||
|
||||
#### 热路径优化
|
||||
- 减少热路径无谓拷贝与重复正则编译，跳过已完成 migration。
|
||||
|
||||
### 内部改进
|
||||
- 删除临时计划文件 `PLAN.md`。
|
||||
- CHANGELOG 补充 LLM 重试机制条目并修正 commit 计数。
|
||||
|
||||
---
|
||||
|
||||
## [0.3.0] - 2026-08-04
|
||||
|
||||
较 [0.2.0] 的 18 个 commit 迭代，聚焦 **Agent 执行与显示层解耦**、**LLM 调用稳定性**、**工程化基线**、**并发持久化稳定性**、**安全加固** 与 **MCP 兼容性** 六大方向。锁屏冻结架构修复经五轮对抗性审查验证。
|
||||
|
||||
### 新增功能
|
||||
|
||||
#### 话题重命名
|
||||
- 新增 `RenameTopic` 命令与 `TopicRenamed` 协议消息，复用存储层已有的 `update_topic_title` 方法。后端响应携带刷新后的完整 topic 列表，前端零额外往返同步侧边栏。
|
||||
- 前端 `TopicList` 侧边栏增加内联编辑入口：悬停显示铅笔图标，Enter 提交 / Esc 取消 / blur 取消，通过 `onMouseDown preventDefault` 防止按钮点击时 input 提前失焦。
|
||||
|
||||
#### 前端自动构建集成
|
||||
- 新增 `build.rs`，`cargo build` 时自动执行 `npm install + npm run build`，消除前后端构建割裂。支持 `SKIP_FRONTEND_BUILD` 环境变量跳过。
|
||||
|
||||
#### 工程化基线
|
||||
- `rustfmt.toml` 固化 `max_width=100` / 4 空格缩进；`Cargo.toml` 配置 `[lints.rust]` 与 `[lints.clippy]` 渐进式规则。
|
||||
- `.github/workflows/ci.yml`：Rust（fmt + clippy + test）+ 前端（eslint + tsc + test + prettier format:check）双平台 CI。
|
||||
- `Makefile` 新增 `check` / `fmt` / `fix` 目标，clippy 对齐 `--all-targets --all-features`。
|
||||
- 前端 eslint flat config + prettier 配置，对 47 个前端文件统一格式化并在 CI 强制 `format:check`。
|
||||
|
||||
#### 默认迭代上限调整
|
||||
- `max_tool_iterations` 默认值调整为 1000，匹配长任务子代理的实际需求。
|
||||
|
||||
#### LLM 主调用可配置重试机制
|
||||
- `ProviderConfig` 增加 `max_retries` 字段（serde default=3，向后兼容），透传至 `AgentRuntimeConfig` 归属 agent 行为层，不进 `ProviderRuntimeConfig` 以保持 provider 构造包纯净。
|
||||
- `agent_loop` 流式 + summary 两个调用点实现重试循环：指数退避 1s/2s/4s，仅对 429/502/503/504/timeout/connection reset 重试；流式仅在未 emit delta 时重试（`AtomicBool` 跟踪），避免重复输出；退避 sleep 期间响应 `cancel_signal`，取消优先。
|
||||
- 前端 `ProviderConfig` 类型和表单增加 `max_retries` 字段。
|
||||
- 7 个单元测试（3 判定 + 4 行为），540 个 lib 测试全绿。
|
||||
|
||||
### 架构修复
|
||||
|
||||
#### Agent 执行与显示层解耦（锁屏冻结根因修复）
|
||||
根因：浏览器锁屏导致 WebSocket 半死，`ws_sender.send().await` 永久阻塞，级联阻塞 dispatcher → MessageBus → Agent Loop，后端停止执行直到解锁。
|
||||
|
||||
基于第一性原理建立"执行-显示解耦"原则：agent 执行只依赖 SQLite 持久化，实时广播是可丢弃的最佳努力通道。
|
||||
|
||||
- `MessageBus::publish_outbound` 由 `send().await` 改为 `try_send()`：bus 满时丢弃消息并告警，agent 不再被显示层阻塞。新增 `BusError::Dropped` 变体。
|
||||
- WebSocket writer task 包裹 `tokio::time::timeout(30s)`：半死连接超时即关闭。使用每连接独立的 `CancellationToken`（非共享的 CliChannel 级 token），避免一个连接超时关闭所有连接。writer 退出时 cancel 通知主 loop 退出，确保 `unregister_connection` 执行。
|
||||
- `CliChannel::send` 由 `send().await` 改为 `try_send()`：dispatcher 单线程顺序处理，原阻塞式发送在 writer 卡住时会阻塞所有连接 37s，改后立即返回。
|
||||
- 全仓 13 处 `publish_outbound` 调用统一区分 `Dropped`（warn，预期背压）/ `Closed`（error，异常）日志级别，消除监控噪音。scheduler 的 3 处 `?` 改为 warn，避免 `Dropped` 触发 misfire 重试风暴。
|
||||
- 前端 WebSocket 添加 25s 客户端 ping + 指数退避重连（3s→6s→12s→24s→60s 封顶，上限 999 次）。重连后区分"重连恢复"与"首次/切换通道"：前者保留断连前 messages 并刷新 topic 列表，后者清空数据避免污染。
|
||||
|
||||
### 安全加固
|
||||
|
||||
#### SSRF 重定向绕过修复
|
||||
- `web_fetch` 工具禁用 HTTP 重定向（`redirect::Policy::none()`）：原 `validate_url` 只校验初始 URL 的 host，跟随 302 跳转可重定向到 `169.254.169.254`（云元数据）或 `127.0.0.1` 等内网地址，绕过 `is_private_host` 的 SSRF 防护。与 `http_request` 工具保持一致。
|
||||
|
||||
#### Safety Guard 正则收紧
|
||||
- `format` 正则收紧为 `\bformat\s+.*[a-z]:`，要求出现盘符才拦截，避免误伤 `dart format`、`buf format`、`pytest --format` 等合法命令。
|
||||
- 按平台分组注入规则，Unix 不再注入 Windows 专用规则；`Remove-Item` 正则改为小写，与 `guard_command` 的大小写处理一致。
|
||||
|
||||
### 稳定性修复
|
||||
|
||||
#### SQLite 并发写入死锁
|
||||
- 7 个写事务从 `BEGIN DEFERRED` 改为 `BEGIN IMMEDIATE`：在事务开始即获取写锁，消除多 sub-agent 并发写入时的死锁路径。`busy_timeout` 从 5s 提升至 30s。
|
||||
- 根因：`BEGIN DEFERRED` 下多个事务可同时读 `MAX(seq)` 不持写锁，提交时互相阻塞，5s timeout 耗尽后返回 `SQLITE_BUSY`。`BEGIN IMMEDIATE` 强制写者串行排队，顺带消除 `MAX(seq)+1` 竞态导致的 UNIQUE 约束冲突。
|
||||
|
||||
#### 测试与代码质量
|
||||
- 修复 `anthropic` 错误链嵌套测试为真正的 `#[source]` 链（原测试是无效的，inner 变量被 `let _` 抑制）。
|
||||
- 补充 `anthropic` provider 15 个纯函数单测（原 396 行零测试），覆盖 data URL 解析、图片过滤、字段过滤、响应反序列化、错误链格式化。
|
||||
- 修复 3 个失败测试：`agent_md_template` 写入模板内容 / `StreamingAccumulator` BTreeMap 保序 / `source_order` 断言补全。
|
||||
- 清理未使用的 dead code 函数与方法。
|
||||
- 新增 `ARCHITECTURE.md` 架构文档，聚焦数据流与 7 个关键设计决策。
|
||||
|
||||
#### CI 覆盖范围扩大
|
||||
- `cargo build --lib` 改为 `cargo build`，确保 `main.rs` 二进制入口被编译验证。
|
||||
- `cargo test --lib` 改为 `cargo test`，纳入 `tests/` 目录集成测试。`test_request_format.rs` 的 8 个序列化测试此前从未在 CI 中运行。
|
||||
|
||||
#### 跨平台构建修复
|
||||
- 修正依赖分类错误：`rmcp` / `schemars` / `http` / `tower-http` / `rust-embed` 均为跨平台 crate，错放在 `[target.'cfg(windows)'.dependencies]` 下导致 Linux 构建失败，移出 target 段。
|
||||
|
||||
#### MCP 工具名清洗
|
||||
- OpenAI 要求 function name 匹配 `^[a-zA-Z0-9_-]+$`，否则整个请求 400。MCP 工具名由 `mcp_{server_key}_{tool_name}` 拼成，两个输入源分别处理：
|
||||
- 后端：`sanitize_tool_name` 替换 `tool_name` 中的非法字符（`.`, `:`, `/` 等）为 `_`，保留 `server_key` 和 `tool_name` 原值用于路由，仅清洗 LLM 可见的 `full_name`，发生清洗时打 warn 日志便于定位。
|
||||
- 前端：MCP 卡片头部改用 `MapEntryHeader` 支持点击重命名，`addMcp` / `renameMcp` 正则校验 server 名称 + toast 提示。顺手修了 `MapEntryHeader` 进入编辑时 `val` 未同步当前 `name` 的 bug。
|
||||
|
||||
## [0.2.0] - 2026-07-31
|
||||
|
||||
较 [0.1.2] 的 47 个 commit 迭代，聚焦 **能力策略**、**模型独立配置**、**话题级并发隔离** 与 **Agent Loop 性能优化** 四大方向。
|
||||
@ -380,9 +113,5 @@
|
||||
- 前端静态文件嵌入二进制。
|
||||
- React Web UI 前端界面。
|
||||
|
||||
[0.3.3]: https://github.com/picobot/picobot/compare/v0.3.2...v0.3.3
|
||||
[0.3.2]: https://github.com/picobot/picobot/compare/v0.3.1...v0.3.2
|
||||
[0.3.1]: https://github.com/picobot/picobot/compare/v0.3.0...v0.3.1
|
||||
[0.3.0]: https://github.com/picobot/picobot/compare/v0.2.0...v0.3.0
|
||||
[0.2.0]: https://github.com/picobot/picobot/compare/v0.1.2...v0.2.0
|
||||
[0.1.2]: https://github.com/picobot/picobot/releases/tag/v0.1.2
|
||||
|
||||
12
rustfmt.toml
12
rustfmt.toml
@ -1,12 +0,0 @@
|
||||
# PicoBot Rust 代码格式化规则
|
||||
#
|
||||
# 设计原则：尽量贴近 rustfmt 默认风格，仅固化少数项目级偏好。
|
||||
# 不追求激进重排，避免一次性产生大量 diff。
|
||||
# 仅使用 stable rustfmt 支持的选项，不依赖 nightly 特性。
|
||||
|
||||
# 行宽：100，现代显示器友好
|
||||
max_width = 100
|
||||
|
||||
# 缩进用 4 空格（Rust 社区主流，与现有代码一致）
|
||||
hard_tabs = false
|
||||
tab_spaces = 4
|
||||
File diff suppressed because it is too large
Load Diff
@ -1,11 +1,11 @@
|
||||
use crate::agent::{AgentError, AgentRuntimeConfig};
|
||||
use crate::bus::{
|
||||
ChatMessage, SYSTEM_CONTEXT_AGENT_PROMPT, SYSTEM_CONTEXT_HISTORY_COMPACTION,
|
||||
SYSTEM_CONTEXT_SCHEDULED_PROMPT,
|
||||
};
|
||||
use crate::config::{CompactionConfig, LLMProviderConfig};
|
||||
use crate::config::LLMProviderConfig;
|
||||
use crate::providers::{ChatCompletionRequest, LLMProvider, Message, create_provider};
|
||||
use crate::text::{char_count, take_prefix_chars};
|
||||
use crate::agent::{AgentError, AgentRuntimeConfig};
|
||||
|
||||
const TOKEN_ESTIMATE_SAFETY_MULTIPLIER: f64 = 1.2;
|
||||
const CJK_CHARS_PER_TOKEN: f64 = 2.0;
|
||||
@ -17,6 +17,13 @@ pub const SYSTEM_CONTEXT_HISTORY_COMPACTION_OLDER: &str = "history_compaction_ol
|
||||
/// System context marker for the light compression (newer segment) summary.
|
||||
pub const SYSTEM_CONTEXT_HISTORY_COMPACTION_NEWER: &str = "history_compaction_newer";
|
||||
|
||||
/// Default threshold ratio: compress when estimated tokens exceed 70% of context window.
|
||||
const DEFAULT_THRESHOLD_RATIO: f64 = 0.7;
|
||||
|
||||
/// Default budget split: older segment gets 30% of summary budget, newer gets 70%.
|
||||
const OLDER_BUDGET_RATIO: f64 = 0.3;
|
||||
const NEWER_BUDGET_RATIO: f64 = 0.7;
|
||||
|
||||
// ============================================================================
|
||||
// HistoryUnit — atomic message units for compression
|
||||
// ============================================================================
|
||||
@ -39,6 +46,22 @@ enum HistoryUnit {
|
||||
AssistantText(ChatMessage),
|
||||
}
|
||||
|
||||
impl HistoryUnit {
|
||||
/// Estimate tokens for this unit alone.
|
||||
fn estimate_tokens(&self) -> usize {
|
||||
match self {
|
||||
HistoryUnit::SystemGuard(msg) | HistoryUnit::UserMessage(msg) | HistoryUnit::AssistantText(msg) => {
|
||||
estimate_tokens(std::slice::from_ref(msg))
|
||||
}
|
||||
HistoryUnit::ToolRound { assistant, results } => {
|
||||
let mut all = vec![assistant.clone()];
|
||||
all.extend(results.clone());
|
||||
estimate_tokens(&all)
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// ============================================================================
|
||||
// Unit parser — one forward pass, O(n)
|
||||
// ============================================================================
|
||||
@ -102,19 +125,9 @@ fn parse_to_units(messages: &[ChatMessage]) -> Vec<HistoryUnit> {
|
||||
units
|
||||
}
|
||||
|
||||
/// Flatten a HistoryUnit back into its constituent ChatMessage(s).
|
||||
/// ToolRound yields [assistant, results...]; others yield a single message.
|
||||
fn unit_to_messages(unit: &HistoryUnit) -> Vec<ChatMessage> {
|
||||
match unit {
|
||||
HistoryUnit::ToolRound { assistant, results } => {
|
||||
let mut msgs = vec![assistant.clone()];
|
||||
msgs.extend(results.clone());
|
||||
msgs
|
||||
}
|
||||
HistoryUnit::AssistantText(msg)
|
||||
| HistoryUnit::UserMessage(msg)
|
||||
| HistoryUnit::SystemGuard(msg) => vec![msg.clone()],
|
||||
}
|
||||
/// Estimate total tokens from a slice of units.
|
||||
fn estimate_tokens_from_units(units: &[HistoryUnit]) -> usize {
|
||||
units.iter().map(|u| u.estimate_tokens()).sum()
|
||||
}
|
||||
|
||||
/// Check if a character is CJK (Chinese, Japanese, Korean)
|
||||
@ -149,8 +162,8 @@ pub fn estimate_tokens(messages: &[ChatMessage]) -> usize {
|
||||
}
|
||||
|
||||
// Weighted token calculation: CJK chars need more tokens per character
|
||||
let content_tokens =
|
||||
(cjk_count as f64 / CJK_CHARS_PER_TOKEN) + (other_count as f64 / OTHER_CHARS_PER_TOKEN);
|
||||
let content_tokens = (cjk_count as f64 / CJK_CHARS_PER_TOKEN)
|
||||
+ (other_count as f64 / OTHER_CHARS_PER_TOKEN);
|
||||
|
||||
// JSON serialization overhead for message structure (fields, brackets, etc.)
|
||||
let json_overhead = messages.len() * JSON_OVERHEAD_PER_MESSAGE;
|
||||
@ -163,72 +176,6 @@ pub fn estimate_tokens(messages: &[ChatMessage]) -> usize {
|
||||
* TOKEN_ESTIMATE_SAFETY_MULTIPLIER) as usize
|
||||
}
|
||||
|
||||
/// 估算纯文本的 token 数（不含消息级 JSON 开销，仅内容）。
|
||||
fn estimate_text_tokens(text: &str) -> usize {
|
||||
let mut cjk = 0usize;
|
||||
let mut other = 0usize;
|
||||
for ch in text.chars() {
|
||||
if is_cjk_char(ch) {
|
||||
cjk += 1;
|
||||
} else {
|
||||
other += 1;
|
||||
}
|
||||
}
|
||||
let content_tokens =
|
||||
(cjk as f64 / CJK_CHARS_PER_TOKEN) + (other as f64 / OTHER_CHARS_PER_TOKEN);
|
||||
(content_tokens * TOKEN_ESTIMATE_SAFETY_MULTIPLIER) as usize
|
||||
}
|
||||
|
||||
/// 将文本截断到约 max_tokens，确保 UTF-8 字符边界安全。
|
||||
/// 截断后追加 "\n...[已截断]" 标记。
|
||||
fn truncate_to_token_limit(text: &str, max_tokens: usize) -> String {
|
||||
let mut token_count = 0.0f64;
|
||||
let mut cutoff_byte = text.len();
|
||||
|
||||
for (byte_idx, ch) in text.char_indices() {
|
||||
let char_tokens = if is_cjk_char(ch) { 0.5 } else { 0.25 };
|
||||
token_count += char_tokens;
|
||||
if token_count >= max_tokens as f64 {
|
||||
cutoff_byte = byte_idx;
|
||||
break;
|
||||
}
|
||||
}
|
||||
|
||||
if cutoff_byte >= text.len() {
|
||||
return text.to_string();
|
||||
}
|
||||
|
||||
let mut result = text[..cutoff_byte].to_string();
|
||||
result.push_str("\n...[已截断]");
|
||||
result
|
||||
}
|
||||
|
||||
/// 工程化压缩：将非子代理的 tool 结果截断到约 max_tokens。
|
||||
///
|
||||
/// 规则：
|
||||
/// - role="tool" 且 tool_name != "task" 的消息，content 截断到约 max_tokens
|
||||
/// - role="tool" 且 tool_name == "task"（子代理返回）的消息，保持原样
|
||||
/// - 其他 role 的消息不受影响
|
||||
/// - 截断时用 char_indices 确保 UTF-8 字符边界安全
|
||||
/// - 截断后追加 "\n...[已截断]" 标记
|
||||
///
|
||||
/// **只修改传入的 messages，不涉及 DB 操作。**
|
||||
pub fn truncate_tool_results_in_place(messages: &mut [ChatMessage], max_tokens: usize) {
|
||||
for msg in messages.iter_mut() {
|
||||
if msg.role != "tool" {
|
||||
continue;
|
||||
}
|
||||
if msg.tool_name.as_deref() == Some("task") {
|
||||
continue;
|
||||
}
|
||||
let estimated = estimate_text_tokens(&msg.content);
|
||||
if estimated <= max_tokens {
|
||||
continue;
|
||||
}
|
||||
msg.content = truncate_to_token_limit(&msg.content, max_tokens);
|
||||
}
|
||||
}
|
||||
|
||||
/// Configuration for context compression.
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct ContextCompressionConfig {
|
||||
@ -267,14 +214,8 @@ impl Default for ContextCompressionConfig {
|
||||
pub struct ContextCompressor {
|
||||
config: ContextCompressionConfig,
|
||||
context_window: usize,
|
||||
/// Threshold ratio to trigger compression (50% of context window).
|
||||
/// Threshold ratio to trigger compression (70% of context window).
|
||||
threshold_ratio: f64,
|
||||
/// LLM 压缩阈值比例（工程化压缩后仍超此比例才调 LLM）
|
||||
llm_compaction_threshold_ratio: f64,
|
||||
/// 三段压缩保留 unit 数（最旧 N + 最新 N）
|
||||
preserve_count: usize,
|
||||
/// 工程化压缩时 tool 结果截断 token 数
|
||||
truncate_max_tokens: usize,
|
||||
}
|
||||
|
||||
impl ContextCompressor {
|
||||
@ -436,6 +377,44 @@ OLDER SEGMENT (events from earlier in the session):
|
||||
|
||||
{}
|
||||
|
||||
"#,
|
||||
target_chars, transcript
|
||||
)
|
||||
}
|
||||
|
||||
/// Prompt for the newer segment — lighter compression, keep more detail.
|
||||
fn build_light_summary_prompt(transcript: &str, target_chars: usize) -> String {
|
||||
format!(
|
||||
r#"You are a conversation compaction engine. Lightly summarize the following RECENT conversation segment. These events happened just before the current moment and will be marked as "近期操作记录（轻度压缩，保留了更多细节）".
|
||||
|
||||
An older segment (already summarized separately) precedes this. Keep enough detail so the model can continue the task seamlessly without re-reading files for data that appears in these recent results.
|
||||
|
||||
=== MUST PRESERVE (keep with more detail than a normal summary) ===
|
||||
- All file paths, URLs, and identifiers
|
||||
- The exact sequence of recent operations (step by step)
|
||||
- Tool parameters (especially file paths, search queries, command strings)
|
||||
- Key outputs and results (shortened but keep the substance)
|
||||
- Error messages (complete, not truncated)
|
||||
- Current task status and what should happen next
|
||||
|
||||
=== SHOULD CONDENSE ===
|
||||
- Tool outputs → truncate to the most meaningful parts (key data, conclusions, not raw output)
|
||||
- Long text → keep the essence but not every word
|
||||
- Repeated similar outputs → note the pattern but keep a representative example
|
||||
|
||||
=== SHOULD DROP ===
|
||||
- Completely irrelevant debug output
|
||||
- Boilerplate text with no informational value
|
||||
- Trivial operations that have no bearing on the task
|
||||
|
||||
Be concise, aim for {} characters or less. Output the summary in Chinese if the original conversation was in Chinese.
|
||||
|
||||
---
|
||||
|
||||
RECENT SEGMENT (events from just before the current moment):
|
||||
|
||||
{}
|
||||
|
||||
"#,
|
||||
target_chars, transcript
|
||||
)
|
||||
@ -518,9 +497,33 @@ OLDER SEGMENT (events from earlier in the session):
|
||||
}
|
||||
|
||||
// =========================================================================
|
||||
// Three-segment compression
|
||||
// Two-segment compression
|
||||
// =========================================================================
|
||||
|
||||
/// Find a safe split point in the unit list. The split ensures:
|
||||
/// - Accumulated tokens up to `ratio` of total are placed in the older segment.
|
||||
/// - The split never lands on a SystemGuard (they're always preserved).
|
||||
/// - The split lands on a unit boundary (ToolRound, AssistantText, or UserMessage).
|
||||
fn find_safe_split_point(&self, units: &[HistoryUnit], ratio: f64) -> usize {
|
||||
let total_tokens = estimate_tokens_from_units(units);
|
||||
let target = (total_tokens as f64 * ratio) as usize;
|
||||
let mut accumulated = 0;
|
||||
|
||||
for (i, unit) in units.iter().enumerate() {
|
||||
// Never split on a SystemGuard — they must be preserved
|
||||
if matches!(unit, HistoryUnit::SystemGuard(_)) {
|
||||
continue;
|
||||
}
|
||||
accumulated += unit.estimate_tokens();
|
||||
if accumulated >= target {
|
||||
return i + 1;
|
||||
}
|
||||
}
|
||||
|
||||
// Fallback: put everything in the older segment
|
||||
units.len()
|
||||
}
|
||||
|
||||
/// Summarize a transcript with a custom prompt builder function.
|
||||
async fn summarize_with_prompt(
|
||||
&self,
|
||||
@ -611,165 +614,206 @@ OLDER SEGMENT (events from earlier in the session):
|
||||
Ok(take_prefix_chars(transcript, target))
|
||||
}
|
||||
|
||||
/// Main entry point for three-segment compression.
|
||||
/// Main entry point for two-segment compression.
|
||||
///
|
||||
/// Preserves the oldest 5 units and newest 5 units in full, then
|
||||
/// summarizes the middle segment with LLM. The result maintains
|
||||
/// ToolRound atomicity and the middle summary is a pure system message
|
||||
/// with no tool_calls — eliminating any risk of API 400 errors
|
||||
/// Splits history into older and newer segments, then summarizes each
|
||||
/// with LLM using different prompts and budget allocations. The result
|
||||
/// contains no tool_calls or tool-result messages — only system summaries
|
||||
/// and the original user message — eliminating any risk of API 400 errors
|
||||
/// from orphaned tool call sequences.
|
||||
pub async fn compress_two_segment(
|
||||
&self,
|
||||
history: &[ChatMessage],
|
||||
provider_config: &LLMProviderConfig,
|
||||
) -> Result<Vec<ChatMessage>, AgentError> {
|
||||
let mut truncated: Vec<ChatMessage> = history.to_vec();
|
||||
truncate_tool_results_in_place(&mut truncated, 100);
|
||||
|
||||
let tokens = estimate_tokens(&truncated);
|
||||
if tokens <= self.threshold() {
|
||||
return Ok(truncated);
|
||||
}
|
||||
|
||||
let provider = create_provider(AgentRuntimeConfig::from(provider_config.clone()).provider)
|
||||
.map_err(|e| AgentError::ProviderCreation(e.to_string()))?;
|
||||
self.compress_two_segment_inner(&truncated, provider.as_ref())
|
||||
.await
|
||||
}
|
||||
|
||||
/// 复用调用方已有的 provider 实例（用于 AgentLoop 内，避免重复创建 provider）。
|
||||
///
|
||||
/// **注意**：此方法假设调用方已对 `history` 做过工程化压缩
|
||||
/// （`truncate_tool_results_in_place`），因此内部不再重复截断。
|
||||
/// 这样可避免二次截断导致 `\n...[已截断]` 标记累积和内容损失。
|
||||
pub async fn compress_two_segment_with_provider(
|
||||
&self,
|
||||
history: &[ChatMessage],
|
||||
provider: &dyn LLMProvider,
|
||||
) -> Result<Vec<ChatMessage>, AgentError> {
|
||||
let tokens = estimate_tokens(history);
|
||||
if tokens <= self.threshold() {
|
||||
tracing::info!(
|
||||
tokens = tokens,
|
||||
threshold = self.threshold(),
|
||||
msg_count = history.len(),
|
||||
"Two-segment compression not needed (under threshold)"
|
||||
);
|
||||
return Ok(history.to_vec());
|
||||
}
|
||||
|
||||
self.compress_two_segment_inner(history, provider).await
|
||||
}
|
||||
|
||||
/// 三段压缩核心逻辑：保留最旧 N + 最新 N unit，中间段用 LLM 压缩。
|
||||
///
|
||||
/// 策略：
|
||||
/// - SystemGuard 永远保留在头部（不计入 N 条配额）
|
||||
/// - 可压缩单元（UserMessage / AssistantText / ToolRound）按时间顺序：
|
||||
/// - 最旧 preserve_count 个 unit 原样保留
|
||||
/// - 最新 preserve_count 个 unit 原样保留
|
||||
/// - 中间段用 LLM 生成摘要（system 消息，无 tool_calls）
|
||||
/// - ToolRound 原子性由 parse_to_units 保证，切分在 unit 边界
|
||||
/// - 中间段摘要为纯文本 system 消息，符合 API 提交要求
|
||||
async fn compress_two_segment_inner(
|
||||
&self,
|
||||
history: &[ChatMessage],
|
||||
provider: &dyn LLMProvider,
|
||||
) -> Result<Vec<ChatMessage>, AgentError> {
|
||||
let preserve_count = self.preserve_count;
|
||||
let tokens = estimate_tokens(history);
|
||||
tracing::info!(
|
||||
tokens = tokens,
|
||||
threshold = self.threshold(),
|
||||
msg_count = history.len(),
|
||||
preserve_count = preserve_count,
|
||||
"Starting three-segment compression"
|
||||
"Starting two-segment compression"
|
||||
);
|
||||
|
||||
// Step 1: Parse into atomic units
|
||||
let units = parse_to_units(history);
|
||||
|
||||
// Step 1: Separate SystemGuard (always preserved, not counted in 5)
|
||||
// Step 2: Separate system guards + user messages from compressible units
|
||||
let mut system_guards: Vec<ChatMessage> = Vec::new();
|
||||
let mut user_messages: Vec<ChatMessage> = Vec::new();
|
||||
let mut compressible: Vec<HistoryUnit> = Vec::new();
|
||||
|
||||
for unit in units {
|
||||
match unit {
|
||||
HistoryUnit::SystemGuard(msg) => system_guards.push(msg),
|
||||
HistoryUnit::UserMessage(msg) => user_messages.push(msg),
|
||||
other => compressible.push(other),
|
||||
}
|
||||
}
|
||||
|
||||
// Step 2: If compressible units are too few, skip LLM compression
|
||||
if compressible.len() <= preserve_count * 2 {
|
||||
tracing::info!(
|
||||
compressible_count = compressible.len(),
|
||||
preserve_threshold = preserve_count * 2,
|
||||
"Too few compressible units, skipping LLM compaction"
|
||||
);
|
||||
let mut result = system_guards;
|
||||
for unit in &compressible {
|
||||
result.extend(unit_to_messages(unit));
|
||||
}
|
||||
return Ok(result);
|
||||
}
|
||||
// Keep only the latest user message in full; older ones go into compression
|
||||
let latest_user_msg = user_messages.pop();
|
||||
|
||||
// Step 3: Three-segment split
|
||||
let split = compressible.len() - preserve_count;
|
||||
let oldest_units = &compressible[..preserve_count];
|
||||
let newest_units = &compressible[split..];
|
||||
let middle_units = &compressible[preserve_count..split];
|
||||
|
||||
// Step 4: Build middle segment messages and transcript
|
||||
let middle_messages: Vec<ChatMessage> = middle_units
|
||||
.iter()
|
||||
.flat_map(unit_to_messages)
|
||||
.collect();
|
||||
let middle_transcript = Self::build_transcript(&middle_messages);
|
||||
|
||||
// Step 5: Summarize middle segment with LLM (heavy prompt)
|
||||
let budget = self.config.summary_max_chars;
|
||||
let middle_summary = if middle_messages.is_empty() {
|
||||
String::new()
|
||||
// Step 3: Find split point in compressible units
|
||||
let total_compressible = estimate_tokens_from_units(&compressible);
|
||||
let split_point = if total_compressible > 0 {
|
||||
self.find_safe_split_point(&compressible, 0.5)
|
||||
} else {
|
||||
self.summarize_units_segment(
|
||||
provider,
|
||||
&middle_messages,
|
||||
&middle_transcript,
|
||||
budget,
|
||||
0
|
||||
};
|
||||
let older_units: Vec<&HistoryUnit> = compressible[..split_point].iter().collect();
|
||||
let newer_units: Vec<&HistoryUnit> = compressible[split_point..].iter().collect();
|
||||
|
||||
// Step 4: Build transcripts
|
||||
let older_messages: Vec<ChatMessage> = older_units
|
||||
.iter()
|
||||
.flat_map(|u| match u {
|
||||
HistoryUnit::ToolRound { assistant, results } => {
|
||||
let mut msgs = vec![assistant.clone()];
|
||||
msgs.extend(results.clone());
|
||||
msgs
|
||||
}
|
||||
HistoryUnit::AssistantText(msg) => vec![msg.clone()],
|
||||
HistoryUnit::UserMessage(msg) => {
|
||||
// Older user messages go into the compressible transcript
|
||||
vec![msg.clone()]
|
||||
}
|
||||
_ => vec![],
|
||||
})
|
||||
.collect();
|
||||
|
||||
let newer_messages: Vec<ChatMessage> = newer_units
|
||||
.iter()
|
||||
.flat_map(|u| match u {
|
||||
HistoryUnit::ToolRound { assistant, results } => {
|
||||
let mut msgs = vec![assistant.clone()];
|
||||
msgs.extend(results.clone());
|
||||
msgs
|
||||
}
|
||||
HistoryUnit::AssistantText(msg) => vec![msg.clone()],
|
||||
_ => vec![],
|
||||
})
|
||||
.collect();
|
||||
|
||||
// Include older user messages in the older transcript
|
||||
let older_user_msgs: Vec<ChatMessage> = user_messages
|
||||
.iter()
|
||||
.map(|m| {
|
||||
let mut msg = m.clone();
|
||||
msg.role = "user".to_string();
|
||||
msg
|
||||
})
|
||||
.collect();
|
||||
let all_older_messages: Vec<ChatMessage> = older_user_msgs
|
||||
.iter()
|
||||
.chain(older_messages.iter())
|
||||
.cloned()
|
||||
.collect();
|
||||
|
||||
let older_transcript = Self::build_transcript(&all_older_messages);
|
||||
let newer_transcript = Self::build_transcript(&newer_messages);
|
||||
|
||||
// Step 5: Budget allocation
|
||||
let total_budget = self.config.summary_max_chars;
|
||||
let older_budget = (total_budget as f64 * OLDER_BUDGET_RATIO) as usize;
|
||||
let newer_budget = (total_budget as f64 * NEWER_BUDGET_RATIO) as usize;
|
||||
|
||||
// Step 6: Create provider and run both summaries (can be parallel)
|
||||
let runtime_config = AgentRuntimeConfig::from(provider_config.clone());
|
||||
let provider = create_provider(runtime_config.provider)
|
||||
.map_err(|e| AgentError::ProviderCreation(e.to_string()))?;
|
||||
|
||||
let (older_result, newer_result) = if older_units.is_empty() && newer_units.is_empty() {
|
||||
(Ok(String::new()), Ok(String::new()))
|
||||
} else if older_units.is_empty() {
|
||||
let result = self
|
||||
.summarize_units_segment(
|
||||
provider.as_ref(),
|
||||
&newer_messages,
|
||||
&newer_transcript,
|
||||
newer_budget,
|
||||
Self::build_light_summary_prompt,
|
||||
)
|
||||
.await;
|
||||
(Ok(String::new()), result)
|
||||
} else if newer_units.is_empty() {
|
||||
let result = self
|
||||
.summarize_units_segment(
|
||||
provider.as_ref(),
|
||||
&all_older_messages,
|
||||
&older_transcript,
|
||||
older_budget,
|
||||
Self::build_heavy_summary_prompt,
|
||||
)
|
||||
.await;
|
||||
(result, Ok(String::new()))
|
||||
} else {
|
||||
let older_fut = self.summarize_units_segment(
|
||||
provider.as_ref(),
|
||||
&all_older_messages,
|
||||
&older_transcript,
|
||||
older_budget,
|
||||
Self::build_heavy_summary_prompt,
|
||||
)
|
||||
.await?
|
||||
);
|
||||
let newer_fut = self.summarize_units_segment(
|
||||
provider.as_ref(),
|
||||
&newer_messages,
|
||||
&newer_transcript,
|
||||
newer_budget,
|
||||
Self::build_light_summary_prompt,
|
||||
);
|
||||
tokio::join!(older_fut, newer_fut)
|
||||
};
|
||||
|
||||
// Step 6: Assemble compressed history
|
||||
// [SystemGuards] + [oldest 5 units raw] + [middle summary] + [newest 5 units raw]
|
||||
let mut compressed: Vec<ChatMessage> = Vec::with_capacity(
|
||||
system_guards.len() + middle_messages.len() + 1 + middle_messages.len(),
|
||||
);
|
||||
let older_summary = older_result?;
|
||||
let newer_summary = newer_result?;
|
||||
|
||||
// Step 7: Assemble compressed history
|
||||
let mut compressed: Vec<ChatMessage> = Vec::with_capacity(4);
|
||||
|
||||
// System guards first
|
||||
compressed.extend(system_guards);
|
||||
|
||||
// Oldest preserve_count units (raw)
|
||||
for unit in oldest_units {
|
||||
compressed.extend(unit_to_messages(unit));
|
||||
// Latest user message
|
||||
if let Some(user_msg) = latest_user_msg {
|
||||
compressed.push(user_msg);
|
||||
}
|
||||
|
||||
// Middle segment summary (system message, no tool_calls)
|
||||
if !middle_summary.is_empty() {
|
||||
// Heavy compression summary (older)
|
||||
if !older_summary.is_empty() {
|
||||
compressed.push(ChatMessage::system_with_context(
|
||||
format!("## 较早的操作记录（已压缩）\n\n{}", middle_summary),
|
||||
format!("## 较早的操作记录（已压缩）\n\n{}", older_summary),
|
||||
Some(SYSTEM_CONTEXT_HISTORY_COMPACTION_OLDER.to_string()),
|
||||
));
|
||||
}
|
||||
|
||||
// Newest preserve_count units (raw)
|
||||
for unit in newest_units {
|
||||
compressed.extend(unit_to_messages(unit));
|
||||
// Light compression summary (newer)
|
||||
if !newer_summary.is_empty() {
|
||||
compressed.push(ChatMessage::system_with_context(
|
||||
format!(
|
||||
"## 近期操作记录（轻度压缩，保留了更多细节）\n\n{}\n\n---\n以上为近期操作记录。如需精确数据，可使用工具重新读取相关文件。",
|
||||
newer_summary
|
||||
),
|
||||
Some(SYSTEM_CONTEXT_HISTORY_COMPACTION_NEWER.to_string()),
|
||||
));
|
||||
}
|
||||
|
||||
tracing::info!(
|
||||
original_tokens = tokens,
|
||||
original_msg_count = history.len(),
|
||||
final_tokens = estimate_tokens(&compressed),
|
||||
final_msg_count = compressed.len(),
|
||||
oldest_units = preserve_count,
|
||||
newest_units = preserve_count,
|
||||
middle_units = middle_units.len(),
|
||||
"Three-segment compression completed"
|
||||
older_units = older_units.len(),
|
||||
newer_units = newer_units.len(),
|
||||
"Two-segment compression completed"
|
||||
);
|
||||
|
||||
Ok(compressed)
|
||||
@ -780,50 +824,39 @@ OLDER SEGMENT (events from earlier in the session):
|
||||
// =========================================================================
|
||||
|
||||
/// Create a new compressor with the given context window size.
|
||||
/// 测试 fallback 路径：所有压缩参数用 CompactionConfig::default()。
|
||||
pub fn new(context_window: usize) -> Self {
|
||||
let default = CompactionConfig::default();
|
||||
Self {
|
||||
config: ContextCompressionConfig::default(),
|
||||
context_window,
|
||||
threshold_ratio: default.threshold_ratio,
|
||||
llm_compaction_threshold_ratio: default.llm_compaction_threshold_ratio,
|
||||
preserve_count: default.preserve_count,
|
||||
truncate_max_tokens: default.truncate_max_tokens,
|
||||
threshold_ratio: DEFAULT_THRESHOLD_RATIO,
|
||||
}
|
||||
}
|
||||
|
||||
/// 从 runtime config + compaction config 构造（AgentFactory 生产路径调用）。
|
||||
/// 所有压缩参数内聚到 ContextCompressor，AgentLoop 零参数耦合。
|
||||
/// 对用户配置做防御性 clamp，避免极端值破坏压缩逻辑：
|
||||
/// - ratio 限制在 [0.1, 1.0]：过低导致每轮压缩，过高导致永不压缩
|
||||
/// - truncate_max_tokens 限制 >= 1：0 会把所有 tool 结果截断成空串
|
||||
/// - preserve_count 限制 >= 1：0 会丢失全部原始上下文
|
||||
pub fn with_compaction_config(
|
||||
context_window: usize,
|
||||
summary_max_chars: usize,
|
||||
compaction: &CompactionConfig,
|
||||
) -> Self {
|
||||
let clamp_ratio = |r: f64| r.clamp(0.1, 1.0);
|
||||
let threshold_ratio = clamp_ratio(compaction.threshold_ratio);
|
||||
let llm_compaction_threshold_ratio = clamp_ratio(compaction.llm_compaction_threshold_ratio);
|
||||
let preserve_count = compaction.preserve_count.max(1);
|
||||
let truncate_max_tokens = compaction.truncate_max_tokens.max(1);
|
||||
Self {
|
||||
config: ContextCompressionConfig {
|
||||
summary_max_chars,
|
||||
pub fn from_provider_config(provider_config: &LLMProviderConfig) -> Self {
|
||||
Self::from_runtime_config(&AgentRuntimeConfig::from(provider_config.clone()))
|
||||
}
|
||||
|
||||
pub fn from_runtime_config(config: &AgentRuntimeConfig) -> Self {
|
||||
Self::with_config(
|
||||
config.context_window_tokens,
|
||||
ContextCompressionConfig {
|
||||
summary_max_chars: config.context_summary_char_budget,
|
||||
..ContextCompressionConfig::default()
|
||||
},
|
||||
)
|
||||
}
|
||||
|
||||
/// Create with custom configuration.
|
||||
pub fn with_config(context_window: usize, config: ContextCompressionConfig) -> Self {
|
||||
Self {
|
||||
config,
|
||||
context_window,
|
||||
threshold_ratio,
|
||||
llm_compaction_threshold_ratio,
|
||||
preserve_count,
|
||||
truncate_max_tokens,
|
||||
threshold_ratio: DEFAULT_THRESHOLD_RATIO,
|
||||
}
|
||||
}
|
||||
|
||||
/// Get the compression threshold in tokens (50% of context window).
|
||||
pub fn threshold(&self) -> usize {
|
||||
/// Get the compression threshold in tokens (70% of context window).
|
||||
fn threshold(&self) -> usize {
|
||||
(self.context_window as f64 * self.threshold_ratio) as usize
|
||||
}
|
||||
|
||||
@ -831,22 +864,6 @@ OLDER SEGMENT (events from earlier in the session):
|
||||
estimate_tokens(history) > self.threshold()
|
||||
}
|
||||
|
||||
/// 触发阈值（50%）：用真实 prompt_tokens 判断是否进入压缩流程。
|
||||
pub fn should_compress_by_usage(&self, prompt_tokens: u32) -> bool {
|
||||
(prompt_tokens as usize) > self.threshold()
|
||||
}
|
||||
|
||||
/// LLM 压缩阈值（30%）：工程化压缩后用 estimate_tokens 判断是否需要 LLM 压缩。
|
||||
pub fn llm_compaction_threshold(&self) -> usize {
|
||||
(self.context_window as f64 * self.llm_compaction_threshold_ratio) as usize
|
||||
}
|
||||
|
||||
/// 工程化压缩：截断非子代理 tool 结果到 self.truncate_max_tokens。
|
||||
/// AgentLoop 调用此方法，不需要知道截断参数细节（参数内聚）。
|
||||
pub fn truncate_tool_results(&self, messages: &mut [ChatMessage]) {
|
||||
truncate_tool_results_in_place(messages, self.truncate_max_tokens);
|
||||
}
|
||||
|
||||
fn user_turn_ranges(&self, history: &[ChatMessage]) -> Vec<UserTurnRange> {
|
||||
let user_indices: Vec<usize> = history
|
||||
.iter()
|
||||
@ -1097,7 +1114,9 @@ mod tests {
|
||||
|
||||
#[test]
|
||||
fn test_estimate_tokens_mixed_content() {
|
||||
let messages = vec![ChatMessage::user("Hello 世界 this is 测试")];
|
||||
let messages = vec![
|
||||
ChatMessage::user("Hello 世界 this is 测试"),
|
||||
];
|
||||
|
||||
let tokens = estimate_tokens(&messages);
|
||||
// Content: 18 English chars + 4 CJK chars
|
||||
@ -1216,7 +1235,7 @@ mod tests {
|
||||
#[test]
|
||||
fn test_threshold() {
|
||||
let compressor = ContextCompressor::new(128_000);
|
||||
assert_eq!(compressor.threshold(), 64_000); // 50% of 128_000
|
||||
assert_eq!(compressor.threshold(), 89_600); // 70% of 128_000
|
||||
}
|
||||
|
||||
#[test]
|
||||
@ -1377,6 +1396,22 @@ mod tests {
|
||||
assert!(units.is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_find_safe_split_point() {
|
||||
let compressor = ContextCompressor::new(100_000);
|
||||
// Build a list of units with known token sizes
|
||||
let messages: Vec<ChatMessage> = (0..10)
|
||||
.map(|i| ChatMessage::assistant(&format!("message content number {}", i)))
|
||||
.collect();
|
||||
let units = parse_to_units(&messages);
|
||||
|
||||
// All units are AssistantText, split at 50% token ratio
|
||||
let split = compressor.find_safe_split_point(&units, 0.5);
|
||||
// Should split somewhere in the middle (not 0, not len())
|
||||
assert!(split > 0 && split < units.len(),
|
||||
"split {} should be between 0 and {}", split, units.len());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_compress_two_segment_no_tool_calls_in_output() {
|
||||
// This test verifies the critical invariant:
|
||||
@ -1403,7 +1438,6 @@ mod tests {
|
||||
extra_headers: std::collections::HashMap::new(),
|
||||
llm_timeout_secs: 120,
|
||||
memory_maintenance_timeout_secs: 300,
|
||||
max_retries: 3,
|
||||
model_id: "test-model".to_string(),
|
||||
temperature: None,
|
||||
max_tokens: None,
|
||||
@ -1424,14 +1458,10 @@ mod tests {
|
||||
assert_eq!(compressed.len(), 3);
|
||||
// Critical invariant: NO tool_calls or tool_call_id anywhere
|
||||
for msg in &compressed {
|
||||
assert!(
|
||||
msg.tool_calls.is_none(),
|
||||
"compress_two_segment output should never contain tool_calls"
|
||||
);
|
||||
assert!(
|
||||
msg.tool_call_id.is_none(),
|
||||
"compress_two_segment output should never contain tool_call_id"
|
||||
);
|
||||
assert!(msg.tool_calls.is_none(),
|
||||
"compress_two_segment output should never contain tool_calls");
|
||||
assert!(msg.tool_call_id.is_none(),
|
||||
"compress_two_segment output should never contain tool_call_id");
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@ -4,12 +4,12 @@ pub mod runtime_config;
|
||||
pub mod system_prompt;
|
||||
|
||||
pub use agent_loop::{
|
||||
AgentError, AgentLoop, AgentProcessResult, CompactionSink, EmittedMessageHandler,
|
||||
AgentError, AgentLoop, AgentProcessResult, EmittedMessageHandler,
|
||||
PersistingEmittedMessageHandler, SkillProvider,
|
||||
};
|
||||
pub use context_compressor::ContextCompressor;
|
||||
pub use runtime_config::AgentRuntimeConfig;
|
||||
pub use system_prompt::{
|
||||
CompositeSystemPromptProvider, SystemPrompt, SystemPromptContext, SystemPromptProvider,
|
||||
generate_system_env_prompt,
|
||||
CompositeSystemPromptProvider, generate_system_env_prompt, SystemPrompt, SystemPromptContext,
|
||||
SystemPromptProvider,
|
||||
};
|
||||
|
||||
@ -12,9 +12,6 @@ pub struct AgentRuntimeConfig {
|
||||
/// 图片上下文限制配置
|
||||
pub max_images_in_context: usize,
|
||||
pub max_image_age_rounds: usize,
|
||||
/// LLM 请求瞬态失败的最大重试次数（仅对 timeout/502/503/504/429 等可恢复错误重试）。
|
||||
/// 0 表示不重试。归属 agent 行为层，不进 ProviderRuntimeConfig（保持 provider 构造包纯净）。
|
||||
pub max_retries: u32,
|
||||
}
|
||||
|
||||
impl From<LLMProviderConfig> for AgentRuntimeConfig {
|
||||
@ -42,7 +39,6 @@ impl From<LLMProviderConfig> for AgentRuntimeConfig {
|
||||
context_tool_result_trim_chars: config.context_tool_result_trim_chars,
|
||||
max_images_in_context: config.max_images_in_context,
|
||||
max_image_age_rounds: config.max_image_age_rounds,
|
||||
max_retries: config.max_retries,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@ -23,4 +23,4 @@ pub fn initialize_process_runtime() {
|
||||
// optionally the RUST_BACKTRACE-based backtrace.
|
||||
default_hook(info);
|
||||
}));
|
||||
}
|
||||
}
|
||||
@ -2,7 +2,6 @@ use serde::{Deserialize, Serialize};
|
||||
use std::collections::{HashMap, HashSet};
|
||||
|
||||
use crate::domain::messages::ToolCall;
|
||||
use crate::utils::current_timestamp;
|
||||
|
||||
pub const SYSTEM_CONTEXT_AGENT_PROMPT: &str = "agent_prompt";
|
||||
pub const SYSTEM_CONTEXT_SCHEDULED_PROMPT: &str = "scheduled_system_prompt";
|
||||
@ -26,7 +25,7 @@ pub struct MediaItem {
|
||||
pub mime_type: Option<String>,
|
||||
pub original_key: Option<String>, // Feishu file_key for download
|
||||
pub content_base64: Option<String>, // Base64-encoded file content for web download
|
||||
pub file_name: Option<String>, // Display file name
|
||||
pub file_name: Option<String>, // Display file name
|
||||
}
|
||||
|
||||
impl MediaItem {
|
||||
@ -67,38 +66,6 @@ pub struct ChatMessage {
|
||||
pub tool_duration_ms: Option<u64>,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub tool_calls: Option<Vec<ToolCall>>,
|
||||
/// LLM 调用 usage（仅 assistant 消息有值，来自 provider 响应）
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub usage: Option<MessageUsage>,
|
||||
}
|
||||
|
||||
/// 单次 LLM 调用的 token 用量
|
||||
#[derive(Debug, Clone, Serialize, Deserialize, Default)]
|
||||
pub struct MessageUsage {
|
||||
pub prompt_tokens: u32,
|
||||
pub completion_tokens: u32,
|
||||
pub total_tokens: u32,
|
||||
/// 本次调用所用模型的上下文窗口大小（来自 AgentRuntimeConfig）。
|
||||
/// 与 prompt_tokens 一起持久化，用于计算上下文占用率。
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub context_window_tokens: Option<u32>,
|
||||
}
|
||||
|
||||
impl MessageUsage {
|
||||
pub fn from_provider_usage(u: crate::providers::Usage) -> Self {
|
||||
Self {
|
||||
prompt_tokens: u.prompt_tokens,
|
||||
completion_tokens: u.completion_tokens,
|
||||
total_tokens: u.total_tokens,
|
||||
context_window_tokens: None,
|
||||
}
|
||||
}
|
||||
|
||||
/// 链式设置 context_window_tokens（来自 AgentRuntimeConfig）
|
||||
pub fn with_context_window(mut self, ctx: usize) -> Self {
|
||||
self.context_window_tokens = Some(ctx as u32);
|
||||
self
|
||||
}
|
||||
}
|
||||
|
||||
impl ChatMessage {
|
||||
@ -116,7 +83,6 @@ impl ChatMessage {
|
||||
tool_duration_ms: None,
|
||||
tool_state: None,
|
||||
tool_calls: None,
|
||||
usage: None,
|
||||
}
|
||||
}
|
||||
|
||||
@ -134,7 +100,6 @@ impl ChatMessage {
|
||||
tool_duration_ms: None,
|
||||
tool_state: None,
|
||||
tool_calls: None,
|
||||
usage: None,
|
||||
}
|
||||
}
|
||||
|
||||
@ -152,7 +117,6 @@ impl ChatMessage {
|
||||
tool_duration_ms: None,
|
||||
tool_state: None,
|
||||
tool_calls: None,
|
||||
usage: None,
|
||||
}
|
||||
}
|
||||
|
||||
@ -182,7 +146,6 @@ impl ChatMessage {
|
||||
tool_duration_ms: None,
|
||||
tool_state: None,
|
||||
tool_calls: Some(tool_calls),
|
||||
usage: None,
|
||||
}
|
||||
}
|
||||
|
||||
@ -217,7 +180,6 @@ impl ChatMessage {
|
||||
tool_duration_ms: None,
|
||||
tool_state: None,
|
||||
tool_calls: None,
|
||||
usage: None,
|
||||
}
|
||||
}
|
||||
|
||||
@ -253,7 +215,6 @@ impl ChatMessage {
|
||||
tool_duration_ms: None,
|
||||
tool_state: Some(tool_state),
|
||||
tool_calls: None,
|
||||
usage: None,
|
||||
}
|
||||
}
|
||||
|
||||
@ -323,13 +284,12 @@ pub(crate) fn sanitize_incomplete_tool_call_sequences(messages: &mut Vec<ChatMes
|
||||
}
|
||||
|
||||
if msg.role == "assistant"
|
||||
&& msg
|
||||
.tool_calls
|
||||
.as_ref()
|
||||
.map_or(false, |calls| !calls.is_empty())
|
||||
&& msg.tool_calls.as_ref().map_or(false, |calls| !calls.is_empty())
|
||||
{
|
||||
let tool_calls = msg.tool_calls.as_ref().unwrap();
|
||||
let all_have_results = tool_calls.iter().all(|tc| resolved_ids.contains(&tc.id));
|
||||
let all_have_results = tool_calls
|
||||
.iter()
|
||||
.all(|tc| resolved_ids.contains(&tc.id));
|
||||
|
||||
if all_have_results {
|
||||
for tc in tool_calls.iter() {
|
||||
@ -399,19 +359,12 @@ pub(crate) fn sanitize_incomplete_tool_call_sequences(messages: &mut Vec<ChatMes
|
||||
}
|
||||
|
||||
if m.role == "assistant"
|
||||
&& m.tool_calls
|
||||
.as_ref()
|
||||
.map_or(false, |calls| !calls.is_empty())
|
||||
&& m.tool_calls.as_ref().map_or(false, |calls| !calls.is_empty())
|
||||
{
|
||||
let already_marked = remove_indices.contains(&i);
|
||||
if !already_marked {
|
||||
pending_tool_ids = m
|
||||
.tool_calls
|
||||
.as_ref()
|
||||
.unwrap()
|
||||
.iter()
|
||||
.map(|tc| tc.id.clone())
|
||||
.collect();
|
||||
pending_tool_ids = m.tool_calls.as_ref().unwrap()
|
||||
.iter().map(|tc| tc.id.clone()).collect();
|
||||
pending_assistant_idx = Some(i);
|
||||
}
|
||||
} else if m.role == "tool" {
|
||||
@ -558,10 +511,7 @@ pub enum OutboundEventKind {
|
||||
|
||||
impl OutboundMessage {
|
||||
pub fn is_stream_delta(&self) -> bool {
|
||||
matches!(
|
||||
self.event_kind,
|
||||
OutboundEventKind::StreamDelta | OutboundEventKind::StreamEnd
|
||||
)
|
||||
matches!(self.event_kind, OutboundEventKind::StreamDelta | OutboundEventKind::StreamEnd)
|
||||
}
|
||||
|
||||
pub fn assistant(
|
||||
@ -598,8 +548,7 @@ impl OutboundMessage {
|
||||
reply_to: Option<String>,
|
||||
metadata: HashMap<String, String>,
|
||||
) -> Self {
|
||||
let mut message =
|
||||
Self::assistant(channel, chat_id, session_id, content, reply_to, metadata);
|
||||
let mut message = Self::assistant(channel, chat_id, session_id, content, reply_to, metadata);
|
||||
message.event_kind = OutboundEventKind::SchedulerNotification;
|
||||
message
|
||||
}
|
||||
@ -612,8 +561,7 @@ impl OutboundMessage {
|
||||
reply_to: Option<String>,
|
||||
metadata: HashMap<String, String>,
|
||||
) -> Self {
|
||||
let mut message =
|
||||
Self::assistant(channel, chat_id, session_id, content, reply_to, metadata);
|
||||
let mut message = Self::assistant(channel, chat_id, session_id, content, reply_to, metadata);
|
||||
message.event_kind = OutboundEventKind::ErrorNotification;
|
||||
message
|
||||
}
|
||||
@ -647,7 +595,7 @@ impl OutboundMessage {
|
||||
message_id: None,
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
pub fn tool_result(
|
||||
channel: impl Into<String>,
|
||||
chat_id: impl Into<String>,
|
||||
@ -678,7 +626,7 @@ impl OutboundMessage {
|
||||
message_id: None,
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
pub fn tool_pending(
|
||||
channel: impl Into<String>,
|
||||
chat_id: impl Into<String>,
|
||||
@ -709,7 +657,7 @@ impl OutboundMessage {
|
||||
message_id: None,
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
/// 构造流式文本增量消息
|
||||
pub fn stream_delta(
|
||||
channel: impl Into<String>,
|
||||
@ -737,7 +685,7 @@ impl OutboundMessage {
|
||||
message_id: None,
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
/// 构造流式结束信号
|
||||
pub fn stream_end(
|
||||
channel: impl Into<String>,
|
||||
@ -801,8 +749,7 @@ impl OutboundMessage {
|
||||
"assistant" => {
|
||||
if let Some(tool_calls) = &message.tool_calls {
|
||||
let mut outbound = Vec::new();
|
||||
let has_content_or_reasoning =
|
||||
!message.content.trim().is_empty() || message.reasoning_content.is_some();
|
||||
let has_content_or_reasoning = !message.content.trim().is_empty() || message.reasoning_content.is_some();
|
||||
if has_content_or_reasoning {
|
||||
let mut resp = Self::assistant(
|
||||
channel.to_string(),
|
||||
@ -819,11 +766,7 @@ impl OutboundMessage {
|
||||
|
||||
// AssistantResponse 已携带 reasoning 时，ToolCall 不再重复；
|
||||
// 只有 AssistantResponse 没发时，ToolCall 才带 reasoning
|
||||
let tc_reasoning = if has_content_or_reasoning {
|
||||
None
|
||||
} else {
|
||||
message.reasoning_content.clone()
|
||||
};
|
||||
let tc_reasoning = if has_content_or_reasoning { None } else { message.reasoning_content.clone() };
|
||||
outbound.extend(tool_calls.iter().map(|tool_call| {
|
||||
let mut tc = Self::tool_call(
|
||||
channel.to_string(),
|
||||
@ -932,6 +875,13 @@ fn format_tool_arguments_json(value: &serde_json::Value) -> String {
|
||||
// Helpers
|
||||
// ============================================================================
|
||||
|
||||
fn current_timestamp() -> i64 {
|
||||
std::time::SystemTime::now()
|
||||
.duration_since(std::time::UNIX_EPOCH)
|
||||
.unwrap()
|
||||
.as_millis() as i64
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::{ChatMessage, OutboundEventKind, OutboundMessage, ToolMessageState};
|
||||
@ -980,7 +930,10 @@ mod tests {
|
||||
"calculator\nargs: {\"expression\":\"1 + 1\"}"
|
||||
);
|
||||
assert_eq!(outbound[1].tool_name.as_deref(), Some("read"));
|
||||
assert_eq!(outbound[1].content, "read\nargs: {\"path\":\"README.md\"}");
|
||||
assert_eq!(
|
||||
outbound[1].content,
|
||||
"read\nargs: {\"path\":\"README.md\"}"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
|
||||
@ -52,26 +52,14 @@ impl MessageBus {
|
||||
Some(msg)
|
||||
}
|
||||
|
||||
/// Publish a message to the outbound queue.
|
||||
///
|
||||
/// Uses `try_send` (non-blocking): if the queue is full, the message is
|
||||
/// dropped immediately with a warning. This ensures the agent loop is never
|
||||
/// blocked by slow or disconnected display consumers. Persistent state is
|
||||
/// unaffected — messages are stored in SQLite independently.
|
||||
/// Publish a message to the outbound queue
|
||||
pub async fn publish_outbound(&self, msg: OutboundMessage) -> Result<(), BusError> {
|
||||
#[cfg(debug_assertions)]
|
||||
tracing::debug!(channel = %msg.channel, chat_id = %msg.chat_id, content_len = %msg.content.len(), "Bus: publishing outbound message");
|
||||
match self.outbound_tx.try_send(msg) {
|
||||
Ok(()) => Ok(()),
|
||||
Err(tokio::sync::mpsc::error::TrySendError::Full(msg)) => {
|
||||
tracing::warn!(
|
||||
channel = %msg.channel,
|
||||
"Outbound bus full, dropping message"
|
||||
);
|
||||
Err(BusError::Dropped)
|
||||
}
|
||||
Err(tokio::sync::mpsc::error::TrySendError::Closed(_)) => Err(BusError::Closed),
|
||||
}
|
||||
self.outbound_tx
|
||||
.send(msg)
|
||||
.await
|
||||
.map_err(|_| BusError::Closed)
|
||||
}
|
||||
|
||||
/// Consume an outbound message from the outbound queue.
|
||||
@ -88,14 +76,12 @@ impl MessageBus {
|
||||
#[derive(Debug)]
|
||||
pub enum BusError {
|
||||
Closed,
|
||||
Dropped,
|
||||
}
|
||||
|
||||
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::Dropped => write!(f, "Bus full, message dropped"),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@ -8,9 +8,6 @@ pub enum ChannelError {
|
||||
ConfigError(String),
|
||||
ConnectionError(String),
|
||||
SendError(String),
|
||||
/// Channel 内部队列已满——不可重试，立即丢弃。
|
||||
/// 重试只会浪费退避时间并阻塞 channel 队列消费。
|
||||
ChannelFull,
|
||||
BusError(String),
|
||||
Other(String),
|
||||
}
|
||||
@ -21,7 +18,6 @@ impl std::fmt::Display for ChannelError {
|
||||
ChannelError::ConfigError(s) => write!(f, "Config error: {}", s),
|
||||
ChannelError::ConnectionError(s) => write!(f, "Connection error: {}", s),
|
||||
ChannelError::SendError(s) => write!(f, "Send error: {}", s),
|
||||
ChannelError::ChannelFull => write!(f, "Channel queue full"),
|
||||
ChannelError::BusError(s) => write!(f, "Bus error: {}", s),
|
||||
ChannelError::Other(s) => write!(f, "Error: {}", s),
|
||||
}
|
||||
|
||||
@ -84,10 +84,7 @@ impl Channel for CliChannel {
|
||||
self.shutdown_token.cancel();
|
||||
let count = self.connections.read().await.len();
|
||||
self.connections.write().await.clear();
|
||||
tracing::info!(
|
||||
connection_count = count,
|
||||
"CliChannel stopped, all connections signaled to close"
|
||||
);
|
||||
tracing::info!(connection_count = count, "CliChannel stopped, all connections signaled to close");
|
||||
Ok(())
|
||||
}
|
||||
|
||||
@ -100,22 +97,12 @@ impl Channel for CliChannel {
|
||||
)));
|
||||
};
|
||||
|
||||
// 使用 try_send 避免阻塞 dispatcher——dispatcher 是单线程顺序处理，
|
||||
// 若 writer task 卡在 ws_sender.send() 上，send().await 会阻塞，
|
||||
// 导致所有连接的实时消息被阻塞。try_send 满时返回 ChannelFull
|
||||
// （不可重试，立即丢弃）；关闭时返回 SendError（可重试）。
|
||||
for outbound in ws_outbound_from_outbound_message(&msg) {
|
||||
match connection.sender.try_send(outbound) {
|
||||
Ok(()) => {}
|
||||
Err(tokio::sync::mpsc::error::TrySendError::Full(_)) => {
|
||||
return Err(ChannelError::ChannelFull);
|
||||
}
|
||||
Err(tokio::sync::mpsc::error::TrySendError::Closed(_)) => {
|
||||
return Err(ChannelError::SendError(
|
||||
"CLI websocket sender closed".to_string(),
|
||||
));
|
||||
}
|
||||
}
|
||||
connection
|
||||
.sender
|
||||
.send(outbound)
|
||||
.await
|
||||
.map_err(|_| ChannelError::SendError("CLI websocket sender closed".to_string()))?;
|
||||
}
|
||||
|
||||
Ok(())
|
||||
|
||||
@ -1,7 +1,6 @@
|
||||
use std::collections::HashMap;
|
||||
use std::path::Path;
|
||||
use std::sync::Arc;
|
||||
use std::sync::LazyLock;
|
||||
use std::time::{Duration, Instant};
|
||||
|
||||
use async_trait::async_trait;
|
||||
@ -11,8 +10,8 @@ use regex::Regex;
|
||||
use serde::Deserialize;
|
||||
use tokio::sync::{RwLock, broadcast};
|
||||
|
||||
use crate::bus::message::OutboundEventKind;
|
||||
use crate::bus::{MediaItem, MessageBus, OutboundMessage};
|
||||
use crate::bus::message::OutboundEventKind;
|
||||
use crate::channels::base::{Channel, ChannelError};
|
||||
use crate::config::{FeishuChannelConfig, LLMProviderConfig};
|
||||
use crate::text::{char_count, truncate_with_ellipsis};
|
||||
@ -549,10 +548,9 @@ impl FeishuChannel {
|
||||
}
|
||||
|
||||
let status = resp.status();
|
||||
let body = resp
|
||||
.text()
|
||||
.await
|
||||
.map_err(|e| ChannelError::Other(format!("Read upload image response error: {}", e)))?;
|
||||
let body = resp.text().await.map_err(|e| {
|
||||
ChannelError::Other(format!("Read upload image response error: {}", e))
|
||||
})?;
|
||||
let result: UploadResp = serde_json::from_str(&body).map_err(|e| {
|
||||
ChannelError::Other(format!(
|
||||
"Parse upload image response error: {} (status={}, body={})",
|
||||
@ -633,10 +631,9 @@ impl FeishuChannel {
|
||||
}
|
||||
|
||||
let status = resp.status();
|
||||
let body = resp
|
||||
.text()
|
||||
.await
|
||||
.map_err(|e| ChannelError::Other(format!("Read upload file response error: {}", e)))?;
|
||||
let body = resp.text().await.map_err(|e| {
|
||||
ChannelError::Other(format!("Read upload file response error: {}", e))
|
||||
})?;
|
||||
let result: UploadResp = serde_json::from_str(&body).map_err(|e| {
|
||||
ChannelError::Other(format!(
|
||||
"Parse upload file response error: {} (status={}, body={})",
|
||||
@ -985,11 +982,9 @@ impl FeishuChannel {
|
||||
reply_to: Option<&str>,
|
||||
) -> Result<(), ChannelError> {
|
||||
if let Some(parent_id) = reply_to {
|
||||
self.reply_to_feishu_message(parent_id, msg_type, content)
|
||||
.await
|
||||
self.reply_to_feishu_message(parent_id, msg_type, content).await
|
||||
} else {
|
||||
self.send_message_to_feishu(receive_id, receive_id_type, msg_type, content)
|
||||
.await
|
||||
self.send_message_to_feishu(receive_id, receive_id_type, msg_type, content).await
|
||||
}
|
||||
}
|
||||
|
||||
@ -1311,7 +1306,6 @@ impl FeishuChannel {
|
||||
let channel = self.clone();
|
||||
let bus = bus.clone();
|
||||
tokio::spawn(async move {
|
||||
#[cfg(debug_assertions)]
|
||||
let media_count = if parsed.media.is_some() { 1 } else { 0 };
|
||||
#[cfg(debug_assertions)]
|
||||
tracing::debug!(open_id = %parsed.open_id, chat_id = %parsed.chat_id, content_len = %parsed.content.len(), media_count = %media_count, "Publishing message to bus");
|
||||
@ -1446,20 +1440,16 @@ fn parse_post_content(content: &str) -> String {
|
||||
}
|
||||
"code_block" => {
|
||||
let lang = el.get("language").and_then(|l| l.as_str()).unwrap_or("");
|
||||
let code_text =
|
||||
if let Some(content_arr) = el.get("content").and_then(|c| c.as_array()) {
|
||||
content_arr
|
||||
.iter()
|
||||
.filter_map(|item| item.get("text").and_then(|t| t.as_str()))
|
||||
.collect::<Vec<_>>()
|
||||
.join("")
|
||||
} else {
|
||||
// Fallback to text field for backwards compatibility
|
||||
el.get("text")
|
||||
.and_then(|t| t.as_str())
|
||||
.unwrap_or("")
|
||||
.to_string()
|
||||
};
|
||||
let code_text = if let Some(content_arr) = el.get("content").and_then(|c| c.as_array()) {
|
||||
content_arr
|
||||
.iter()
|
||||
.filter_map(|item| item.get("text").and_then(|t| t.as_str()))
|
||||
.collect::<Vec<_>>()
|
||||
.join("")
|
||||
} else {
|
||||
// Fallback to text field for backwards compatibility
|
||||
el.get("text").and_then(|t| t.as_str()).unwrap_or("").to_string()
|
||||
};
|
||||
out.push(format!("\n```{}\n{}\n```\n", lang, code_text));
|
||||
}
|
||||
_ => {
|
||||
@ -1865,13 +1855,10 @@ impl Default for MdPatterns {
|
||||
}
|
||||
}
|
||||
|
||||
/// 全局唯一的正则模式实例，仅编译一次。
|
||||
static MD_PATTERNS: LazyLock<MdPatterns> = LazyLock::new(MdPatterns::new);
|
||||
|
||||
impl FeishuChannel {
|
||||
/// Determine the optimal Feishu message format for content.
|
||||
fn detect_msg_format(content: &str) -> MsgFormat {
|
||||
let patterns = &MD_PATTERNS;
|
||||
let patterns = MdPatterns::new();
|
||||
let stripped = content.trim();
|
||||
|
||||
// Tables and headings are not supported by post `md` nodes, so use cards.
|
||||
@ -1907,7 +1894,7 @@ impl FeishuChannel {
|
||||
|
||||
/// Strip markdown formatting markers from text for plain display.
|
||||
fn strip_md_formatting(text: &str) -> String {
|
||||
let patterns = &MD_PATTERNS;
|
||||
let patterns = MdPatterns::new();
|
||||
let mut result = text.to_string();
|
||||
|
||||
// Remove bold markers
|
||||
@ -1986,7 +1973,7 @@ impl FeishuChannel {
|
||||
|
||||
/// Split content by headings, converting headings to div elements.
|
||||
fn split_headings(content: &str) -> Vec<serde_json::Value> {
|
||||
let patterns = &MD_PATTERNS;
|
||||
let patterns = MdPatterns::new();
|
||||
let mut protected = content.to_string();
|
||||
|
||||
// Protect code blocks by replacing them with placeholders
|
||||
@ -2092,7 +2079,7 @@ impl FeishuChannel {
|
||||
|
||||
/// Build content into card elements (div/markdown + table).
|
||||
fn build_card_elements(content: &str) -> Vec<serde_json::Value> {
|
||||
let patterns = &MD_PATTERNS;
|
||||
let patterns = MdPatterns::new();
|
||||
let mut elements: Vec<serde_json::Value> = Vec::new();
|
||||
let mut last_end = 0;
|
||||
|
||||
@ -2393,8 +2380,7 @@ mod tests {
|
||||
#[test]
|
||||
fn parse_post_content_handles_empty_code_block() {
|
||||
// Test code_block with empty content
|
||||
let post_json =
|
||||
r#"{"post":{"zh_cn":{"content":[[{"tag":"code_block","language":"go"}]]}}}"#;
|
||||
let post_json = r#"{"post":{"zh_cn":{"content":[[{"tag":"code_block","language":"go"}]]}}}"#;
|
||||
let result = parse_post_content(post_json);
|
||||
assert!(result.contains("```go"));
|
||||
}
|
||||
@ -2475,18 +2461,8 @@ impl Channel for FeishuChannel {
|
||||
}
|
||||
|
||||
async fn send(&self, msg: OutboundMessage) -> Result<(), ChannelError> {
|
||||
if matches!(
|
||||
msg.event_kind,
|
||||
OutboundEventKind::ToolResult
|
||||
| OutboundEventKind::ToolPending
|
||||
| OutboundEventKind::StreamDelta
|
||||
| OutboundEventKind::StreamEnd
|
||||
| OutboundEventKind::ExecutionCompleted
|
||||
) || msg
|
||||
.metadata
|
||||
.get("is_subagent_event")
|
||||
.map(|v| v == "true")
|
||||
.unwrap_or(false)
|
||||
if matches!(msg.event_kind, OutboundEventKind::ToolResult | OutboundEventKind::ToolPending | OutboundEventKind::StreamDelta | OutboundEventKind::StreamEnd | OutboundEventKind::ExecutionCompleted)
|
||||
|| msg.metadata.get("is_subagent_event").map(|v| v == "true").unwrap_or(false)
|
||||
{
|
||||
return Ok(());
|
||||
}
|
||||
@ -2509,15 +2485,6 @@ impl Channel for FeishuChannel {
|
||||
"open_id"
|
||||
};
|
||||
|
||||
// reaction 语义:仅"终态消息"成功发送后才移除 reaction。
|
||||
// 终态 = AssistantResponse(最终响应)或 ErrorNotification(agent 异常终止的错误通知)。
|
||||
// ToolCall 等中间过程事件不触碰 reaction——让 reaction 真实反映用户是否已收到
|
||||
// agent 的最终产出(reaction 持续 = 尚未收到响应或错误;reaction 消失 = 已收到)。
|
||||
// 发送失败时保留 reaction 作为异常信号;空内容不视为送达,保留 reaction。
|
||||
let is_terminal = matches!(
|
||||
msg.event_kind,
|
||||
OutboundEventKind::AssistantResponse | OutboundEventKind::ErrorNotification
|
||||
);
|
||||
let remove_reaction = async {
|
||||
self.remove_reaction_from_metadata(&msg.metadata).await;
|
||||
};
|
||||
@ -2528,14 +2495,7 @@ impl Channel for FeishuChannel {
|
||||
|
||||
// Empty content
|
||||
if content.is_empty() {
|
||||
// 空最终响应是异常:保留 reaction 让用户察觉,仅记录 warn。
|
||||
// 非最终响应(如 ToolCall 空内容)本来就不触碰 reaction。
|
||||
if is_terminal {
|
||||
tracing::warn!(
|
||||
chat_id = %msg.chat_id,
|
||||
"Final response has empty content, keeping reaction as anomaly signal"
|
||||
);
|
||||
}
|
||||
remove_reaction.await;
|
||||
return Ok(());
|
||||
}
|
||||
|
||||
@ -2547,10 +2507,7 @@ impl Channel for FeishuChannel {
|
||||
let result = self
|
||||
.dispatch_send(receive_id, receive_id_type, "text", content, reply_to)
|
||||
.await;
|
||||
// 仅最终响应且发送成功才移除 reaction
|
||||
if is_terminal && result.is_ok() {
|
||||
remove_reaction.await;
|
||||
}
|
||||
remove_reaction.await;
|
||||
return result;
|
||||
}
|
||||
MsgFormat::Post => {
|
||||
@ -2559,9 +2516,7 @@ impl Channel for FeishuChannel {
|
||||
let result = self
|
||||
.dispatch_send(receive_id, receive_id_type, "post", &post_body, reply_to)
|
||||
.await;
|
||||
if is_terminal && result.is_ok() {
|
||||
remove_reaction.await;
|
||||
}
|
||||
remove_reaction.await;
|
||||
return result;
|
||||
}
|
||||
MsgFormat::Interactive => {
|
||||
@ -2587,31 +2542,19 @@ impl Channel for FeishuChannel {
|
||||
reply_to,
|
||||
)
|
||||
.await;
|
||||
// 回退成功才移除 reaction;回退失败则保留作为异常信号
|
||||
if is_terminal && result.is_ok() {
|
||||
remove_reaction.await;
|
||||
}
|
||||
remove_reaction.await;
|
||||
return result;
|
||||
}
|
||||
}
|
||||
// 所有 chunk 成功:仅最终响应移除 reaction
|
||||
if is_terminal {
|
||||
remove_reaction.await;
|
||||
}
|
||||
remove_reaction.await;
|
||||
return Ok(());
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if !msg.content.trim().is_empty() {
|
||||
self.dispatch_send(
|
||||
receive_id,
|
||||
receive_id_type,
|
||||
"text",
|
||||
msg.content.trim(),
|
||||
reply_to,
|
||||
)
|
||||
.await?;
|
||||
self.dispatch_send(receive_id, receive_id_type, "text", msg.content.trim(), reply_to)
|
||||
.await?;
|
||||
}
|
||||
|
||||
let mut sent_media = 0usize;
|
||||
@ -2650,23 +2593,19 @@ impl Channel for FeishuChannel {
|
||||
Ok(()) => sent_media += 1,
|
||||
Err(error) => {
|
||||
tracing::warn!(error = %error, path = %path, media_type = %media_item.media_type, "Failed to send media message to Feishu");
|
||||
// 媒体失败:保留 reaction 作为异常信号(不调用 remove_reaction)
|
||||
return Err(error);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if msg.content.trim().is_empty() && sent_media == 0 {
|
||||
// 无内容无媒体:保留 reaction 作为异常信号
|
||||
remove_reaction.await;
|
||||
return Err(ChannelError::Other(
|
||||
"No supported media items were sent to Feishu".to_string(),
|
||||
));
|
||||
}
|
||||
|
||||
// 全部成功:仅最终响应移除 reaction
|
||||
if is_terminal {
|
||||
remove_reaction.await;
|
||||
}
|
||||
remove_reaction.await;
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
@ -26,7 +26,7 @@ impl ChannelManager {
|
||||
|
||||
Self {
|
||||
channels: Arc::new(RwLock::new(channels)),
|
||||
bus: MessageBus::new(256),
|
||||
bus: MessageBus::new(100),
|
||||
websocket_channel,
|
||||
}
|
||||
}
|
||||
@ -48,9 +48,7 @@ impl ChannelManager {
|
||||
) -> Result<(), ChannelError> {
|
||||
for (name, channel_config) in &config.channels {
|
||||
match channel_config {
|
||||
crate::config::ChannelConfig::Tagged(TaggedChannelConfig::Feishu(
|
||||
feishu_config,
|
||||
))
|
||||
crate::config::ChannelConfig::Tagged(TaggedChannelConfig::Feishu(feishu_config))
|
||||
| crate::config::ChannelConfig::LegacyFeishu(feishu_config) => {
|
||||
if feishu_config.enabled {
|
||||
let channel = FeishuChannel::new(
|
||||
@ -74,9 +72,7 @@ impl ChannelManager {
|
||||
tracing::info!(channel = %name, kind = channel_config.kind(), "Channel disabled in config");
|
||||
}
|
||||
}
|
||||
crate::config::ChannelConfig::Tagged(TaggedChannelConfig::Wechat(
|
||||
wechat_config,
|
||||
)) => {
|
||||
crate::config::ChannelConfig::Tagged(TaggedChannelConfig::Wechat(wechat_config)) => {
|
||||
if wechat_config.enabled {
|
||||
let channel = WechatChannel::new(
|
||||
name.clone(),
|
||||
@ -257,14 +253,8 @@ mod tests {
|
||||
names.sort();
|
||||
|
||||
assert_eq!(names, vec!["backup", "primary", "websocket"]);
|
||||
assert_eq!(
|
||||
manager.get_channel("primary").await.unwrap().name(),
|
||||
"primary"
|
||||
);
|
||||
assert_eq!(
|
||||
manager.get_channel("backup").await.unwrap().name(),
|
||||
"backup"
|
||||
);
|
||||
assert_eq!(manager.get_channel("primary").await.unwrap().name(), "primary");
|
||||
assert_eq!(manager.get_channel("backup").await.unwrap().name(), "backup");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
@ -305,8 +295,7 @@ mod tests {
|
||||
"cred_path": "<CRED_PATH>"
|
||||
}
|
||||
}
|
||||
}"#
|
||||
.replace("<CRED_PATH>", &cred_path_json),
|
||||
}"#.replace("<CRED_PATH>", &cred_path_json),
|
||||
)
|
||||
.unwrap();
|
||||
|
||||
@ -325,9 +314,6 @@ mod tests {
|
||||
names.sort();
|
||||
|
||||
assert_eq!(names, vec!["websocket", "wechat_main"]);
|
||||
assert_eq!(
|
||||
manager.get_channel("wechat_main").await.unwrap().name(),
|
||||
"wechat_main"
|
||||
);
|
||||
assert_eq!(manager.get_channel("wechat_main").await.unwrap().name(), "wechat_main");
|
||||
}
|
||||
}
|
||||
|
||||
@ -13,8 +13,8 @@ use tokio::sync::RwLock;
|
||||
use tokio::task::JoinHandle;
|
||||
use wechatbot::{BotOptions, SendContent, WeChatBot};
|
||||
|
||||
use crate::bus::message::OutboundEventKind;
|
||||
use crate::bus::{InboundMessage, MediaItem, MessageBus, OutboundMessage};
|
||||
use crate::bus::message::OutboundEventKind;
|
||||
use crate::channels::base::{Channel, ChannelError};
|
||||
use crate::config::{LLMProviderConfig, WechatChannelConfig};
|
||||
|
||||
@ -55,29 +55,19 @@ impl WechatChannel {
|
||||
}
|
||||
|
||||
fn sender_allowed(&self, sender_id: &str) -> bool {
|
||||
self.config
|
||||
.allow_from
|
||||
.iter()
|
||||
.any(|pattern| pattern == "*" || pattern == sender_id)
|
||||
self.config.allow_from.iter().any(|pattern| pattern == "*" || pattern == sender_id)
|
||||
}
|
||||
|
||||
async fn media_to_send_content(
|
||||
fn media_to_send_content(
|
||||
media: &MediaItem,
|
||||
caption: Option<String>,
|
||||
) -> Result<SendContent, ChannelError> {
|
||||
// 媒体文件读取是阻塞 IO，放到 blocking 线程池避免阻塞 async worker。
|
||||
let path = media.path.clone();
|
||||
let data = tokio::task::spawn_blocking(move || std::fs::read(&path))
|
||||
.await
|
||||
.map_err(|e| {
|
||||
ChannelError::SendError(format!("WeChat media read task failed: {}", e))
|
||||
})?
|
||||
.map_err(|error| {
|
||||
ChannelError::SendError(format!(
|
||||
"WeChat media read failed for '{}': {}",
|
||||
media.path, error
|
||||
))
|
||||
})?;
|
||||
let data = std::fs::read(&media.path).map_err(|error| {
|
||||
ChannelError::SendError(format!(
|
||||
"WeChat media read failed for '{}': {}",
|
||||
media.path, error
|
||||
))
|
||||
})?;
|
||||
|
||||
if data.is_empty() {
|
||||
return Err(ChannelError::SendError(format!(
|
||||
@ -104,7 +94,7 @@ impl WechatChannel {
|
||||
}
|
||||
|
||||
fn default_media_dir() -> PathBuf {
|
||||
let home = crate::platform::picobot_home_dir();
|
||||
let home = dirs::home_dir().unwrap_or_else(|| PathBuf::from("."));
|
||||
home.join(".picobot").join("media").join("wechat")
|
||||
}
|
||||
|
||||
@ -142,17 +132,14 @@ impl WechatChannel {
|
||||
) -> Result<Vec<MediaItem>, ChannelError> {
|
||||
let Some(downloaded) = bot.download(&msg).await.map_err(|error| {
|
||||
ChannelError::Other(format!("WeChat media download failed: {}", error))
|
||||
})?
|
||||
else {
|
||||
})? else {
|
||||
return Ok(Vec::new());
|
||||
};
|
||||
|
||||
let media_dir = Self::default_media_dir();
|
||||
tokio::fs::create_dir_all(&media_dir)
|
||||
.await
|
||||
.map_err(|error| {
|
||||
ChannelError::Other(format!("Failed to create WeChat media dir: {}", error))
|
||||
})?;
|
||||
.map_err(|error| ChannelError::Other(format!("Failed to create WeChat media dir: {}", error)))?;
|
||||
|
||||
let filename = Self::build_download_filename(
|
||||
&downloaded.media_type,
|
||||
@ -162,9 +149,7 @@ impl WechatChannel {
|
||||
let file_path = media_dir.join(&filename);
|
||||
tokio::fs::write(&file_path, downloaded.data)
|
||||
.await
|
||||
.map_err(|error| {
|
||||
ChannelError::Other(format!("Failed to write WeChat media file: {}", error))
|
||||
})?;
|
||||
.map_err(|error| ChannelError::Other(format!("Failed to write WeChat media file: {}", error)))?;
|
||||
|
||||
tracing::info!(filename = %filename, media_type = %downloaded.media_type, "Downloaded WeChat media");
|
||||
|
||||
@ -331,11 +316,7 @@ impl Channel for WechatChannel {
|
||||
| OutboundEventKind::StreamDelta
|
||||
| OutboundEventKind::StreamEnd
|
||||
| OutboundEventKind::ExecutionCompleted
|
||||
) || msg
|
||||
.metadata
|
||||
.get("is_subagent_event")
|
||||
.map(|v| v == "true")
|
||||
.unwrap_or(false)
|
||||
) || msg.metadata.get("is_subagent_event").map(|v| v == "true").unwrap_or(false)
|
||||
{
|
||||
return Ok(());
|
||||
}
|
||||
@ -362,13 +343,10 @@ impl Channel for WechatChannel {
|
||||
} else {
|
||||
None
|
||||
};
|
||||
let content = Self::media_to_send_content(media, caption).await?;
|
||||
self.bot
|
||||
.send_media(&msg.chat_id, content)
|
||||
.await
|
||||
.map_err(|error| {
|
||||
ChannelError::SendError(format!("WeChat media send failed: {}", error))
|
||||
})?;
|
||||
let content = Self::media_to_send_content(media, caption)?;
|
||||
self.bot.send_media(&msg.chat_id, content).await.map_err(|error| {
|
||||
ChannelError::SendError(format!("WeChat media send failed: {}", error))
|
||||
})?;
|
||||
tracing::info!(
|
||||
channel = %self.name,
|
||||
chat_id = %msg.chat_id,
|
||||
@ -410,33 +388,34 @@ mod tests {
|
||||
assert!(filename.ends_with(".silk"));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn media_to_send_content_maps_image() {
|
||||
#[test]
|
||||
fn media_to_send_content_maps_image() {
|
||||
let file = NamedTempFile::new().unwrap();
|
||||
std::fs::write(file.path(), b"demo-image").unwrap();
|
||||
let image_path = file.path().with_extension("png");
|
||||
std::fs::rename(file.path(), &image_path).unwrap();
|
||||
|
||||
let media = MediaItem::new(image_path.to_string_lossy().to_string(), "image");
|
||||
let content = WechatChannel::media_to_send_content(&media, None).await.unwrap();
|
||||
let content = WechatChannel::media_to_send_content(&media, None).unwrap();
|
||||
|
||||
assert!(matches!(content, SendContent::Image { .. }));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn media_to_send_content_maps_generic_file() {
|
||||
#[test]
|
||||
fn media_to_send_content_maps_generic_file() {
|
||||
let file = NamedTempFile::new().unwrap();
|
||||
std::fs::write(file.path(), b"hello").unwrap();
|
||||
let doc_path = file.path().with_extension("md");
|
||||
std::fs::rename(file.path(), &doc_path).unwrap();
|
||||
|
||||
let media = MediaItem::new(doc_path.to_string_lossy().to_string(), "file");
|
||||
let content =
|
||||
WechatChannel::media_to_send_content(&media, Some("note".to_string())).await.unwrap();
|
||||
let content = WechatChannel::media_to_send_content(&media, Some("note".to_string())).unwrap();
|
||||
|
||||
match content {
|
||||
SendContent::File {
|
||||
file_name, caption, ..
|
||||
file_name,
|
||||
caption,
|
||||
..
|
||||
} => {
|
||||
assert_eq!(file_name, doc_path.file_name().unwrap().to_string_lossy());
|
||||
assert_eq!(caption.as_deref(), Some("note"));
|
||||
@ -444,4 +423,4 @@ mod tests {
|
||||
_ => panic!("expected file send content"),
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
@ -18,7 +18,7 @@ pub struct InitWizard {
|
||||
|
||||
impl InitWizard {
|
||||
pub fn new() -> Self {
|
||||
let home = crate::platform::picobot_home_dir();
|
||||
let home = dirs::home_dir().unwrap_or_else(|| PathBuf::from("."));
|
||||
Self {
|
||||
read: BufReader::new(tokio::io::stdin()),
|
||||
write: tokio::io::stdout(),
|
||||
@ -81,7 +81,6 @@ impl InitWizard {
|
||||
image_context: crate::config::ImageContextConfig::default(),
|
||||
subagents: crate::config::SubagentsConfig::default(),
|
||||
experts: crate::config::ExpertsConfig::default(),
|
||||
compaction: crate::config::CompactionConfig::default(),
|
||||
}
|
||||
}
|
||||
|
||||
@ -118,11 +117,7 @@ impl InitWizard {
|
||||
}
|
||||
|
||||
let input = line.trim().to_string();
|
||||
Ok(if input.is_empty() {
|
||||
default.to_string()
|
||||
} else {
|
||||
input
|
||||
})
|
||||
Ok(if input.is_empty() { default.to_string() } else { input })
|
||||
}
|
||||
|
||||
async fn prompt_required(&mut self, label: &str) -> Result<String, InitError> {
|
||||
@ -159,9 +154,9 @@ impl InitWizard {
|
||||
let default_str = (default + 1).to_string();
|
||||
let input = self.prompt_with_default(label, &default_str).await?;
|
||||
|
||||
let selected: usize = input
|
||||
.parse()
|
||||
.map_err(|_| InitError::InputError(format!("Invalid selection: {}", input)))?;
|
||||
let selected: usize = input.parse().map_err(|_| {
|
||||
InitError::InputError(format!("Invalid selection: {}", input))
|
||||
})?;
|
||||
|
||||
if selected == 0 || selected > options.len() {
|
||||
return Err(InitError::InputError(format!(
|
||||
@ -200,7 +195,9 @@ impl InitWizard {
|
||||
println!(" 4. Skip");
|
||||
println!();
|
||||
|
||||
let choice = self.prompt_with_default("Select option", "1").await?;
|
||||
let choice = self
|
||||
.prompt_with_default("Select option", "1")
|
||||
.await?;
|
||||
|
||||
match choice.as_str() {
|
||||
"1" => return self.add_provider(existing).await,
|
||||
@ -246,7 +243,9 @@ impl InitWizard {
|
||||
&mut self,
|
||||
existing: &Config,
|
||||
) -> Result<HashMap<String, ProviderConfig>, InitError> {
|
||||
let provider_name = self.prompt_with_default("Provider name", "default").await?;
|
||||
let provider_name = self
|
||||
.prompt_with_default("Provider name", "default")
|
||||
.await?;
|
||||
|
||||
println!("Provider type:");
|
||||
println!(" 1. openai");
|
||||
@ -272,7 +271,6 @@ impl InitWizard {
|
||||
extra_headers: HashMap::new(),
|
||||
llm_timeout_secs: 120,
|
||||
memory_maintenance_timeout_secs: 600,
|
||||
max_retries: 3,
|
||||
};
|
||||
|
||||
let mut providers = existing.providers.clone();
|
||||
@ -309,9 +307,7 @@ impl InitWizard {
|
||||
};
|
||||
let type_options = vec!["openai".to_string(), "anthropic".to_string()];
|
||||
println!("Provider type:");
|
||||
let type_idx = self
|
||||
.prompt_select("", &type_options, current_type_idx)
|
||||
.await?;
|
||||
let type_idx = self.prompt_select("", &type_options, current_type_idx).await?;
|
||||
let provider_type = &type_options[type_idx];
|
||||
|
||||
let base_url = self
|
||||
@ -331,7 +327,6 @@ impl InitWizard {
|
||||
extra_headers: current_provider.extra_headers.clone(),
|
||||
llm_timeout_secs: current_provider.llm_timeout_secs,
|
||||
memory_maintenance_timeout_secs: current_provider.memory_maintenance_timeout_secs,
|
||||
max_retries: current_provider.max_retries,
|
||||
};
|
||||
|
||||
let mut providers = existing.providers.clone();
|
||||
@ -541,7 +536,9 @@ impl InitWizard {
|
||||
providers: &HashMap<String, ProviderConfig>,
|
||||
models: &HashMap<String, ModelConfig>,
|
||||
) -> Result<HashMap<String, AgentConfig>, InitError> {
|
||||
let agent_name = self.prompt_with_default("Agent name", "default").await?;
|
||||
let agent_name = self
|
||||
.prompt_with_default("Agent name", "default")
|
||||
.await?;
|
||||
|
||||
// Select provider
|
||||
let provider_names: Vec<String> = providers.keys().cloned().collect();
|
||||
@ -603,9 +600,7 @@ impl InitWizard {
|
||||
.position(|p| p == ¤t_agent.provider)
|
||||
.unwrap_or(0);
|
||||
println!("Select provider:");
|
||||
let provider_idx = self
|
||||
.prompt_select("", &provider_names, current_provider_idx)
|
||||
.await?;
|
||||
let provider_idx = self.prompt_select("", &provider_names, current_provider_idx).await?;
|
||||
let selected_provider = &provider_names[provider_idx];
|
||||
|
||||
// Select new model
|
||||
@ -616,9 +611,7 @@ impl InitWizard {
|
||||
.unwrap_or(0);
|
||||
println!();
|
||||
println!("Select model:");
|
||||
let model_idx = self
|
||||
.prompt_select("", &model_names, current_model_idx)
|
||||
.await?;
|
||||
let model_idx = self.prompt_select("", &model_names, current_model_idx).await?;
|
||||
let selected_model = &model_names[model_idx];
|
||||
|
||||
let agent = AgentConfig {
|
||||
@ -657,11 +650,7 @@ impl InitWizard {
|
||||
if !existing.channels.is_empty() {
|
||||
println!("Existing channels:");
|
||||
for (name, config) in &existing.channels {
|
||||
let status = if config.enabled() {
|
||||
"enabled"
|
||||
} else {
|
||||
"disabled"
|
||||
};
|
||||
let status = if config.enabled() { "enabled" } else { "disabled" };
|
||||
println!(" - {} ({})", name, status);
|
||||
}
|
||||
println!();
|
||||
@ -711,7 +700,9 @@ impl InitWizard {
|
||||
println!("Configuring Feishu channel...");
|
||||
println!();
|
||||
|
||||
let channel_name = self.prompt_with_default("Channel name", "feishu").await?;
|
||||
let channel_name = self
|
||||
.prompt_with_default("Channel name", "feishu")
|
||||
.await?;
|
||||
|
||||
let _existing_config = existing.get(&channel_name).and_then(|c| c.as_feishu());
|
||||
|
||||
@ -775,8 +766,7 @@ impl InitWizard {
|
||||
println!();
|
||||
println!("Starting WeChat login...");
|
||||
|
||||
self.do_wechat_login(base_url, &Self::default_wechat_cred_path())
|
||||
.await?;
|
||||
self.do_wechat_login(base_url, &Self::default_wechat_cred_path()).await?;
|
||||
|
||||
println!();
|
||||
println!("WeChat login successful! Credentials saved.");
|
||||
@ -806,10 +796,9 @@ impl InitWizard {
|
||||
})),
|
||||
});
|
||||
|
||||
let creds = bot
|
||||
.login(true)
|
||||
.await
|
||||
.map_err(|e| InitError::WeChatError(format!("WeChat login failed: {}", e)))?;
|
||||
let creds = bot.login(true).await.map_err(|e| {
|
||||
InitError::WeChatError(format!("WeChat login failed: {}", e))
|
||||
})?;
|
||||
|
||||
println!();
|
||||
println!(
|
||||
@ -846,7 +835,6 @@ impl InitWizard {
|
||||
image_context: existing.image_context.clone(),
|
||||
subagents: existing.subagents.clone(),
|
||||
experts: existing.experts.clone(),
|
||||
compaction: existing.compaction.clone(),
|
||||
}
|
||||
}
|
||||
|
||||
@ -892,14 +880,14 @@ impl InitWizard {
|
||||
}
|
||||
|
||||
fn default_feishu_media_dir() -> String {
|
||||
let home = crate::platform::picobot_home_dir();
|
||||
let home = dirs::home_dir().unwrap_or_else(|| PathBuf::from("."));
|
||||
home.join(".picobot/media/feishu")
|
||||
.to_string_lossy()
|
||||
.to_string()
|
||||
}
|
||||
|
||||
fn default_wechat_cred_path() -> String {
|
||||
let home = crate::platform::picobot_home_dir();
|
||||
let home = dirs::home_dir().unwrap_or_else(|| PathBuf::from("."));
|
||||
home.join(".picobot/wechat/credentials.json")
|
||||
.to_string_lossy()
|
||||
.to_string()
|
||||
@ -937,4 +925,4 @@ impl From<std::io::Error> for InitError {
|
||||
fn from(e: std::io::Error) -> Self {
|
||||
InitError::IoError(e.to_string())
|
||||
}
|
||||
}
|
||||
}
|
||||
@ -1,7 +1,7 @@
|
||||
pub mod channel;
|
||||
pub mod init;
|
||||
pub mod input;
|
||||
pub mod init;
|
||||
|
||||
pub use channel::CliChannel;
|
||||
pub use init::InitWizard;
|
||||
pub use input::{InputCommand, InputEvent, InputHandler};
|
||||
pub use init::InitWizard;
|
||||
|
||||
@ -26,11 +26,8 @@ pub async fn run(gateway_url: &str) -> Result<(), Box<dyn std::error::Error>> {
|
||||
let mut current_session_id: Option<String> = None;
|
||||
// Track message IDs that were already streamed so we can skip
|
||||
// the duplicate AssistantResponse that arrives afterwards.
|
||||
let mut streamed_message_ids: std::collections::HashSet<String> =
|
||||
std::collections::HashSet::new();
|
||||
input
|
||||
.write_output("picobot CLI - Commands: /new [title], /save [filepath], /quit\n")
|
||||
.await?;
|
||||
let mut streamed_message_ids: std::collections::HashSet<String> = std::collections::HashSet::new();
|
||||
input.write_output("picobot CLI - Commands: /new [title], /save [filepath], /quit\n").await?;
|
||||
|
||||
// Main loop: poll both stdin and WebSocket
|
||||
loop {
|
||||
|
||||
@ -1,6 +1,6 @@
|
||||
use crate::command::Command;
|
||||
use crate::command::context::AdapterContext;
|
||||
use crate::command::response::{CommandError, CommandResponse};
|
||||
use crate::command::Command;
|
||||
|
||||
/// 输入适配器：将渠道特定输入转换为 Command
|
||||
///
|
||||
@ -13,7 +13,11 @@ pub trait InputAdapter: Send + Sync {
|
||||
/// - `Ok(Some(Command))`：成功解析为命令
|
||||
/// - `Ok(None)`：不是命令（如普通聊天消息）
|
||||
/// - `Err(CommandError)`：解析错误（如缺少参数）
|
||||
fn try_parse(&self, input: &str, ctx: AdapterContext) -> Result<Option<Command>, AdapterError>;
|
||||
fn try_parse(
|
||||
&self,
|
||||
input: &str,
|
||||
ctx: AdapterContext,
|
||||
) -> Result<Option<Command>, AdapterError>;
|
||||
}
|
||||
|
||||
/// 输出适配器：将 CommandResponse 转换为渠道特定输出
|
||||
|
||||
@ -1,6 +1,6 @@
|
||||
use crate::command::Command;
|
||||
use crate::command::adapter::{AdapterError, InputAdapter};
|
||||
use crate::command::context::AdapterContext;
|
||||
use crate::command::Command;
|
||||
|
||||
/// Channel 输入适配器
|
||||
///
|
||||
|
||||
@ -1,7 +1,7 @@
|
||||
use crate::command::Command;
|
||||
use crate::command::adapter::{AdapterError, InputAdapter, OutputAdapter};
|
||||
use crate::command::context::AdapterContext;
|
||||
use crate::command::response::{CommandResponse, MessageKind};
|
||||
use crate::command::Command;
|
||||
|
||||
/// CLI 输入适配器
|
||||
///
|
||||
@ -313,14 +313,7 @@ mod tests {
|
||||
|
||||
assert!(result.is_some());
|
||||
let cmd = result.unwrap();
|
||||
assert!(matches!(
|
||||
cmd,
|
||||
Command::SaveSession {
|
||||
filepath: None,
|
||||
include_all: false,
|
||||
..
|
||||
}
|
||||
));
|
||||
assert!(matches!(cmd, Command::SaveSession { filepath: None, include_all: false, .. }));
|
||||
}
|
||||
|
||||
#[test]
|
||||
@ -328,9 +321,7 @@ mod tests {
|
||||
let adapter = CliInputAdapter::new();
|
||||
let ctx = AdapterContext::new("test");
|
||||
|
||||
let result = adapter
|
||||
.try_parse("/save-session ./debug/session.md", ctx)
|
||||
.unwrap();
|
||||
let result = adapter.try_parse("/save-session ./debug/session.md", ctx).unwrap();
|
||||
|
||||
assert!(result.is_some());
|
||||
let cmd = result.unwrap();
|
||||
@ -353,14 +344,7 @@ mod tests {
|
||||
|
||||
assert!(result.is_some());
|
||||
let cmd = result.unwrap();
|
||||
assert!(matches!(
|
||||
cmd,
|
||||
Command::SaveSession {
|
||||
filepath: None,
|
||||
include_all: true,
|
||||
..
|
||||
}
|
||||
));
|
||||
assert!(matches!(cmd, Command::SaveSession { filepath: None, include_all: true, .. }));
|
||||
}
|
||||
|
||||
#[test]
|
||||
@ -368,9 +352,7 @@ mod tests {
|
||||
let adapter = CliInputAdapter::new();
|
||||
let ctx = AdapterContext::new("test");
|
||||
|
||||
let result = adapter
|
||||
.try_parse("/save-session all ./debug/session.md", ctx)
|
||||
.unwrap();
|
||||
let result = adapter.try_parse("/save-session all ./debug/session.md", ctx).unwrap();
|
||||
|
||||
assert!(result.is_some());
|
||||
let cmd = result.unwrap();
|
||||
|
||||
@ -1,7 +1,7 @@
|
||||
use crate::command::Command;
|
||||
use crate::command::adapter::{AdapterError, InputAdapter, OutputAdapter};
|
||||
use crate::command::context::AdapterContext;
|
||||
use crate::command::response::{CommandResponse, MessageKind};
|
||||
use crate::command::Command;
|
||||
use crate::protocol::WsOutbound;
|
||||
|
||||
/// WebSocket 输入适配器
|
||||
@ -79,12 +79,8 @@ impl OutputAdapter for WebSocketOutputAdapter {
|
||||
id: response.request_id.to_string(),
|
||||
content: msg.content.clone(),
|
||||
role: "assistant".to_string(),
|
||||
attachments: Vec::new(),
|
||||
subagent_task_id: None,
|
||||
topic_id: None,
|
||||
timestamp: Some(crate::protocol::now_timestamp()),
|
||||
reasoning_content: None,
|
||||
user_message_id: None,
|
||||
attachments: Vec::new(), subagent_task_id: None, topic_id: None, timestamp: Some(crate::protocol::now_timestamp()),
|
||||
reasoning_content: None, user_message_id: None,
|
||||
},
|
||||
MessageKind::Notification => {
|
||||
// 根据元数据判断具体类型
|
||||
@ -94,13 +90,9 @@ impl OutputAdapter for WebSocketOutputAdapter {
|
||||
response.metadata.get("topic_id"),
|
||||
response.metadata.get("title"),
|
||||
) {
|
||||
match serde_json::from_str::<Vec<crate::protocol::TopicSummary>>(
|
||||
topics_json,
|
||||
) {
|
||||
match serde_json::from_str::<Vec<crate::protocol::TopicSummary>>(topics_json) {
|
||||
Ok(topics) => {
|
||||
let session_id = response
|
||||
.metadata
|
||||
.get("session_id")
|
||||
let session_id = response.metadata.get("session_id")
|
||||
.cloned()
|
||||
.unwrap_or_default();
|
||||
WsOutbound::TopicRenamed {
|
||||
@ -114,37 +106,28 @@ impl OutputAdapter for WebSocketOutputAdapter {
|
||||
id: response.request_id.to_string(),
|
||||
content: msg.content.clone(),
|
||||
role: "assistant".to_string(),
|
||||
attachments: Vec::new(),
|
||||
subagent_task_id: None,
|
||||
topic_id: None,
|
||||
timestamp: Some(crate::protocol::now_timestamp()),
|
||||
reasoning_content: None,
|
||||
user_message_id: None,
|
||||
attachments: Vec::new(), subagent_task_id: None, topic_id: None, timestamp: Some(crate::protocol::now_timestamp()),
|
||||
reasoning_content: None, user_message_id: None,
|
||||
},
|
||||
}
|
||||
} else if let Some(topics_json) = response.metadata.get("topics") {
|
||||
// Topic 列表响应 - 优先检查 topics
|
||||
match serde_json::from_str::<Vec<crate::protocol::TopicSummary>>(
|
||||
topics_json,
|
||||
) {
|
||||
match serde_json::from_str::<Vec<crate::protocol::TopicSummary>>(topics_json) {
|
||||
Ok(topics) => {
|
||||
let session_id = response
|
||||
.metadata
|
||||
.get("session_id")
|
||||
let session_id = response.metadata.get("session_id")
|
||||
.cloned()
|
||||
.unwrap_or_default();
|
||||
WsOutbound::TopicList { topics, session_id }
|
||||
WsOutbound::TopicList {
|
||||
topics,
|
||||
session_id,
|
||||
}
|
||||
}
|
||||
Err(_) => WsOutbound::AssistantResponse {
|
||||
id: response.request_id.to_string(),
|
||||
content: msg.content.clone(),
|
||||
role: "assistant".to_string(),
|
||||
attachments: Vec::new(),
|
||||
subagent_task_id: None,
|
||||
topic_id: None,
|
||||
timestamp: Some(crate::protocol::now_timestamp()),
|
||||
reasoning_content: None,
|
||||
user_message_id: None,
|
||||
attachments: Vec::new(), subagent_task_id: None, topic_id: None, timestamp: Some(crate::protocol::now_timestamp()),
|
||||
reasoning_content: None, user_message_id: None,
|
||||
},
|
||||
}
|
||||
} else if let Some(session_id) = response.metadata.get("session_id") {
|
||||
@ -156,9 +139,7 @@ impl OutputAdapter for WebSocketOutputAdapter {
|
||||
}
|
||||
} else {
|
||||
// 加载会话
|
||||
let message_count = response
|
||||
.metadata
|
||||
.get("message_count")
|
||||
let message_count = response.metadata.get("message_count")
|
||||
.and_then(|s| s.parse().ok())
|
||||
.unwrap_or(0);
|
||||
WsOutbound::SessionLoaded {
|
||||
@ -169,9 +150,7 @@ impl OutputAdapter for WebSocketOutputAdapter {
|
||||
}
|
||||
} else if let Some(topic_id) = response.metadata.get("topic_id") {
|
||||
// 只有 topic_id，可能是加载话题
|
||||
let message_count = response
|
||||
.metadata
|
||||
.get("message_count")
|
||||
let message_count = response.metadata.get("message_count")
|
||||
.and_then(|s| s.parse().ok())
|
||||
.unwrap_or(0);
|
||||
WsOutbound::SessionLoaded {
|
||||
@ -187,19 +166,13 @@ impl OutputAdapter for WebSocketOutputAdapter {
|
||||
id: response.request_id.to_string(),
|
||||
content: msg.content.clone(),
|
||||
role: "assistant".to_string(),
|
||||
attachments: Vec::new(),
|
||||
subagent_task_id: None,
|
||||
topic_id: None,
|
||||
timestamp: Some(crate::protocol::now_timestamp()),
|
||||
reasoning_content: None,
|
||||
user_message_id: None,
|
||||
attachments: Vec::new(), subagent_task_id: None, topic_id: None, timestamp: Some(crate::protocol::now_timestamp()),
|
||||
reasoning_content: None, user_message_id: None,
|
||||
},
|
||||
}
|
||||
} else if let Some(sessions_json) = response.metadata.get("sessions") {
|
||||
// 会话列表响应
|
||||
match serde_json::from_str::<Vec<crate::protocol::SessionSummary>>(
|
||||
sessions_json,
|
||||
) {
|
||||
match serde_json::from_str::<Vec<crate::protocol::SessionSummary>>(sessions_json) {
|
||||
Ok(sessions) => {
|
||||
let channel_name = response.metadata.get("channel_name").cloned();
|
||||
WsOutbound::SessionList {
|
||||
@ -212,37 +185,28 @@ impl OutputAdapter for WebSocketOutputAdapter {
|
||||
id: response.request_id.to_string(),
|
||||
content: msg.content.clone(),
|
||||
role: "assistant".to_string(),
|
||||
attachments: Vec::new(),
|
||||
subagent_task_id: None,
|
||||
topic_id: None,
|
||||
timestamp: Some(crate::protocol::now_timestamp()),
|
||||
reasoning_content: None,
|
||||
user_message_id: None,
|
||||
attachments: Vec::new(), subagent_task_id: None, topic_id: None, timestamp: Some(crate::protocol::now_timestamp()),
|
||||
reasoning_content: None, user_message_id: None,
|
||||
},
|
||||
}
|
||||
} else if let Some(topics_json) = response.metadata.get("topics") {
|
||||
// Topic 列表响应
|
||||
match serde_json::from_str::<Vec<crate::protocol::TopicSummary>>(
|
||||
topics_json,
|
||||
) {
|
||||
match serde_json::from_str::<Vec<crate::protocol::TopicSummary>>(topics_json) {
|
||||
Ok(topics) => {
|
||||
let session_id = response
|
||||
.metadata
|
||||
.get("session_id")
|
||||
let session_id = response.metadata.get("session_id")
|
||||
.cloned()
|
||||
.unwrap_or_default();
|
||||
WsOutbound::TopicList { topics, session_id }
|
||||
WsOutbound::TopicList {
|
||||
topics,
|
||||
session_id,
|
||||
}
|
||||
}
|
||||
Err(_) => WsOutbound::AssistantResponse {
|
||||
id: response.request_id.to_string(),
|
||||
content: msg.content.clone(),
|
||||
role: "assistant".to_string(),
|
||||
attachments: Vec::new(),
|
||||
subagent_task_id: None,
|
||||
topic_id: None,
|
||||
timestamp: Some(crate::protocol::now_timestamp()),
|
||||
reasoning_content: None,
|
||||
user_message_id: None,
|
||||
attachments: Vec::new(), subagent_task_id: None, topic_id: None, timestamp: Some(crate::protocol::now_timestamp()),
|
||||
reasoning_content: None, user_message_id: None,
|
||||
},
|
||||
}
|
||||
} else {
|
||||
@ -251,12 +215,8 @@ impl OutputAdapter for WebSocketOutputAdapter {
|
||||
id: response.request_id.to_string(),
|
||||
content: msg.content.clone(),
|
||||
role: "assistant".to_string(),
|
||||
attachments: Vec::new(),
|
||||
subagent_task_id: None,
|
||||
topic_id: None,
|
||||
timestamp: Some(crate::protocol::now_timestamp()),
|
||||
reasoning_content: None,
|
||||
user_message_id: None,
|
||||
attachments: Vec::new(), subagent_task_id: None, topic_id: None, timestamp: Some(crate::protocol::now_timestamp()),
|
||||
reasoning_content: None, user_message_id: None,
|
||||
}
|
||||
}
|
||||
}
|
||||
@ -270,12 +230,8 @@ impl OutputAdapter for WebSocketOutputAdapter {
|
||||
id: response.request_id.to_string(),
|
||||
content: msg.content.clone(),
|
||||
role: "assistant".to_string(),
|
||||
attachments: Vec::new(),
|
||||
subagent_task_id: None,
|
||||
topic_id: None,
|
||||
timestamp: Some(crate::protocol::now_timestamp()),
|
||||
reasoning_content: None,
|
||||
user_message_id: None,
|
||||
attachments: Vec::new(), subagent_task_id: None, topic_id: None, timestamp: Some(crate::protocol::now_timestamp()),
|
||||
reasoning_content: None, user_message_id: None,
|
||||
},
|
||||
};
|
||||
outbounds.push(outbound);
|
||||
|
||||
@ -1,8 +1,8 @@
|
||||
use crate::agent::AgentError;
|
||||
use crate::bus::InboundMessage;
|
||||
use crate::command::Command;
|
||||
use crate::command::context::CommandContext;
|
||||
use crate::command::response::{CommandError, CommandResponse};
|
||||
use crate::command::Command;
|
||||
use crate::agent::AgentError;
|
||||
use crate::gateway::session::SessionManager;
|
||||
use async_trait::async_trait;
|
||||
use std::sync::Arc;
|
||||
@ -78,7 +78,7 @@ pub trait InChatCommandHandler: Send + Sync {
|
||||
/// 负责将命令分发到合适的处理器
|
||||
pub struct CommandRouter {
|
||||
handlers: Vec<Box<dyn CommandHandler>>,
|
||||
metadata: Arc<parking_lot::Mutex<Vec<CommandMetadata>>>,
|
||||
metadata: Arc<std::sync::Mutex<Vec<CommandMetadata>>>,
|
||||
}
|
||||
|
||||
impl CommandRouter {
|
||||
@ -86,7 +86,7 @@ impl CommandRouter {
|
||||
pub fn new() -> Self {
|
||||
Self {
|
||||
handlers: Vec::new(),
|
||||
metadata: Arc::new(parking_lot::Mutex::new(Vec::new())),
|
||||
metadata: Arc::new(std::sync::Mutex::new(Vec::new())),
|
||||
}
|
||||
}
|
||||
|
||||
@ -96,13 +96,13 @@ impl CommandRouter {
|
||||
/// * `handler` - 要注册的处理器
|
||||
pub fn register(&mut self, handler: Box<dyn CommandHandler>) {
|
||||
if let Some(meta) = handler.metadata() {
|
||||
self.metadata.lock().push(meta);
|
||||
self.metadata.lock().unwrap().push(meta);
|
||||
}
|
||||
self.handlers.push(handler);
|
||||
}
|
||||
|
||||
/// 获取已注册命令的元数据列表（用于 Help 命令）
|
||||
pub fn metadata_arc(&self) -> Arc<parking_lot::Mutex<Vec<CommandMetadata>>> {
|
||||
pub fn metadata_arc(&self) -> Arc<std::sync::Mutex<Vec<CommandMetadata>>> {
|
||||
self.metadata.clone()
|
||||
}
|
||||
|
||||
|
||||
@ -1,8 +1,8 @@
|
||||
use crate::command::Command;
|
||||
use crate::command::context::CommandContext;
|
||||
use crate::command::handler::{CommandHandler, CommandMetadata};
|
||||
use crate::command::handlers::list_topics::build_topic_summaries;
|
||||
use crate::command::handlers::list_topics::TopicSummary;
|
||||
use crate::command::response::{CommandError, CommandResponse, MessageKind};
|
||||
use crate::command::Command;
|
||||
use crate::gateway::session::SessionManager;
|
||||
use crate::storage::SessionStore;
|
||||
use async_trait::async_trait;
|
||||
@ -87,10 +87,22 @@ async fn handle_delete_topic(
|
||||
.list_topics(session_id)
|
||||
.map_err(|e| CommandError::new("LIST_TOPICS_ERROR", e.to_string()))?;
|
||||
|
||||
let topic_summaries = build_topic_summaries(handler.store.as_ref(), topics)?;
|
||||
let topic_summaries: Vec<TopicSummary> = topics
|
||||
.into_iter()
|
||||
.map(|t| TopicSummary {
|
||||
topic_id: t.id,
|
||||
session_id: t.session_id,
|
||||
title: t.title,
|
||||
description: t.description.filter(|d| !d.is_empty()),
|
||||
message_count: t.message_count,
|
||||
created_at: t.created_at,
|
||||
last_active_at: t.last_active_at,
|
||||
})
|
||||
.collect();
|
||||
|
||||
let topics_json = serde_json::to_string(&topic_summaries)
|
||||
.map_err(|e| CommandError::new("SERIALIZE_ERROR", e.to_string()))?;
|
||||
let topics_json =
|
||||
serde_json::to_string(&topic_summaries)
|
||||
.map_err(|e| CommandError::new("SERIALIZE_ERROR", e.to_string()))?;
|
||||
|
||||
let message = format!("✓ 已删除话题: {}", topic_title);
|
||||
|
||||
@ -118,7 +130,9 @@ mod tests {
|
||||
|
||||
// 先创建 session 和 topic
|
||||
let session = store.create_session("test_channel", Some("test")).unwrap();
|
||||
let topic = store.create_topic(&session.id, "test topic", None).unwrap();
|
||||
let topic = store
|
||||
.create_topic(&session.id, "test topic", None)
|
||||
.unwrap();
|
||||
|
||||
let ctx = CommandContext::new("test", "test_channel")
|
||||
.with_session_id(&session.id)
|
||||
|
||||
@ -1,9 +1,9 @@
|
||||
use crate::agent::context_compressor::estimate_tokens;
|
||||
use crate::agent::{SystemPromptContext, SystemPromptProvider};
|
||||
use crate::command::Command;
|
||||
use crate::command::context::CommandContext;
|
||||
use crate::command::handler::{CommandHandler, CommandMetadata};
|
||||
use crate::command::response::{CommandError, CommandResponse, MessageKind};
|
||||
use crate::command::Command;
|
||||
use crate::storage::SessionStore;
|
||||
use async_trait::async_trait;
|
||||
use std::sync::Arc;
|
||||
@ -58,28 +58,22 @@ async fn handle_get_current_session(
|
||||
handler: &GetCurrentSessionCommandHandler,
|
||||
ctx: CommandContext,
|
||||
) -> Result<CommandResponse, CommandError> {
|
||||
let topic_id = ctx
|
||||
.topic_id
|
||||
.as_deref()
|
||||
let topic_id = ctx.topic_id.as_deref()
|
||||
.ok_or_else(|| CommandError::new("NO_CURRENT_TOPIC", "No current topic"))?;
|
||||
|
||||
let chat_id = ctx
|
||||
.chat_id
|
||||
.as_deref()
|
||||
let chat_id = ctx.chat_id.as_deref()
|
||||
.ok_or_else(|| CommandError::new("NO_CHAT_ID", "No chat id".to_string()))?;
|
||||
|
||||
let topic = handler
|
||||
.store
|
||||
.get_topic(topic_id)
|
||||
.map_err(|e| CommandError::new("GET_TOPIC_ERROR", e.to_string()))?
|
||||
.ok_or_else(|| {
|
||||
CommandError::new("TOPIC_NOT_FOUND", format!("Topic not found: {}", topic_id))
|
||||
})?;
|
||||
.ok_or_else(|| CommandError::new("TOPIC_NOT_FOUND", format!("Topic not found: {}", topic_id)))?;
|
||||
|
||||
// 直读 DB 按话题加载消息（不依赖 SessionManager 内存，重启后也能正确获取历史）
|
||||
let messages = handler
|
||||
.store
|
||||
.load_messages_for_topic_full(topic_id, Some(&topic.session_id))
|
||||
.load_messages_for_topic(topic_id, Some(&topic.session_id))
|
||||
.map_err(|e| CommandError::new("LOAD_MESSAGES_ERROR", e.to_string()))?;
|
||||
|
||||
let actual_message_count = messages.len();
|
||||
@ -94,8 +88,7 @@ async fn handle_get_current_session(
|
||||
user_message_count,
|
||||
};
|
||||
|
||||
provider
|
||||
.build(&system_prompt_context)
|
||||
provider.build(&system_prompt_context)
|
||||
.map(|sp| {
|
||||
use crate::bus::ChatMessage;
|
||||
let system_msg = ChatMessage::system(&sp.content);
|
||||
@ -162,4 +155,4 @@ fn format_time_ago(timestamp_ms: i64) -> String {
|
||||
} else {
|
||||
format!("{} days ago", diff_secs / 86400)
|
||||
}
|
||||
}
|
||||
}
|
||||
@ -1,10 +1,9 @@
|
||||
use crate::command::Command;
|
||||
use crate::command::context::CommandContext;
|
||||
use crate::command::handler::{CommandHandler, CommandMetadata};
|
||||
use crate::command::response::{CommandError, CommandResponse, MessageKind};
|
||||
use crate::command::Command;
|
||||
use async_trait::async_trait;
|
||||
use parking_lot::Mutex;
|
||||
use std::sync::Arc;
|
||||
use std::sync::{Arc, Mutex};
|
||||
|
||||
/// Help 命令处理器
|
||||
///
|
||||
@ -42,10 +41,11 @@ impl CommandHandler for HelpCommandHandler {
|
||||
_cmd: Command,
|
||||
ctx: CommandContext,
|
||||
) -> Result<CommandResponse, CommandError> {
|
||||
let metadata = self.metadata.lock();
|
||||
let metadata = self.metadata.lock().unwrap();
|
||||
let help_text = format_help(&metadata);
|
||||
|
||||
Ok(CommandResponse::success(ctx.request_id).with_message(MessageKind::Text, &help_text))
|
||||
Ok(CommandResponse::success(ctx.request_id)
|
||||
.with_message(MessageKind::Text, &help_text))
|
||||
}
|
||||
}
|
||||
|
||||
@ -58,4 +58,4 @@ fn format_help(commands: &[CommandMetadata]) -> String {
|
||||
}
|
||||
|
||||
output
|
||||
}
|
||||
}
|
||||
@ -1,8 +1,8 @@
|
||||
use crate::channels::manager::ChannelManager;
|
||||
use crate::command::Command;
|
||||
use crate::command::context::CommandContext;
|
||||
use crate::command::handler::{CommandHandler, CommandMetadata};
|
||||
use crate::command::response::{CommandError, CommandResponse, MessageKind};
|
||||
use crate::command::Command;
|
||||
use async_trait::async_trait;
|
||||
use std::sync::Arc;
|
||||
|
||||
|
||||
@ -1,7 +1,7 @@
|
||||
use crate::command::Command;
|
||||
use crate::command::context::CommandContext;
|
||||
use crate::command::handler::{CommandHandler, CommandMetadata};
|
||||
use crate::command::response::{CommandError, CommandResponse};
|
||||
use crate::command::Command;
|
||||
use crate::storage::SessionStore;
|
||||
use async_trait::async_trait;
|
||||
use serde::{Deserialize, Serialize};
|
||||
@ -68,6 +68,7 @@ impl CommandHandler for ListMemoriesCommandHandler {
|
||||
let memories_json = serde_json::to_string(&summaries)
|
||||
.map_err(|e| CommandError::new("SERIALIZE_ERROR", e.to_string()))?;
|
||||
|
||||
Ok(CommandResponse::success(ctx.request_id).with_metadata("memories", &memories_json))
|
||||
Ok(CommandResponse::success(ctx.request_id)
|
||||
.with_metadata("memories", &memories_json))
|
||||
}
|
||||
}
|
||||
|
||||
@ -1,7 +1,7 @@
|
||||
use crate::command::Command;
|
||||
use crate::command::context::CommandContext;
|
||||
use crate::command::handler::{CommandHandler, CommandMetadata};
|
||||
use crate::command::response::{CommandError, CommandResponse};
|
||||
use crate::command::Command;
|
||||
use crate::protocol::{SchedulerJobSessionLookup, SchedulerJobSummary};
|
||||
use crate::storage::SessionStore;
|
||||
use async_trait::async_trait;
|
||||
@ -67,7 +67,8 @@ impl CommandHandler for ListSchedulerJobsCommandHandler {
|
||||
let jobs_json = serde_json::to_string(&summaries)
|
||||
.map_err(|e| CommandError::new("SERIALIZE_ERROR", e.to_string()))?;
|
||||
|
||||
Ok(CommandResponse::success(ctx.request_id).with_metadata("scheduler_jobs", &jobs_json))
|
||||
Ok(CommandResponse::success(ctx.request_id)
|
||||
.with_metadata("scheduler_jobs", &jobs_json))
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@ -1,7 +1,7 @@
|
||||
use crate::command::Command;
|
||||
use crate::command::context::CommandContext;
|
||||
use crate::command::handler::{CommandHandler, CommandMetadata};
|
||||
use crate::command::response::{CommandError, CommandResponse, MessageKind};
|
||||
use crate::command::Command;
|
||||
use crate::storage::SessionStore;
|
||||
use async_trait::async_trait;
|
||||
use std::sync::Arc;
|
||||
@ -50,9 +50,7 @@ async fn handle_list_sessions(
|
||||
_include_archived: bool,
|
||||
ctx: CommandContext,
|
||||
) -> Result<CommandResponse, CommandError> {
|
||||
let session_id = ctx
|
||||
.session_id
|
||||
.as_deref()
|
||||
let session_id = ctx.session_id.as_deref()
|
||||
.ok_or_else(|| CommandError::new("NO_SESSION", "No active session"))?;
|
||||
|
||||
let topics = handler
|
||||
@ -74,8 +72,7 @@ async fn handle_list_sessions(
|
||||
let marker = if is_current { " *" } else { "" };
|
||||
|
||||
// 使用辅助方法获取消息数量
|
||||
let msg_count = handler
|
||||
.store
|
||||
let msg_count = handler.store
|
||||
.get_topic_message_count(&topic.id)
|
||||
.unwrap_or(0);
|
||||
|
||||
@ -107,4 +104,4 @@ async fn handle_list_sessions(
|
||||
.with_metadata("topics", &topics_json)
|
||||
.with_metadata("count", &topics.len().to_string())
|
||||
.with_metadata("current_topic_id", current_topic_id))
|
||||
}
|
||||
}
|
||||
@ -1,7 +1,7 @@
|
||||
use crate::command::Command;
|
||||
use crate::command::context::CommandContext;
|
||||
use crate::command::handler::{CommandHandler, CommandMetadata};
|
||||
use crate::command::response::{CommandError, CommandResponse, MessageKind};
|
||||
use crate::command::Command;
|
||||
use crate::protocol::SessionSummary;
|
||||
use crate::storage::SessionStore;
|
||||
use async_trait::async_trait;
|
||||
|
||||
@ -1,7 +1,7 @@
|
||||
use crate::command::Command;
|
||||
use crate::command::context::CommandContext;
|
||||
use crate::command::handler::{CommandHandler, CommandMetadata};
|
||||
use crate::command::response::{CommandError, CommandResponse};
|
||||
use crate::command::Command;
|
||||
use crate::skills::SkillRuntime;
|
||||
use async_trait::async_trait;
|
||||
use serde::{Deserialize, Serialize};
|
||||
@ -58,6 +58,7 @@ impl CommandHandler for ListSkillsCommandHandler {
|
||||
let skills_json = serde_json::to_string(&summaries)
|
||||
.map_err(|e| CommandError::new("SERIALIZE_ERROR", e.to_string()))?;
|
||||
|
||||
Ok(CommandResponse::success(ctx.request_id).with_metadata("skills", &skills_json))
|
||||
Ok(CommandResponse::success(ctx.request_id)
|
||||
.with_metadata("skills", &skills_json))
|
||||
}
|
||||
}
|
||||
|
||||
@ -1,7 +1,7 @@
|
||||
use crate::command::Command;
|
||||
use crate::command::context::CommandContext;
|
||||
use crate::command::handler::{CommandHandler, CommandMetadata};
|
||||
use crate::command::response::{CommandError, CommandResponse};
|
||||
use crate::command::Command;
|
||||
use crate::protocol::TodoItemSummary;
|
||||
use crate::storage::SessionStore;
|
||||
use async_trait::async_trait;
|
||||
@ -64,7 +64,7 @@ impl CommandHandler for ListTodosCommandHandler {
|
||||
.list_todos(&scope_key)
|
||||
.map_err(|e| CommandError::new("LIST_TODOS_ERROR", e.to_string()))?;
|
||||
|
||||
tracing::debug!(
|
||||
tracing::info!(
|
||||
scope_key = %scope_key,
|
||||
record_count = records.len(),
|
||||
"list_todos handler: reading from store"
|
||||
|
||||
@ -1,28 +1,12 @@
|
||||
use crate::command::Command;
|
||||
use crate::command::context::CommandContext;
|
||||
use crate::command::handler::{CommandHandler, CommandMetadata};
|
||||
use crate::command::response::{CommandError, CommandResponse, MessageKind};
|
||||
use crate::storage::{SessionStore, SessionTokenStats, TopicRecord};
|
||||
use crate::command::Command;
|
||||
use crate::storage::SessionStore;
|
||||
use async_trait::async_trait;
|
||||
use serde::{Deserialize, Serialize};
|
||||
use std::collections::HashMap;
|
||||
use std::sync::Arc;
|
||||
|
||||
/// Topic 维度的 token 统计（cost 累计 + context 瞬时）。
|
||||
///
|
||||
/// - `prompt_tokens` / `completion_tokens` / `total_tokens`：累计求和（cost 维度）
|
||||
/// - `last_prompt_tokens`：最后一条 assistant 消息的 prompt_tokens（context 占用瞬时值）
|
||||
/// - `context_window_tokens`：当前 session 使用的模型上下文窗口上限（来自配置）；
|
||||
/// 为 0 表示未配置，前端不显示百分比。
|
||||
#[derive(Debug, Clone, Default, Serialize, Deserialize)]
|
||||
pub struct TopicTokenStats {
|
||||
pub prompt_tokens: u64,
|
||||
pub completion_tokens: u64,
|
||||
pub total_tokens: u64,
|
||||
pub last_prompt_tokens: Option<u32>,
|
||||
pub context_window_tokens: u32,
|
||||
}
|
||||
|
||||
/// Topic 摘要信息
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct TopicSummary {
|
||||
@ -33,65 +17,6 @@ pub struct TopicSummary {
|
||||
pub message_count: i64,
|
||||
pub created_at: i64,
|
||||
pub last_active_at: i64,
|
||||
/// Token 用量统计。老 topic 或无 LLM 调用时为 None，前端不显示 token 标签。
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub token_stats: Option<TopicTokenStats>,
|
||||
}
|
||||
|
||||
/// 构建 TopicSummary 列表的公共函数。
|
||||
///
|
||||
/// 一次批量查询所有 topic 的 token 统计，按 `topic_id` 聚合，
|
||||
/// 避免按 session_id 聚合时同 session 下多个 topic 共享同一总和。
|
||||
/// 子代理天然隔离：子代理消息的 topic_id 属于子代理自身的 topic，
|
||||
/// 不在主 topic 列表中。
|
||||
///
|
||||
/// `context_window_tokens` 来自最新 assistant 消息记录（LLM 调用时持久化），
|
||||
/// 无需从配置链路注入，保持存储层与配置解耦。
|
||||
pub fn build_topic_summaries(
|
||||
store: &SessionStore,
|
||||
topics: Vec<TopicRecord>,
|
||||
) -> Result<Vec<TopicSummary>, CommandError> {
|
||||
if topics.is_empty() {
|
||||
return Ok(Vec::new());
|
||||
}
|
||||
|
||||
// 收集所有 topic_id（去重）
|
||||
let mut topic_ids: Vec<String> = Vec::new();
|
||||
for t in &topics {
|
||||
if !topic_ids.contains(&t.id) {
|
||||
topic_ids.push(t.id.clone());
|
||||
}
|
||||
}
|
||||
let topic_id_refs: Vec<&str> = topic_ids.iter().map(|s| s.as_str()).collect();
|
||||
|
||||
// 一次批量查询 token 统计（按 topic_id 聚合）
|
||||
let stats_map: HashMap<String, SessionTokenStats> = store
|
||||
.batch_topic_token_stats(&topic_id_refs)
|
||||
.map_err(|e| CommandError::new("TOKEN_STATS_ERROR", e.to_string()))?;
|
||||
|
||||
let summaries = topics
|
||||
.into_iter()
|
||||
.map(|t| {
|
||||
let token_stats = stats_map.get(&t.id).map(|s| TopicTokenStats {
|
||||
prompt_tokens: s.prompt_tokens,
|
||||
completion_tokens: s.completion_tokens,
|
||||
total_tokens: s.total_tokens,
|
||||
last_prompt_tokens: s.last_prompt_tokens,
|
||||
context_window_tokens: s.context_window_tokens.unwrap_or(0),
|
||||
});
|
||||
TopicSummary {
|
||||
topic_id: t.id,
|
||||
session_id: t.session_id,
|
||||
title: t.title,
|
||||
description: t.description.filter(|d| !d.is_empty()),
|
||||
message_count: t.message_count,
|
||||
created_at: t.created_at,
|
||||
last_active_at: t.last_active_at,
|
||||
token_stats,
|
||||
}
|
||||
})
|
||||
.collect();
|
||||
Ok(summaries)
|
||||
}
|
||||
|
||||
/// 列出 Session 的 Topics 命令处理器
|
||||
@ -125,7 +50,9 @@ impl CommandHandler for ListTopicsCommandHandler {
|
||||
ctx: CommandContext,
|
||||
) -> Result<CommandResponse, CommandError> {
|
||||
match cmd {
|
||||
Command::ListTopics { session_id } => handle_list_topics(self, session_id, ctx).await,
|
||||
Command::ListTopics { session_id } => {
|
||||
handle_list_topics(self, session_id, ctx).await
|
||||
}
|
||||
_ => unreachable!(),
|
||||
}
|
||||
}
|
||||
@ -141,7 +68,18 @@ async fn handle_list_topics(
|
||||
.list_topics(&session_id)
|
||||
.map_err(|e| CommandError::new("LIST_TOPICS_ERROR", e.to_string()))?;
|
||||
|
||||
let summaries = build_topic_summaries(handler.store.as_ref(), topics)?;
|
||||
let summaries: Vec<TopicSummary> = topics
|
||||
.into_iter()
|
||||
.map(|t| TopicSummary {
|
||||
topic_id: t.id,
|
||||
session_id: t.session_id,
|
||||
title: t.title,
|
||||
description: t.description.filter(|d| !d.is_empty()),
|
||||
message_count: t.message_count,
|
||||
created_at: t.created_at,
|
||||
last_active_at: t.last_active_at,
|
||||
})
|
||||
.collect();
|
||||
|
||||
let topics_json = serde_json::to_string(&summaries)
|
||||
.map_err(|e| CommandError::new("SERIALIZE_ERROR", e.to_string()))?;
|
||||
|
||||
@ -1,7 +1,7 @@
|
||||
use crate::command::Command;
|
||||
use crate::command::context::CommandContext;
|
||||
use crate::command::handler::{CommandHandler, CommandMetadata};
|
||||
use crate::command::response::{CommandError, CommandResponse};
|
||||
use crate::command::Command;
|
||||
use async_trait::async_trait;
|
||||
|
||||
/// 加载指定 channel + chat_id 的对话消息。
|
||||
|
||||
@ -1,8 +1,7 @@
|
||||
use crate::command::Command;
|
||||
use crate::command::context::CommandContext;
|
||||
use crate::command::handler::{CommandHandler, CommandMetadata};
|
||||
use crate::command::handlers::list_topics::TopicTokenStats;
|
||||
use crate::command::response::{CommandError, CommandResponse};
|
||||
use crate::command::Command;
|
||||
use crate::storage::SessionStore;
|
||||
use crate::tools::task::repository::TaskRepository;
|
||||
use crate::tools::task::types::{TaskSession, TaskSessionState};
|
||||
@ -15,11 +14,11 @@ pub struct LoadTaskMessagesCommandHandler {
|
||||
}
|
||||
|
||||
impl LoadTaskMessagesCommandHandler {
|
||||
pub fn new(task_repository: Arc<dyn TaskRepository>, store: Arc<SessionStore>) -> Self {
|
||||
Self {
|
||||
task_repository,
|
||||
store,
|
||||
}
|
||||
pub fn new(
|
||||
task_repository: Arc<dyn TaskRepository>,
|
||||
store: Arc<SessionStore>,
|
||||
) -> Self {
|
||||
Self { task_repository, store }
|
||||
}
|
||||
}
|
||||
|
||||
@ -63,7 +62,11 @@ async fn handle_load_task_messages(
|
||||
);
|
||||
|
||||
// 1. Try in-memory repository first
|
||||
let task = match handler.task_repository.load_task_session(&task_id).await {
|
||||
let task = match handler
|
||||
.task_repository
|
||||
.load_task_session(&task_id)
|
||||
.await
|
||||
{
|
||||
Ok(Some(task)) => {
|
||||
tracing::info!(
|
||||
task_id = %task.id,
|
||||
@ -101,19 +104,6 @@ async fn handle_load_task_messages(
|
||||
|
||||
let status = format!("{:?}", task.state).to_lowercase();
|
||||
|
||||
// 查询子代理 session 的 token 统计（按 session_id 精确匹配，不过滤 sub:%）
|
||||
let token_stats = handler
|
||||
.store
|
||||
.get_session_token_stats(&task.session_id)
|
||||
.map_err(|e| CommandError::new("TOKEN_STATS_ERROR", e.to_string()))?
|
||||
.map(|s| TopicTokenStats {
|
||||
prompt_tokens: s.prompt_tokens,
|
||||
completion_tokens: s.completion_tokens,
|
||||
total_tokens: s.total_tokens,
|
||||
last_prompt_tokens: s.last_prompt_tokens,
|
||||
context_window_tokens: s.context_window_tokens.unwrap_or(0),
|
||||
});
|
||||
|
||||
let mut response = CommandResponse::success(ctx.request_id)
|
||||
.with_metadata("task_session_id", &task.session_id)
|
||||
.with_metadata("task_id", &task.id)
|
||||
@ -125,12 +115,6 @@ async fn handle_load_task_messages(
|
||||
response = response.with_metadata("task_summary", summary);
|
||||
}
|
||||
|
||||
if let Some(ref stats) = token_stats {
|
||||
let stats_json = serde_json::to_string(stats)
|
||||
.map_err(|e| CommandError::new("SERIALIZE_ERROR", e.to_string()))?;
|
||||
response = response.with_metadata("task_token_stats", &stats_json);
|
||||
}
|
||||
|
||||
Ok(response)
|
||||
}
|
||||
|
||||
@ -188,7 +172,6 @@ fn reconstruct_task_from_db(
|
||||
updated_at: now,
|
||||
summary: None,
|
||||
error: None,
|
||||
tool_call_id: None,
|
||||
}))
|
||||
}
|
||||
|
||||
@ -203,9 +186,6 @@ fn parse_subagent_title(title: &str) -> (String, String) {
|
||||
return (agent_type, desc);
|
||||
}
|
||||
}
|
||||
let desc = title
|
||||
.strip_prefix("Subagent: ")
|
||||
.unwrap_or(title)
|
||||
.to_string();
|
||||
let desc = title.strip_prefix("Subagent: ").unwrap_or(title).to_string();
|
||||
("general".to_string(), desc)
|
||||
}
|
||||
|
||||
@ -1,7 +1,7 @@
|
||||
use crate::command::Command;
|
||||
use crate::command::context::CommandContext;
|
||||
use crate::command::handler::{CommandHandler, CommandMetadata};
|
||||
use crate::command::response::{CommandError, CommandResponse, MessageKind};
|
||||
use crate::command::Command;
|
||||
use crate::storage::SessionStore;
|
||||
use async_trait::async_trait;
|
||||
use std::sync::Arc;
|
||||
@ -37,7 +37,9 @@ impl CommandHandler for LoadTopicCommandHandler {
|
||||
ctx: CommandContext,
|
||||
) -> Result<CommandResponse, CommandError> {
|
||||
match cmd {
|
||||
Command::LoadTopic { topic_id } => handle_load_topic(self, topic_id, ctx).await,
|
||||
Command::LoadTopic { topic_id } => {
|
||||
handle_load_topic(self, topic_id, ctx).await
|
||||
}
|
||||
_ => unreachable!(),
|
||||
}
|
||||
}
|
||||
@ -52,9 +54,7 @@ async fn handle_load_topic(
|
||||
.store
|
||||
.get_topic(&topic_id)
|
||||
.map_err(|e| CommandError::new("LOAD_TOPIC_ERROR", e.to_string()))?
|
||||
.ok_or_else(|| {
|
||||
CommandError::new("TOPIC_NOT_FOUND", format!("Topic not found: {}", topic_id))
|
||||
})?;
|
||||
.ok_or_else(|| CommandError::new("TOPIC_NOT_FOUND", format!("Topic not found: {}", topic_id)))?;
|
||||
|
||||
Ok(CommandResponse::success(ctx.request_id)
|
||||
.with_message(MessageKind::Notification, &topic.title)
|
||||
|
||||
@ -1,8 +1,8 @@
|
||||
use crate::command::Command;
|
||||
use crate::command::context::CommandContext;
|
||||
use crate::command::handler::{CommandHandler, CommandMetadata};
|
||||
use crate::command::response::{CommandError, CommandResponse};
|
||||
use crate::storage::{GLOBAL_SCOPE_KEY, MemoryUpsert, SessionStore};
|
||||
use crate::command::Command;
|
||||
use crate::storage::{MemoryUpsert, SessionStore, GLOBAL_SCOPE_KEY};
|
||||
use async_trait::async_trait;
|
||||
use std::sync::Arc;
|
||||
|
||||
@ -17,7 +17,10 @@ impl MemoryCrudCommandHandler {
|
||||
}
|
||||
|
||||
/// 通过 ID 查找记忆的 namespace 和 memory_key
|
||||
fn find_by_id(store: &SessionStore, id: &str) -> Result<Option<(String, String)>, CommandError> {
|
||||
fn find_by_id(
|
||||
store: &SessionStore,
|
||||
id: &str,
|
||||
) -> Result<Option<(String, String)>, CommandError> {
|
||||
let records = store
|
||||
.list_memories_for_scope("user", GLOBAL_SCOPE_KEY)
|
||||
.map_err(|e| CommandError::new("LIST_ERROR", e.to_string()))?;
|
||||
@ -32,9 +35,7 @@ impl CommandHandler for MemoryCrudCommandHandler {
|
||||
fn can_handle(&self, cmd: &Command) -> bool {
|
||||
matches!(
|
||||
cmd,
|
||||
Command::CreateMemory { .. }
|
||||
| Command::UpdateMemory { .. }
|
||||
| Command::DeleteMemory { .. }
|
||||
Command::CreateMemory { .. } | Command::UpdateMemory { .. } | Command::DeleteMemory { .. }
|
||||
)
|
||||
}
|
||||
|
||||
@ -111,6 +112,7 @@ impl CommandHandler for MemoryCrudCommandHandler {
|
||||
_ => unreachable!(),
|
||||
}
|
||||
|
||||
Ok(CommandResponse::success(ctx.request_id).with_metadata("memory_updated", "true"))
|
||||
Ok(CommandResponse::success(ctx.request_id)
|
||||
.with_metadata("memory_updated", "true"))
|
||||
}
|
||||
}
|
||||
|
||||
@ -4,15 +4,15 @@ pub mod help;
|
||||
pub mod list_channels;
|
||||
pub mod list_memories;
|
||||
pub mod list_scheduler_jobs;
|
||||
pub mod list_sessions;
|
||||
pub mod list_sessions_by_channel;
|
||||
pub mod list_skills;
|
||||
pub mod list_todos;
|
||||
pub mod memory_crud;
|
||||
pub mod list_sessions;
|
||||
pub mod list_sessions_by_channel;
|
||||
pub mod list_topics;
|
||||
pub mod load_chat_messages;
|
||||
pub mod load_task_messages;
|
||||
pub mod load_topic;
|
||||
pub mod memory_crud;
|
||||
pub mod rename_topic;
|
||||
pub mod save_session;
|
||||
pub mod save_topic;
|
||||
@ -22,7 +22,8 @@ pub mod switch_topic;
|
||||
|
||||
// 导出公共函数供其他模块复用
|
||||
pub use save_session::{
|
||||
SubagentTaskData, escape_yaml_string, format_message_content, format_timestamp,
|
||||
generate_messages_markdown, generate_subagent_tasks_markdown, generate_system_prompt_markdown,
|
||||
load_subagent_data,
|
||||
escape_yaml_string, format_message_content, format_timestamp,
|
||||
generate_messages_markdown, generate_system_prompt_markdown,
|
||||
generate_subagent_tasks_markdown, load_subagent_data, SubagentTaskData,
|
||||
};
|
||||
|
||||
|
||||
@ -1,8 +1,8 @@
|
||||
use crate::command::Command;
|
||||
use crate::command::context::CommandContext;
|
||||
use crate::command::handler::{CommandHandler, CommandMetadata};
|
||||
use crate::command::handlers::list_topics::build_topic_summaries;
|
||||
use crate::command::handlers::list_topics::TopicSummary;
|
||||
use crate::command::response::{CommandError, CommandResponse, MessageKind};
|
||||
use crate::command::Command;
|
||||
use crate::storage::SessionStore;
|
||||
use async_trait::async_trait;
|
||||
use std::sync::Arc;
|
||||
@ -75,38 +75,28 @@ async fn handle_rename_topic(
|
||||
CommandError::new("TOPIC_NOT_FOUND", format!("Topic not found: {}", topic_id))
|
||||
})?;
|
||||
|
||||
let old_display = topic
|
||||
.description
|
||||
.as_deref()
|
||||
.filter(|s| !s.is_empty())
|
||||
.unwrap_or(&topic.title)
|
||||
.to_string();
|
||||
let old_title = topic.title.clone();
|
||||
|
||||
// 标题未变化时直接返回当前列表，避免无意义写入
|
||||
if old_display == trimmed_title {
|
||||
if old_title == trimmed_title {
|
||||
let topics = handler
|
||||
.store
|
||||
.list_topics(session_id)
|
||||
.map_err(|e| CommandError::new("LIST_TOPICS_ERROR", e.to_string()))?;
|
||||
let topic_summaries = build_topic_summaries(handler.store.as_ref(), topics)?;
|
||||
let topic_summaries_json = serde_json::to_string(&topic_summaries)
|
||||
.map_err(|e| CommandError::new("SERIALIZE_ERROR", e.to_string()))?;
|
||||
let topic_summaries = serialize_summaries(&topics);
|
||||
|
||||
return Ok(CommandResponse::success(ctx.request_id)
|
||||
.with_message(
|
||||
MessageKind::Notification,
|
||||
&format!("✓ 话题标题未变化: {}", trimmed_title),
|
||||
)
|
||||
.with_metadata("topics", &topic_summaries_json)
|
||||
.with_message(MessageKind::Notification, &format!("✓ 话题标题未变化: {}", trimmed_title))
|
||||
.with_metadata("topics", &topic_summaries)
|
||||
.with_metadata("topic_id", &topic_id)
|
||||
.with_metadata("title", trimmed_title)
|
||||
.with_metadata("session_id", session_id));
|
||||
}
|
||||
|
||||
// 执行重命名：更新 description（显示字段），保留 title 作为内部标识
|
||||
// 执行重命名（存储层方法已存在）
|
||||
handler
|
||||
.store
|
||||
.update_topic_description(&topic_id, trimmed_title)
|
||||
.update_topic_title(&topic_id, trimmed_title)
|
||||
.map_err(|e| CommandError::new("RENAME_TOPIC_ERROR", e.to_string()))?;
|
||||
|
||||
// 查询更新后的话题列表，返回给前端刷新侧边栏
|
||||
@ -115,20 +105,34 @@ async fn handle_rename_topic(
|
||||
.list_topics(session_id)
|
||||
.map_err(|e| CommandError::new("LIST_TOPICS_ERROR", e.to_string()))?;
|
||||
|
||||
let topic_summaries = build_topic_summaries(handler.store.as_ref(), topics)?;
|
||||
let topic_summaries_json = serde_json::to_string(&topic_summaries)
|
||||
.map_err(|e| CommandError::new("SERIALIZE_ERROR", e.to_string()))?;
|
||||
let topic_summaries = serialize_summaries(&topics);
|
||||
|
||||
let message = format!("✓ 已重命名话题: {} → {}", old_display, trimmed_title);
|
||||
let message = format!("✓ 已重命名话题: {} → {}", old_title, trimmed_title);
|
||||
|
||||
Ok(CommandResponse::success(ctx.request_id)
|
||||
.with_message(MessageKind::Notification, &message)
|
||||
.with_metadata("topics", &topic_summaries_json)
|
||||
.with_metadata("topics", &topic_summaries)
|
||||
.with_metadata("topic_id", &topic_id)
|
||||
.with_metadata("title", trimmed_title)
|
||||
.with_metadata("session_id", session_id))
|
||||
}
|
||||
|
||||
fn serialize_summaries(topics: &[crate::storage::TopicRecord]) -> String {
|
||||
let summaries: Vec<TopicSummary> = topics
|
||||
.iter()
|
||||
.map(|t| TopicSummary {
|
||||
topic_id: t.id.clone(),
|
||||
session_id: t.session_id.clone(),
|
||||
title: t.title.clone(),
|
||||
description: t.description.clone().filter(|d| !d.is_empty()),
|
||||
message_count: t.message_count,
|
||||
created_at: t.created_at,
|
||||
last_active_at: t.last_active_at,
|
||||
})
|
||||
.collect();
|
||||
serde_json::to_string(&summaries).unwrap_or_else(|_| "[]".to_string())
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
@ -145,7 +149,9 @@ mod tests {
|
||||
let store = handler.store.clone();
|
||||
|
||||
let session = store.create_session("test_channel", Some("test")).unwrap();
|
||||
let topic = store.create_topic(&session.id, "old title", None).unwrap();
|
||||
let topic = store
|
||||
.create_topic(&session.id, "old title", None)
|
||||
.unwrap();
|
||||
|
||||
let ctx = CommandContext::new("test", "test_channel")
|
||||
.with_session_id(&session.id)
|
||||
@ -160,20 +166,13 @@ mod tests {
|
||||
|
||||
let resp = result.unwrap();
|
||||
assert!(resp.success);
|
||||
assert_eq!(
|
||||
resp.metadata.get("title").map(String::as_str),
|
||||
Some("new title")
|
||||
);
|
||||
assert_eq!(
|
||||
resp.metadata.get("topic_id").map(String::as_str),
|
||||
Some(topic.id.as_str())
|
||||
);
|
||||
assert_eq!(resp.metadata.get("title").map(String::as_str), Some("new title"));
|
||||
assert_eq!(resp.metadata.get("topic_id").map(String::as_str), Some(topic.id.as_str()));
|
||||
assert!(resp.metadata.contains_key("topics"));
|
||||
|
||||
// 验证存储层已更新 description（显示字段），title 保持原内部标识
|
||||
// 验证存储层已更新
|
||||
let updated = store.get_topic(&topic.id).unwrap().unwrap();
|
||||
assert_eq!(updated.title, "old title");
|
||||
assert_eq!(updated.description.as_deref(), Some("new title"));
|
||||
assert_eq!(updated.title, "new title");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
@ -182,7 +181,9 @@ mod tests {
|
||||
let store = handler.store.clone();
|
||||
|
||||
let session = store.create_session("test_channel", Some("test")).unwrap();
|
||||
let topic = store.create_topic(&session.id, "old title", None).unwrap();
|
||||
let topic = store
|
||||
.create_topic(&session.id, "old title", None)
|
||||
.unwrap();
|
||||
|
||||
let ctx = CommandContext::new("test", "test_channel")
|
||||
.with_session_id(&session.id)
|
||||
@ -228,8 +229,9 @@ mod tests {
|
||||
let store = handler.store.clone();
|
||||
|
||||
let session = store.create_session("test_channel", Some("test")).unwrap();
|
||||
// description=None，当前显示值 fallback 到 title
|
||||
let topic = store.create_topic(&session.id, "same title", None).unwrap();
|
||||
let topic = store
|
||||
.create_topic(&session.id, "same title", None)
|
||||
.unwrap();
|
||||
let original_updated_at = store.get_topic(&topic.id).unwrap().unwrap().updated_at;
|
||||
|
||||
// 等待一秒确保 updated_at 会变化（如果真的写入）
|
||||
@ -251,38 +253,6 @@ mod tests {
|
||||
assert_eq!(after.updated_at, original_updated_at);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_rename_topic_compares_against_description_when_present() {
|
||||
// 有 description 时，"未变化"比较应基于 description 而非 title
|
||||
let handler = create_test_handler();
|
||||
let store = handler.store.clone();
|
||||
|
||||
let session = store.create_session("test_channel", Some("test")).unwrap();
|
||||
// title 是内部标识，description 是显示值
|
||||
let topic = store
|
||||
.create_topic(&session.id, "topic_internal_id", Some("AI 生成的摘要"))
|
||||
.unwrap();
|
||||
let original_updated_at = store.get_topic(&topic.id).unwrap().unwrap().updated_at;
|
||||
|
||||
tokio::time::sleep(std::time::Duration::from_secs(1)).await;
|
||||
|
||||
let ctx = CommandContext::new("test", "test_channel")
|
||||
.with_session_id(&session.id)
|
||||
.with_chat_id(&session.id);
|
||||
// 提交与 description 相同的值（不是 title），应跳过写入
|
||||
let cmd = Command::RenameTopic {
|
||||
topic_id: topic.id.clone(),
|
||||
title: "AI 生成的摘要".to_string(),
|
||||
};
|
||||
|
||||
let result = handler.handle(cmd, ctx).await;
|
||||
assert!(result.is_ok());
|
||||
|
||||
let after = store.get_topic(&topic.id).unwrap().unwrap();
|
||||
assert_eq!(after.updated_at, original_updated_at);
|
||||
assert_eq!(after.title, "topic_internal_id");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_can_handle() {
|
||||
let handler = create_test_handler();
|
||||
|
||||
@ -1,12 +1,12 @@
|
||||
use crate::agent::AgentError;
|
||||
use crate::agent::{SystemPrompt, SystemPromptContext, SystemPromptProvider};
|
||||
use crate::bus::InboundMessage;
|
||||
use crate::command::Command;
|
||||
use crate::command::context::CommandContext;
|
||||
use crate::command::handler::{CommandHandler, CommandMetadata, InChatCommandHandler};
|
||||
use crate::command::response::{CommandError, CommandResponse, MessageKind};
|
||||
use crate::command::Command;
|
||||
use crate::storage::{SessionRecord, SessionStore};
|
||||
use crate::tools::task::repository::TaskRepository;
|
||||
use crate::agent::AgentError;
|
||||
use async_trait::async_trait;
|
||||
use chrono::{Local, TimeZone};
|
||||
use std::path::PathBuf;
|
||||
@ -65,8 +65,7 @@ pub async fn save_session_to_file(
|
||||
let system_prompt = build_system_prompt(system_prompt_provider, &record, user_message_count);
|
||||
|
||||
// 生成 Markdown 内容
|
||||
let markdown =
|
||||
generate_markdown_with_subagents(&record, &system_prompt, &messages, &subagent_data);
|
||||
let markdown = generate_markdown_with_subagents(&record, &system_prompt, &messages, &subagent_data);
|
||||
|
||||
// 确定输出路径
|
||||
let output_path = resolve_filepath(filepath, &record);
|
||||
@ -80,7 +79,8 @@ pub async fn save_session_to_file(
|
||||
}
|
||||
|
||||
// 写入文件
|
||||
std::fs::write(&output_path, markdown).map_err(|e| format!("Failed to write file: {}", e))?;
|
||||
std::fs::write(&output_path, markdown)
|
||||
.map_err(|e| format!("Failed to write file: {}", e))?;
|
||||
|
||||
Ok(output_path)
|
||||
}
|
||||
@ -134,11 +134,9 @@ impl CommandHandler for SaveSessionCommandHandler {
|
||||
ctx: CommandContext,
|
||||
) -> Result<CommandResponse, CommandError> {
|
||||
match cmd {
|
||||
Command::SaveSession {
|
||||
filepath,
|
||||
include_all,
|
||||
include_subagents,
|
||||
} => handle_save_session(self, filepath, include_all, include_subagents, ctx).await,
|
||||
Command::SaveSession { filepath, include_all, include_subagents } => {
|
||||
handle_save_session(self, filepath, include_all, include_subagents, ctx).await
|
||||
}
|
||||
_ => unreachable!(),
|
||||
}
|
||||
}
|
||||
@ -201,9 +199,13 @@ async fn handle_save_session(
|
||||
|
||||
// 根据 include_all 获取消息数量
|
||||
let message_count = if include_all {
|
||||
handler.store.load_all_messages(session_id)
|
||||
handler
|
||||
.store
|
||||
.load_all_messages(session_id)
|
||||
} else {
|
||||
handler.store.load_messages(session_id)
|
||||
handler
|
||||
.store
|
||||
.load_messages(session_id)
|
||||
}
|
||||
.map_err(|e| CommandError::new("LOAD_MESSAGES_ERROR", e.to_string()))?
|
||||
.len();
|
||||
@ -213,15 +215,9 @@ async fn handle_save_session(
|
||||
MessageKind::Notification,
|
||||
// 路径中的反斜杠在 Markdown 渲染时会被当作转义符吃掉，
|
||||
// 统一转换为正斜杠以保证显示完整（跨平台兼容）
|
||||
&format!(
|
||||
"Session saved to: {}",
|
||||
output_path.display().to_string().replace('\\', "/")
|
||||
),
|
||||
)
|
||||
.with_metadata(
|
||||
"filepath",
|
||||
&output_path.display().to_string().replace('\\', "/"),
|
||||
&format!("Session saved to: {}", output_path.display().to_string().replace('\\', "/")),
|
||||
)
|
||||
.with_metadata("filepath", &output_path.display().to_string().replace('\\', "/"))
|
||||
.with_metadata("message_count", &message_count.to_string()))
|
||||
}
|
||||
|
||||
@ -351,18 +347,12 @@ pub fn generate_subagent_tasks_markdown(subagent_data: &[SubagentTaskData]) -> S
|
||||
output.push_str("# Subagent Tasks\n\n");
|
||||
|
||||
for task in subagent_data {
|
||||
output.push_str(&format!(
|
||||
"## Task: {} ({})",
|
||||
task.description, task.subagent_type
|
||||
));
|
||||
output.push_str(&format!("## Task: {} ({})", task.description, task.subagent_type));
|
||||
output.push('\n');
|
||||
output.push_str(&format!("**Task ID:** `{}`\n\n", task.task_id));
|
||||
output.push_str(&format!("**Session ID:** `{}`\n\n", task.session_id));
|
||||
output.push_str(&format!("**Status:** {}\n\n", task.state));
|
||||
output.push_str(&format!(
|
||||
"**Created:** {}\n\n",
|
||||
format_timestamp(task.created_at)
|
||||
));
|
||||
output.push_str(&format!("**Created:** {}\n\n", format_timestamp(task.created_at)));
|
||||
output.push_str(&format!("**Message Count:** {}\n\n", task.messages.len()));
|
||||
|
||||
// 子智能体消息
|
||||
@ -371,10 +361,7 @@ pub fn generate_subagent_tasks_markdown(subagent_data: &[SubagentTaskData]) -> S
|
||||
for (idx, msg) in task.messages.iter().enumerate() {
|
||||
output.push_str(&format!("#### Message {}\n\n", idx + 1));
|
||||
output.push_str(&format!("**Role:** {}\n\n", msg.role));
|
||||
output.push_str(&format!(
|
||||
"**Time:** {}\n\n",
|
||||
format_timestamp(msg.timestamp)
|
||||
));
|
||||
output.push_str(&format!("**Time:** {}\n\n", format_timestamp(msg.timestamp)));
|
||||
|
||||
if let Some(ref reasoning) = msg.reasoning_content {
|
||||
output.push_str("**Reasoning:**\n");
|
||||
@ -689,12 +676,7 @@ impl InChatCommandHandler for SaveSessionInChatHandler {
|
||||
inbound: &InboundMessage,
|
||||
session_manager: &crate::gateway::session::SessionManager,
|
||||
) -> Result<Option<String>, AgentError> {
|
||||
let Command::SaveSession {
|
||||
filepath,
|
||||
include_all,
|
||||
include_subagents,
|
||||
} = cmd
|
||||
else {
|
||||
let Command::SaveSession { filepath, include_all, include_subagents } = cmd else {
|
||||
return Ok(None);
|
||||
};
|
||||
|
||||
@ -725,10 +707,7 @@ impl InChatCommandHandler for SaveSessionInChatHandler {
|
||||
// 返回成功或失败消息
|
||||
match result {
|
||||
Ok(output_path) => {
|
||||
let msg = format!(
|
||||
"Session saved to: {}",
|
||||
output_path.display().to_string().replace('\\', "/")
|
||||
);
|
||||
let msg = format!("Session saved to: {}", output_path.display().to_string().replace('\\', "/"));
|
||||
tracing::info!("{}", msg);
|
||||
Ok(Some(msg))
|
||||
}
|
||||
@ -795,10 +774,7 @@ mod tests {
|
||||
fn test_escape_yaml_string() {
|
||||
assert_eq!(escape_yaml_string("simple"), "simple");
|
||||
assert_eq!(escape_yaml_string("with: colon"), "\"with: colon\"");
|
||||
assert_eq!(
|
||||
escape_yaml_string("with \"quote\""),
|
||||
"\"with \\\"quote\\\"\""
|
||||
);
|
||||
assert_eq!(escape_yaml_string("with \"quote\""), "\"with \\\"quote\\\"\"");
|
||||
}
|
||||
|
||||
#[test]
|
||||
@ -859,26 +835,14 @@ mod tests {
|
||||
#[test]
|
||||
fn test_can_handle() {
|
||||
let store = Arc::new(SessionStore::in_memory().unwrap());
|
||||
let task_repository =
|
||||
Arc::new(crate::tools::task::repository::InMemoryTaskRepository::new());
|
||||
let task_repository = Arc::new(crate::tools::task::repository::InMemoryTaskRepository::new());
|
||||
let provider = Arc::new(TestSystemPromptProvider);
|
||||
let handler = SaveSessionCommandHandler::new(store, task_repository, provider);
|
||||
|
||||
assert!(handler.can_handle(&Command::SaveSession {
|
||||
filepath: None,
|
||||
include_all: false,
|
||||
include_subagents: false
|
||||
}));
|
||||
assert!(handler.can_handle(&Command::SaveSession {
|
||||
filepath: None,
|
||||
include_all: true,
|
||||
include_subagents: false
|
||||
}));
|
||||
assert!(handler.can_handle(&Command::SaveSession { filepath: None, include_all: false, include_subagents: false }));
|
||||
assert!(handler.can_handle(&Command::SaveSession { filepath: None, include_all: true, include_subagents: false }));
|
||||
assert!(!handler.can_handle(&Command::CreateSession { title: None }));
|
||||
assert!(!handler.can_handle(&Command::SaveTopic {
|
||||
filepath: None,
|
||||
include_subagents: false
|
||||
}));
|
||||
assert!(!handler.can_handle(&Command::SaveTopic { filepath: None, include_subagents: false }));
|
||||
}
|
||||
|
||||
/// 测试用的系统提示词提供者
|
||||
|
||||
@ -1,13 +1,14 @@
|
||||
use crate::agent::{SystemPrompt, SystemPromptContext, SystemPromptProvider};
|
||||
use crate::bus::ChatMessage;
|
||||
use crate::command::Command;
|
||||
use crate::command::context::CommandContext;
|
||||
use crate::command::handler::{CommandHandler, CommandMetadata};
|
||||
use crate::command::handlers::{
|
||||
SubagentTaskData, escape_yaml_string, format_timestamp, generate_messages_markdown,
|
||||
generate_subagent_tasks_markdown, generate_system_prompt_markdown, load_subagent_data,
|
||||
escape_yaml_string, format_timestamp, generate_messages_markdown,
|
||||
generate_subagent_tasks_markdown, generate_system_prompt_markdown,
|
||||
load_subagent_data, SubagentTaskData,
|
||||
};
|
||||
use crate::command::response::{CommandError, CommandResponse, MessageKind};
|
||||
use crate::command::Command;
|
||||
use crate::storage::{SessionStore, TopicRecord};
|
||||
use crate::tools::task::repository::TaskRepository;
|
||||
use async_trait::async_trait;
|
||||
@ -62,7 +63,8 @@ pub async fn save_topic_to_file(
|
||||
}
|
||||
|
||||
// 写入文件
|
||||
std::fs::write(&output_path, markdown).map_err(|e| format!("Failed to write file: {}", e))?;
|
||||
std::fs::write(&output_path, markdown)
|
||||
.map_err(|e| format!("Failed to write file: {}", e))?;
|
||||
|
||||
Ok(output_path)
|
||||
}
|
||||
@ -208,10 +210,9 @@ impl CommandHandler for SaveTopicCommandHandler {
|
||||
ctx: CommandContext,
|
||||
) -> Result<CommandResponse, CommandError> {
|
||||
match cmd {
|
||||
Command::SaveTopic {
|
||||
filepath,
|
||||
include_subagents,
|
||||
} => handle_save_topic(self, filepath, include_subagents, ctx).await,
|
||||
Command::SaveTopic { filepath, include_subagents } => {
|
||||
handle_save_topic(self, filepath, include_subagents, ctx).await
|
||||
}
|
||||
_ => unreachable!(),
|
||||
}
|
||||
}
|
||||
@ -248,19 +249,14 @@ async fn handle_save_topic(
|
||||
.store
|
||||
.get_topic(topic_id)
|
||||
.map_err(|e| CommandError::new("GET_TOPIC_ERROR", e.to_string()))?
|
||||
.ok_or_else(|| {
|
||||
CommandError::new("TOPIC_NOT_FOUND", format!("Topic not found: {}", topic_id))
|
||||
})?;
|
||||
.ok_or_else(|| CommandError::new("TOPIC_NOT_FOUND", format!("Topic not found: {}", topic_id)))?;
|
||||
|
||||
let messages = handler
|
||||
.store
|
||||
.load_messages_for_topic_full(topic_id, Some(&topic_record.session_id))
|
||||
.load_messages_for_topic(topic_id, Some(&topic_record.session_id))
|
||||
.map_err(|e| CommandError::new("LOAD_MESSAGES_ERROR", e.to_string()))?;
|
||||
|
||||
tracing::debug!(
|
||||
message_count = messages.len(),
|
||||
"Loaded messages from DB for topic"
|
||||
);
|
||||
tracing::debug!(message_count = messages.len(), "Loaded messages from DB for topic");
|
||||
|
||||
// 调用保存函数
|
||||
let output_path = save_topic_to_file(
|
||||
@ -282,14 +278,8 @@ async fn handle_save_topic(
|
||||
MessageKind::Notification,
|
||||
// 路径中的反斜杠在 Markdown 渲染时会被当作转义符吃掉，
|
||||
// 统一转换为正斜杠以保证显示完整（跨平台兼容）
|
||||
&format!(
|
||||
"Topic saved to: {}",
|
||||
output_path.display().to_string().replace('\\', "/")
|
||||
),
|
||||
)
|
||||
.with_metadata(
|
||||
"filepath",
|
||||
&output_path.display().to_string().replace('\\', "/"),
|
||||
&format!("Topic saved to: {}", output_path.display().to_string().replace('\\', "/")),
|
||||
)
|
||||
.with_metadata("filepath", &output_path.display().to_string().replace('\\', "/"))
|
||||
.with_metadata("message_count", &message_count.to_string()))
|
||||
}
|
||||
}
|
||||
@ -1,8 +1,8 @@
|
||||
use crate::command::Command;
|
||||
use crate::command::context::CommandContext;
|
||||
use crate::command::handler::{CommandHandler, CommandMetadata};
|
||||
use crate::command::handlers::list_topics::build_topic_summaries;
|
||||
use crate::command::handlers::list_topics::TopicSummary;
|
||||
use crate::command::response::{CommandError, CommandResponse, MessageKind};
|
||||
use crate::command::Command;
|
||||
use crate::gateway::session::SessionManager;
|
||||
use crate::storage::SessionStore;
|
||||
use async_trait::async_trait;
|
||||
@ -56,9 +56,7 @@ impl CommandHandler for SessionCommandHandler {
|
||||
) -> Result<CommandResponse, CommandError> {
|
||||
match cmd {
|
||||
Command::CreateSession { title } => handle_create_session(self, title, ctx).await,
|
||||
Command::SaveSession { .. } => {
|
||||
unreachable!("SaveSession should be handled by SaveSessionCommandHandler")
|
||||
}
|
||||
Command::SaveSession { .. } => unreachable!("SaveSession should be handled by SaveSessionCommandHandler"),
|
||||
_ => unreachable!("Other commands should be handled by other handlers"),
|
||||
}
|
||||
}
|
||||
@ -71,16 +69,13 @@ async fn handle_create_session(
|
||||
ctx: CommandContext,
|
||||
) -> Result<CommandResponse, CommandError> {
|
||||
// 获取当前 session_id，如果没有则报错
|
||||
let session_id = ctx.session_id.as_deref().ok_or_else(|| {
|
||||
CommandError::new(
|
||||
"NO_SESSION",
|
||||
"No active session. Please ensure a session exists first.",
|
||||
)
|
||||
})?;
|
||||
let session_id = ctx.session_id.as_deref()
|
||||
.ok_or_else(|| CommandError::new("NO_SESSION", "No active session. Please ensure a session exists first."))?;
|
||||
|
||||
// 创建新话题（在同一个 Session 内）
|
||||
let topic_title =
|
||||
title.unwrap_or_else(|| format!("Topic {}", &uuid::Uuid::new_v4().to_string()[..8]));
|
||||
let topic_title = title.unwrap_or_else(|| {
|
||||
format!("Topic {}", &uuid::Uuid::new_v4().to_string()[..8])
|
||||
});
|
||||
|
||||
let topic = handler
|
||||
.store
|
||||
@ -88,17 +83,14 @@ async fn handle_create_session(
|
||||
.map_err(|e| CommandError::new("CREATE_TOPIC_ERROR", e.to_string()))?;
|
||||
|
||||
// 获取 chat_id
|
||||
let chat_id = ctx
|
||||
.chat_id
|
||||
.as_deref()
|
||||
let chat_id = ctx.chat_id.as_deref()
|
||||
.ok_or_else(|| CommandError::new("NO_CHAT_ID", "No chat_id in context"))?;
|
||||
|
||||
// 如果有 SessionManager，自动切换到新话题
|
||||
if let Some(ref session_manager) = handler.session_manager {
|
||||
if let Some(session) = session_manager.get(&ctx.channel_name).await {
|
||||
let mut session_guard = session.lock().await;
|
||||
session_guard
|
||||
.switch_topic(chat_id, &topic.id)
|
||||
session_guard.switch_topic(chat_id, &topic.id)
|
||||
.map_err(|e| CommandError::new("SWITCH_TOPIC_ERROR", e.to_string()))?;
|
||||
}
|
||||
}
|
||||
@ -109,7 +101,18 @@ async fn handle_create_session(
|
||||
.list_topics(session_id)
|
||||
.map_err(|e| CommandError::new("LIST_TOPICS_ERROR", e.to_string()))?;
|
||||
|
||||
let topic_summaries = build_topic_summaries(handler.store.as_ref(), topics)?;
|
||||
let topic_summaries: Vec<TopicSummary> = topics
|
||||
.into_iter()
|
||||
.map(|t| TopicSummary {
|
||||
topic_id: t.id,
|
||||
session_id: t.session_id,
|
||||
title: t.title,
|
||||
description: t.description.filter(|d| !d.is_empty()),
|
||||
message_count: t.message_count,
|
||||
created_at: t.created_at,
|
||||
last_active_at: t.last_active_at,
|
||||
})
|
||||
.collect();
|
||||
|
||||
let topics_json = serde_json::to_string(&topic_summaries)
|
||||
.map_err(|e| CommandError::new("SERIALIZE_ERROR", e.to_string()))?;
|
||||
|
||||
@ -1,9 +1,9 @@
|
||||
use async_trait::async_trait;
|
||||
|
||||
use crate::command::Command;
|
||||
use crate::command::context::CommandContext;
|
||||
use crate::command::handler::{CommandHandler, CommandMetadata};
|
||||
use crate::command::response::{CommandError, CommandResponse, MessageKind};
|
||||
use crate::command::Command;
|
||||
use crate::gateway::cancel_manager::CancelManager;
|
||||
use crate::gateway::session::SessionManager;
|
||||
|
||||
@ -15,10 +15,7 @@ pub struct StopExecutionCommandHandler {
|
||||
|
||||
impl StopExecutionCommandHandler {
|
||||
pub fn new(cancel_manager: CancelManager, session_manager: SessionManager) -> Self {
|
||||
Self {
|
||||
cancel_manager,
|
||||
session_manager,
|
||||
}
|
||||
Self { cancel_manager, session_manager }
|
||||
}
|
||||
}
|
||||
|
||||
@ -56,11 +53,7 @@ impl CommandHandler for StopExecutionCommandHandler {
|
||||
None => {
|
||||
// 从 SessionManager 获取真实的 current topic
|
||||
let chat_id = ctx.chat_id.as_deref().unwrap_or("");
|
||||
match self
|
||||
.session_manager
|
||||
.get_current_topic(&ctx.channel_name, chat_id)
|
||||
.await
|
||||
{
|
||||
match self.session_manager.get_current_topic(&ctx.channel_name, chat_id).await {
|
||||
Ok(Some(id)) => {
|
||||
tracing::info!(
|
||||
channel = %ctx.channel_name,
|
||||
@ -72,16 +65,12 @@ impl CommandHandler for StopExecutionCommandHandler {
|
||||
id
|
||||
}
|
||||
Ok(None) => {
|
||||
return Ok(CommandResponse::success(ctx.request_id).with_message(
|
||||
MessageKind::Notification,
|
||||
"当前没有活跃的话题，无法停止",
|
||||
));
|
||||
return Ok(CommandResponse::success(ctx.request_id)
|
||||
.with_message(MessageKind::Notification, "当前没有活跃的话题，无法停止"));
|
||||
}
|
||||
Err(e) => {
|
||||
return Ok(CommandResponse::error(
|
||||
ctx.request_id,
|
||||
CommandError::new("QUERY_TOPIC_ERROR", e.to_string()),
|
||||
));
|
||||
return Ok(CommandResponse::error(ctx.request_id,
|
||||
CommandError::new("QUERY_TOPIC_ERROR", e.to_string())));
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@ -1,7 +1,7 @@
|
||||
use crate::command::Command;
|
||||
use crate::command::context::CommandContext;
|
||||
use crate::command::handler::{CommandHandler, CommandMetadata};
|
||||
use crate::command::response::{CommandError, CommandResponse, MessageKind};
|
||||
use crate::command::Command;
|
||||
use crate::gateway::session::SessionManager;
|
||||
use crate::storage::SessionStore;
|
||||
use async_trait::async_trait;
|
||||
@ -15,10 +15,7 @@ pub struct SwitchTopicCommandHandler {
|
||||
|
||||
impl SwitchTopicCommandHandler {
|
||||
pub fn new(store: Arc<SessionStore>) -> Self {
|
||||
Self {
|
||||
store,
|
||||
session_manager: None,
|
||||
}
|
||||
Self { store, session_manager: None }
|
||||
}
|
||||
|
||||
pub fn with_session_manager(mut self, session_manager: SessionManager) -> Self {
|
||||
@ -47,7 +44,9 @@ impl CommandHandler for SwitchTopicCommandHandler {
|
||||
ctx: CommandContext,
|
||||
) -> Result<CommandResponse, CommandError> {
|
||||
match cmd {
|
||||
Command::SwitchTopic { topic_id } => handle_switch_topic(self, topic_id, ctx).await,
|
||||
Command::SwitchTopic { topic_id } => {
|
||||
handle_switch_topic(self, topic_id, ctx).await
|
||||
}
|
||||
_ => unreachable!(),
|
||||
}
|
||||
}
|
||||
@ -58,13 +57,9 @@ async fn handle_switch_topic(
|
||||
topic_id: String,
|
||||
ctx: CommandContext,
|
||||
) -> Result<CommandResponse, CommandError> {
|
||||
let session_id = ctx
|
||||
.session_id
|
||||
.as_deref()
|
||||
let session_id = ctx.session_id.as_deref()
|
||||
.ok_or_else(|| CommandError::new("NO_SESSION", "No active session"))?;
|
||||
let chat_id = ctx
|
||||
.chat_id
|
||||
.as_deref()
|
||||
let chat_id = ctx.chat_id.as_deref()
|
||||
.ok_or_else(|| CommandError::new("NO_CHAT_ID", "No chat_id in context"))?;
|
||||
|
||||
// 尝试解析为序号
|
||||
@ -78,11 +73,7 @@ async fn handle_switch_topic(
|
||||
if index >= topics.len() {
|
||||
return Err(CommandError::new(
|
||||
"INVALID_TOPIC_INDEX",
|
||||
format!(
|
||||
"Topic index {} is out of range (1-{})",
|
||||
index + 1,
|
||||
topics.len()
|
||||
),
|
||||
format!("Topic index {} is out of range (1-{})", index + 1, topics.len())
|
||||
));
|
||||
}
|
||||
topics[index].id.clone()
|
||||
@ -95,26 +86,19 @@ async fn handle_switch_topic(
|
||||
.store
|
||||
.get_topic(&target_topic_id)
|
||||
.map_err(|e| CommandError::new("SWITCH_TOPIC_ERROR", e.to_string()))?
|
||||
.ok_or_else(|| {
|
||||
CommandError::new(
|
||||
"TOPIC_NOT_FOUND",
|
||||
format!("Topic not found: {}", target_topic_id),
|
||||
)
|
||||
})?;
|
||||
.ok_or_else(|| CommandError::new("TOPIC_NOT_FOUND", format!("Topic not found: {}", target_topic_id)))?;
|
||||
|
||||
// 如果有 SessionManager，实际切换话题历史
|
||||
if let Some(ref session_manager) = handler.session_manager {
|
||||
if let Some(session) = session_manager.get(&ctx.channel_name).await {
|
||||
let mut session_guard = session.lock().await;
|
||||
session_guard
|
||||
.switch_topic(chat_id, &target_topic_id)
|
||||
session_guard.switch_topic(chat_id, &target_topic_id)
|
||||
.map_err(|e| CommandError::new("SWITCH_TOPIC_ERROR", e.to_string()))?;
|
||||
}
|
||||
}
|
||||
|
||||
// 使用辅助方法获取消息数量
|
||||
let msg_count = handler
|
||||
.store
|
||||
let msg_count = handler.store
|
||||
.get_topic_message_count(&target_topic_id)
|
||||
.unwrap_or(0);
|
||||
|
||||
|
||||
@ -48,7 +48,10 @@ pub enum Command {
|
||||
/// 列出所有定时任务
|
||||
ListSchedulerJobs,
|
||||
/// 加载指定 channel + chat_id 的对话消息
|
||||
LoadChatMessages { channel: String, chat_id: String },
|
||||
LoadChatMessages {
|
||||
channel: String,
|
||||
chat_id: String,
|
||||
},
|
||||
/// 删除指定话题
|
||||
DeleteTopic { topic_id: String },
|
||||
/// 重命名指定话题
|
||||
@ -64,7 +67,10 @@ pub enum Command {
|
||||
content: String,
|
||||
},
|
||||
/// 更新已有记忆
|
||||
UpdateMemory { id: String, content: String },
|
||||
UpdateMemory {
|
||||
id: String,
|
||||
content: String,
|
||||
},
|
||||
/// 删除记忆
|
||||
DeleteMemory { id: String },
|
||||
/// 列出所有技能
|
||||
|
||||
@ -40,8 +40,6 @@ pub struct Config {
|
||||
pub subagents: SubagentsConfig,
|
||||
#[serde(default)]
|
||||
pub experts: ExpertsConfig,
|
||||
#[serde(default)]
|
||||
pub compaction: CompactionConfig,
|
||||
}
|
||||
|
||||
/// 图片上下文限制配置
|
||||
@ -74,54 +72,6 @@ impl Default for ImageContextConfig {
|
||||
}
|
||||
}
|
||||
|
||||
/// 上下文压缩算法配置（全局，所有 agent 共享）
|
||||
#[derive(Debug, Clone, Deserialize, Serialize)]
|
||||
pub struct CompactionConfig {
|
||||
/// 工程化压缩触发阈值（占 context_window 的比例，0.0-1.0）
|
||||
/// 当 LLM 返回的 prompt_tokens 超过 context_window × 此比例时，触发压缩流程
|
||||
#[serde(default = "default_threshold_ratio")]
|
||||
pub threshold_ratio: f64,
|
||||
|
||||
/// LLM 压缩触发阈值（占 context_window 的比例，0.0-1.0）
|
||||
/// 工程化压缩（截断 tool 结果）后，若估算 token 仍超过此比例才调 LLM 压缩
|
||||
#[serde(default = "default_llm_compaction_threshold_ratio")]
|
||||
pub llm_compaction_threshold_ratio: f64,
|
||||
|
||||
/// 工程化压缩时 tool 结果截断到的 token 数
|
||||
/// 子代理返回（tool_name="task"）不受此限制
|
||||
#[serde(default = "default_truncate_max_tokens")]
|
||||
pub truncate_max_tokens: usize,
|
||||
|
||||
/// LLM 三段压缩保留的最旧/最新 unit 数
|
||||
/// 压缩后保留最旧 N 个 unit + 中间段摘要 + 最新 N 个 unit
|
||||
#[serde(default = "default_preserve_count")]
|
||||
pub preserve_count: usize,
|
||||
}
|
||||
|
||||
fn default_threshold_ratio() -> f64 {
|
||||
0.5
|
||||
}
|
||||
fn default_llm_compaction_threshold_ratio() -> f64 {
|
||||
0.3
|
||||
}
|
||||
fn default_truncate_max_tokens() -> usize {
|
||||
100
|
||||
}
|
||||
fn default_preserve_count() -> usize {
|
||||
5
|
||||
}
|
||||
|
||||
impl Default for CompactionConfig {
|
||||
fn default() -> Self {
|
||||
Self {
|
||||
threshold_ratio: default_threshold_ratio(),
|
||||
llm_compaction_threshold_ratio: default_llm_compaction_threshold_ratio(),
|
||||
truncate_max_tokens: default_truncate_max_tokens(),
|
||||
preserve_count: default_preserve_count(),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Deserialize, Serialize)]
|
||||
pub struct TimeConfig {
|
||||
#[serde(default = "default_timezone")]
|
||||
@ -311,7 +261,7 @@ fn default_task_enabled() -> bool {
|
||||
}
|
||||
|
||||
fn default_task_max_execution_secs() -> u64 {
|
||||
3600 // 60分钟
|
||||
3600 // 60分钟
|
||||
}
|
||||
|
||||
fn default_task_ttl_hours() -> u64 {
|
||||
@ -457,7 +407,7 @@ fn default_allow_from() -> Vec<String> {
|
||||
}
|
||||
|
||||
fn default_media_dir() -> String {
|
||||
let home = crate::platform::picobot_home_dir();
|
||||
let home = dirs::home_dir().unwrap_or_else(|| std::path::PathBuf::from("."));
|
||||
home.join(".picobot/media/feishu")
|
||||
.to_string_lossy()
|
||||
.to_string()
|
||||
@ -471,7 +421,7 @@ fn default_wechat_base_url() -> String {
|
||||
}
|
||||
|
||||
fn default_wechat_cred_path() -> String {
|
||||
let home = crate::platform::picobot_home_dir();
|
||||
let home = dirs::home_dir().unwrap_or_else(|| std::path::PathBuf::from("."));
|
||||
home.join(".picobot/wechat/credentials.json")
|
||||
.to_string_lossy()
|
||||
.to_string()
|
||||
@ -501,10 +451,6 @@ pub struct ProviderConfig {
|
||||
pub llm_timeout_secs: u64,
|
||||
#[serde(default = "default_memory_maintenance_timeout_secs")]
|
||||
pub memory_maintenance_timeout_secs: u64,
|
||||
/// LLM 请求瞬态失败（timeout/502/503/504/429 等）的最大重试次数。
|
||||
/// 0 表示不重试。默认 3。
|
||||
#[serde(default = "default_max_retries")]
|
||||
pub max_retries: u32,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Deserialize, Serialize)]
|
||||
@ -552,10 +498,6 @@ fn default_memory_maintenance_timeout_secs() -> u64 {
|
||||
600
|
||||
}
|
||||
|
||||
fn default_max_retries() -> u32 {
|
||||
3
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Deserialize, Serialize)]
|
||||
pub struct GatewayConfig {
|
||||
#[serde(default = "default_gateway_host")]
|
||||
@ -573,11 +515,6 @@ pub struct GatewayConfig {
|
||||
pub max_concurrent_requests: usize,
|
||||
#[serde(default, rename = "session_ttl_hours")]
|
||||
pub session_ttl_hours: Option<u64>,
|
||||
/// 网关认证 token。当绑定到非 loopback 地址时必须配置，否则启动会报错。
|
||||
/// 绑定到 loopback 时可不配置（本地访问免认证）。
|
||||
/// 前端通过 Authorization: Bearer <token> 头或 WS query param ?token=<token> 携带。
|
||||
#[serde(default, rename = "auth_token")]
|
||||
pub auth_token: Option<String>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Deserialize, Serialize)]
|
||||
@ -887,7 +824,6 @@ impl Default for GatewayConfig {
|
||||
agent_prompt_reinject_every: default_agent_prompt_reinject_every(),
|
||||
max_concurrent_requests: default_max_concurrent_requests(),
|
||||
session_ttl_hours: Some(24),
|
||||
auth_token: None,
|
||||
}
|
||||
}
|
||||
}
|
||||
@ -921,8 +857,6 @@ pub struct LLMProviderConfig {
|
||||
pub extra_headers: HashMap<String, String>,
|
||||
pub llm_timeout_secs: u64,
|
||||
pub memory_maintenance_timeout_secs: u64,
|
||||
/// LLM 请求瞬态失败的最大重试次数（透传自 ProviderConfig）。
|
||||
pub max_retries: u32,
|
||||
pub model_id: String,
|
||||
pub temperature: Option<f32>,
|
||||
pub max_tokens: Option<u32>,
|
||||
@ -955,7 +889,7 @@ impl LLMProviderConfig {
|
||||
}
|
||||
|
||||
pub(crate) fn get_default_config_path() -> PathBuf {
|
||||
let home = crate::platform::picobot_home_dir();
|
||||
let home = dirs::home_dir().unwrap_or_else(|| PathBuf::from("."));
|
||||
home.join(".picobot").join("config.json")
|
||||
}
|
||||
|
||||
@ -1070,7 +1004,6 @@ impl Config {
|
||||
extra_headers: provider.extra_headers.clone(),
|
||||
llm_timeout_secs: provider.llm_timeout_secs,
|
||||
memory_maintenance_timeout_secs: provider.memory_maintenance_timeout_secs,
|
||||
max_retries: provider.max_retries,
|
||||
model_id: model.model_id.clone(),
|
||||
temperature: model.temperature,
|
||||
max_tokens: model.max_tokens,
|
||||
@ -1128,10 +1061,7 @@ pub struct ModelResolver {
|
||||
}
|
||||
|
||||
impl ModelResolver {
|
||||
pub fn new(
|
||||
providers: HashMap<String, ProviderConfig>,
|
||||
models: HashMap<String, ModelConfig>,
|
||||
) -> Self {
|
||||
pub fn new(providers: HashMap<String, ProviderConfig>, models: HashMap<String, ModelConfig>) -> Self {
|
||||
Self { providers, models }
|
||||
}
|
||||
|
||||
@ -1165,7 +1095,6 @@ impl ModelResolver {
|
||||
result.extra_headers = provider.extra_headers.clone();
|
||||
result.llm_timeout_secs = provider.llm_timeout_secs;
|
||||
result.memory_maintenance_timeout_secs = provider.memory_maintenance_timeout_secs;
|
||||
result.max_retries = provider.max_retries;
|
||||
}
|
||||
|
||||
if let Some(name) = model_name.map(str::trim).filter(|s| !s.is_empty()) {
|
||||
@ -1222,17 +1151,16 @@ fn resolve_env_placeholders(content: &str) -> String {
|
||||
let re_braces = Regex::new(r"\$\{([A-Z_][A-Z0-9_]*)\}").expect("invalid regex");
|
||||
let re_angle = Regex::new(r"<([A-Z_]+)>").expect("invalid regex");
|
||||
|
||||
let content = re_braces.replace_all(content, |caps: ®ex::Captures<'_>| {
|
||||
let content = re_braces.replace_all(content, |caps: ®ex::Captures| {
|
||||
let var_name = &caps[1];
|
||||
env::var(var_name).unwrap_or_else(|_| caps[0].to_string())
|
||||
});
|
||||
|
||||
re_angle
|
||||
.replace_all(&content, |caps: ®ex::Captures<'_>| {
|
||||
let var_name = &caps[1];
|
||||
env::var(var_name).unwrap_or_else(|_| caps[0].to_string())
|
||||
})
|
||||
.to_string()
|
||||
re_angle.replace_all(&content, |caps: ®ex::Captures| {
|
||||
let var_name = &caps[1];
|
||||
env::var(var_name).unwrap_or_else(|_| caps[0].to_string())
|
||||
})
|
||||
.to_string()
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
@ -1810,8 +1738,7 @@ mod tests {
|
||||
"allow_from": ["wxid_1"]
|
||||
}
|
||||
}
|
||||
}"#
|
||||
.replace("<CRED_PATH>", &cred_path_json),
|
||||
}"#.replace("<CRED_PATH>", &cred_path_json),
|
||||
)
|
||||
.unwrap();
|
||||
|
||||
@ -1931,7 +1858,7 @@ mod tests {
|
||||
timezone: "Asia/Shanghai".to_string(),
|
||||
});
|
||||
assert_eq!(effective_jobs.len(), 3); // 2个内置 + 1个自定义
|
||||
// 第一个作业：内存维护（被覆盖为禁用）
|
||||
// 第一个作业：内存维护（被覆盖为禁用）
|
||||
assert_eq!(effective_jobs[0].id, BUILTIN_MEMORY_MAINTENANCE_JOB_ID);
|
||||
assert!(!effective_jobs[0].enabled);
|
||||
assert_eq!(
|
||||
@ -2252,25 +2179,33 @@ mod tests {
|
||||
|
||||
#[test]
|
||||
fn test_scheduler_schedule_validation_rejects_invalid_values() {
|
||||
assert!(SchedulerSchedule::Delay { seconds: 0 }
|
||||
.validate("delay.job")
|
||||
.is_err());
|
||||
assert!(SchedulerSchedule::Interval {
|
||||
seconds: 0,
|
||||
startup_delay_secs: 0,
|
||||
}
|
||||
.validate("interval.job")
|
||||
.is_err());
|
||||
assert!(SchedulerSchedule::At {
|
||||
timestamp: "bad timestamp".to_string(),
|
||||
}
|
||||
.validate("at.job")
|
||||
.is_err());
|
||||
assert!(SchedulerSchedule::Cron {
|
||||
expression: "bad cron".to_string(),
|
||||
}
|
||||
.validate("cron.job")
|
||||
.is_err());
|
||||
assert!(
|
||||
SchedulerSchedule::Delay { seconds: 0 }
|
||||
.validate("delay.job")
|
||||
.is_err()
|
||||
);
|
||||
assert!(
|
||||
SchedulerSchedule::Interval {
|
||||
seconds: 0,
|
||||
startup_delay_secs: 0,
|
||||
}
|
||||
.validate("interval.job")
|
||||
.is_err()
|
||||
);
|
||||
assert!(
|
||||
SchedulerSchedule::At {
|
||||
timestamp: "bad timestamp".to_string(),
|
||||
}
|
||||
.validate("at.job")
|
||||
.is_err()
|
||||
);
|
||||
assert!(
|
||||
SchedulerSchedule::Cron {
|
||||
expression: "bad cron".to_string(),
|
||||
}
|
||||
.validate("cron.job")
|
||||
.is_err()
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
|
||||
@ -5,17 +5,14 @@ use serde::{Deserialize, Serialize};
|
||||
use std::collections::{HashMap, HashSet};
|
||||
use std::fs;
|
||||
use std::path::{Path, PathBuf};
|
||||
use std::sync::Arc;
|
||||
use parking_lot::RwLock;
|
||||
use std::sync::{Arc, RwLock};
|
||||
|
||||
#[cfg(test)]
|
||||
static EXPERT_TEST_ENV_LOCK: std::sync::Mutex<()> = std::sync::Mutex::new(());
|
||||
|
||||
#[cfg(test)]
|
||||
pub(crate) fn acquire_expert_test_env_lock() -> std::sync::MutexGuard<'static, ()> {
|
||||
EXPERT_TEST_ENV_LOCK
|
||||
.lock()
|
||||
.unwrap_or_else(|err| err.into_inner())
|
||||
EXPERT_TEST_ENV_LOCK.lock().unwrap_or_else(|err| err.into_inner())
|
||||
}
|
||||
|
||||
/// A discovered expert definition.
|
||||
@ -294,18 +291,13 @@ impl ExpertRuntime {
|
||||
|
||||
/// Re-discover experts from the filesystem.
|
||||
pub fn reload(&self) -> Result<ExpertCatalog, String> {
|
||||
let config = self
|
||||
.config
|
||||
.read()
|
||||
.clone();
|
||||
let config = self.config.read().expect("experts config rwlock poisoned").clone();
|
||||
let catalog = ExpertCatalog::discover_with_state(
|
||||
&config,
|
||||
&self.cwd,
|
||||
Some(&load_expert_disable_state(&self.cwd)),
|
||||
);
|
||||
let mut guard = self
|
||||
.catalog
|
||||
.write();
|
||||
let mut guard = self.catalog.write().expect("experts catalog rwlock poisoned");
|
||||
*guard = catalog.clone();
|
||||
Ok(catalog)
|
||||
}
|
||||
@ -314,7 +306,7 @@ impl ExpertRuntime {
|
||||
/// 用于前端保存配置后即时生效，无需重启网关。
|
||||
pub fn update_config(&self, new_config: ExpertsConfig) -> Result<(), String> {
|
||||
{
|
||||
let mut guard = self.config.write();
|
||||
let mut guard = self.config.write().expect("experts config rwlock poisoned");
|
||||
*guard = new_config;
|
||||
}
|
||||
self.reload()?;
|
||||
@ -325,16 +317,14 @@ impl ExpertRuntime {
|
||||
pub fn list_experts(&self) -> Vec<Expert> {
|
||||
self.catalog
|
||||
.read()
|
||||
.expect("experts catalog rwlock poisoned")
|
||||
.experts
|
||||
.clone()
|
||||
}
|
||||
|
||||
/// List all discovered experts including disabled ones, with their disabled scopes.
|
||||
pub fn list_experts_with_status(&self) -> Vec<ExpertWithStatus> {
|
||||
let config = self
|
||||
.config
|
||||
.read()
|
||||
.clone();
|
||||
let config = self.config.read().expect("experts config rwlock poisoned").clone();
|
||||
let catalog = ExpertCatalog::discover_without_state(&config, &self.cwd);
|
||||
let disable_state = load_expert_disable_state(&self.cwd);
|
||||
|
||||
@ -363,6 +353,7 @@ impl ExpertRuntime {
|
||||
pub fn get_expert(&self, name: &str) -> Option<Expert> {
|
||||
self.catalog
|
||||
.read()
|
||||
.expect("experts catalog rwlock poisoned")
|
||||
.find_expert(name)
|
||||
.cloned()
|
||||
}
|
||||
@ -420,15 +411,7 @@ impl ExpertRuntime {
|
||||
let next_provider = provider.cloned().unwrap_or(existing.provider);
|
||||
let next_model = model.cloned().unwrap_or(existing.model);
|
||||
|
||||
write_expert_file(
|
||||
&path,
|
||||
name,
|
||||
next_description,
|
||||
next_body,
|
||||
&next_capability,
|
||||
&next_provider,
|
||||
&next_model,
|
||||
)?;
|
||||
write_expert_file(&path, name, next_description, next_body, &next_capability, &next_provider, &next_model)?;
|
||||
let expert = parse_expert_file(&path, scope.into())?;
|
||||
if reload {
|
||||
let _ = self.reload()?;
|
||||
@ -474,10 +457,7 @@ impl ExpertRuntime {
|
||||
|
||||
pub fn has_expert_definition(&self, name: &str) -> Result<bool, String> {
|
||||
validate_expert_name(name)?;
|
||||
let config = self
|
||||
.config
|
||||
.read()
|
||||
.clone();
|
||||
let config = self.config.read().expect("experts config rwlock poisoned").clone();
|
||||
let catalog = ExpertCatalog::discover_without_state(&config, &self.cwd);
|
||||
Ok(catalog.find_expert(name).is_some())
|
||||
}
|
||||
@ -511,7 +491,8 @@ impl ExpertRuntime {
|
||||
{
|
||||
let mut state = self
|
||||
.disable_state
|
||||
.write();
|
||||
.write()
|
||||
.expect("experts disable_state rwlock poisoned");
|
||||
match scope {
|
||||
ExpertScope::User => {
|
||||
if enabled {
|
||||
@ -535,7 +516,8 @@ impl ExpertRuntime {
|
||||
|
||||
let state = self
|
||||
.disable_state
|
||||
.read();
|
||||
.read()
|
||||
.expect("experts disable_state rwlock poisoned");
|
||||
let disabled_in_scopes = state.disabled_scopes_for(name);
|
||||
|
||||
Ok(ExpertAvailabilityChange {
|
||||
@ -554,13 +536,17 @@ impl ExpertRuntime {
|
||||
}
|
||||
// The expert must exist (and not be disabled) for selection to be meaningful.
|
||||
if self.get_expert(expert_name).is_none() {
|
||||
return Err(format!("expert '{}' not found or disabled", expert_name));
|
||||
return Err(format!(
|
||||
"expert '{}' not found or disabled",
|
||||
expert_name
|
||||
));
|
||||
}
|
||||
|
||||
{
|
||||
let mut sessions = self
|
||||
.session_experts
|
||||
.write();
|
||||
.write()
|
||||
.expect("experts session_experts rwlock poisoned");
|
||||
sessions.insert(session_id.to_string(), expert_name.to_string());
|
||||
}
|
||||
persist_session_experts(&self.cwd, |state| {
|
||||
@ -575,7 +561,8 @@ impl ExpertRuntime {
|
||||
{
|
||||
let mut sessions = self
|
||||
.session_experts
|
||||
.write();
|
||||
.write()
|
||||
.expect("experts session_experts rwlock poisoned");
|
||||
sessions.remove(session_id);
|
||||
}
|
||||
persist_session_experts(&self.cwd, |state| {
|
||||
@ -588,14 +575,16 @@ impl ExpertRuntime {
|
||||
let name = {
|
||||
let sessions = self
|
||||
.session_experts
|
||||
.read();
|
||||
.read()
|
||||
.expect("experts session_experts rwlock poisoned");
|
||||
sessions.get(session_id).cloned()
|
||||
}?;
|
||||
|
||||
// Filter out disabled experts.
|
||||
let state = self
|
||||
.disable_state
|
||||
.read();
|
||||
.read()
|
||||
.expect("experts disable_state rwlock poisoned");
|
||||
if state.is_disabled(&name) {
|
||||
return None;
|
||||
}
|
||||
@ -635,7 +624,10 @@ impl SystemPromptProvider for ExpertPromptProvider {
|
||||
|
||||
let content = if expert.body.trim().is_empty() {
|
||||
// Empty body is OK; inject a header so the LLM still knows the role.
|
||||
format!("# 专家角色: {}\n\n{}", expert.name, expert.description)
|
||||
format!(
|
||||
"# 专家角色: {}\n\n{}",
|
||||
expert.name, expert.description
|
||||
)
|
||||
} else {
|
||||
expert.body.clone()
|
||||
};
|
||||
@ -675,9 +667,8 @@ fn expert_state_path(scope: ExpertScope, cwd: &Path) -> PathBuf {
|
||||
|
||||
fn root_for_scope(scope: ExpertScope, cwd: &Path) -> Result<PathBuf, String> {
|
||||
match scope {
|
||||
ExpertScope::User => {
|
||||
user_experts_root().ok_or_else(|| "failed to resolve home directory".to_string())
|
||||
}
|
||||
ExpertScope::User => user_experts_root()
|
||||
.ok_or_else(|| "failed to resolve home directory".to_string()),
|
||||
ExpertScope::Project => Ok(project_experts_root(cwd)),
|
||||
}
|
||||
}
|
||||
@ -1029,10 +1020,7 @@ fn load_project_session_experts(cwd: &Path) -> HashMap<String, String> {
|
||||
|
||||
/// Persist a mutation to the project-scope state file's session_experts while
|
||||
/// preserving the existing disabled_experts field.
|
||||
fn persist_session_experts<F: FnOnce(&mut ExpertStateFile)>(
|
||||
cwd: &Path,
|
||||
mutate: F,
|
||||
) -> Result<(), String> {
|
||||
fn persist_session_experts<F: FnOnce(&mut ExpertStateFile)>(cwd: &Path, mutate: F) -> Result<(), String> {
|
||||
let path = project_expert_state_path(cwd);
|
||||
let mut state = load_expert_state_file(&path)?;
|
||||
mutate(&mut state);
|
||||
@ -1109,11 +1097,7 @@ mod tests {
|
||||
let expert_dir = dir.path().join("demo");
|
||||
fs::create_dir_all(&expert_dir).unwrap();
|
||||
let expert_md = expert_dir.join("EXPERT.md");
|
||||
fs::write(
|
||||
&expert_md,
|
||||
"---\r\ndescription: demo expert\r\n---\r\nStep A\r\nStep B",
|
||||
)
|
||||
.unwrap();
|
||||
fs::write(&expert_md, "---\r\ndescription: demo expert\r\n---\r\nStep A\r\nStep B").unwrap();
|
||||
|
||||
let expert = parse_expert_file(&expert_md, ExpertSource::Project).unwrap();
|
||||
assert_eq!(expert.name, "demo");
|
||||
@ -1153,49 +1137,33 @@ mod tests {
|
||||
|
||||
#[test]
|
||||
fn test_render_expert_file_requires_description() {
|
||||
let err = render_expert_file(
|
||||
"demo",
|
||||
" ",
|
||||
"body",
|
||||
&CapabilityPolicy::default(),
|
||||
&None,
|
||||
&None,
|
||||
)
|
||||
.unwrap_err();
|
||||
let err = render_expert_file("demo", " ", "body", &CapabilityPolicy::default(), &None, &None).unwrap_err();
|
||||
assert!(err.contains("description"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_capability_policy_is_empty_helpers() {
|
||||
assert!(CapabilityPolicy::default().is_empty());
|
||||
assert!(
|
||||
!CapabilityPolicy {
|
||||
allowed_skills: Some(vec!["a".to_string()]),
|
||||
..Default::default()
|
||||
}
|
||||
.is_empty()
|
||||
);
|
||||
assert!(
|
||||
CapabilityPolicy {
|
||||
denied_tools: vec![],
|
||||
..Default::default()
|
||||
}
|
||||
.is_empty()
|
||||
);
|
||||
assert!(
|
||||
CapabilityPolicy {
|
||||
allowed_tools: Some(vec![]),
|
||||
..Default::default()
|
||||
}
|
||||
.has_tool_policy()
|
||||
);
|
||||
assert!(
|
||||
CapabilityPolicy {
|
||||
denied_skills: vec!["x".to_string()],
|
||||
..Default::default()
|
||||
}
|
||||
.has_skill_policy()
|
||||
);
|
||||
assert!(!CapabilityPolicy {
|
||||
allowed_skills: Some(vec!["a".to_string()]),
|
||||
..Default::default()
|
||||
}
|
||||
.is_empty());
|
||||
assert!(CapabilityPolicy {
|
||||
denied_tools: vec![],
|
||||
..Default::default()
|
||||
}
|
||||
.is_empty());
|
||||
assert!(CapabilityPolicy {
|
||||
allowed_tools: Some(vec![]),
|
||||
..Default::default()
|
||||
}
|
||||
.has_tool_policy());
|
||||
assert!(CapabilityPolicy {
|
||||
denied_skills: vec!["x".to_string()],
|
||||
..Default::default()
|
||||
}
|
||||
.has_skill_policy());
|
||||
}
|
||||
|
||||
#[test]
|
||||
@ -1228,15 +1196,7 @@ mod tests {
|
||||
#[test]
|
||||
fn test_empty_capability_omits_keys() {
|
||||
// 空策略不应输出多余 frontmatter 键，保持旧文件格式兼容
|
||||
let rendered = render_expert_file(
|
||||
"plain",
|
||||
"desc",
|
||||
"body",
|
||||
&CapabilityPolicy::default(),
|
||||
&None,
|
||||
&None,
|
||||
)
|
||||
.unwrap();
|
||||
let rendered = render_expert_file("plain", "desc", "body", &CapabilityPolicy::default(), &None, &None).unwrap();
|
||||
assert!(!rendered.contains("allowed_skills"));
|
||||
assert!(!rendered.contains("denied_skills"));
|
||||
assert!(!rendered.contains("allowed_tools"));
|
||||
@ -1264,7 +1224,10 @@ mod tests {
|
||||
.unwrap();
|
||||
|
||||
// project scope (overrides user)
|
||||
let project_dir_expert = project_dir.join(".picobot").join("experts").join("demo");
|
||||
let project_dir_expert = project_dir
|
||||
.join(".picobot")
|
||||
.join("experts")
|
||||
.join("demo");
|
||||
fs::create_dir_all(&project_dir_expert).unwrap();
|
||||
fs::write(
|
||||
project_dir_expert.join("EXPERT.md"),
|
||||
@ -1347,16 +1310,7 @@ mod tests {
|
||||
|
||||
// update with None preserves fields
|
||||
let updated_none = runtime
|
||||
.update_expert(
|
||||
ExpertScope::Project,
|
||||
"translator",
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
true,
|
||||
)
|
||||
.update_expert(ExpertScope::Project, "translator", None, None, None, None, None, true)
|
||||
.unwrap();
|
||||
assert_eq!(updated_none.description, "更新翻译专家");
|
||||
assert_eq!(updated_none.body, "你是一名中文教师。");
|
||||
@ -1555,21 +1509,13 @@ mod tests {
|
||||
.unwrap();
|
||||
|
||||
let items = runtime.list_experts_with_status();
|
||||
assert_eq!(
|
||||
items.len(),
|
||||
1,
|
||||
"list_experts_with_status should include disabled experts"
|
||||
);
|
||||
assert_eq!(items.len(), 1, "list_experts_with_status should include disabled experts");
|
||||
assert_eq!(items[0].name, "planner");
|
||||
assert_eq!(items[0].disabled_in_scopes, vec!["project".to_string()]);
|
||||
|
||||
// list_experts (filtered) should be empty
|
||||
let active = runtime.list_experts();
|
||||
assert_eq!(
|
||||
active.len(),
|
||||
0,
|
||||
"list_experts should filter out disabled experts"
|
||||
);
|
||||
assert_eq!(active.len(), 0, "list_experts should filter out disabled experts");
|
||||
}
|
||||
|
||||
#[test]
|
||||
@ -1580,9 +1526,7 @@ mod tests {
|
||||
session_experts: HashMap::new(),
|
||||
disabled_experts: vec!["demo".to_string()],
|
||||
};
|
||||
state
|
||||
.session_experts
|
||||
.insert("sess-1".to_string(), "demo".to_string());
|
||||
state.session_experts.insert("sess-1".to_string(), "demo".to_string());
|
||||
|
||||
save_expert_state_file(&path, &state).unwrap();
|
||||
|
||||
|
||||
@ -1,5 +1,5 @@
|
||||
use gray_matter::Matter;
|
||||
use gray_matter::engine::YAML;
|
||||
use gray_matter::Matter;
|
||||
use serde::de::DeserializeOwned;
|
||||
|
||||
/// Parse a markdown document with YAML frontmatter into `(frontmatter, body)`.
|
||||
@ -45,13 +45,7 @@ mod tests {
|
||||
fn parses_lf_endings() {
|
||||
let input = "---\ndescription: demo\n---\nbody text";
|
||||
let (fm, body) = parse::<FrontMatter>(input).unwrap();
|
||||
assert_eq!(
|
||||
fm,
|
||||
FrontMatter {
|
||||
description: "demo".to_string(),
|
||||
name: None
|
||||
}
|
||||
);
|
||||
assert_eq!(fm, FrontMatter { description: "demo".to_string(), name: None });
|
||||
assert_eq!(body, "body text");
|
||||
}
|
||||
|
||||
|
||||
@ -1,8 +1,7 @@
|
||||
use std::sync::Arc;
|
||||
|
||||
use crate::agent::context_compressor::ContextCompressor;
|
||||
use crate::agent::{AgentError, AgentLoop, AgentRuntimeConfig, CompositeSystemPromptProvider, SystemPromptProvider};
|
||||
use crate::config::{CompactionConfig, LLMProviderConfig, ModelResolver};
|
||||
use crate::agent::{AgentError, AgentLoop, CompositeSystemPromptProvider, SystemPromptProvider};
|
||||
use crate::config::{LLMProviderConfig, ModelResolver};
|
||||
use crate::domain::CapabilityPolicy;
|
||||
use crate::experts::ExpertPromptProvider;
|
||||
use crate::experts::ExpertRuntime;
|
||||
@ -10,8 +9,8 @@ use crate::gateway::agent_prompt_provider::AgentPromptProvider;
|
||||
use crate::gateway::model_selection::ModelSelectionStore;
|
||||
use crate::gateway::tool_prompt_provider::ToolPromptProvider;
|
||||
use crate::skills::{SkillPromptProvider, SkillRuntime};
|
||||
use crate::storage::PromptInjectionRepository;
|
||||
use crate::storage::persistent_session_id;
|
||||
use crate::storage::PromptInjectionRepository;
|
||||
use crate::tools::task::runtime::{SubagentPromptProvider, SubagentRuntime};
|
||||
use crate::tools::{ToolContext, ToolRegistry};
|
||||
|
||||
@ -54,8 +53,6 @@ pub(crate) struct AgentFactory {
|
||||
model_resolver: Arc<ModelResolver>,
|
||||
/// per-session 的用户模型选择（最高优先级，覆盖专家配置）
|
||||
model_selections: Arc<ModelSelectionStore>,
|
||||
/// 上下文压缩算法配置（所有 agent 共享）
|
||||
compaction_config: CompactionConfig,
|
||||
/// 实例创建时间戳（用于区分新旧 AgentFactory 实例）
|
||||
instance_id: u64,
|
||||
}
|
||||
@ -83,7 +80,6 @@ impl AgentFactory {
|
||||
prompt_repository: Arc<dyn PromptInjectionRepository>,
|
||||
model_resolver: Arc<ModelResolver>,
|
||||
model_selections: Arc<ModelSelectionStore>,
|
||||
compaction_config: CompactionConfig,
|
||||
) -> Self {
|
||||
// 使用 Arc 指针地址作为实例标识符，用于区分新旧 AgentFactory 实例
|
||||
let instance_id = Arc::as_ptr(&tools) as u64;
|
||||
@ -101,22 +97,10 @@ impl AgentFactory {
|
||||
prompt_repository,
|
||||
model_resolver,
|
||||
model_selections,
|
||||
compaction_config,
|
||||
instance_id,
|
||||
}
|
||||
}
|
||||
|
||||
/// 构造 ContextCompressor（参数内聚到 ContextCompressor，CompactionConfig 注入）。
|
||||
/// AgentLoop（in-loop 压缩）和 Session（sync 兜底压缩）共用此方法，
|
||||
/// 确保两条压缩路径使用同一套用户配置的压缩参数。
|
||||
pub(crate) fn build_compressor(&self, runtime_config: &AgentRuntimeConfig) -> ContextCompressor {
|
||||
ContextCompressor::with_compaction_config(
|
||||
runtime_config.context_window_tokens,
|
||||
runtime_config.context_summary_char_budget,
|
||||
&self.compaction_config,
|
||||
)
|
||||
}
|
||||
|
||||
pub(crate) fn create(&self, request: AgentBuildRequest<'_>) -> Result<AgentLoop, AgentError> {
|
||||
let session_id = persistent_session_id(request.channel_name, request.session_chat_id);
|
||||
|
||||
@ -128,14 +112,12 @@ impl AgentFactory {
|
||||
// 引用不存在的 provider/model 名时报错并阻止会话（用户主动选择的角色，配置错误应明确反馈）。
|
||||
let expert_provider_config = match &expert {
|
||||
Some(e) if e.provider.is_some() || e.model.is_some() => {
|
||||
let resolved = self
|
||||
.model_resolver
|
||||
.resolve(
|
||||
e.provider.as_deref(),
|
||||
e.model.as_deref(),
|
||||
&request.provider_config,
|
||||
)
|
||||
.map_err(|e| AgentError::Other(e.to_string()))?;
|
||||
let resolved = self.model_resolver.resolve(
|
||||
e.provider.as_deref(),
|
||||
e.model.as_deref(),
|
||||
&request.provider_config,
|
||||
)
|
||||
.map_err(|e| AgentError::Other(e.to_string()))?;
|
||||
tracing::info!(
|
||||
instance_id = self.instance_id,
|
||||
session_id = %session_id,
|
||||
@ -151,29 +133,30 @@ impl AgentFactory {
|
||||
|
||||
// 按用户手动选择的 provider/model 覆盖（最高优先级，覆盖专家配置）。
|
||||
// 引用不存在的 provider/model 名时报错并阻止会话（用户主动选择，配置错误应明确反馈）。
|
||||
let effective_provider_config = match self.model_selections.get(&session_id) {
|
||||
Some((user_provider, user_model))
|
||||
if user_provider.is_some() || user_model.is_some() =>
|
||||
{
|
||||
let resolved = self
|
||||
.model_resolver
|
||||
.resolve(
|
||||
user_provider.as_deref(),
|
||||
user_model.as_deref(),
|
||||
&expert_provider_config,
|
||||
)
|
||||
.map_err(|e| AgentError::Other(e.to_string()))?;
|
||||
tracing::info!(
|
||||
instance_id = self.instance_id,
|
||||
session_id = %session_id,
|
||||
provider = %resolved.name,
|
||||
model_id = %resolved.model_id,
|
||||
"AgentFactory: applied user model override"
|
||||
);
|
||||
resolved
|
||||
}
|
||||
_ => expert_provider_config,
|
||||
};
|
||||
let effective_provider_config =
|
||||
match self.model_selections.get(&session_id) {
|
||||
Some((user_provider, user_model))
|
||||
if user_provider.is_some() || user_model.is_some() =>
|
||||
{
|
||||
let resolved = self
|
||||
.model_resolver
|
||||
.resolve(
|
||||
user_provider.as_deref(),
|
||||
user_model.as_deref(),
|
||||
&expert_provider_config,
|
||||
)
|
||||
.map_err(|e| AgentError::Other(e.to_string()))?;
|
||||
tracing::info!(
|
||||
instance_id = self.instance_id,
|
||||
session_id = %session_id,
|
||||
provider = %resolved.name,
|
||||
model_id = %resolved.model_id,
|
||||
"AgentFactory: applied user model override"
|
||||
);
|
||||
resolved
|
||||
}
|
||||
_ => expert_provider_config,
|
||||
};
|
||||
|
||||
// 诊断日志：记录 agent 实际使用的配置和实例 ID
|
||||
tracing::info!(
|
||||
@ -216,7 +199,7 @@ impl AgentFactory {
|
||||
};
|
||||
|
||||
AgentLoop::with_tools_and_system_prompt_provider(
|
||||
effective_provider_config.clone(),
|
||||
effective_provider_config,
|
||||
tools,
|
||||
system_prompt_provider,
|
||||
Some(self.skills.clone()),
|
||||
@ -226,27 +209,22 @@ impl AgentFactory {
|
||||
let tool_chat_id = request
|
||||
.notification_chat_id
|
||||
.unwrap_or(request.session_chat_id);
|
||||
// 构建上下文压缩器（参数内聚到 ContextCompressor，CompactionConfig 注入）
|
||||
let runtime_config = AgentRuntimeConfig::from(effective_provider_config.clone());
|
||||
let compressor = Arc::new(self.build_compressor(&runtime_config));
|
||||
let mut agent = agent
|
||||
.with_tool_context(ToolContext {
|
||||
channel_name: Some(request.channel_name.to_string()),
|
||||
sender_id: request.sender_id.map(str::to_string),
|
||||
chat_id: Some(tool_chat_id.to_string()),
|
||||
session_id: Some(session_id),
|
||||
topic_id: request.topic_id.clone(),
|
||||
message_id: request.message_id.map(str::to_string),
|
||||
message_seq: None,
|
||||
subagent_description: None,
|
||||
nesting_depth: 0,
|
||||
task_id: None,
|
||||
parent_task_id: None,
|
||||
tool_call_id: None,
|
||||
// 注入专家 capability，TaskTool 据此强制校验子代理白/黑名单
|
||||
parent_capability: expert_capability.clone(),
|
||||
})
|
||||
.with_compressor(Some(compressor));
|
||||
let mut agent = agent.with_tool_context(ToolContext {
|
||||
channel_name: Some(request.channel_name.to_string()),
|
||||
sender_id: request.sender_id.map(str::to_string),
|
||||
chat_id: Some(tool_chat_id.to_string()),
|
||||
session_id: Some(session_id),
|
||||
topic_id: request.topic_id.clone(),
|
||||
message_id: request.message_id.map(str::to_string),
|
||||
message_seq: None,
|
||||
subagent_description: None,
|
||||
nesting_depth: 0,
|
||||
task_id: None,
|
||||
parent_task_id: None,
|
||||
tool_call_id: None,
|
||||
// 注入专家 capability，TaskTool 据此强制校验子代理白/黑名单
|
||||
parent_capability: expert_capability.clone(),
|
||||
});
|
||||
// 如果有取消信号接收端，注入 Agent
|
||||
if let Some(token) = request.cancel_token {
|
||||
agent = agent.with_cancel_token(token);
|
||||
|
||||
@ -129,7 +129,6 @@ mod tests {
|
||||
extra_headers: HashMap::new(),
|
||||
llm_timeout_secs: 120,
|
||||
memory_maintenance_timeout_secs: 600,
|
||||
max_retries: 3,
|
||||
model_id: "test-model".to_string(),
|
||||
temperature: Some(0.0),
|
||||
max_tokens: Some(32),
|
||||
|
||||
@ -40,13 +40,7 @@ impl AgentTaskExecutor {
|
||||
options: ScheduledAgentTaskOptions,
|
||||
) -> Result<Vec<OutboundMessage>, AgentError> {
|
||||
self.session_manager
|
||||
.run_silent_agent_task(
|
||||
channel_name,
|
||||
session_chat_id,
|
||||
notification_chat_id,
|
||||
prompt,
|
||||
options,
|
||||
)
|
||||
.run_silent_agent_task(channel_name, session_chat_id, notification_chat_id, prompt, options)
|
||||
.await
|
||||
}
|
||||
}
|
||||
@ -99,12 +93,8 @@ impl SchedulerMaintenanceService {
|
||||
self.session_manager.cleanup_expired_sessions().await
|
||||
}
|
||||
|
||||
async fn run_memory_maintenance(
|
||||
&self,
|
||||
) -> Result<Option<MemoryMaintenanceScopeResult>, AgentError> {
|
||||
self.session_manager
|
||||
.run_memory_maintenance_for_all_scopes()
|
||||
.await
|
||||
async fn run_memory_maintenance(&self) -> Result<Option<MemoryMaintenanceScopeResult>, AgentError> {
|
||||
self.session_manager.run_memory_maintenance_for_all_scopes().await
|
||||
}
|
||||
}
|
||||
|
||||
@ -114,9 +104,7 @@ impl MaintenanceExecutor for SchedulerMaintenanceService {
|
||||
self.cleanup_sessions().await
|
||||
}
|
||||
|
||||
async fn run_memory_maintenance_for_all_scopes(
|
||||
&self,
|
||||
) -> anyhow::Result<Vec<MaintenanceRunSummary>> {
|
||||
async fn run_memory_maintenance_for_all_scopes(&self) -> anyhow::Result<Vec<MaintenanceRunSummary>> {
|
||||
self.run_memory_maintenance()
|
||||
.await
|
||||
.map(|results| {
|
||||
|
||||
@ -1,243 +0,0 @@
|
||||
//! 网关认证与访问控制。
|
||||
//!
|
||||
//! 设计目标（第一性原理）：
|
||||
//! - 本地单机部署（host 为 loopback）：免认证，仅靠 CORS 防御 DNS rebinding / CSRF。
|
||||
//! - 远程访问（host 非 loopback）：必须配置 `auth_token`，所有 `/api/*` 与 `/ws` 强制校验。
|
||||
//! - token 通过 `Authorization: Bearer <token>`（HTTP）或 `?token=<token>`（WS）传递。
|
||||
//! - 校验使用常量时间比较，避免计时侧信道。
|
||||
|
||||
use axum::extract::Request;
|
||||
use axum::http::{HeaderMap, StatusCode};
|
||||
use axum::middleware::Next;
|
||||
use axum::response::{IntoResponse, Response};
|
||||
use axum::Json;
|
||||
use serde_json::json;
|
||||
use subtle::ConstantTimeEq;
|
||||
|
||||
/// 判定给定 host 是否为 loopback 地址。
|
||||
/// 通过 `std::net::IpAddr` 解析，覆盖 IPv4/IPv6 的所有等价表示。
|
||||
/// 也接受 `localhost` 主机名。
|
||||
pub fn is_loopback_host(host: &str) -> bool {
|
||||
let h = host.trim();
|
||||
if h.eq_ignore_ascii_case("localhost") {
|
||||
return true;
|
||||
}
|
||||
// 去除 IPv6 方括号（如 `[::1]`）
|
||||
let bare = h
|
||||
.strip_prefix('[')
|
||||
.and_then(|s| s.strip_suffix(']'))
|
||||
.unwrap_or(h);
|
||||
// 尝试解析为 IpAddr
|
||||
match bare.parse::<std::net::IpAddr>() {
|
||||
Ok(std::net::IpAddr::V4(v4)) => v4.is_loopback(),
|
||||
Ok(std::net::IpAddr::V6(v6)) => v6.is_loopback(),
|
||||
Err(_) => {
|
||||
// 解析失败（如域名），保守判定为非 loopback，强制认证
|
||||
false
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// 判定当前配置是否需要强制认证。
|
||||
/// - host 非 loopback：必须认证（且 auth_token 必须存在，否则启动会报错）
|
||||
/// - host 为 loopback 但显式配置了 auth_token：也启用认证（用户主动加固）
|
||||
pub fn requires_auth(host: &str, auth_token: &Option<String>) -> bool {
|
||||
!is_loopback_host(host) || auth_token.is_some()
|
||||
}
|
||||
|
||||
/// 校验请求携带的 token 是否与配置的 auth_token 匹配（常量时间）。
|
||||
/// 返回 true 表示通过（含「未配置 token 则放行」的兜底，调用方应先用 requires_auth 判定）。
|
||||
pub fn token_matches(provided: Option<&str>, expected: &Option<String>) -> bool {
|
||||
match expected {
|
||||
// 未配置 token：放行（调用方已通过 requires_auth 保证只在 loopback 下到达此处）
|
||||
None => true,
|
||||
Some(expected_str) => match provided {
|
||||
None => false,
|
||||
Some(provided_str) => {
|
||||
// subtle::ConstantTimeEq 要求两端等长；长度不等时 ct_eq 返回 0（false），
|
||||
// 但为避免长度差异导致的提前退出计时泄露，我们确保比较路径不因长度而分支提前返回。
|
||||
let p = provided_str.as_bytes();
|
||||
let e = expected_str.as_bytes();
|
||||
// ct_eq 内部在长度不等时仍会遍历较短的一侧，返回 0，无提前退出
|
||||
p.ct_eq(e).unwrap_u8() == 1
|
||||
}
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
/// 从 Authorization 头提取 Bearer token（RFC 6750，scheme 不区分大小写）。
|
||||
pub fn extract_bearer_token(headers: &HeaderMap) -> Option<&str> {
|
||||
headers
|
||||
.get(axum::http::header::AUTHORIZATION)
|
||||
.and_then(|v| v.to_str().ok())
|
||||
.and_then(|s| {
|
||||
// RFC 6750: scheme 名不区分大小写（Bearer / bearer / BEARER 均合法）
|
||||
// 找到第一个空格分隔 scheme 与 token，比较 scheme 部分忽略大小写
|
||||
let trimmed = s.trim_start();
|
||||
let split = trimmed.find(char::is_whitespace)?;
|
||||
let (scheme, rest) = trimmed.split_at(split);
|
||||
if scheme.eq_ignore_ascii_case("Bearer") {
|
||||
Some(rest.trim())
|
||||
} else {
|
||||
None
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
/// axum 中间件：对 `/api/*` 路由强制 Bearer token 校验。
|
||||
/// 仅在 `requires_auth` 为 true 时挂载。
|
||||
/// `/health`、`/ws`、静态资源放行；`/ws` 的 token 校验在 ws_handler 内完成。
|
||||
pub async fn require_bearer_auth(
|
||||
headers: HeaderMap,
|
||||
request: Request,
|
||||
next: Next,
|
||||
) -> Response {
|
||||
let path = request.uri().path();
|
||||
|
||||
// 仅对 /api/ 前缀的请求强制认证
|
||||
if !path.starts_with("/api/") {
|
||||
return next.run(request).await;
|
||||
}
|
||||
|
||||
// expected_token 通过 extension 注入（见 mod.rs 装配处）
|
||||
let expected = request
|
||||
.extensions()
|
||||
.get::<AuthConfig>()
|
||||
.map(|c| c.token.clone())
|
||||
.unwrap_or(None);
|
||||
|
||||
let provided = extract_bearer_token(&headers);
|
||||
if token_matches(provided, &expected) {
|
||||
next.run(request).await
|
||||
} else {
|
||||
(
|
||||
StatusCode::UNAUTHORIZED,
|
||||
Json(json!({ "error": "unauthorized", "message": "missing or invalid token" })),
|
||||
)
|
||||
.into_response()
|
||||
}
|
||||
}
|
||||
|
||||
/// 通过 extension 注入到 Router 的认证配置。
|
||||
#[derive(Clone, Debug)]
|
||||
pub struct AuthConfig {
|
||||
pub token: Option<String>,
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use axum::http::HeaderMap;
|
||||
|
||||
#[test]
|
||||
fn loopback_detection() {
|
||||
// IPv4 loopback 整个 127.0.0.0/8 段
|
||||
assert!(is_loopback_host("127.0.0.1"));
|
||||
assert!(is_loopback_host("127.0.0.2"));
|
||||
assert!(is_loopback_host("127.1.2.3"));
|
||||
assert!(is_loopback_host("127.255.255.255"));
|
||||
// IPv6 loopback 的各种等价表示
|
||||
assert!(is_loopback_host("::1"));
|
||||
assert!(is_loopback_host("[::1]"));
|
||||
assert!(is_loopback_host("0:0:0:0:0:0:0:1"));
|
||||
// localhost 主机名（大小写不敏感）
|
||||
assert!(is_loopback_host("localhost"));
|
||||
assert!(is_loopback_host("LOCALHOST"));
|
||||
assert!(is_loopback_host("Localhost"));
|
||||
// 空白容错
|
||||
assert!(is_loopback_host(" 127.0.0.1 "));
|
||||
assert!(is_loopback_host(" localhost "));
|
||||
// 非 loopback
|
||||
assert!(!is_loopback_host("0.0.0.0"));
|
||||
assert!(!is_loopback_host("192.168.1.1"));
|
||||
assert!(!is_loopback_host("10.0.0.1"));
|
||||
assert!(!is_loopback_host("172.16.0.1"));
|
||||
assert!(!is_loopback_host("example.com"));
|
||||
assert!(!is_loopback_host("picobot.local"));
|
||||
assert!(!is_loopback_host("::"));
|
||||
assert!(!is_loopback_host(""));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn requires_auth_logic() {
|
||||
// 非 loopback 无论 token 是否配置都要认证
|
||||
assert!(requires_auth("0.0.0.0", &None));
|
||||
assert!(requires_auth("192.168.1.1", &None));
|
||||
assert!(requires_auth("0.0.0.0", &Some("t".into())));
|
||||
// loopback 无 token：免认证
|
||||
assert!(!requires_auth("127.0.0.1", &None));
|
||||
assert!(!requires_auth("::1", &None));
|
||||
assert!(!requires_auth("localhost", &None));
|
||||
// loopback 有 token：认证（用户主动加固）
|
||||
assert!(requires_auth("127.0.0.1", &Some("t".into())));
|
||||
// 域名始终需认证
|
||||
assert!(requires_auth("myserver.com", &None));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn token_matching() {
|
||||
// 未配置 token：始终放行
|
||||
assert!(token_matches(None, &None));
|
||||
assert!(token_matches(Some("anything"), &None));
|
||||
// 配置了 token
|
||||
let expected = Some("s3cr3t".to_string());
|
||||
assert!(token_matches(Some("s3cr3t"), &expected));
|
||||
assert!(!token_matches(Some("wrong"), &expected));
|
||||
assert!(!token_matches(None, &expected));
|
||||
// 长度不同
|
||||
assert!(!token_matches(Some("s3cr3t-extra"), &expected));
|
||||
assert!(!token_matches(Some("s3"), &expected));
|
||||
assert!(!token_matches(Some(""), &expected));
|
||||
// 空字符串 token（配置了但为空 — 等同于未配置的放行语义由 requires_auth 控制）
|
||||
let empty_expected = Some(String::new());
|
||||
assert!(token_matches(Some(""), &empty_expected));
|
||||
assert!(!token_matches(Some("x"), &empty_expected));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn extract_bearer_token_rfc6750() {
|
||||
fn make_auth_header(value: &str) -> HeaderMap {
|
||||
let mut h = HeaderMap::new();
|
||||
h.insert(
|
||||
axum::http::header::AUTHORIZATION,
|
||||
axum::http::HeaderValue::from_str(value).unwrap(),
|
||||
);
|
||||
h
|
||||
}
|
||||
|
||||
// 标准格式
|
||||
assert_eq!(
|
||||
extract_bearer_token(&make_auth_header("Bearer abc123")),
|
||||
Some("abc123")
|
||||
);
|
||||
// scheme 大小写不敏感
|
||||
assert_eq!(
|
||||
extract_bearer_token(&make_auth_header("bearer abc123")),
|
||||
Some("abc123")
|
||||
);
|
||||
assert_eq!(
|
||||
extract_bearer_token(&make_auth_header("BEARER abc123")),
|
||||
Some("abc123")
|
||||
);
|
||||
// 多空格容忍
|
||||
assert_eq!(
|
||||
extract_bearer_token(&make_auth_header("Bearer abc123")),
|
||||
Some("abc123")
|
||||
);
|
||||
assert_eq!(
|
||||
extract_bearer_token(&make_auth_header("Bearer\tabc123")),
|
||||
Some("abc123")
|
||||
);
|
||||
// 非 Bearer scheme
|
||||
assert_eq!(
|
||||
extract_bearer_token(&make_auth_header("Basic dXNlcjpwYXNz")),
|
||||
None
|
||||
);
|
||||
// 无 Authorization 头
|
||||
let empty = HeaderMap::new();
|
||||
assert_eq!(extract_bearer_token(&empty), None);
|
||||
// 缺少 token 部分
|
||||
assert_eq!(extract_bearer_token(&make_auth_header("Bearer")), None);
|
||||
assert_eq!(extract_bearer_token(&make_auth_header("Bearer ")), Some(""));
|
||||
}
|
||||
}
|
||||
@ -58,14 +58,6 @@ impl CancelManager {
|
||||
self.tokens.lock().await.len()
|
||||
}
|
||||
|
||||
/// 返回当前正在执行的 Agent 的 topic_id 列表。
|
||||
///
|
||||
/// 用于前端重连时对账执行状态：前端通过此 API 判断断连期间
|
||||
/// 哪些话题的智能体仍在运行、哪些已完成。
|
||||
pub async fn list_active_topic_ids(&self) -> Vec<String> {
|
||||
self.tokens.lock().await.keys().cloned().collect()
|
||||
}
|
||||
|
||||
/// 取消所有正在运行的 Agent 并清空注册表。
|
||||
///
|
||||
/// 用于 graceful shutdown / restart 场景。
|
||||
|
||||
@ -52,17 +52,17 @@ pub(crate) async fn schedule_background_history_compaction(
|
||||
.compress_two_segment(&history, &provider_config)
|
||||
.await?;
|
||||
|
||||
// 保留原始消息（标记 is_compacted=1）+ 插入压缩摘要，不删除原消息，
|
||||
// 从而让前端仍能展示完整原始对话，LLM 只看压缩后的精简历史。
|
||||
// Replace only this topic's history in DB (not the entire session).
|
||||
// This avoids clobbering other topics' messages during compaction.
|
||||
store
|
||||
.compact_topic_history(&session_id, &topic_id, &compressed)
|
||||
.map_err(|e| AgentError::Other(format!("compact_topic_history error: {}", e)))?;
|
||||
.replace_topic_history(&session_id, &topic_id, &compressed)
|
||||
.map_err(|e| AgentError::Other(format!("replace_topic_history error: {}", e)))?;
|
||||
|
||||
tracing::info!(
|
||||
chat_id = %chat_id,
|
||||
topic_id = %topic_id,
|
||||
compressed_msg_count = compressed.len(),
|
||||
"Two-segment compression committed (original messages retained)"
|
||||
"Two-segment compression committed"
|
||||
);
|
||||
|
||||
session_guard.reload_topic_history(&chat_id, &topic_id)?;
|
||||
|
||||
@ -1,15 +1,12 @@
|
||||
use std::collections::HashMap;
|
||||
use std::sync::Arc;
|
||||
|
||||
use crate::agent::{
|
||||
AgentError, AgentProcessResult, CompactionSink, EmittedMessageHandler,
|
||||
PersistingEmittedMessageHandler, SystemPromptContext,
|
||||
};
|
||||
use async_trait::async_trait;
|
||||
use crate::agent::{AgentError, AgentProcessResult, EmittedMessageHandler, PersistingEmittedMessageHandler, SystemPromptContext};
|
||||
use crate::bus::message::ToolMessageState;
|
||||
use crate::bus::{ChatMessage, MediaItem, OutboundMessage, SYSTEM_CONTEXT_SCHEDULED_PROMPT};
|
||||
use crate::config::LLMProviderConfig;
|
||||
use crate::storage::{ConversationRepository, persistent_session_id};
|
||||
use async_trait::async_trait;
|
||||
use crate::storage::{persistent_session_id, ConversationRepository};
|
||||
use tokio::sync::Mutex;
|
||||
|
||||
use super::compaction::schedule_background_history_compaction;
|
||||
@ -24,51 +21,6 @@ impl EmittedMessageHandler for NoOpEmittedMessageHandler {
|
||||
async fn handle(&self, _message: ChatMessage) {}
|
||||
}
|
||||
|
||||
/// CompactionSink 实现：在 AgentLoop 内部触发 LLM 压缩时，
|
||||
/// 把压缩后的消息写回 DB（标记原消息 is_compacted=1 + 插入摘要）。
|
||||
///
|
||||
/// 不在此处 reload 内存历史——process() 仍在使用局部 messages 变量，
|
||||
/// 内存历史的刷新由 finalize_result 在 process 返回后统一处理。
|
||||
pub(crate) struct CompactionSinkImpl {
|
||||
session: Arc<Mutex<Session>>,
|
||||
chat_id: String,
|
||||
topic_id: String,
|
||||
}
|
||||
|
||||
impl CompactionSinkImpl {
|
||||
pub(crate) fn new(session: Arc<Mutex<Session>>, chat_id: String, topic_id: String) -> Self {
|
||||
Self {
|
||||
session,
|
||||
chat_id,
|
||||
topic_id,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl CompactionSink for CompactionSinkImpl {
|
||||
async fn compact(&self, compressed: &[ChatMessage]) -> Result<(), AgentError> {
|
||||
let mut session_guard = self.session.lock().await;
|
||||
session_guard.ensure_persistent_session(&self.chat_id)?;
|
||||
session_guard.ensure_chat_loaded(&self.chat_id, Some(&self.topic_id))?;
|
||||
|
||||
let store = session_guard.store();
|
||||
let session_id = session_guard.persistent_session_id(&self.chat_id);
|
||||
|
||||
store
|
||||
.compact_topic_history(&session_id, &self.topic_id, compressed)
|
||||
.map_err(|e| AgentError::Other(format!("compact_topic_history error: {}", e)))?;
|
||||
|
||||
tracing::info!(
|
||||
chat_id = %self.chat_id,
|
||||
topic_id = %self.topic_id,
|
||||
compressed_msg_count = compressed.len(),
|
||||
"In-loop LLM compaction committed to DB (original messages retained as is_compacted=1)"
|
||||
);
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
const SCHEDULED_TASK_EXECUTION_SYSTEM_PROMPT: &str = "系统说明：当前输入来自一次已经触发的定时任务执行。你现在需要执行任务内容本身，而不是创建、修改、恢复、暂停或查询新的定时任务。除非当前任务内容明确要求管理调度器，否则不要调用任何定时任务管理工具；像“每小时”、“每天”、“cron”、“定时”等词，只应视为任务背景，不应再解释为新的建任务请求。";
|
||||
|
||||
pub(crate) fn compose_scheduled_task_system_prompt(system_prompt: Option<&str>) -> String {
|
||||
@ -148,12 +100,9 @@ impl AgentExecutionService {
|
||||
};
|
||||
|
||||
if !is_current_turn {
|
||||
let (latest_user_id, latest_user_preview, compression_in_flight, history_len) = session
|
||||
.stale_result_diagnostics(
|
||||
request
|
||||
.original_topic_id
|
||||
.as_deref()
|
||||
.unwrap_or(request.chat_id),
|
||||
let (latest_user_id, latest_user_preview, compression_in_flight, history_len) =
|
||||
session.stale_result_diagnostics(
|
||||
request.original_topic_id.as_deref().unwrap_or(request.chat_id),
|
||||
);
|
||||
tracing::info!(
|
||||
channel = %request.channel_name,
|
||||
@ -173,31 +122,14 @@ impl AgentExecutionService {
|
||||
// 始终使用执行开始时捕获的 original_topic_id，避免从共享状态重复读取竞态
|
||||
let target_topic_id = request.original_topic_id.as_deref();
|
||||
|
||||
// 如果 AgentLoop 内部已触发 LLM 压缩，DB 已被 CompactionSink 更新
|
||||
// （原消息标记 is_compacted=1 + 压缩摘要已插入）。
|
||||
// 此时 emitted_messages 早已通过 handler 持久化，且作为"最新5个 unit"
|
||||
// 包含在压缩输出中。直接 append 到内存历史会产生重复，因此从 DB 重新加载。
|
||||
if request.result.compaction_performed && is_current_turn {
|
||||
let reload_topic = target_topic_id.unwrap_or(request.chat_id);
|
||||
if let Err(err) = session.reload_topic_history(request.chat_id, reload_topic) {
|
||||
tracing::error!(
|
||||
error = %err,
|
||||
chat_id = %request.chat_id,
|
||||
topic_id = %reload_topic,
|
||||
"Failed to reload topic history after in-loop compaction"
|
||||
);
|
||||
}
|
||||
tracing::info!(
|
||||
chat_id = %request.chat_id,
|
||||
topic_id = %reload_topic,
|
||||
"In-loop compaction was performed; reloaded topic history from DB"
|
||||
);
|
||||
} else if let Some(topic_id) = target_topic_id {
|
||||
// 将结果消息保存到确定的话题
|
||||
if let Some(topic_id) = target_topic_id {
|
||||
if is_current_turn {
|
||||
// 话题未切换（current_topic == original_topic_id），安全更新内存历史
|
||||
if let Err(err) = session
|
||||
.append_persisted_messages(topic_id, request.result.emitted_messages.clone())
|
||||
{
|
||||
if let Err(err) = session.append_persisted_messages(
|
||||
topic_id,
|
||||
request.result.emitted_messages.clone(),
|
||||
) {
|
||||
tracing::error!(
|
||||
error = %err,
|
||||
topic_id = %topic_id,
|
||||
@ -221,9 +153,10 @@ impl AgentExecutionService {
|
||||
} else if is_current_turn {
|
||||
// 没有话题：直接更新内存历史（append_persisted_messages 会处理持久化）
|
||||
// 无 topic 场景用 chat_id 作为 topic_histories 的回退 key
|
||||
if let Err(err) = session
|
||||
.append_persisted_messages(request.chat_id, request.result.emitted_messages.clone())
|
||||
{
|
||||
if let Err(err) = session.append_persisted_messages(
|
||||
request.chat_id,
|
||||
request.result.emitted_messages.clone(),
|
||||
) {
|
||||
tracing::error!(
|
||||
error = %err,
|
||||
chat_id = %request.chat_id,
|
||||
@ -259,13 +192,8 @@ impl AgentExecutionService {
|
||||
Vec::new()
|
||||
};
|
||||
|
||||
// 只有当是最新回合且未在 loop 内触发过任何压缩（工程化或 LLM）时，
|
||||
// 才触发兜底历史压缩。in-loop 已做工程化压缩时跳过——因为兜底基于
|
||||
// 未压缩历史的 estimate_tokens 判断会不准确，可能冗余触发 LLM 压缩，
|
||||
// 违背"in-loop 已判断工程化压缩足够则不 LLM 压缩"的意图。
|
||||
let should_schedule_compaction = is_current_turn
|
||||
&& !request.result.compaction_performed
|
||||
&& !request.result.engineering_compaction_applied;
|
||||
// 只有当是最新回合时才触发历史压缩
|
||||
let should_schedule_compaction = is_current_turn;
|
||||
|
||||
Ok(FinalizedAgentResult {
|
||||
outbound_messages,
|
||||
@ -346,13 +274,7 @@ impl AgentExecutionService {
|
||||
agent = agent.with_emitted_message_handler(handler);
|
||||
}
|
||||
|
||||
(
|
||||
history,
|
||||
agent,
|
||||
user_message,
|
||||
user_message_count,
|
||||
original_topic_id,
|
||||
)
|
||||
(history, agent, user_message, user_message_count, original_topic_id)
|
||||
};
|
||||
|
||||
// 构建系统提示词上下文
|
||||
@ -362,20 +284,7 @@ impl AgentExecutionService {
|
||||
user_message_count,
|
||||
};
|
||||
|
||||
// 构建 CompactionSink：在 AgentLoop 内部触发 LLM 压缩时把结果写回 DB。
|
||||
// topic_id 退化为 chat_id（与 history_key 一致）。
|
||||
let compaction_topic_id = original_topic_id
|
||||
.clone()
|
||||
.unwrap_or_else(|| request.chat_id.to_string());
|
||||
let compaction_sink = CompactionSinkImpl::new(
|
||||
request.session.clone(),
|
||||
request.chat_id.to_string(),
|
||||
compaction_topic_id,
|
||||
);
|
||||
|
||||
let result = agent
|
||||
.process(history, Some(&system_prompt_context), Some(&compaction_sink))
|
||||
.await?;
|
||||
let result = agent.process(history, Some(&system_prompt_context)).await?;
|
||||
let mut metadata = HashMap::new();
|
||||
// 把用户消息的 UUID 回传给前端，前端用此更新本地消息 ID，使 todo 点击跳转能匹配
|
||||
metadata.insert("user_message_id".to_string(), user_message.id.clone());
|
||||
@ -415,15 +324,7 @@ impl AgentExecutionService {
|
||||
// 等待该 topic 的前一条消息处理完成（含压缩）
|
||||
let _serial_guard = serial_lock.lock().await;
|
||||
|
||||
let (
|
||||
history,
|
||||
mut agent,
|
||||
user_message,
|
||||
user_message_count,
|
||||
original_topic_id,
|
||||
store,
|
||||
session_id,
|
||||
) = {
|
||||
let (history, mut agent, user_message, user_message_count, original_topic_id, store, session_id) = {
|
||||
let mut session_guard = request.session.lock().await;
|
||||
|
||||
session_guard.ensure_persistent_session(request.chat_id)?;
|
||||
@ -481,18 +382,12 @@ impl AgentExecutionService {
|
||||
|
||||
// 获取 store 和 session_id，用于构造消息持久化 handler
|
||||
let store = session_guard.store();
|
||||
let session_id =
|
||||
crate::storage::persistent_session_id(request.channel_name, request.chat_id);
|
||||
let session_id = crate::storage::persistent_session_id(
|
||||
request.channel_name,
|
||||
request.chat_id,
|
||||
);
|
||||
|
||||
(
|
||||
history,
|
||||
agent,
|
||||
user_message,
|
||||
user_message_count,
|
||||
original_topic_id,
|
||||
store,
|
||||
session_id,
|
||||
)
|
||||
(history, agent, user_message, user_message_count, original_topic_id, store, session_id)
|
||||
};
|
||||
|
||||
// 定时任务没有 live_emitter，需要 PersistingEmittedMessageHandler 来持久化消息
|
||||
@ -513,35 +408,22 @@ impl AgentExecutionService {
|
||||
user_message_count,
|
||||
};
|
||||
|
||||
// 构建 CompactionSink：在 AgentLoop 内部触发 LLM 压缩时把结果写回 DB。
|
||||
let compaction_topic_id = original_topic_id
|
||||
.clone()
|
||||
.unwrap_or_else(|| request.chat_id.to_string());
|
||||
let compaction_sink = CompactionSinkImpl::new(
|
||||
let result = agent.process(history, Some(&system_prompt_context)).await?;
|
||||
|
||||
let outbound_messages = self.finalize_result_and_schedule_compaction(
|
||||
request.session.clone(),
|
||||
request.chat_id.to_string(),
|
||||
compaction_topic_id,
|
||||
);
|
||||
|
||||
let result = agent
|
||||
.process(history, Some(&system_prompt_context), Some(&compaction_sink))
|
||||
.await?;
|
||||
|
||||
let outbound_messages = self
|
||||
.finalize_result_and_schedule_compaction(
|
||||
request.session.clone(),
|
||||
FinalizeAgentResultRequest {
|
||||
channel_name: request.channel_name,
|
||||
chat_id: request.chat_id,
|
||||
user_message: &user_message,
|
||||
result,
|
||||
metadata: request.metadata,
|
||||
suppress_live_tool_calls: false,
|
||||
execution_kind: "scheduled_task",
|
||||
original_topic_id: original_topic_id.clone(),
|
||||
},
|
||||
)
|
||||
.await?;
|
||||
FinalizeAgentResultRequest {
|
||||
channel_name: request.channel_name,
|
||||
chat_id: request.chat_id,
|
||||
user_message: &user_message,
|
||||
result,
|
||||
metadata: request.metadata,
|
||||
suppress_live_tool_calls: false,
|
||||
execution_kind: "scheduled_task",
|
||||
original_topic_id: original_topic_id.clone(),
|
||||
},
|
||||
)
|
||||
.await?;
|
||||
|
||||
// 清理内存历史，释放内存（数据库历史保留）
|
||||
{
|
||||
@ -661,7 +543,11 @@ mod tests {
|
||||
let _guard1 = lock.lock().await;
|
||||
|
||||
// 第二次获取应阻塞，1ms 超时验证
|
||||
let result = tokio::time::timeout(std::time::Duration::from_millis(1), lock.lock()).await;
|
||||
let result = tokio::time::timeout(
|
||||
std::time::Duration::from_millis(1),
|
||||
lock.lock(),
|
||||
)
|
||||
.await;
|
||||
|
||||
assert!(result.is_err(), "第二次获取同一锁应阻塞");
|
||||
}
|
||||
@ -675,8 +561,11 @@ mod tests {
|
||||
let _guard_a = lock_a.lock().await;
|
||||
|
||||
// 不同锁应立即可获取
|
||||
let result =
|
||||
tokio::time::timeout(std::time::Duration::from_millis(100), lock_b.lock()).await;
|
||||
let result = tokio::time::timeout(
|
||||
std::time::Duration::from_millis(100),
|
||||
lock_b.lock(),
|
||||
)
|
||||
.await;
|
||||
|
||||
assert!(result.is_ok(), "不同 topic 的锁应互不影响");
|
||||
}
|
||||
@ -696,7 +585,11 @@ mod tests {
|
||||
}
|
||||
|
||||
// 锁应已释放，可再次获取
|
||||
let result = tokio::time::timeout(std::time::Duration::from_millis(100), lock.lock()).await;
|
||||
let result = tokio::time::timeout(
|
||||
std::time::Duration::from_millis(100),
|
||||
lock.lock(),
|
||||
)
|
||||
.await;
|
||||
|
||||
assert!(result.is_ok(), "错误返回后锁应已释放");
|
||||
}
|
||||
|
||||
@ -1,8 +1,5 @@
|
||||
use axum::{Json, extract::{Query, State}};
|
||||
use axum::http::StatusCode;
|
||||
use axum::{
|
||||
Json,
|
||||
extract::{Query, State},
|
||||
};
|
||||
use serde::{Deserialize, Serialize};
|
||||
use std::sync::Arc;
|
||||
|
||||
@ -72,13 +69,6 @@ fn mask_config(config: &Config) -> Config {
|
||||
}
|
||||
}
|
||||
}
|
||||
// 掩码网关认证 token（避免通过 /api/config 泄露）
|
||||
if let Some(ref token) = masked.gateway.auth_token {
|
||||
if !token.is_empty() {
|
||||
let visible: String = token.chars().take(4).collect();
|
||||
masked.gateway.auth_token = Some(format!("{}{}", visible, API_KEY_MASK));
|
||||
}
|
||||
}
|
||||
masked
|
||||
}
|
||||
|
||||
@ -100,7 +90,9 @@ pub struct SaveConfigResponse {
|
||||
}
|
||||
|
||||
/// GET /api/config — Return current config with masked sensitive fields
|
||||
pub async fn get_config(State(state): State<Arc<GatewayState>>) -> Json<Config> {
|
||||
pub async fn get_config(
|
||||
State(state): State<Arc<GatewayState>>,
|
||||
) -> Json<Config> {
|
||||
Json(mask_config(&*state.config.read().await))
|
||||
}
|
||||
|
||||
@ -133,12 +125,6 @@ pub async fn save_config(
|
||||
}
|
||||
}
|
||||
}
|
||||
// 保留原始 auth_token（若提交的是掩码值）
|
||||
if let Some(ref submitted) = new_config.gateway.auth_token {
|
||||
if is_masked_key(submitted) {
|
||||
new_config.gateway.auth_token = cfg.gateway.auth_token.clone();
|
||||
}
|
||||
}
|
||||
} // read lock released here
|
||||
|
||||
// Validate timezone
|
||||
@ -152,19 +138,11 @@ pub async fn save_config(
|
||||
.unwrap_or_else(|_| get_default_config_path());
|
||||
|
||||
// Serialize and write to disk (no lock held)
|
||||
let json = serde_json::to_string_pretty(&new_config).map_err(|e| {
|
||||
(
|
||||
StatusCode::INTERNAL_SERVER_ERROR,
|
||||
format!("Serialize error: {}", e),
|
||||
)
|
||||
})?;
|
||||
let json = serde_json::to_string_pretty(&new_config)
|
||||
.map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, format!("Serialize error: {}", e)))?;
|
||||
|
||||
std::fs::write(&config_path, &json).map_err(|e| {
|
||||
(
|
||||
StatusCode::INTERNAL_SERVER_ERROR,
|
||||
format!("Write error: {}", e),
|
||||
)
|
||||
})?;
|
||||
std::fs::write(&config_path, &json)
|
||||
.map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, format!("Write error: {}", e)))?;
|
||||
|
||||
// Update in-memory config (write lock, held only for assignment)
|
||||
{
|
||||
@ -225,21 +203,6 @@ pub async fn restart(
|
||||
}))
|
||||
}
|
||||
|
||||
#[derive(Serialize)]
|
||||
pub struct ExecutionsResponse {
|
||||
/// 当前正在执行的 Agent 的 topic_id 列表
|
||||
pub topic_ids: Vec<String>,
|
||||
}
|
||||
|
||||
/// GET /api/executions — 返回当前正在执行的 Agent 的 topic_id 列表
|
||||
///
|
||||
/// 供前端重连时对账执行状态：前端据此判断断连期间哪些话题的
|
||||
/// 智能体仍在运行（需保持禁用）、哪些已完成（应解锁）。
|
||||
pub async fn list_executions(State(state): State<Arc<GatewayState>>) -> Json<ExecutionsResponse> {
|
||||
let topic_ids = state.cancel_manager.list_active_topic_ids().await;
|
||||
Json(ExecutionsResponse { topic_ids })
|
||||
}
|
||||
|
||||
/// GET /api/mcp/status — Return MCP server connection status
|
||||
pub async fn mcp_status(
|
||||
State(state): State<Arc<GatewayState>>,
|
||||
@ -263,7 +226,9 @@ pub async fn mcp_status(
|
||||
}
|
||||
|
||||
/// GET /api/skills — Return all discovered skills with their disabled status
|
||||
pub async fn skills_list(State(state): State<Arc<GatewayState>>) -> Json<SkillListResponse> {
|
||||
pub async fn skills_list(
|
||||
State(state): State<Arc<GatewayState>>,
|
||||
) -> Json<SkillListResponse> {
|
||||
let skills_enabled = state.config.read().await.skills.enabled;
|
||||
|
||||
if !skills_enabled {
|
||||
@ -316,7 +281,9 @@ pub struct CurrentModel {
|
||||
|
||||
/// GET /api/tools — Return all registered tools (builtin + MCP) with name/description/source.
|
||||
/// 通过 SessionManager::tools() 只读访问 ToolRegistry，不修改状态。
|
||||
pub async fn tools_list(State(state): State<Arc<GatewayState>>) -> Json<ToolsListResponse> {
|
||||
pub async fn tools_list(
|
||||
State(state): State<Arc<GatewayState>>,
|
||||
) -> Json<ToolsListResponse> {
|
||||
let registry = state.session_manager.tools();
|
||||
let tools: Vec<ToolInfo> = registry
|
||||
.get_definitions()
|
||||
@ -345,7 +312,9 @@ pub async fn tools_list(State(state): State<Arc<GatewayState>>) -> Json<ToolsLis
|
||||
}
|
||||
|
||||
/// GET /api/model-options — 返回 config.json 中配置的 provider/model 名列表。
|
||||
pub async fn model_options(State(state): State<Arc<GatewayState>>) -> Json<ModelOptionsResponse> {
|
||||
pub async fn model_options(
|
||||
State(state): State<Arc<GatewayState>>,
|
||||
) -> Json<ModelOptionsResponse> {
|
||||
let config = state.config.read().await;
|
||||
let resolver = crate::config::ModelResolver::from_config(&config);
|
||||
// 当前默认 agent 的 provider/model 名（直接引用 providers/models 表的 key）
|
||||
@ -402,11 +371,7 @@ pub async fn skills_toggle(
|
||||
changed: Some(change.changed),
|
||||
available: Some(change.available),
|
||||
disabled_in_scopes: Some(
|
||||
change
|
||||
.disabled_in_scopes
|
||||
.iter()
|
||||
.map(|s| s.as_str().to_string())
|
||||
.collect(),
|
||||
change.disabled_in_scopes.iter().map(|s| s.as_str().to_string()).collect(),
|
||||
),
|
||||
error: None,
|
||||
}),
|
||||
@ -459,7 +424,9 @@ pub struct SubagentListResponse {
|
||||
}
|
||||
|
||||
/// GET /api/subagents — Return all discovered subagents with their disabled status
|
||||
pub async fn subagents_list(State(state): State<Arc<GatewayState>>) -> Json<SubagentListResponse> {
|
||||
pub async fn subagents_list(
|
||||
State(state): State<Arc<GatewayState>>,
|
||||
) -> Json<SubagentListResponse> {
|
||||
let subagents_enabled = state.config.read().await.subagents.enabled;
|
||||
|
||||
if !subagents_enabled {
|
||||
@ -606,7 +573,6 @@ pub async fn subagents_update(
|
||||
capability: updated.capability.clone(),
|
||||
provider: updated.provider.clone(),
|
||||
model: updated.model.clone(),
|
||||
body: updated.body.clone(),
|
||||
});
|
||||
|
||||
Ok(Json(SubagentUpdateResponse {
|
||||
@ -616,124 +582,6 @@ pub async fn subagents_update(
|
||||
}))
|
||||
}
|
||||
|
||||
#[derive(Deserialize)]
|
||||
pub struct SubagentCreateRequest {
|
||||
pub name: String,
|
||||
pub description: String,
|
||||
#[serde(default)]
|
||||
pub body: String,
|
||||
pub scope: String,
|
||||
#[serde(default)]
|
||||
pub capability: CapabilityPolicy,
|
||||
#[serde(default)]
|
||||
pub provider: Option<String>,
|
||||
#[serde(default)]
|
||||
pub model: Option<String>,
|
||||
}
|
||||
|
||||
#[derive(Deserialize)]
|
||||
pub struct SubagentDeleteRequest {
|
||||
pub name: String,
|
||||
}
|
||||
|
||||
#[derive(Serialize)]
|
||||
pub struct SubagentDeleteResponse {
|
||||
pub success: bool,
|
||||
pub path: String,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub error: Option<String>,
|
||||
}
|
||||
|
||||
/// POST /api/subagents/create — Create a new subagent (writes SUBAGENT.md)
|
||||
pub async fn subagents_create(
|
||||
State(state): State<Arc<GatewayState>>,
|
||||
Json(req): Json<SubagentCreateRequest>,
|
||||
) -> Result<Json<SubagentUpdateResponse>, (StatusCode, String)> {
|
||||
let scope = SubagentScope::parse(&req.scope).ok_or_else(|| {
|
||||
(
|
||||
StatusCode::BAD_REQUEST,
|
||||
format!("invalid scope: {}", req.scope),
|
||||
)
|
||||
})?;
|
||||
|
||||
let created = state
|
||||
.subagent_runtime
|
||||
.create_subagent(
|
||||
scope,
|
||||
&req.name,
|
||||
&req.description,
|
||||
&req.body,
|
||||
&req.capability,
|
||||
&req.provider,
|
||||
&req.model,
|
||||
true,
|
||||
)
|
||||
.map_err(|err| {
|
||||
let status = if err.contains("already exists") {
|
||||
StatusCode::CONFLICT
|
||||
} else {
|
||||
StatusCode::BAD_REQUEST
|
||||
};
|
||||
(status, err)
|
||||
})?;
|
||||
|
||||
// 返回创建后的状态（含 disabled_in_scopes）
|
||||
let status = state
|
||||
.subagent_runtime
|
||||
.list_with_status()
|
||||
.into_iter()
|
||||
.find(|s| s.name == created.name)
|
||||
.unwrap_or_else(|| SubagentWithStatus {
|
||||
name: created.name.clone(),
|
||||
description: created.description.clone(),
|
||||
source: created.source.as_str().to_string(),
|
||||
disabled_in_scopes: vec![],
|
||||
capability: created.capability.clone(),
|
||||
provider: created.provider.clone(),
|
||||
model: created.model.clone(),
|
||||
body: created.body.clone(),
|
||||
});
|
||||
|
||||
Ok(Json(SubagentUpdateResponse {
|
||||
success: true,
|
||||
subagent: Some(status),
|
||||
error: None,
|
||||
}))
|
||||
}
|
||||
|
||||
/// DELETE /api/subagents/delete?name= — Delete a subagent (removes SUBAGENT.md)
|
||||
pub async fn subagents_delete(
|
||||
State(state): State<Arc<GatewayState>>,
|
||||
Query(req): Query<SubagentDeleteRequest>,
|
||||
) -> Result<Json<SubagentDeleteResponse>, (StatusCode, Json<SubagentDeleteResponse>)> {
|
||||
let path = state
|
||||
.subagent_runtime
|
||||
.delete_subagent(&req.name, true)
|
||||
.map_err(|err| {
|
||||
let status = if err.contains("not found") {
|
||||
StatusCode::NOT_FOUND
|
||||
} else if err.contains("builtin") {
|
||||
StatusCode::BAD_REQUEST
|
||||
} else {
|
||||
StatusCode::INTERNAL_SERVER_ERROR
|
||||
};
|
||||
(
|
||||
status,
|
||||
Json(SubagentDeleteResponse {
|
||||
success: false,
|
||||
path: String::new(),
|
||||
error: Some(err),
|
||||
}),
|
||||
)
|
||||
})?;
|
||||
|
||||
Ok(Json(SubagentDeleteResponse {
|
||||
success: true,
|
||||
path: path.display().to_string(),
|
||||
error: None,
|
||||
}))
|
||||
}
|
||||
|
||||
// ===================== Experts =====================
|
||||
|
||||
#[derive(Deserialize)]
|
||||
@ -877,7 +725,9 @@ pub struct ExpertDeleteResponse {
|
||||
}
|
||||
|
||||
/// GET /api/experts — Return all discovered experts with their disabled status
|
||||
pub async fn experts_list(State(state): State<Arc<GatewayState>>) -> Json<ExpertListResponse> {
|
||||
pub async fn experts_list(
|
||||
State(state): State<Arc<GatewayState>>,
|
||||
) -> Json<ExpertListResponse> {
|
||||
let experts_enabled = state.config.read().await.experts.enabled;
|
||||
|
||||
if !experts_enabled {
|
||||
@ -967,12 +817,8 @@ pub async fn experts_create(
|
||||
State(state): State<Arc<GatewayState>>,
|
||||
Json(req): Json<ExpertCreateRequest>,
|
||||
) -> Result<Json<ExpertResponse>, (StatusCode, String)> {
|
||||
let scope = ExpertScope::parse(&req.scope).ok_or_else(|| {
|
||||
(
|
||||
StatusCode::BAD_REQUEST,
|
||||
format!("invalid scope: {}", req.scope),
|
||||
)
|
||||
})?;
|
||||
let scope = ExpertScope::parse(&req.scope)
|
||||
.ok_or_else(|| (StatusCode::BAD_REQUEST, format!("invalid scope: {}", req.scope)))?;
|
||||
|
||||
let expert = state
|
||||
.experts
|
||||
@ -1003,12 +849,8 @@ pub async fn experts_update(
|
||||
State(state): State<Arc<GatewayState>>,
|
||||
Json(req): Json<ExpertUpdateRequest>,
|
||||
) -> Result<Json<ExpertResponse>, (StatusCode, String)> {
|
||||
let scope = ExpertScope::parse(&req.scope).ok_or_else(|| {
|
||||
(
|
||||
StatusCode::BAD_REQUEST,
|
||||
format!("invalid scope: {}", req.scope),
|
||||
)
|
||||
})?;
|
||||
let scope = ExpertScope::parse(&req.scope)
|
||||
.ok_or_else(|| (StatusCode::BAD_REQUEST, format!("invalid scope: {}", req.scope)))?;
|
||||
|
||||
let expert = state
|
||||
.experts
|
||||
@ -1169,7 +1011,9 @@ pub async fn session_select_model(
|
||||
}
|
||||
drop(config);
|
||||
|
||||
state.model_selections.set(&req.session_id, provider, model);
|
||||
state
|
||||
.model_selections
|
||||
.set(&req.session_id, provider, model);
|
||||
(
|
||||
StatusCode::OK,
|
||||
Json(SelectModelResponse {
|
||||
|
||||
@ -26,7 +26,7 @@ pub(crate) struct MemoryMaintenanceCandidate {
|
||||
pub(crate) namespace: String,
|
||||
pub(crate) key: String,
|
||||
pub(crate) content: String,
|
||||
pub(crate) updated_at: i64, // 记忆更新时间（Unix timestamp）
|
||||
pub(crate) updated_at: i64, // 记忆更新时间（Unix timestamp）
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Default, Serialize, Deserialize, PartialEq, Eq)]
|
||||
@ -227,8 +227,8 @@ impl MemoryMaintenanceService {
|
||||
Ok(parsed) => return Ok(parsed),
|
||||
Err(err) => {
|
||||
let error_msg = err.to_string();
|
||||
let is_truncated =
|
||||
error_msg.contains("EOF while parsing") || error_msg.contains("expected");
|
||||
let is_truncated = error_msg.contains("EOF while parsing")
|
||||
|| error_msg.contains("expected");
|
||||
|
||||
let should_retry = delay_ms.is_some() && is_truncated;
|
||||
last_error = Some(error_msg.clone());
|
||||
@ -369,10 +369,9 @@ impl MemoryMaintenanceService {
|
||||
pub(crate) async fn run_for_all_scopes(
|
||||
&self,
|
||||
) -> Result<Option<MemoryMaintenanceScopeResult>, AgentError> {
|
||||
let scope_keys = self
|
||||
.store
|
||||
.list_memory_scope_keys("user")
|
||||
.map_err(|err| AgentError::Other(format!("list memory scope keys error: {}", err)))?;
|
||||
let scope_keys = self.store.list_memory_scope_keys("user").map_err(|err| {
|
||||
AgentError::Other(format!("list memory scope keys error: {}", err))
|
||||
})?;
|
||||
|
||||
if scope_keys.is_empty() {
|
||||
return Ok(None);
|
||||
@ -419,8 +418,7 @@ impl MemoryMaintenanceService {
|
||||
let managed_markdown = if all_remaining_memories.is_empty() {
|
||||
String::new()
|
||||
} else {
|
||||
self.generate_summary("all", &all_remaining_memories)
|
||||
.await?
|
||||
self.generate_summary("all", &all_remaining_memories).await?
|
||||
};
|
||||
|
||||
if !managed_markdown.is_empty() {
|
||||
@ -680,29 +678,24 @@ pub(crate) fn validate_memory_maintenance_output(
|
||||
}
|
||||
|
||||
// 验证 2: 跨 namespace 合并检测（完全禁止）
|
||||
let candidates_by_id: HashMap<&str, &MemoryMaintenanceCandidate> =
|
||||
plan.candidates.iter().map(|c| (c.id.as_str(), c)).collect();
|
||||
let candidates_by_id: HashMap<&str, &MemoryMaintenanceCandidate> = plan
|
||||
.candidates
|
||||
.iter()
|
||||
.map(|c| (c.id.as_str(), c))
|
||||
.collect();
|
||||
|
||||
for merge in &output.merges {
|
||||
let source_namespaces: HashSet<&str> = merge
|
||||
.source_ids
|
||||
.iter()
|
||||
.filter_map(|id| {
|
||||
candidates_by_id
|
||||
.get(id.as_str())
|
||||
.map(|c| c.namespace.as_str())
|
||||
})
|
||||
.filter_map(|id| candidates_by_id.get(id.as_str()).map(|c| c.namespace.as_str()))
|
||||
.collect();
|
||||
|
||||
// 检查是否跨越多个 namespace
|
||||
if source_namespaces.len() > 1 {
|
||||
return Err(format!(
|
||||
"跨 namespace 合并被禁止: 源来自 {}",
|
||||
source_namespaces
|
||||
.iter()
|
||||
.cloned()
|
||||
.collect::<Vec<_>>()
|
||||
.join(", ")
|
||||
source_namespaces.iter().cloned().collect::<Vec<_>>().join(", ")
|
||||
));
|
||||
}
|
||||
|
||||
@ -725,7 +718,11 @@ pub(crate) fn validate_memory_maintenance_output(
|
||||
.map(|s| s.as_str())
|
||||
.collect();
|
||||
|
||||
let deleted_ids: HashSet<&str> = output.low_value_ids.iter().map(|s| s.as_str()).collect();
|
||||
let deleted_ids: HashSet<&str> = output
|
||||
.low_value_ids
|
||||
.iter()
|
||||
.map(|s| s.as_str())
|
||||
.collect();
|
||||
|
||||
let affected = merged_ids.len() + deleted_ids.len();
|
||||
let max_allowed = (total as f32 * max_merge_ratio).ceil() as usize;
|
||||
@ -761,14 +758,8 @@ pub(crate) fn apply_memory_maintenance_output(
|
||||
max_merge_per_group: usize,
|
||||
) -> Result<(), AgentError> {
|
||||
// 新增: 验证合并输出
|
||||
validate_memory_maintenance_output(
|
||||
plan,
|
||||
output,
|
||||
max_merge_ratio,
|
||||
min_memories_to_keep,
|
||||
max_merge_per_group,
|
||||
)
|
||||
.map_err(|e| AgentError::Other(e))?;
|
||||
validate_memory_maintenance_output(plan, output, max_merge_ratio, min_memories_to_keep, max_merge_per_group)
|
||||
.map_err(|e| AgentError::Other(e))?;
|
||||
|
||||
let all_candidates = plan.candidates.clone();
|
||||
|
||||
|
||||
@ -1,7 +1,6 @@
|
||||
pub mod agent_factory;
|
||||
pub mod agent_prompt_provider;
|
||||
pub mod agent_task_executor;
|
||||
pub mod auth;
|
||||
pub mod cancel_manager;
|
||||
pub mod cli_session;
|
||||
pub mod command;
|
||||
@ -26,16 +25,15 @@ pub mod session_message_sender;
|
||||
pub mod session_message_service;
|
||||
pub mod session_pool;
|
||||
pub mod static_files;
|
||||
pub mod tool_prompt_provider;
|
||||
pub mod tool_registry_factory;
|
||||
pub mod tool_prompt_provider;
|
||||
pub mod ws;
|
||||
|
||||
use axum::{Router, middleware, routing};
|
||||
use axum::{Router, routing};
|
||||
use std::collections::HashMap;
|
||||
use std::sync::Arc;
|
||||
use tokio::net::TcpSocket;
|
||||
use tokio::sync::Semaphore;
|
||||
use tower_http::cors::{Any, CorsLayer};
|
||||
use tower_http::services::ServeDir;
|
||||
|
||||
use crate::bus::MessageBus;
|
||||
@ -52,11 +50,11 @@ use cancel_manager::CancelManager;
|
||||
use outbound_dispatcher::OutboundDispatcher;
|
||||
use processor::InboundProcessor;
|
||||
use runtime::build_session_manager_with_sender;
|
||||
use session::SessionManager;
|
||||
use session_message_sender::BusSessionMessageSender;
|
||||
use session::SessionManager;
|
||||
use static_files::static_handler;
|
||||
|
||||
use tokio::sync::{RwLock, watch};
|
||||
use tokio::sync::{watch, RwLock};
|
||||
|
||||
pub struct GatewayState {
|
||||
pub config: Arc<RwLock<Config>>,
|
||||
@ -75,10 +73,7 @@ pub struct GatewayState {
|
||||
}
|
||||
|
||||
impl GatewayState {
|
||||
pub fn from_config(
|
||||
config: Config,
|
||||
restart_tx: watch::Sender<bool>,
|
||||
) -> Result<Self, Box<dyn std::error::Error>> {
|
||||
pub fn from_config(config: Config, restart_tx: watch::Sender<bool>) -> Result<Self, Box<dyn std::error::Error>> {
|
||||
// Get provider config for SessionManager
|
||||
let provider_config = config.get_provider_config("default")?;
|
||||
let mut provider_configs = HashMap::<String, LLMProviderConfig>::new();
|
||||
@ -92,9 +87,7 @@ impl GatewayState {
|
||||
let session_ttl_hours = config.gateway.session_ttl_hours;
|
||||
|
||||
let skills = Arc::new(SkillRuntime::from_config(config.skills.clone()));
|
||||
let experts = Arc::new(crate::experts::ExpertRuntime::from_config(
|
||||
config.experts.clone(),
|
||||
));
|
||||
let experts = Arc::new(crate::experts::ExpertRuntime::from_config(config.experts.clone()));
|
||||
let channel_manager = ChannelManager::new();
|
||||
let bus = channel_manager.bus();
|
||||
|
||||
@ -102,26 +95,24 @@ impl GatewayState {
|
||||
mcp_servers: config.mcp_servers.clone(),
|
||||
};
|
||||
|
||||
let (session_manager, task_repository, mcp_manager, subagent_runtime, model_selections) =
|
||||
build_session_manager_with_sender(
|
||||
agent_prompt_reinject_every,
|
||||
show_tool_results,
|
||||
config.time.timezone.clone(),
|
||||
provider_config,
|
||||
provider_configs,
|
||||
skills.clone(),
|
||||
experts.clone(),
|
||||
Arc::new(BusSessionMessageSender::new(bus.clone())),
|
||||
std::collections::HashSet::new(),
|
||||
config.tools.task.clone(),
|
||||
config.subagents.clone(),
|
||||
config.memory_maintenance.clone(),
|
||||
session_ttl_hours,
|
||||
mcp_config,
|
||||
Some(bus.clone()),
|
||||
Arc::new(crate::config::ModelResolver::from_config(&config)),
|
||||
config.compaction.clone(),
|
||||
)?;
|
||||
let (session_manager, task_repository, mcp_manager, subagent_runtime, model_selections) = build_session_manager_with_sender(
|
||||
agent_prompt_reinject_every,
|
||||
show_tool_results,
|
||||
config.time.timezone.clone(),
|
||||
provider_config,
|
||||
provider_configs,
|
||||
skills.clone(),
|
||||
experts.clone(),
|
||||
Arc::new(BusSessionMessageSender::new(bus.clone())),
|
||||
std::collections::HashSet::new(),
|
||||
config.tools.task.clone(),
|
||||
config.subagents.clone(),
|
||||
config.memory_maintenance.clone(),
|
||||
session_ttl_hours,
|
||||
mcp_config,
|
||||
Some(bus.clone()),
|
||||
Arc::new(crate::config::ModelResolver::from_config(&config)),
|
||||
)?;
|
||||
|
||||
// 诊断日志：记录新 GatewayState 的创建（用于排查重启后是否使用了新状态）
|
||||
tracing::info!(
|
||||
@ -164,13 +155,8 @@ impl GatewayState {
|
||||
drop(cfg); // release read lock before spawning long-running tasks
|
||||
|
||||
let semaphore = Arc::new(Semaphore::new(max_concurrent));
|
||||
let inbound_processor = InboundProcessor::new(
|
||||
self.bus.clone(),
|
||||
self.session_manager.clone(),
|
||||
semaphore,
|
||||
provider_config,
|
||||
self.cancel_manager.clone(),
|
||||
);
|
||||
let inbound_processor =
|
||||
InboundProcessor::new(self.bus.clone(), self.session_manager.clone(), semaphore, provider_config, self.cancel_manager.clone());
|
||||
tokio::spawn(inbound_processor.run());
|
||||
|
||||
// Spawn outbound dispatcher
|
||||
@ -241,132 +227,70 @@ pub async fn run(
|
||||
}
|
||||
|
||||
// CLI args override config file values
|
||||
let (bind_host, bind_port, auth_token) = {
|
||||
let (bind_host, bind_port) = {
|
||||
let cfg = state.config.read().await;
|
||||
let h = host.unwrap_or_else(|| cfg.gateway.host.clone());
|
||||
let p = port.unwrap_or(cfg.gateway.port);
|
||||
(h, p, cfg.gateway.auth_token.clone())
|
||||
(h, p)
|
||||
};
|
||||
|
||||
// 安全校验：绑定到非 loopback 地址时必须配置 auth_token
|
||||
if !auth::is_loopback_host(&bind_host) && auth_token.is_none() {
|
||||
return Err(format!(
|
||||
"Gateway is bound to non-loopback address '{}' but no `gateway.auth_token` is configured. \
|
||||
Remote access requires authentication. \
|
||||
Please set `auth_token` in the `gateway` section of config.json.",
|
||||
bind_host
|
||||
)
|
||||
.into());
|
||||
}
|
||||
|
||||
let auth_required = auth::requires_auth(&bind_host, &auth_token);
|
||||
let auth_config = auth::AuthConfig {
|
||||
token: auth_token.clone(),
|
||||
};
|
||||
|
||||
if auth_required {
|
||||
tracing::info!(
|
||||
host = %bind_host,
|
||||
has_token = auth_token.is_some(),
|
||||
"Authentication enabled for gateway"
|
||||
);
|
||||
} else {
|
||||
tracing::info!(
|
||||
host = %bind_host,
|
||||
"Authentication disabled (loopback binding without explicit token)"
|
||||
);
|
||||
}
|
||||
|
||||
// 使用嵌入的静态文件（编译时打包进二进制）
|
||||
// 开发模式下可通过 STATIC_DIR 环境变量使用磁盘文件
|
||||
let use_embedded = std::env::var("STATIC_DIR").is_err();
|
||||
|
||||
// 公共路由：生产/开发两种模式共享，避免重复注册导致漏配。
|
||||
// 仅 fallback（嵌入 vs 磁盘）与 state 绑定按模式区分。
|
||||
let app = Router::new()
|
||||
.route("/health", routing::get(http::health))
|
||||
.route(
|
||||
"/api/config",
|
||||
routing::get(http::get_config).put(http::save_config),
|
||||
)
|
||||
.route("/api/restart", routing::post(http::restart))
|
||||
.route("/api/executions", routing::get(http::list_executions))
|
||||
.route("/api/mcp/status", routing::get(http::mcp_status))
|
||||
.route("/api/skills", routing::get(http::skills_list))
|
||||
.route("/api/skills/toggle", routing::post(http::skills_toggle))
|
||||
.route("/api/tools", routing::get(http::tools_list))
|
||||
.route("/api/model-options", routing::get(http::model_options))
|
||||
.route("/api/subagents", routing::get(http::subagents_list))
|
||||
.route(
|
||||
"/api/subagents/toggle",
|
||||
routing::post(http::subagents_toggle),
|
||||
)
|
||||
.route(
|
||||
"/api/subagents/update",
|
||||
routing::put(http::subagents_update),
|
||||
)
|
||||
.route(
|
||||
"/api/subagents/create",
|
||||
routing::post(http::subagents_create),
|
||||
)
|
||||
.route(
|
||||
"/api/subagents/delete",
|
||||
routing::delete(http::subagents_delete),
|
||||
)
|
||||
.route("/api/experts", routing::get(http::experts_list))
|
||||
.route("/api/experts/toggle", routing::post(http::experts_toggle))
|
||||
.route("/api/experts/create", routing::post(http::experts_create))
|
||||
.route("/api/experts/update", routing::put(http::experts_update))
|
||||
.route("/api/experts/delete", routing::delete(http::experts_delete))
|
||||
.route(
|
||||
"/api/experts/selected",
|
||||
routing::get(http::experts_selected),
|
||||
)
|
||||
.route("/api/experts/select", routing::post(http::experts_select))
|
||||
.route(
|
||||
"/api/session/select-model",
|
||||
routing::post(http::session_select_model),
|
||||
)
|
||||
.route(
|
||||
"/api/session/selected-model",
|
||||
routing::get(http::session_selected_model),
|
||||
)
|
||||
.route("/ws", routing::get(ws::ws_handler));
|
||||
|
||||
// 仅 fallback 按模式区分：嵌入资源 vs 磁盘目录。
|
||||
// fallback 必须在 with_state 之前调用，否则 handler 的 State 类型无法推断。
|
||||
let app = if use_embedded {
|
||||
app.fallback(static_handler)
|
||||
Router::new()
|
||||
.route("/health", routing::get(http::health))
|
||||
.route("/api/config", routing::get(http::get_config).put(http::save_config))
|
||||
.route("/api/restart", routing::post(http::restart))
|
||||
.route("/api/mcp/status", routing::get(http::mcp_status))
|
||||
.route("/api/skills", routing::get(http::skills_list))
|
||||
.route("/api/skills/toggle", routing::post(http::skills_toggle))
|
||||
.route("/api/tools", routing::get(http::tools_list))
|
||||
.route("/api/model-options", routing::get(http::model_options))
|
||||
.route("/api/subagents", routing::get(http::subagents_list))
|
||||
.route("/api/subagents/toggle", routing::post(http::subagents_toggle))
|
||||
.route("/api/subagents/update", routing::put(http::subagents_update))
|
||||
.route("/api/experts", routing::get(http::experts_list))
|
||||
.route("/api/experts/toggle", routing::post(http::experts_toggle))
|
||||
.route("/api/experts/create", routing::post(http::experts_create))
|
||||
.route("/api/experts/update", routing::put(http::experts_update))
|
||||
.route("/api/experts/delete", routing::delete(http::experts_delete))
|
||||
.route("/api/experts/selected", routing::get(http::experts_selected))
|
||||
.route("/api/experts/select", routing::post(http::experts_select))
|
||||
.route("/api/session/select-model", routing::post(http::session_select_model))
|
||||
.route("/api/session/selected-model", routing::get(http::session_selected_model))
|
||||
.route("/ws", routing::get(ws::ws_handler))
|
||||
.fallback(static_handler)
|
||||
.with_state(state.clone())
|
||||
} else {
|
||||
let static_dir = std::env::var("STATIC_DIR").unwrap_or_else(|_| "static".to_string());
|
||||
app.fallback_service(ServeDir::new(&static_dir))
|
||||
}
|
||||
.with_state(state.clone());
|
||||
|
||||
// 条件性挂载认证中间件：仅在需要认证时启用。
|
||||
// 中间件内部按 path 前缀判断，仅 /api/* 需要校验；
|
||||
// /health、/ws（WS 在 handler 内单独校验）和静态资源放行。
|
||||
let app = if auth_required {
|
||||
app.layer(axum::Extension(auth_config))
|
||||
.layer(middleware::from_fn(auth::require_bearer_auth))
|
||||
} else {
|
||||
app
|
||||
Router::new()
|
||||
.route("/health", routing::get(http::health))
|
||||
.route("/api/config", routing::get(http::get_config).put(http::save_config))
|
||||
.route("/api/restart", routing::post(http::restart))
|
||||
.route("/api/mcp/status", routing::get(http::mcp_status))
|
||||
.route("/api/skills", routing::get(http::skills_list))
|
||||
.route("/api/skills/toggle", routing::post(http::skills_toggle))
|
||||
.route("/api/tools", routing::get(http::tools_list))
|
||||
.route("/api/model-options", routing::get(http::model_options))
|
||||
.route("/api/subagents", routing::get(http::subagents_list))
|
||||
.route("/api/subagents/toggle", routing::post(http::subagents_toggle))
|
||||
.route("/api/subagents/update", routing::put(http::subagents_update))
|
||||
.route("/api/experts", routing::get(http::experts_list))
|
||||
.route("/api/experts/toggle", routing::post(http::experts_toggle))
|
||||
.route("/api/experts/create", routing::post(http::experts_create))
|
||||
.route("/api/experts/update", routing::put(http::experts_update))
|
||||
.route("/api/experts/delete", routing::delete(http::experts_delete))
|
||||
.route("/api/experts/selected", routing::get(http::experts_selected))
|
||||
.route("/api/experts/select", routing::post(http::experts_select))
|
||||
.route("/api/session/select-model", routing::post(http::session_select_model))
|
||||
.route("/api/session/selected-model", routing::get(http::session_selected_model))
|
||||
.route("/ws", routing::get(ws::ws_handler))
|
||||
.fallback_service(ServeDir::new(&static_dir))
|
||||
.with_state(state.clone())
|
||||
};
|
||||
|
||||
// CORS：loopback 下宽松（仅同源）；非 loopback 下允许任意来源（由 auth_token 保护）。
|
||||
// 不论哪种情况都显式设置以避免浏览器默认行为差异。
|
||||
let cors = if auth::is_loopback_host(&bind_host) {
|
||||
// 本地开发：同源即可，阻止跨域（防 DNS rebinding）
|
||||
CorsLayer::new()
|
||||
.allow_origin(tower_http::cors::AllowOrigin::mirror_request())
|
||||
.allow_methods(Any)
|
||||
.allow_headers(Any)
|
||||
} else {
|
||||
// 远程访问：允许跨域，但由 token 保护
|
||||
CorsLayer::permissive()
|
||||
};
|
||||
let app = app.layer(cors);
|
||||
|
||||
let addr: std::net::SocketAddr = format!("{}:{}", bind_host, bind_port).parse()?;
|
||||
let listener = {
|
||||
let socket = match addr {
|
||||
|
||||
@ -1,5 +1,5 @@
|
||||
use std::collections::HashMap;
|
||||
use parking_lot::RwLock;
|
||||
use std::sync::RwLock;
|
||||
|
||||
/// per-session 的用户模型覆盖选择存储。
|
||||
///
|
||||
@ -16,10 +16,16 @@ impl ModelSelectionStore {
|
||||
}
|
||||
|
||||
/// 设置 session 的用户模型覆盖。provider 和 model 均为 None 时清除该 session 的选择。
|
||||
pub fn set(&self, session_id: &str, provider: Option<String>, model: Option<String>) {
|
||||
pub fn set(
|
||||
&self,
|
||||
session_id: &str,
|
||||
provider: Option<String>,
|
||||
model: Option<String>,
|
||||
) {
|
||||
let mut selections = self
|
||||
.selections
|
||||
.write();
|
||||
.write()
|
||||
.expect("model selections rwlock poisoned");
|
||||
if provider.is_none() && model.is_none() {
|
||||
selections.remove(session_id);
|
||||
} else {
|
||||
@ -31,6 +37,7 @@ impl ModelSelectionStore {
|
||||
pub fn get(&self, session_id: &str) -> Option<(Option<String>, Option<String>)> {
|
||||
self.selections
|
||||
.read()
|
||||
.expect("model selections rwlock poisoned")
|
||||
.get(session_id)
|
||||
.cloned()
|
||||
}
|
||||
@ -69,6 +76,9 @@ mod tests {
|
||||
fn set_only_provider_keeps_entry() {
|
||||
let store = ModelSelectionStore::new();
|
||||
store.set("s1", Some("p1".to_string()), None);
|
||||
assert_eq!(store.get("s1"), Some((Some("p1".to_string()), None)));
|
||||
assert_eq!(
|
||||
store.get("s1"),
|
||||
Some((Some("p1".to_string()), None))
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
@ -1,69 +1,20 @@
|
||||
use std::collections::HashMap;
|
||||
use std::sync::Arc;
|
||||
|
||||
use tokio::sync::{RwLock, mpsc};
|
||||
use tokio_util::sync::CancellationToken;
|
||||
use tokio::sync::RwLock;
|
||||
|
||||
use crate::bus::message::OutboundEventKind;
|
||||
use crate::bus::{MessageBus, OutboundMessage};
|
||||
use crate::channels::base::{Channel, ChannelError};
|
||||
|
||||
/// 低优先级队列容量（ToolCall / ToolResult / StreamDelta 等中间过程事件）。
|
||||
///
|
||||
/// 略小于 MessageBus 的容量，确保 bus 的 `try_send` 丢消息防线仍有效——
|
||||
/// channel 队列满时 dispatcher 立即丢弃该消息，不阻塞路由循环影响其他 channel。
|
||||
const LOW_PRIORITY_QUEUE_CAPACITY: usize = 64;
|
||||
|
||||
/// 高优先级队列容量（仅 AssistantResponse 最终响应）。
|
||||
///
|
||||
/// 单次 agent 执行仅产生 1 条最终响应；32 足够缓冲多 topic 并发的最终响应，
|
||||
/// 极罕见满。最终响应是 agent 与用户的契约，必达。
|
||||
const HIGH_PRIORITY_QUEUE_CAPACITY: usize = 32;
|
||||
|
||||
/// 低优先级事件的发送重试间隔（秒）。
|
||||
const RETRY_DELAYS_SECS: [u64; 3] = [1, 2, 4];
|
||||
|
||||
/// 高优先级事件（最终响应）的发送重试间隔（秒）——更多次、更长退避，
|
||||
/// 尽最大努力送达最终响应。
|
||||
const EXTENDED_RETRY_DELAYS_SECS: [u64; 5] = [1, 2, 4, 8, 16];
|
||||
/// Consumes outbound messages from MessageBus and dispatches them to channels.
|
||||
pub struct OutboundDispatcher {
|
||||
bus: Arc<MessageBus>,
|
||||
channels: Arc<RwLock<HashMap<String, Arc<dyn Channel + Send + Sync>>>>,
|
||||
}
|
||||
|
||||
/// Prefix for virtual scheduler chat IDs that should not be sent to external channels.
|
||||
const SCHEDULER_VIRTUAL_CHAT_ID_PREFIX: &str = "scheduler/";
|
||||
|
||||
/// 判断消息是否为高优先级（必达）。
|
||||
///
|
||||
/// `AssistantResponse`（最终响应）和 `ErrorNotification`（agent 异常终止的错误通知）
|
||||
/// 都是 agent 与用户的终态契约，必须送达；其他事件（工具调用进度、流式增量、
|
||||
/// 执行完成信号等）是中间过程，可丢。
|
||||
fn is_high_priority(msg: &OutboundMessage) -> bool {
|
||||
matches!(
|
||||
msg.event_kind,
|
||||
OutboundEventKind::AssistantResponse | OutboundEventKind::ErrorNotification
|
||||
)
|
||||
}
|
||||
|
||||
/// 单个 channel 的发送上下文：双优先级 mpsc 队列 + sender task。
|
||||
///
|
||||
/// dispatcher 按消息优先级 `try_send` 到 `high_tx`（最终响应，必达）或
|
||||
/// `low_tx`（中间过程，可丢）；`sender_task` 优先消费 high 队列，再消费
|
||||
/// low 队列，串行调用 `Channel::send`（含重试）。channel 之间完全隔离——
|
||||
/// 某个 channel 的慢发送或重试 sleep 不会阻塞其他 channel 的消息投递。
|
||||
#[derive(Clone)]
|
||||
struct ChannelSink {
|
||||
high_tx: mpsc::Sender<OutboundMessage>,
|
||||
low_tx: mpsc::Sender<OutboundMessage>,
|
||||
cancel: CancellationToken,
|
||||
}
|
||||
|
||||
/// Consumes outbound messages from MessageBus and dispatches them to channels.
|
||||
///
|
||||
/// 架构：dispatcher 主循环只负责路由（O(1) try_send），不参与发送。
|
||||
/// 每个 channel 拥有独立的 sender task 和有界队列，实现 channel 级隔离。
|
||||
pub struct OutboundDispatcher {
|
||||
bus: Arc<MessageBus>,
|
||||
channels: Arc<RwLock<HashMap<String, ChannelSink>>>,
|
||||
}
|
||||
|
||||
impl OutboundDispatcher {
|
||||
pub fn new(bus: Arc<MessageBus>) -> Self {
|
||||
Self {
|
||||
@ -72,145 +23,11 @@ impl OutboundDispatcher {
|
||||
}
|
||||
}
|
||||
|
||||
/// 注册 channel 并启动其独立 sender task。
|
||||
///
|
||||
/// sender task 生命周期与 dispatcher 一致：dispatcher `run()` 退出时
|
||||
/// 通过 cancel token 终止所有 sender task。
|
||||
pub async fn register_channel(&self, name: &str, channel: Arc<dyn Channel + Send + Sync>) {
|
||||
let (high_tx, high_rx) =
|
||||
mpsc::channel::<OutboundMessage>(HIGH_PRIORITY_QUEUE_CAPACITY);
|
||||
let (low_tx, low_rx) =
|
||||
mpsc::channel::<OutboundMessage>(LOW_PRIORITY_QUEUE_CAPACITY);
|
||||
let cancel = CancellationToken::new();
|
||||
|
||||
let channel_name = name.to_string();
|
||||
let cancel_for_task = cancel.clone();
|
||||
|
||||
tokio::spawn(async move {
|
||||
Self::run_sender_task(
|
||||
&channel_name,
|
||||
channel,
|
||||
high_rx,
|
||||
low_rx,
|
||||
cancel_for_task,
|
||||
)
|
||||
.await;
|
||||
});
|
||||
|
||||
self.channels
|
||||
.write()
|
||||
.await
|
||||
.insert(
|
||||
name.to_string(),
|
||||
ChannelSink {
|
||||
high_tx,
|
||||
low_tx,
|
||||
cancel,
|
||||
},
|
||||
);
|
||||
}
|
||||
|
||||
/// sender task：优先消费 high 队列（最终响应），再消费 low 队列（中间过程），
|
||||
/// 串行调用 `Channel::send` 并重试。
|
||||
///
|
||||
/// 重试 sleep 只阻塞当前 channel 的 task，不影响其他 channel。
|
||||
async fn run_sender_task(
|
||||
channel_name: &str,
|
||||
channel: Arc<dyn Channel + Send + Sync>,
|
||||
mut high_rx: mpsc::Receiver<OutboundMessage>,
|
||||
mut low_rx: mpsc::Receiver<OutboundMessage>,
|
||||
cancel: CancellationToken,
|
||||
) {
|
||||
tracing::info!(channel = %channel_name, "Channel sender task started");
|
||||
|
||||
loop {
|
||||
// 优先消费 high 队列：保证最终响应优先发送，不被中间过程阻塞。
|
||||
// try_recv 非阻塞，有则立即处理，无则进入 select! 等待。
|
||||
if let Ok(msg) = high_rx.try_recv() {
|
||||
Self::send_one(&*channel, channel_name, msg).await;
|
||||
continue;
|
||||
}
|
||||
|
||||
tokio::select! {
|
||||
// high 优先：一旦有最终响应立即处理
|
||||
msg = high_rx.recv() => {
|
||||
match msg {
|
||||
Some(msg) => Self::send_one(&*channel, channel_name, msg).await,
|
||||
None => {
|
||||
// high 关闭：仅消费 low 残留后退出
|
||||
tracing::debug!(channel = %channel_name, "High-priority queue closed, draining low queue");
|
||||
Self::drain_low(&*channel, channel_name, &mut low_rx).await;
|
||||
break;
|
||||
}
|
||||
}
|
||||
}
|
||||
// low：high 空时消费中间过程
|
||||
msg = low_rx.recv() => {
|
||||
match msg {
|
||||
Some(msg) => Self::send_one(&*channel, channel_name, msg).await,
|
||||
None => {
|
||||
// low 关闭：仅消费 high 残留后退出
|
||||
tracing::debug!(channel = %channel_name, "Low-priority queue closed, draining high queue");
|
||||
Self::drain_high(&*channel, channel_name, &mut high_rx).await;
|
||||
break;
|
||||
}
|
||||
}
|
||||
}
|
||||
// dispatcher 退出时取消所有 sender task
|
||||
_ = cancel.cancelled() => {
|
||||
tracing::info!(channel = %channel_name, "Sender task cancelled, stopping");
|
||||
break;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// 发送单条消息，处理重试结果日志。
|
||||
async fn send_one(
|
||||
channel: &dyn Channel,
|
||||
channel_name: &str,
|
||||
msg: OutboundMessage,
|
||||
) {
|
||||
match Self::send_with_retry(channel, msg).await {
|
||||
Ok(()) => {}
|
||||
Err(ChannelError::ChannelFull) => {
|
||||
// 队列满是不可重试的——send_with_retry 已跳过重试。
|
||||
// 记 warn 而非 error：这是预期的背压丢弃。
|
||||
tracing::warn!(
|
||||
channel = %channel_name,
|
||||
"Message dropped: channel queue full"
|
||||
);
|
||||
}
|
||||
Err(error) => {
|
||||
tracing::error!(
|
||||
channel = %channel_name,
|
||||
error = %error,
|
||||
"Failed to send message after retries"
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// 排空 high 队列残留消息后返回。
|
||||
async fn drain_high(
|
||||
channel: &dyn Channel,
|
||||
channel_name: &str,
|
||||
high_rx: &mut mpsc::Receiver<OutboundMessage>,
|
||||
) {
|
||||
while let Ok(msg) = high_rx.try_recv() {
|
||||
Self::send_one(channel, channel_name, msg).await;
|
||||
}
|
||||
}
|
||||
|
||||
/// 排空 low 队列残留消息后返回。
|
||||
async fn drain_low(
|
||||
channel: &dyn Channel,
|
||||
channel_name: &str,
|
||||
low_rx: &mut mpsc::Receiver<OutboundMessage>,
|
||||
) {
|
||||
while let Ok(msg) = low_rx.try_recv() {
|
||||
Self::send_one(channel, channel_name, msg).await;
|
||||
}
|
||||
.insert(name.to_string(), channel);
|
||||
}
|
||||
|
||||
pub async fn run(&self) {
|
||||
@ -224,7 +41,6 @@ impl OutboundDispatcher {
|
||||
break;
|
||||
}
|
||||
};
|
||||
|
||||
#[cfg(debug_assertions)]
|
||||
tracing::debug!(
|
||||
channel = %msg.channel,
|
||||
@ -236,7 +52,6 @@ impl OutboundDispatcher {
|
||||
// Skip messages with virtual scheduler chat IDs (e.g., "scheduler/job_id")
|
||||
// These are internal messages from SilentAgentTask that should not be sent externally
|
||||
if msg.chat_id.starts_with(SCHEDULER_VIRTUAL_CHAT_ID_PREFIX) {
|
||||
#[cfg(debug_assertions)]
|
||||
tracing::debug!(
|
||||
channel = %msg.channel,
|
||||
chat_id = %msg.chat_id,
|
||||
@ -246,35 +61,12 @@ impl OutboundDispatcher {
|
||||
}
|
||||
|
||||
let channel_name = msg.channel.clone();
|
||||
let sink = self.channels.read().await.get(&channel_name).cloned();
|
||||
let channel = self.channels.read().await.get(&channel_name).cloned();
|
||||
|
||||
match sink {
|
||||
Some(sink) => {
|
||||
// try_send 保证 dispatcher 永不阻塞：队列满时立即丢弃该消息，
|
||||
// 不影响其他 channel 的投递。与 bus.publish_outbound 策略一致。
|
||||
// 高优先级（最终响应）投递到 high 队列——容量 32 且仅最终响应，
|
||||
// 极罕见满；低优先级（中间过程）投递到 low 队列，满则丢弃。
|
||||
let (queue, priority_label) = if is_high_priority(&msg) {
|
||||
(&sink.high_tx, "high")
|
||||
} else {
|
||||
(&sink.low_tx, "low")
|
||||
};
|
||||
match queue.try_send(msg) {
|
||||
Ok(()) => {}
|
||||
Err(mpsc::error::TrySendError::Full(_)) => {
|
||||
tracing::warn!(
|
||||
channel = %channel_name,
|
||||
priority = priority_label,
|
||||
"Channel queue full, dropping message"
|
||||
);
|
||||
}
|
||||
Err(mpsc::error::TrySendError::Closed(_)) => {
|
||||
tracing::warn!(
|
||||
channel = %channel_name,
|
||||
priority = priority_label,
|
||||
"Channel queue closed, dropping message"
|
||||
);
|
||||
}
|
||||
match channel {
|
||||
Some(ch) => {
|
||||
if let Err(error) = self.send_with_retry(&*ch, msg).await {
|
||||
tracing::error!(channel = %channel_name, error = %error, "Failed to send message after retries");
|
||||
}
|
||||
}
|
||||
None => {
|
||||
@ -282,44 +74,22 @@ impl OutboundDispatcher {
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// 通知所有 sender task 退出
|
||||
let sinks = self.channels.write().await;
|
||||
for (name, sink) in sinks.iter() {
|
||||
sink.cancel.cancel();
|
||||
tracing::debug!(channel = %name, "Cancelled sender task");
|
||||
}
|
||||
}
|
||||
|
||||
/// 发送消息，失败时按重试间隔重试。
|
||||
///
|
||||
/// 高优先级消息（最终响应）使用 `EXTENDED_RETRY_DELAYS_SECS`（5 次，最长 16 秒），
|
||||
/// 尽最大努力送达；低优先级消息（中间过程）使用 `RETRY_DELAYS_SECS`（3 次）。
|
||||
///
|
||||
/// `ChannelFull` 不可重试——队列满时重试只会浪费退避时间并阻塞
|
||||
/// channel 队列消费。立即返回让 sender task 尽快处理下一条消息。
|
||||
///
|
||||
/// 仅在单个 channel 的 sender task 内执行——重试 sleep 只阻塞
|
||||
/// 该 channel 的发送，不影响其他 channel。
|
||||
async fn send_with_retry(
|
||||
&self,
|
||||
channel: &dyn Channel,
|
||||
msg: OutboundMessage,
|
||||
) -> Result<(), ChannelError> {
|
||||
let delays: &[u64] = if is_high_priority(&msg) {
|
||||
&EXTENDED_RETRY_DELAYS_SECS
|
||||
} else {
|
||||
&RETRY_DELAYS_SECS
|
||||
};
|
||||
for (attempt_index, delay) in delays.iter().enumerate() {
|
||||
const DELAYS: [u64; 3] = [1, 2, 4];
|
||||
|
||||
for (attempt_index, delay) in DELAYS.iter().enumerate() {
|
||||
match channel.send(msg.clone()).await {
|
||||
Ok(()) => return Ok(()),
|
||||
// 队列满：不可重试，立即返回
|
||||
Err(ChannelError::ChannelFull) => return Err(ChannelError::ChannelFull),
|
||||
Err(error) if attempt_index < delays.len() - 1 => {
|
||||
Err(error) if attempt_index < DELAYS.len() - 1 => {
|
||||
tracing::warn!(
|
||||
attempt = attempt_index + 1,
|
||||
delay = delay,
|
||||
high_priority = is_high_priority(&msg),
|
||||
error = %error,
|
||||
"Send failed, retrying"
|
||||
);
|
||||
@ -329,573 +99,6 @@ impl OutboundDispatcher {
|
||||
}
|
||||
}
|
||||
|
||||
// 防御性兜底：正常情况下循环最后一次迭代会 return Err(error)。
|
||||
// 若 delays 被改为空数组，循环体不执行，返回错误而非 panic。
|
||||
Err(ChannelError::SendError(
|
||||
"send_with_retry exhausted with no retry attempts configured".into(),
|
||||
))
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use crate::bus::OutboundMessage;
|
||||
use async_trait::async_trait;
|
||||
use std::sync::atomic::{AtomicU32, Ordering};
|
||||
use std::time::Duration;
|
||||
|
||||
/// 测试用 channel：记录所有收到的消息内容，可配置人为延迟和失败。
|
||||
struct TestChannel {
|
||||
name: String,
|
||||
received: Arc<AtomicU32>,
|
||||
/// 按发送顺序记录每条成功发送消息的 content，用于断言优先级顺序。
|
||||
received_contents: Arc<std::sync::Mutex<Vec<String>>>,
|
||||
delay_ms: u64,
|
||||
fail_first_n: u32,
|
||||
/// 始终返回 `ChannelFull`，用于验证 `send_with_retry` 不重试队列满错误。
|
||||
always_full: bool,
|
||||
call_count: Arc<AtomicU32>,
|
||||
}
|
||||
|
||||
impl TestChannel {
|
||||
fn new(name: &str) -> Self {
|
||||
Self {
|
||||
name: name.to_string(),
|
||||
received: Arc::new(AtomicU32::new(0)),
|
||||
received_contents: Arc::new(std::sync::Mutex::new(Vec::new())),
|
||||
delay_ms: 0,
|
||||
fail_first_n: 0,
|
||||
always_full: false,
|
||||
call_count: Arc::new(AtomicU32::new(0)),
|
||||
}
|
||||
}
|
||||
|
||||
fn with_delay(mut self, ms: u64) -> Self {
|
||||
self.delay_ms = ms;
|
||||
self
|
||||
}
|
||||
|
||||
fn with_fail_first_n(mut self, n: u32) -> Self {
|
||||
self.fail_first_n = n;
|
||||
self
|
||||
}
|
||||
|
||||
fn with_channel_full(mut self) -> Self {
|
||||
self.always_full = true;
|
||||
self
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl Channel for TestChannel {
|
||||
fn name(&self) -> &str {
|
||||
&self.name
|
||||
}
|
||||
|
||||
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> {
|
||||
let count = self.call_count.fetch_add(1, Ordering::SeqCst);
|
||||
|
||||
// 队列满：不可重试错误，用于验证 send_with_retry 立即返回
|
||||
if self.always_full {
|
||||
return Err(ChannelError::ChannelFull);
|
||||
}
|
||||
|
||||
if (count as u32) < self.fail_first_n {
|
||||
return Err(ChannelError::SendError("simulated failure".to_string()));
|
||||
}
|
||||
|
||||
if self.delay_ms > 0 {
|
||||
tokio::time::sleep(Duration::from_millis(self.delay_ms)).await;
|
||||
}
|
||||
|
||||
self.received.fetch_add(1, Ordering::SeqCst);
|
||||
if let Ok(mut guard) = self.received_contents.lock() {
|
||||
guard.push(msg.content.clone());
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
fn make_message(channel: &str, chat_id: &str, content: &str) -> OutboundMessage {
|
||||
OutboundMessage::assistant(
|
||||
channel,
|
||||
chat_id,
|
||||
None,
|
||||
content,
|
||||
None,
|
||||
std::collections::HashMap::new(),
|
||||
)
|
||||
}
|
||||
|
||||
/// 构造低优先级消息（ToolCall），用于验证 low 队列的丢弃与优先级行为。
|
||||
fn make_low_message(channel: &str, chat_id: &str, content: &str) -> OutboundMessage {
|
||||
OutboundMessage::tool_call(
|
||||
channel,
|
||||
chat_id,
|
||||
None,
|
||||
"msg-id",
|
||||
content,
|
||||
serde_json::json!({}),
|
||||
None,
|
||||
std::collections::HashMap::new(),
|
||||
)
|
||||
}
|
||||
|
||||
/// 构造错误通知（agent 异常终止），用于验证 ErrorNotification 走高优队列。
|
||||
fn make_error_message(channel: &str, chat_id: &str, content: &str) -> OutboundMessage {
|
||||
OutboundMessage::error_notification(
|
||||
channel,
|
||||
chat_id,
|
||||
None,
|
||||
content,
|
||||
None,
|
||||
std::collections::HashMap::new(),
|
||||
)
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_is_high_priority_classifies_terminal_events() {
|
||||
// AssistantResponse 和 ErrorNotification 都是终态,必须走高优队列必达。
|
||||
let assistant = make_message("c", "chat", "final");
|
||||
let error = make_error_message("c", "chat", "agent failed");
|
||||
let tool_call = make_low_message("c", "chat", "calling tool");
|
||||
let tool_result = OutboundMessage::tool_result(
|
||||
"c", "chat", None, "id", "tool", "result", None,
|
||||
std::collections::HashMap::new(),
|
||||
);
|
||||
let exec_done = OutboundMessage::execution_completed(
|
||||
"c", "chat", None,
|
||||
std::collections::HashMap::new(),
|
||||
);
|
||||
|
||||
assert!(is_high_priority(&assistant), "AssistantResponse should be high priority");
|
||||
assert!(is_high_priority(&error), "ErrorNotification should be high priority");
|
||||
assert!(!is_high_priority(&tool_call), "ToolCall should be low priority");
|
||||
assert!(!is_high_priority(&tool_result), "ToolResult should be low priority");
|
||||
assert!(!is_high_priority(&exec_done), "ExecutionCompleted should be low priority");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_fast_channel_not_blocked_by_slow_channel() {
|
||||
// 验证核心目标：channel A 慢发送不应阻塞 channel B 的消息投递
|
||||
let bus = MessageBus::new(16);
|
||||
let dispatcher = OutboundDispatcher::new(bus.clone());
|
||||
|
||||
let slow = Arc::new(TestChannel::new("slow").with_delay(500));
|
||||
let fast = Arc::new(TestChannel::new("fast"));
|
||||
|
||||
let slow_received = slow.received.clone();
|
||||
let fast_received = fast.received.clone();
|
||||
|
||||
dispatcher.register_channel("slow", slow).await;
|
||||
dispatcher.register_channel("fast", fast).await;
|
||||
|
||||
let dispatcher_handle = tokio::spawn(async move {
|
||||
dispatcher.run().await;
|
||||
});
|
||||
|
||||
// 先发一条 slow（500ms 延迟），紧接着发一条 fast
|
||||
bus.publish_outbound(make_message("slow", "chat-1", "slow-msg")).await.unwrap();
|
||||
bus.publish_outbound(make_message("fast", "chat-2", "fast-msg")).await.unwrap();
|
||||
|
||||
// 等待 fast 消息被投递（远早于 slow 完成）
|
||||
tokio::time::timeout(Duration::from_millis(200), async {
|
||||
while fast_received.load(Ordering::SeqCst) == 0 {
|
||||
tokio::time::sleep(Duration::from_millis(10)).await;
|
||||
}
|
||||
})
|
||||
.await
|
||||
.expect("fast channel should receive message within 200ms, but was blocked by slow channel");
|
||||
|
||||
// 等待 slow 消息完成
|
||||
tokio::time::timeout(Duration::from_secs(2), async {
|
||||
while slow_received.load(Ordering::SeqCst) == 0 {
|
||||
tokio::time::sleep(Duration::from_millis(10)).await;
|
||||
}
|
||||
})
|
||||
.await
|
||||
.expect("slow channel should eventually receive message");
|
||||
|
||||
assert_eq!(fast_received.load(Ordering::SeqCst), 1);
|
||||
assert_eq!(slow_received.load(Ordering::SeqCst), 1);
|
||||
|
||||
// 关闭 bus 让 dispatcher 退出
|
||||
// dispatcher 持有 Arc<MessageBus> 克隆，drop(bus) 不会关闭 bus。
|
||||
// 直接 abort dispatcher 及其 sender task 即可清理。
|
||||
dispatcher_handle.abort();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_retry_does_not_block_other_channel() {
|
||||
// 验证：channel A 重试 sleep（1+2=3秒）期间，channel B 正常投递
|
||||
let bus = MessageBus::new(16);
|
||||
let dispatcher = OutboundDispatcher::new(bus.clone());
|
||||
|
||||
let flaky = Arc::new(
|
||||
TestChannel::new("flaky")
|
||||
.with_fail_first_n(2) // 前 2 次失败，触发 1+2 秒重试
|
||||
.with_delay(0),
|
||||
);
|
||||
let stable = Arc::new(TestChannel::new("stable"));
|
||||
|
||||
let flaky_received = flaky.received.clone();
|
||||
let stable_received = stable.received.clone();
|
||||
|
||||
dispatcher.register_channel("flaky", flaky).await;
|
||||
dispatcher.register_channel("stable", stable).await;
|
||||
|
||||
let dispatcher_handle = tokio::spawn(async move {
|
||||
dispatcher.run().await;
|
||||
});
|
||||
|
||||
// 先发 flaky（会重试 3 秒），紧接着发 stable
|
||||
bus.publish_outbound(make_message("flaky", "chat-1", "flaky-msg")).await.unwrap();
|
||||
bus.publish_outbound(make_message("stable", "chat-2", "stable-msg")).await.unwrap();
|
||||
|
||||
// stable 应在 200ms 内收到，远早于 flaky 的 3 秒重试完成
|
||||
tokio::time::timeout(Duration::from_millis(200), async {
|
||||
while stable_received.load(Ordering::SeqCst) == 0 {
|
||||
tokio::time::sleep(Duration::from_millis(10)).await;
|
||||
}
|
||||
})
|
||||
.await
|
||||
.expect("stable channel should not be blocked by flaky channel's retry sleep");
|
||||
|
||||
// 等待 flaky 重试成功（第 3 次尝试）
|
||||
tokio::time::timeout(Duration::from_secs(5), async {
|
||||
while flaky_received.load(Ordering::SeqCst) == 0 {
|
||||
tokio::time::sleep(Duration::from_millis(50)).await;
|
||||
}
|
||||
})
|
||||
.await
|
||||
.expect("flaky channel should eventually succeed after retries");
|
||||
|
||||
assert_eq!(stable_received.load(Ordering::SeqCst), 1);
|
||||
assert_eq!(flaky_received.load(Ordering::SeqCst), 1);
|
||||
|
||||
// dispatcher 持有 Arc<MessageBus> 克隆，drop(bus) 不会关闭 bus。
|
||||
// 直接 abort dispatcher 及其 sender task 即可清理。
|
||||
dispatcher_handle.abort();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_scheduler_virtual_chat_id_skipped() {
|
||||
// 验证：scheduler/ 前缀的 chat_id 不被投递到任何 channel
|
||||
let bus = MessageBus::new(16);
|
||||
let dispatcher = OutboundDispatcher::new(bus.clone());
|
||||
|
||||
let channel = Arc::new(TestChannel::new("test"));
|
||||
let received = channel.received.clone();
|
||||
|
||||
dispatcher.register_channel("test", channel).await;
|
||||
|
||||
let dispatcher_handle = tokio::spawn(async move {
|
||||
dispatcher.run().await;
|
||||
});
|
||||
|
||||
// scheduler 虚拟消息应被跳过
|
||||
bus.publish_outbound(make_message("test", "scheduler/job-1", "internal"))
|
||||
.await
|
||||
.unwrap();
|
||||
// 正常消息应被投递
|
||||
bus.publish_outbound(make_message("test", "chat-1", "normal"))
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
tokio::time::timeout(Duration::from_secs(2), async {
|
||||
while received.load(Ordering::SeqCst) == 0 {
|
||||
tokio::time::sleep(Duration::from_millis(10)).await;
|
||||
}
|
||||
})
|
||||
.await
|
||||
.expect("normal message should be delivered");
|
||||
|
||||
// 只收到 1 条（scheduler 虚拟消息被跳过）
|
||||
assert_eq!(received.load(Ordering::SeqCst), 1);
|
||||
|
||||
// dispatcher 持有 Arc<MessageBus> 克隆，drop(bus) 不会关闭 bus。
|
||||
// 直接 abort dispatcher 及其 sender task 即可清理。
|
||||
dispatcher_handle.abort();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_unknown_channel_warns_and_continues() {
|
||||
// 验证：未知 channel 的消息被跳过，不影响后续消息投递
|
||||
let bus = MessageBus::new(16);
|
||||
let dispatcher = OutboundDispatcher::new(bus.clone());
|
||||
|
||||
let channel = Arc::new(TestChannel::new("known"));
|
||||
let received = channel.received.clone();
|
||||
|
||||
dispatcher.register_channel("known", channel).await;
|
||||
|
||||
let dispatcher_handle = tokio::spawn(async move {
|
||||
dispatcher.run().await;
|
||||
});
|
||||
|
||||
// 发往未知 channel 的消息
|
||||
bus.publish_outbound(make_message("unknown", "chat-1", "lost"))
|
||||
.await
|
||||
.unwrap();
|
||||
// 发往已知 channel 的消息
|
||||
bus.publish_outbound(make_message("known", "chat-2", "delivered"))
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
tokio::time::timeout(Duration::from_secs(2), async {
|
||||
while received.load(Ordering::SeqCst) == 0 {
|
||||
tokio::time::sleep(Duration::from_millis(10)).await;
|
||||
}
|
||||
})
|
||||
.await
|
||||
.expect("known channel should receive message despite preceding unknown channel message");
|
||||
|
||||
assert_eq!(received.load(Ordering::SeqCst), 1);
|
||||
|
||||
// dispatcher 持有 Arc<MessageBus> 克隆，drop(bus) 不会关闭 bus。
|
||||
// 直接 abort dispatcher 及其 sender task 即可清理。
|
||||
dispatcher_handle.abort();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_send_with_retry_no_retry_on_channel_full() {
|
||||
// 验证：send_with_retry 遇到 ChannelFull 应立即返回，不重试。
|
||||
// 若错误地重试，会 sleep 1+2+4=7 秒，测试将在超时阈值内失败。
|
||||
let channel = TestChannel::new("full").with_channel_full();
|
||||
let call_count = channel.call_count.clone();
|
||||
|
||||
let msg = make_message("full", "chat-1", "dropped");
|
||||
|
||||
// 500ms 阈值：远小于首次重试间隔 1s，足以区分"立即返回"与"至少一次重试"
|
||||
let result = tokio::time::timeout(Duration::from_millis(500), async {
|
||||
OutboundDispatcher::send_with_retry(&channel, msg).await
|
||||
})
|
||||
.await
|
||||
.expect("send_with_retry should return immediately on ChannelFull, not retry");
|
||||
|
||||
// 必须返回 ChannelFull 错误
|
||||
assert!(
|
||||
matches!(result, Err(ChannelError::ChannelFull)),
|
||||
"expected ChannelFull error, got {:?}",
|
||||
result
|
||||
);
|
||||
|
||||
// send 只被调用一次——证明没有重试
|
||||
assert_eq!(
|
||||
call_count.load(Ordering::SeqCst),
|
||||
1,
|
||||
"send should be called exactly once when ChannelFull is returned"
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_high_priority_not_dropped_when_low_full() {
|
||||
// 验证:low 队列被大量低优消息填满时,高优(AssistantResponse)仍被发送。
|
||||
// 这是本次修复的核心目标——最终响应必达,不被中间过程挤占丢弃。
|
||||
let bus = MessageBus::new(256);
|
||||
let dispatcher = OutboundDispatcher::new(bus.clone());
|
||||
|
||||
// 慢 channel:50ms/条,确保 low 队列持续积压
|
||||
let channel = Arc::new(TestChannel::new("test").with_delay(50));
|
||||
let contents = channel.received_contents.clone();
|
||||
dispatcher.register_channel("test", channel).await;
|
||||
|
||||
let dispatcher_handle = tokio::spawn(async move {
|
||||
dispatcher.run().await;
|
||||
});
|
||||
|
||||
// 投递大量低优消息(超过 low 容量 64),填满 low 队列
|
||||
for i in 0..80 {
|
||||
bus.publish_outbound(make_low_message("test", "chat-1", &format!("low-{i}")))
|
||||
.await
|
||||
.unwrap();
|
||||
}
|
||||
// 投递高优消息(最终响应)
|
||||
bus.publish_outbound(make_message("test", "chat-1", "HIGH-FINAL"))
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
// 高优应在 3 秒内被发送(独立 high 队列 + sender 优先消费)
|
||||
tokio::time::timeout(Duration::from_secs(3), async {
|
||||
loop {
|
||||
let sent = contents.lock().unwrap().clone();
|
||||
if sent.iter().any(|c| c == "HIGH-FINAL") {
|
||||
break;
|
||||
}
|
||||
tokio::time::sleep(Duration::from_millis(20)).await;
|
||||
}
|
||||
})
|
||||
.await
|
||||
.expect("high priority message should be sent even when low queue is full");
|
||||
|
||||
dispatcher_handle.abort();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_low_priority_dropped_when_full() {
|
||||
// 验证:low 队列满时,后续低优消息被丢弃(不进入 high 队列)。
|
||||
// 低优是中间过程,可丢——这是与高优必达的对比行为。
|
||||
let bus = MessageBus::new(256);
|
||||
let dispatcher = OutboundDispatcher::new(bus.clone());
|
||||
|
||||
// 慢 channel:50ms/条,确保 low 队列积压触发丢弃
|
||||
let channel = Arc::new(TestChannel::new("test").with_delay(50));
|
||||
let received = channel.received.clone();
|
||||
dispatcher.register_channel("test", channel).await;
|
||||
|
||||
let dispatcher_handle = tokio::spawn(async move {
|
||||
dispatcher.run().await;
|
||||
});
|
||||
|
||||
// 投递远超 low 容量的低优消息(200 条),确保触发丢弃
|
||||
let total_sent: u32 = 200;
|
||||
for i in 0..total_sent {
|
||||
bus.publish_outbound(make_low_message("test", "chat-1", &format!("low-{i}")))
|
||||
.await
|
||||
.unwrap();
|
||||
}
|
||||
|
||||
// 等 sender 消费完 low 队列内的消息(至少 LOW_PRIORITY_QUEUE_CAPACITY 条)
|
||||
tokio::time::timeout(Duration::from_secs(10), async {
|
||||
while received.load(Ordering::SeqCst) < LOW_PRIORITY_QUEUE_CAPACITY as u32 {
|
||||
tokio::time::sleep(Duration::from_millis(20)).await;
|
||||
}
|
||||
})
|
||||
.await
|
||||
.expect("should receive at least LOW_PRIORITY_QUEUE_CAPACITY messages");
|
||||
|
||||
// 额外等待确认 sender 已消费完 low 队列残留(64 * 50ms ≈ 3.2s,给足 1s 余量)
|
||||
tokio::time::sleep(Duration::from_secs(1)).await;
|
||||
let final_received = received.load(Ordering::SeqCst);
|
||||
|
||||
// 断言:有丢弃发生(received < 投递数),且 low 队列曾被填满(received >= 容量)。
|
||||
// 不断言精确数量——sender 与 dispatcher 的并发竞态会使进入 low 的条数略多于容量。
|
||||
assert!(
|
||||
final_received < total_sent,
|
||||
"some low-priority messages should be dropped when queue full, got {}/{}",
|
||||
final_received,
|
||||
total_sent
|
||||
);
|
||||
assert!(
|
||||
final_received >= LOW_PRIORITY_QUEUE_CAPACITY as u32,
|
||||
"should receive at least {} messages (queue was filled), got {}",
|
||||
LOW_PRIORITY_QUEUE_CAPACITY,
|
||||
final_received
|
||||
);
|
||||
|
||||
dispatcher_handle.abort();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_high_priority_uses_extended_retry() {
|
||||
// 验证:高优消息(AssistantResponse)使用 EXTENDED_RETRY_DELAYS(5次)。
|
||||
// fail_first_n(3):前 3 次失败,第 4 次成功——证明高优在 3 次后仍继续重试。
|
||||
let channel = TestChannel::new("flaky").with_fail_first_n(3);
|
||||
let call_count = channel.call_count.clone();
|
||||
let msg = make_message("flaky", "chat-1", "final"); // 高优
|
||||
|
||||
// EXTENDED_RETRY_DELAYS = [1,2,4,8,16];前 3 次失败后第 4 次成功,耗时 1+2+4=7s
|
||||
let result = tokio::time::timeout(Duration::from_secs(15), async {
|
||||
OutboundDispatcher::send_with_retry(&channel, msg).await
|
||||
})
|
||||
.await
|
||||
.expect("high priority should succeed within extended retry budget");
|
||||
|
||||
assert!(result.is_ok(), "high priority should succeed after 4 attempts");
|
||||
assert_eq!(
|
||||
call_count.load(Ordering::SeqCst),
|
||||
4,
|
||||
"high priority should attempt 4 times (extended retry)"
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_low_priority_gives_up_after_standard_retries() {
|
||||
// 验证:低优消息(ToolCall)使用 RETRY_DELAYS(3次),前 3 次失败后放弃。
|
||||
// 与高优的 5 次重试形成对比——低优是中间过程,放弃可接受。
|
||||
let channel = TestChannel::new("flaky").with_fail_first_n(3);
|
||||
let call_count = channel.call_count.clone();
|
||||
let msg = make_low_message("flaky", "chat-1", "tool"); // 低优
|
||||
|
||||
// RETRY_DELAYS = [1,2,4];3 次都失败,耗时 1+2=3s
|
||||
let result = tokio::time::timeout(Duration::from_secs(10), async {
|
||||
OutboundDispatcher::send_with_retry(&channel, msg).await
|
||||
})
|
||||
.await
|
||||
.expect("low priority should give up within standard retry budget");
|
||||
|
||||
assert!(result.is_err(), "low priority should fail after 3 attempts");
|
||||
assert_eq!(
|
||||
call_count.load(Ordering::SeqCst),
|
||||
3,
|
||||
"low priority should attempt 3 times (standard retry)"
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_sender_consumes_high_first() {
|
||||
// 验证:high 队列优先于 low 队列被消费。
|
||||
// 投递多条低优后再投递高优,高优应在大部分低优之前被发送。
|
||||
let bus = MessageBus::new(64);
|
||||
let dispatcher = OutboundDispatcher::new(bus.clone());
|
||||
|
||||
let channel = Arc::new(TestChannel::new("test").with_delay(30));
|
||||
let contents = channel.received_contents.clone();
|
||||
let received = channel.received.clone();
|
||||
dispatcher.register_channel("test", channel).await;
|
||||
|
||||
let dispatcher_handle = tokio::spawn(async move {
|
||||
dispatcher.run().await;
|
||||
});
|
||||
|
||||
// 投递 5 条低优
|
||||
for i in 0..5 {
|
||||
bus.publish_outbound(make_low_message("test", "chat-1", &format!("low-{i}")))
|
||||
.await
|
||||
.unwrap();
|
||||
}
|
||||
// 等待低优入 low 队列
|
||||
tokio::time::sleep(Duration::from_millis(10)).await;
|
||||
// 投递 1 条高优
|
||||
bus.publish_outbound(make_message("test", "chat-1", "HIGH"))
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
// 等待 6 条全部发送完成
|
||||
tokio::time::timeout(Duration::from_secs(3), async {
|
||||
while received.load(Ordering::SeqCst) < 6 {
|
||||
tokio::time::sleep(Duration::from_millis(10)).await;
|
||||
}
|
||||
})
|
||||
.await
|
||||
.expect("all 6 messages should be sent within 3s");
|
||||
|
||||
let sent = contents.lock().unwrap().clone();
|
||||
let high_index = sent
|
||||
.iter()
|
||||
.position(|c| c == "HIGH")
|
||||
.expect("HIGH should be in sent list");
|
||||
// HIGH 应在前 3 条内:sender 完成当前低优后立即取 high,优先于剩余低优
|
||||
assert!(
|
||||
high_index < 3,
|
||||
"HIGH should be sent before most low-priority messages, got index {} in {:?}",
|
||||
high_index,
|
||||
sent
|
||||
);
|
||||
|
||||
dispatcher_handle.abort();
|
||||
unreachable!()
|
||||
}
|
||||
}
|
||||
|
||||
@ -1,6 +1,5 @@
|
||||
use std::collections::HashSet;
|
||||
use std::sync::Arc;
|
||||
use parking_lot::Mutex;
|
||||
use std::sync::{Arc, Mutex};
|
||||
|
||||
use tokio::sync::Semaphore;
|
||||
|
||||
@ -23,7 +22,7 @@ use crate::command::handlers::switch_topic::SwitchTopicCommandHandler;
|
||||
use crate::config::LLMProviderConfig;
|
||||
use crate::gateway::agent_factory::build_system_prompt_provider;
|
||||
use crate::gateway::cancel_manager::CancelManager;
|
||||
use crate::providers::{ProviderRuntimeConfig, create_provider};
|
||||
use crate::providers::{create_provider, ProviderRuntimeConfig};
|
||||
use crate::storage::persistent_session_id;
|
||||
use crate::topic_description::generate_topic_description;
|
||||
|
||||
@ -53,8 +52,8 @@ impl InboundProcessor {
|
||||
let store = session_manager.store();
|
||||
|
||||
// 注册 Session 处理器
|
||||
let session_handler =
|
||||
SessionCommandHandler::new(store.clone()).with_session_manager(session_manager.clone());
|
||||
let session_handler = SessionCommandHandler::new(store.clone())
|
||||
.with_session_manager(session_manager.clone());
|
||||
command_router.register(Box::new(session_handler));
|
||||
|
||||
// 注册 list_sessions 处理器
|
||||
@ -80,7 +79,7 @@ impl InboundProcessor {
|
||||
// 注册 get_current 处理器
|
||||
command_router.register(Box::new(
|
||||
GetCurrentSessionCommandHandler::new(store.clone())
|
||||
.with_system_prompt_provider(system_prompt_provider.clone()),
|
||||
.with_system_prompt_provider(system_prompt_provider.clone())
|
||||
));
|
||||
|
||||
// 注册 load_topic 处理器
|
||||
@ -186,8 +185,7 @@ impl InboundProcessor {
|
||||
let session_id = persistent_session_id(&inbound.channel, &inbound.chat_id);
|
||||
|
||||
// 获取当前话题（封装了 session 创建逻辑）
|
||||
let current_topic = self
|
||||
.session_manager
|
||||
let current_topic = self.session_manager
|
||||
.get_current_topic(&inbound.channel, &inbound.chat_id)
|
||||
.await?;
|
||||
|
||||
@ -198,19 +196,15 @@ impl InboundProcessor {
|
||||
|
||||
if let Ok(Some(cmd)) = adapter.try_parse(&inbound.content, ctx) {
|
||||
// 使用命令路由器处理
|
||||
let mut cmd_ctx =
|
||||
crate::command::context::CommandContext::new(&inbound.channel, &inbound.channel)
|
||||
.with_session_id(&session_id)
|
||||
.with_chat_id(&inbound.chat_id);
|
||||
let mut cmd_ctx = crate::command::context::CommandContext::new(&inbound.channel, &inbound.channel)
|
||||
.with_session_id(&session_id)
|
||||
.with_chat_id(&inbound.chat_id);
|
||||
// 只在有话题时才设置 topic_id
|
||||
if let Some(ref topic_id) = current_topic {
|
||||
cmd_ctx = cmd_ctx.with_topic_id(topic_id.as_str());
|
||||
}
|
||||
|
||||
let response = self
|
||||
.command_router
|
||||
.dispatch_with_response(cmd, cmd_ctx)
|
||||
.await;
|
||||
let response = self.command_router.dispatch_with_response(cmd, cmd_ctx).await;
|
||||
|
||||
// 发送响应给用户
|
||||
if response.success {
|
||||
@ -230,14 +224,7 @@ impl InboundProcessor {
|
||||
))
|
||||
.await
|
||||
{
|
||||
match error {
|
||||
crate::bus::BusError::Dropped => {
|
||||
tracing::warn!(error = %error, "Outbound dropped (bus full)");
|
||||
}
|
||||
crate::bus::BusError::Closed => {
|
||||
tracing::error!(error = %error, "Failed to publish command response");
|
||||
}
|
||||
}
|
||||
tracing::error!(error = %error, "Failed to publish command response");
|
||||
}
|
||||
}
|
||||
} else if let Some(error) = response.error {
|
||||
@ -253,14 +240,7 @@ impl InboundProcessor {
|
||||
))
|
||||
.await
|
||||
{
|
||||
match e {
|
||||
crate::bus::BusError::Dropped => {
|
||||
tracing::warn!(error = %e, "Outbound dropped (bus full)");
|
||||
}
|
||||
crate::bus::BusError::Closed => {
|
||||
tracing::error!(error = %e, "Failed to publish error response");
|
||||
}
|
||||
}
|
||||
tracing::error!(error = %e, "Failed to publish error response");
|
||||
}
|
||||
}
|
||||
return Ok(());
|
||||
@ -315,19 +295,10 @@ impl InboundProcessor {
|
||||
outbound.metadata.extend(inbound.forwarded_metadata.clone());
|
||||
// 注入 topic_id 到 outbound metadata，用于前端按话题隔离消息
|
||||
if let Some(ref topic_id) = current_topic {
|
||||
outbound
|
||||
.metadata
|
||||
.insert("topic_id".to_string(), topic_id.clone());
|
||||
outbound.metadata.insert("topic_id".to_string(), topic_id.clone());
|
||||
}
|
||||
if let Err(error) = self.bus.publish_outbound(outbound).await {
|
||||
match error {
|
||||
crate::bus::BusError::Dropped => {
|
||||
tracing::warn!(error = %error, "Outbound dropped (bus full)");
|
||||
}
|
||||
crate::bus::BusError::Closed => {
|
||||
tracing::error!(error = %error, "Failed to publish outbound");
|
||||
}
|
||||
}
|
||||
tracing::error!(error = %error, "Failed to publish outbound");
|
||||
}
|
||||
}
|
||||
|
||||
@ -335,17 +306,10 @@ impl InboundProcessor {
|
||||
if let Some(ref topic_id) = current_topic {
|
||||
let store = self.session_manager.store();
|
||||
if let Ok(Some(topic)) = store.get_topic(topic_id) {
|
||||
if topic.description.is_none()
|
||||
|| topic
|
||||
.description
|
||||
.as_ref()
|
||||
.map(|d| d.is_empty())
|
||||
.unwrap_or(true)
|
||||
{
|
||||
if topic.description.is_none() || topic.description.as_ref().map(|d| d.is_empty()).unwrap_or(true) {
|
||||
// 检查并设置"生成中"守卫，防止竞态条件导致重复生成
|
||||
let should_generate = {
|
||||
let mut in_flight =
|
||||
self.description_generation_in_flight.lock();
|
||||
let mut in_flight = self.description_generation_in_flight.lock().unwrap();
|
||||
if in_flight.contains(topic_id) {
|
||||
false
|
||||
} else {
|
||||
@ -363,38 +327,25 @@ impl InboundProcessor {
|
||||
tokio::spawn(async move {
|
||||
// 从 DB 查询该 topic 的第一条用户消息作为描述生成的依据
|
||||
let first_user_message = store_clone
|
||||
.load_messages_for_topic_full(&topic_id_clone, None)
|
||||
.load_messages_for_topic(&topic_id_clone, None)
|
||||
.ok()
|
||||
.and_then(|msgs| {
|
||||
msgs.into_iter().find(|m| m.role == "user")
|
||||
})
|
||||
.and_then(|msgs| msgs.into_iter().find(|m| m.role == "user"))
|
||||
.map(|m| m.content);
|
||||
|
||||
let message_content = match first_user_message {
|
||||
Some(content) => content,
|
||||
None => {
|
||||
tracing::warn!(topic_id = %topic_id_clone, "No user message found for topic, skipping description generation");
|
||||
in_flight.lock().remove(&topic_id_clone);
|
||||
in_flight.lock().unwrap().remove(&topic_id_clone);
|
||||
return;
|
||||
}
|
||||
};
|
||||
|
||||
let runtime_config: ProviderRuntimeConfig =
|
||||
provider_config.into();
|
||||
let runtime_config: ProviderRuntimeConfig = provider_config.into();
|
||||
if let Ok(provider) = create_provider(runtime_config) {
|
||||
match generate_topic_description(
|
||||
provider.as_ref(),
|
||||
&message_content,
|
||||
)
|
||||
.await
|
||||
{
|
||||
match generate_topic_description(provider.as_ref(), &message_content).await {
|
||||
Ok(description) => {
|
||||
if let Err(e) = store_clone
|
||||
.update_topic_description(
|
||||
&topic_id_clone,
|
||||
&description,
|
||||
)
|
||||
{
|
||||
if let Err(e) = store_clone.update_topic_description(&topic_id_clone, &description) {
|
||||
tracing::error!(error = %e, topic_id = %topic_id_clone, "Failed to update topic description");
|
||||
} else {
|
||||
tracing::info!(topic_id = %topic_id_clone, description = %description, "Topic description generated");
|
||||
@ -406,7 +357,7 @@ impl InboundProcessor {
|
||||
}
|
||||
}
|
||||
// 无论成功失败，释放生成守卫
|
||||
in_flight.lock().remove(&topic_id_clone);
|
||||
in_flight.lock().unwrap().remove(&topic_id_clone);
|
||||
});
|
||||
}
|
||||
}
|
||||
@ -429,14 +380,7 @@ impl InboundProcessor {
|
||||
))
|
||||
.await
|
||||
{
|
||||
match publish_error {
|
||||
crate::bus::BusError::Dropped => {
|
||||
tracing::warn!(error = %publish_error, "Outbound dropped (bus full)");
|
||||
}
|
||||
crate::bus::BusError::Closed => {
|
||||
tracing::error!(error = %publish_error, "Failed to publish execution error outbound");
|
||||
}
|
||||
}
|
||||
tracing::error!(error = %publish_error, "Failed to publish execution error outbound");
|
||||
}
|
||||
}
|
||||
}
|
||||
@ -462,14 +406,7 @@ impl InboundProcessor {
|
||||
))
|
||||
.await
|
||||
{
|
||||
match error {
|
||||
crate::bus::BusError::Dropped => {
|
||||
tracing::warn!(error = %error, "Outbound dropped (bus full)");
|
||||
}
|
||||
crate::bus::BusError::Closed => {
|
||||
tracing::error!(error = %error, "Failed to publish execution_completed");
|
||||
}
|
||||
}
|
||||
tracing::error!(error = %error, "Failed to publish execution_completed");
|
||||
}
|
||||
|
||||
Ok(())
|
||||
|
||||
@ -57,9 +57,8 @@ fn load_prompt_from_sources(sources: &[PromptSource]) -> Result<Option<String>,
|
||||
ensure_parent_dir(path)?;
|
||||
// 文件不存在时创建空白模板
|
||||
if !path.exists() {
|
||||
fs::write(path, template).map_err(|err| {
|
||||
AgentError::Other(format!("create AGENT.md template error: {}", err))
|
||||
})?;
|
||||
fs::write(path, template)
|
||||
.map_err(|err| AgentError::Other(format!("create AGENT.md template error: {}", err)))?;
|
||||
}
|
||||
// 读取内容，仅当非空（去除注释后）时注入
|
||||
let content = fs::read_to_string(path)
|
||||
@ -71,9 +70,8 @@ fn load_prompt_from_sources(sources: &[PromptSource]) -> Result<Option<String>,
|
||||
}
|
||||
PromptSource::AutoGenerated(path) => {
|
||||
if path.exists() {
|
||||
let content = fs::read_to_string(path).map_err(|err| {
|
||||
AgentError::Other(format!("read MEMORY_SUMMARY.md error: {}", err))
|
||||
})?;
|
||||
let content = fs::read_to_string(path)
|
||||
.map_err(|err| AgentError::Other(format!("read MEMORY_SUMMARY.md error: {}", err)))?;
|
||||
let without_comments = strip_comments_and_whitespace(&content);
|
||||
if !without_comments.is_empty() {
|
||||
fragments.push(without_comments);
|
||||
@ -339,9 +337,6 @@ mod tests {
|
||||
|
||||
persist_memory_summary(&memory_path, "\n## 用户记忆摘要\n- 偏好简洁\n\n").unwrap();
|
||||
|
||||
assert_eq!(
|
||||
fs::read_to_string(&memory_path).unwrap(),
|
||||
"## 用户记忆摘要\n- 偏好简洁\n"
|
||||
);
|
||||
assert_eq!(fs::read_to_string(&memory_path).unwrap(), "## 用户记忆摘要\n- 偏好简洁\n");
|
||||
}
|
||||
}
|
||||
|
||||
@ -59,7 +59,6 @@ mod tests {
|
||||
extra_headers: HashMap::new(),
|
||||
llm_timeout_secs: 120,
|
||||
memory_maintenance_timeout_secs: 600,
|
||||
max_retries: 3,
|
||||
model_id: model_id.to_string(),
|
||||
temperature: Some(0.0),
|
||||
max_tokens: Some(32),
|
||||
|
||||
@ -8,10 +8,7 @@ use tokio::sync::RwLock;
|
||||
|
||||
use crate::agent::AgentError;
|
||||
use crate::bus::MessageBus;
|
||||
use crate::config::{
|
||||
CompactionConfig, LLMProviderConfig, MemoryMaintenanceConfig, ModelResolver, SubagentsConfig,
|
||||
TaskConfig,
|
||||
};
|
||||
use crate::config::{LLMProviderConfig, MemoryMaintenanceConfig, ModelResolver, SubagentsConfig, TaskConfig};
|
||||
use crate::gateway::model_selection::ModelSelectionStore;
|
||||
use crate::gateway::tool_registry_factory::ToolRegistryFactory;
|
||||
use crate::mcp::McpInitializer;
|
||||
@ -21,13 +18,13 @@ use crate::storage::{
|
||||
ConversationRepository, MemoryRepository, PromptInjectionRepository, SchedulerJobRepository,
|
||||
SessionStore, SkillEventRepository, TodoRepository,
|
||||
};
|
||||
use crate::tools::task::repository::TaskRepository;
|
||||
use crate::tools::task::runtime::SubagentRuntime;
|
||||
use crate::tools::todo_write::TodoItem;
|
||||
use crate::tools::{
|
||||
DefaultSubAgentRuntime, InMemoryTaskRepository, NoopSessionMessageSender, SessionMessageSender,
|
||||
SubAgentRuntimeConfig, SubagentCatalog, TaskTool, ToolRegistry,
|
||||
DefaultSubAgentRuntime, InMemoryTaskRepository, NoopSessionMessageSender,
|
||||
SessionMessageSender, SubAgentRuntimeConfig, SubagentCatalog, TaskTool, ToolRegistry,
|
||||
};
|
||||
use crate::tools::task::repository::TaskRepository;
|
||||
use crate::tools::todo_write::TodoItem;
|
||||
|
||||
use super::agent_factory::AgentFactory;
|
||||
use super::cli_session::CliSessionService;
|
||||
@ -58,17 +55,7 @@ pub(crate) fn build_session_manager(
|
||||
mcp_config: crate::mcp::McpConfig,
|
||||
bus: Option<Arc<MessageBus>>,
|
||||
model_resolver: Arc<ModelResolver>,
|
||||
compaction_config: CompactionConfig,
|
||||
) -> Result<
|
||||
(
|
||||
SessionManager,
|
||||
Arc<dyn TaskRepository>,
|
||||
Option<Arc<McpClientManager>>,
|
||||
Arc<SubagentRuntime>,
|
||||
Arc<ModelSelectionStore>,
|
||||
),
|
||||
AgentError,
|
||||
> {
|
||||
) -> Result<(SessionManager, Arc<dyn TaskRepository>, Option<Arc<McpClientManager>>, Arc<SubagentRuntime>, Arc<ModelSelectionStore>), AgentError> {
|
||||
build_session_manager_with_sender(
|
||||
agent_prompt_reinject_every,
|
||||
show_tool_results,
|
||||
@ -86,7 +73,6 @@ pub(crate) fn build_session_manager(
|
||||
mcp_config,
|
||||
bus,
|
||||
model_resolver,
|
||||
compaction_config,
|
||||
)
|
||||
}
|
||||
|
||||
@ -108,17 +94,7 @@ pub(crate) fn build_session_manager_with_sender(
|
||||
mcp_config: crate::mcp::McpConfig,
|
||||
bus: Option<Arc<MessageBus>>,
|
||||
model_resolver: Arc<ModelResolver>,
|
||||
compaction_config: CompactionConfig,
|
||||
) -> Result<
|
||||
(
|
||||
SessionManager,
|
||||
Arc<dyn TaskRepository>,
|
||||
Option<Arc<McpClientManager>>,
|
||||
Arc<SubagentRuntime>,
|
||||
Arc<ModelSelectionStore>,
|
||||
),
|
||||
AgentError,
|
||||
> {
|
||||
) -> Result<(SessionManager, Arc<dyn TaskRepository>, Option<Arc<McpClientManager>>, Arc<SubagentRuntime>, Arc<ModelSelectionStore>), AgentError> {
|
||||
let store = Arc::new(
|
||||
SessionStore::new()
|
||||
.map_err(|err| AgentError::Other(format!("session store init error: {}", err)))?,
|
||||
@ -205,20 +181,18 @@ pub(crate) fn build_session_manager_with_sender(
|
||||
}
|
||||
|
||||
// Create SubAgentRuntime (if task tool is enabled)
|
||||
let (factory, task_repository, subagent_runtime): (
|
||||
_,
|
||||
Arc<dyn TaskRepository>,
|
||||
Arc<SubagentRuntime>,
|
||||
) = if task_config.enabled {
|
||||
let (factory, task_repository, subagent_runtime): (_, Arc<dyn TaskRepository>, Arc<SubagentRuntime>) = if task_config.enabled {
|
||||
let task_repository = Arc::new(InMemoryTaskRepository::new());
|
||||
// Build subagent tools with MCP tools (task tool registered separately below)
|
||||
let subagent_tools = Arc::new(factory.build_subagent_tools(
|
||||
if mcp_tools_for_subagents.is_empty() {
|
||||
None
|
||||
} else {
|
||||
Some(mcp_tools_for_subagents.clone())
|
||||
},
|
||||
));
|
||||
let subagent_tools = Arc::new(
|
||||
factory.build_subagent_tools(
|
||||
if mcp_tools_for_subagents.is_empty() {
|
||||
None
|
||||
} else {
|
||||
Some(mcp_tools_for_subagents.clone())
|
||||
}
|
||||
)
|
||||
);
|
||||
|
||||
// Create subagent catalog with discovery, wrap in SubagentRuntime
|
||||
let catalog = SubagentCatalog::discover(&subagents_config);
|
||||
@ -256,19 +230,11 @@ pub(crate) fn build_session_manager_with_sender(
|
||||
));
|
||||
}
|
||||
|
||||
(
|
||||
factory.with_subagent_runtime(default_subagent_runtime),
|
||||
task_repository,
|
||||
subagent_runtime,
|
||||
)
|
||||
(factory.with_subagent_runtime(default_subagent_runtime), task_repository, subagent_runtime)
|
||||
} else {
|
||||
// task_config 未启用时仍创建 subagent_runtime（供 API 使用）
|
||||
let subagent_runtime = Arc::new(SubagentRuntime::from_config(subagents_config.clone()));
|
||||
(
|
||||
factory,
|
||||
Arc::new(InMemoryTaskRepository::new()),
|
||||
subagent_runtime,
|
||||
)
|
||||
(factory, Arc::new(InMemoryTaskRepository::new()), subagent_runtime)
|
||||
};
|
||||
|
||||
// Build base tools
|
||||
@ -317,7 +283,6 @@ pub(crate) fn build_session_manager_with_sender(
|
||||
prompt_repository.clone(),
|
||||
model_resolver.clone(),
|
||||
model_selections.clone(),
|
||||
compaction_config,
|
||||
);
|
||||
let session_factory = SessionFactory::new(
|
||||
provider_config.clone(),
|
||||
@ -341,24 +306,18 @@ pub(crate) fn build_session_manager_with_sender(
|
||||
// Extract MCP manager for lifecycle management (e.g., disconnect on restart)
|
||||
let mcp_manager = mcp_initializer.manager();
|
||||
|
||||
Ok((
|
||||
SessionManager::from_services(SessionManagerServices {
|
||||
tools: tools as Arc<ToolRegistry>,
|
||||
skills,
|
||||
experts,
|
||||
subagent_runtime: subagent_runtime.clone(),
|
||||
store,
|
||||
show_tool_results,
|
||||
lifecycle,
|
||||
cli_sessions,
|
||||
messages,
|
||||
scheduled_tasks,
|
||||
memory_maintenance,
|
||||
task_repository: task_repository.clone(),
|
||||
}),
|
||||
task_repository,
|
||||
mcp_manager,
|
||||
subagent_runtime,
|
||||
model_selections,
|
||||
))
|
||||
Ok((SessionManager::from_services(SessionManagerServices {
|
||||
tools: tools as Arc<ToolRegistry>,
|
||||
skills,
|
||||
experts,
|
||||
subagent_runtime: subagent_runtime.clone(),
|
||||
store,
|
||||
show_tool_results,
|
||||
lifecycle,
|
||||
cli_sessions,
|
||||
messages,
|
||||
scheduled_tasks,
|
||||
memory_maintenance,
|
||||
task_repository: task_repository.clone(),
|
||||
}), task_repository, mcp_manager, subagent_runtime, model_selections))
|
||||
}
|
||||
|
||||
@ -37,10 +37,7 @@ impl ScheduledAgentTaskService {
|
||||
// 根据 chat_id 自动选择 Session：
|
||||
// - scheduler/ 开头：使用定时任务专用 Session（独立实例，不与用户消息竞争锁）
|
||||
// - 其他：使用主 Session
|
||||
let session = self
|
||||
.lifecycle
|
||||
.active_session_for_chat_id(channel_name, chat_id)
|
||||
.await?;
|
||||
let session = self.lifecycle.active_session_for_chat_id(channel_name, chat_id).await?;
|
||||
let sender_id = options
|
||||
.sender_id
|
||||
.clone()
|
||||
|
||||
@ -1,16 +1,13 @@
|
||||
use crate::agent::{AgentError, AgentLoop, AgentRuntimeConfig, ContextCompressor, EmittedMessageHandler};
|
||||
use crate::agent::{AgentError, AgentLoop, ContextCompressor, EmittedMessageHandler};
|
||||
#[cfg(test)]
|
||||
use crate::bus::SYSTEM_CONTEXT_SCHEDULED_PROMPT;
|
||||
use crate::bus::{ChatMessage, MessageBus, OutboundMessage};
|
||||
use crate::providers::StreamDelta;
|
||||
use crate::config::LLMProviderConfig;
|
||||
use crate::protocol::WsOutbound;
|
||||
use crate::providers::StreamDelta;
|
||||
use crate::scheduler::ScheduledAgentTaskOptions;
|
||||
use crate::skills::SkillRuntime;
|
||||
use crate::storage::{
|
||||
ConversationRepository, PromptInjectionRepository, SessionRecord, SessionStore,
|
||||
SkillEventRepository,
|
||||
};
|
||||
use crate::storage::{ConversationRepository, PromptInjectionRepository, SessionRecord, SessionStore, SkillEventRepository};
|
||||
use crate::tools::ToolRegistry;
|
||||
use crate::tools::task::repository::TaskRepository;
|
||||
use crate::tools::task::runtime::SubagentRuntime;
|
||||
@ -27,7 +24,8 @@ use super::execution::should_display_message_to_user;
|
||||
#[cfg(test)]
|
||||
use super::memory_maintenance::{
|
||||
MemoryMaintenanceMerge, apply_memory_maintenance_output, build_memory_maintenance_plan,
|
||||
extract_json_object, is_recoverable_maintenance_llm_error, strip_json_code_fence,
|
||||
extract_json_object, is_recoverable_maintenance_llm_error,
|
||||
strip_json_code_fence,
|
||||
};
|
||||
use super::memory_maintenance::{MemoryMaintenanceScopeResult, MemoryOrganizationOutput};
|
||||
use super::memory_maintenance_coordinator::MemoryMaintenanceCoordinator;
|
||||
@ -60,7 +58,7 @@ pub struct BusToolCallEmitter {
|
||||
chat_id: String,
|
||||
metadata: HashMap<String, String>,
|
||||
store: Arc<SessionStore>,
|
||||
stream_message_id: parking_lot::Mutex<Option<String>>,
|
||||
stream_message_id: std::sync::Mutex<Option<String>>,
|
||||
}
|
||||
|
||||
impl BusToolCallEmitter {
|
||||
@ -77,7 +75,7 @@ impl BusToolCallEmitter {
|
||||
chat_id: chat_id.into(),
|
||||
metadata,
|
||||
store,
|
||||
stream_message_id: parking_lot::Mutex::new(None),
|
||||
stream_message_id: std::sync::Mutex::new(None),
|
||||
}
|
||||
}
|
||||
}
|
||||
@ -94,14 +92,7 @@ impl EmittedMessageHandler for BusToolCallEmitter {
|
||||
&message,
|
||||
) {
|
||||
if let Err(error) = self.bus.publish_outbound(outbound).await {
|
||||
match error {
|
||||
crate::bus::BusError::Dropped => {
|
||||
tracing::warn!(error = %error, channel = %self.channel_name, chat_id = %self.chat_id, "Outbound dropped (bus full)");
|
||||
}
|
||||
crate::bus::BusError::Closed => {
|
||||
tracing::error!(error = %error, channel = %self.channel_name, chat_id = %self.chat_id, "Failed to publish live outbound tool call");
|
||||
}
|
||||
}
|
||||
tracing::error!(error = %error, channel = %self.channel_name, chat_id = %self.chat_id, "Failed to publish live outbound tool call");
|
||||
}
|
||||
}
|
||||
}
|
||||
@ -120,14 +111,7 @@ impl EmittedMessageHandler for BusToolCallEmitter {
|
||||
&message,
|
||||
) {
|
||||
if let Err(error) = self.bus.publish_outbound(outbound).await {
|
||||
match error {
|
||||
crate::bus::BusError::Dropped => {
|
||||
tracing::warn!(error = %error, channel = %self.channel_name, chat_id = %self.chat_id, "Outbound dropped (bus full)");
|
||||
}
|
||||
crate::bus::BusError::Closed => {
|
||||
tracing::error!(error = %error, channel = %self.channel_name, chat_id = %self.chat_id, "Failed to publish live outbound tool call");
|
||||
}
|
||||
}
|
||||
tracing::error!(error = %error, channel = %self.channel_name, chat_id = %self.chat_id, "Failed to publish live outbound tool call");
|
||||
}
|
||||
}
|
||||
|
||||
@ -140,10 +124,8 @@ impl EmittedMessageHandler for BusToolCallEmitter {
|
||||
async fn handle_stream_delta(&self, delta: &StreamDelta) {
|
||||
// Get or create the stream message ID
|
||||
let message_id = {
|
||||
let mut guard = self.stream_message_id.lock();
|
||||
guard
|
||||
.get_or_insert_with(|| Uuid::new_v4().to_string())
|
||||
.clone()
|
||||
let mut guard = self.stream_message_id.lock().unwrap();
|
||||
guard.get_or_insert_with(|| Uuid::new_v4().to_string()).clone()
|
||||
};
|
||||
|
||||
// Empty content + no reasoning = stream end signal
|
||||
@ -168,19 +150,12 @@ impl EmittedMessageHandler for BusToolCallEmitter {
|
||||
};
|
||||
|
||||
if let Err(error) = self.bus.publish_outbound(outbound).await {
|
||||
match error {
|
||||
crate::bus::BusError::Dropped => {
|
||||
tracing::warn!(error = %error, channel = %self.channel_name, "Outbound dropped (bus full)");
|
||||
}
|
||||
crate::bus::BusError::Closed => {
|
||||
tracing::error!(error = %error, channel = %self.channel_name, "Failed to publish stream delta");
|
||||
}
|
||||
}
|
||||
tracing::error!(error = %error, channel = %self.channel_name, "Failed to publish stream delta");
|
||||
}
|
||||
}
|
||||
|
||||
async fn set_stream_message_id(&self, id: &str) {
|
||||
*self.stream_message_id.lock() = Some(id.to_string());
|
||||
*self.stream_message_id.lock().unwrap() = Some(id.to_string());
|
||||
}
|
||||
}
|
||||
|
||||
@ -205,11 +180,7 @@ impl BusToolCallEmitter {
|
||||
.cloned()
|
||||
.unwrap_or_else(|| session_id.clone());
|
||||
|
||||
let topic_id = self
|
||||
.metadata
|
||||
.get("topic_id")
|
||||
.filter(|t| !t.is_empty())
|
||||
.cloned();
|
||||
let topic_id = self.metadata.get("topic_id").filter(|t| !t.is_empty()).cloned();
|
||||
|
||||
let now = std::time::SystemTime::now()
|
||||
.duration_since(std::time::UNIX_EPOCH)
|
||||
@ -301,7 +272,6 @@ impl Session {
|
||||
prompt_repository.clone(),
|
||||
model_resolver,
|
||||
Arc::new(super::model_selection::ModelSelectionStore::new()),
|
||||
crate::config::CompactionConfig::default(),
|
||||
);
|
||||
Self::with_factories(
|
||||
channel_name,
|
||||
@ -326,10 +296,6 @@ impl Session {
|
||||
skill_events: Arc<dyn SkillEventRepository>,
|
||||
store: Arc<SessionStore>,
|
||||
) -> Result<Self, AgentError> {
|
||||
// Session 的 compressor 用于 sync 兜底压缩路径（compaction.rs），
|
||||
// 必须与 AgentLoop 的 compressor 共用同一套用户配置的压缩参数。
|
||||
let runtime_config = AgentRuntimeConfig::from(provider_config.clone());
|
||||
let compressor = agent_factory.build_compressor(&runtime_config);
|
||||
Ok(Self {
|
||||
id: Uuid::new_v4(),
|
||||
channel_name: channel_name.clone(),
|
||||
@ -337,8 +303,12 @@ impl Session {
|
||||
provider_config: provider_config.clone(),
|
||||
skills,
|
||||
agent_factory,
|
||||
compressor,
|
||||
history: SessionHistory::new(channel_name, conversations, skill_events),
|
||||
compressor: ContextCompressor::from_provider_config(&provider_config),
|
||||
history: SessionHistory::new(
|
||||
channel_name,
|
||||
conversations,
|
||||
skill_events,
|
||||
),
|
||||
store,
|
||||
pending_cancel_tokens: HashMap::new(),
|
||||
})
|
||||
@ -616,7 +586,7 @@ impl Session {
|
||||
) -> Result<AgentLoop, AgentError> {
|
||||
self.create_agent_with_provider_config(
|
||||
chat_id,
|
||||
None, // notification_chat_id = None，使用 session_chat_id
|
||||
None, // notification_chat_id = None，使用 session_chat_id
|
||||
sender_id,
|
||||
message_id,
|
||||
self.provider_config.clone(),
|
||||
@ -641,9 +611,7 @@ impl Session {
|
||||
// 消费 pending 的取消信号接收端（如果存在）
|
||||
// 优先按 topic_id 查找；无 topic 时回退 chat_id
|
||||
let cancel_token = match &topic_id {
|
||||
Some(tid) => self
|
||||
.pending_cancel_tokens
|
||||
.remove(tid)
|
||||
Some(tid) => self.pending_cancel_tokens.remove(tid)
|
||||
.or_else(|| self.pending_cancel_tokens.remove(session_chat_id)),
|
||||
None => self.pending_cancel_tokens.remove(session_chat_id),
|
||||
};
|
||||
@ -749,7 +717,6 @@ impl SessionManager {
|
||||
mcp_config,
|
||||
None,
|
||||
model_resolver,
|
||||
crate::config::CompactionConfig::default(),
|
||||
)
|
||||
.map(|(session_manager, _, _, _, _)| session_manager)
|
||||
}
|
||||
@ -813,11 +780,7 @@ impl SessionManager {
|
||||
}
|
||||
|
||||
/// 获取指定 chat 的当前话题（确保 session 存在，自动从数据库恢复）
|
||||
pub async fn get_current_topic(
|
||||
&self,
|
||||
channel_name: &str,
|
||||
chat_id: &str,
|
||||
) -> Result<Option<String>, AgentError> {
|
||||
pub async fn get_current_topic(&self, channel_name: &str, chat_id: &str) -> Result<Option<String>, AgentError> {
|
||||
self.ensure_session(channel_name).await?;
|
||||
if let Some(session) = self.get(channel_name).await {
|
||||
let mut guard = session.lock().await;
|
||||
@ -825,9 +788,7 @@ impl SessionManager {
|
||||
// 如果内存中没有当前话题，从数据库恢复最近活跃的话题
|
||||
if guard.current_topic(chat_id).is_none() {
|
||||
let session_id = guard.persistent_session_id(chat_id);
|
||||
let topics = self
|
||||
.store
|
||||
.list_topics(&session_id)
|
||||
let topics = self.store.list_topics(&session_id)
|
||||
.map_err(|e| AgentError::Other(format!("Failed to list topics: {}", e)))?;
|
||||
|
||||
if let Some(latest_topic) = topics.first() {
|
||||
@ -841,7 +802,10 @@ impl SessionManager {
|
||||
);
|
||||
} else {
|
||||
// 数据库中也没有话题，自动创建默认话题
|
||||
let title = format!("话题 {}", chrono::Local::now().format("%m/%d %H:%M"));
|
||||
let title = format!(
|
||||
"话题 {}",
|
||||
chrono::Local::now().format("%m/%d %H:%M")
|
||||
);
|
||||
match self.store.create_topic(&session_id, &title, None) {
|
||||
Ok(topic) => {
|
||||
guard.set_current_topic(chat_id, Some(topic.id.clone()));
|
||||
@ -881,10 +845,7 @@ impl SessionManager {
|
||||
token: tokio::sync::watch::Receiver<()>,
|
||||
) {
|
||||
if let Some(session) = self.get(channel_name).await {
|
||||
session
|
||||
.lock()
|
||||
.await
|
||||
.set_cancel_receiver(chat_id, topic_id, token);
|
||||
session.lock().await.set_cancel_receiver(chat_id, topic_id, token);
|
||||
}
|
||||
}
|
||||
|
||||
@ -943,13 +904,7 @@ impl SessionManager {
|
||||
options: ScheduledAgentTaskOptions,
|
||||
) -> Result<Vec<OutboundMessage>, AgentError> {
|
||||
self.scheduled_tasks
|
||||
.run(
|
||||
channel_name,
|
||||
session_chat_id,
|
||||
notification_chat_id,
|
||||
prompt,
|
||||
options,
|
||||
)
|
||||
.run(channel_name, session_chat_id, notification_chat_id, prompt, options)
|
||||
.await
|
||||
}
|
||||
|
||||
@ -991,7 +946,6 @@ mod tests {
|
||||
extra_headers: HashMap::new(),
|
||||
llm_timeout_secs: 120,
|
||||
memory_maintenance_timeout_secs: 600,
|
||||
max_retries: 3,
|
||||
model_id: "test-model".to_string(),
|
||||
temperature: Some(0.0),
|
||||
max_tokens: Some(32),
|
||||
@ -1018,9 +972,9 @@ mod tests {
|
||||
store.clone(),
|
||||
store.clone(),
|
||||
Arc::new(NoopSessionMessageSender),
|
||||
HashSet::new(),
|
||||
HashSet::new(),
|
||||
"Asia/Shanghai".to_string(),
|
||||
HashSet::new(),
|
||||
HashSet::new(),
|
||||
Default::default(),
|
||||
)
|
||||
.build(),
|
||||
@ -1046,16 +1000,12 @@ mod tests {
|
||||
|
||||
let first = session.create_user_message("first", Vec::new());
|
||||
let first_id = first.id.clone();
|
||||
session
|
||||
.append_persisted_message("chat-1", Some(&topic_id), first)
|
||||
.unwrap();
|
||||
session.append_persisted_message("chat-1", Some(&topic_id), first).unwrap();
|
||||
assert!(session.is_latest_user_message(&topic_id, &first_id));
|
||||
|
||||
let second = session.create_user_message("second", Vec::new());
|
||||
let second_id = second.id.clone();
|
||||
session
|
||||
.append_persisted_message("chat-1", Some(&topic_id), second)
|
||||
.unwrap();
|
||||
session.append_persisted_message("chat-1", Some(&topic_id), second).unwrap();
|
||||
|
||||
assert!(!session.is_latest_user_message(&topic_id, &first_id));
|
||||
assert!(session.is_latest_user_message(&topic_id, &second_id));
|
||||
@ -1074,9 +1024,9 @@ mod tests {
|
||||
store.clone(),
|
||||
store.clone(),
|
||||
Arc::new(NoopSessionMessageSender),
|
||||
HashSet::new(),
|
||||
HashSet::new(),
|
||||
"Asia/Shanghai".to_string(),
|
||||
HashSet::new(),
|
||||
HashSet::new(),
|
||||
Default::default(),
|
||||
)
|
||||
.build(),
|
||||
@ -1102,15 +1052,9 @@ mod tests {
|
||||
|
||||
let first = session.create_user_message("first", Vec::new());
|
||||
let first_id = first.id.clone();
|
||||
session.append_persisted_message("chat-1", Some(&topic_id), first).unwrap();
|
||||
session
|
||||
.append_persisted_message("chat-1", Some(&topic_id), first)
|
||||
.unwrap();
|
||||
session
|
||||
.append_persisted_message(
|
||||
"chat-1",
|
||||
Some(&topic_id),
|
||||
ChatMessage::assistant("answer-1"),
|
||||
)
|
||||
.append_persisted_message("chat-1", Some(&topic_id), ChatMessage::assistant("answer-1"))
|
||||
.unwrap();
|
||||
|
||||
let second = session.create_user_message("second", Vec::new());
|
||||
@ -1118,11 +1062,7 @@ mod tests {
|
||||
.append_persisted_message("chat-1", Some(&topic_id), second.clone())
|
||||
.unwrap();
|
||||
session
|
||||
.append_persisted_message(
|
||||
"chat-1",
|
||||
Some(&topic_id),
|
||||
ChatMessage::assistant("answer-2"),
|
||||
)
|
||||
.append_persisted_message("chat-1", Some(&topic_id), ChatMessage::assistant("answer-2"))
|
||||
.unwrap();
|
||||
|
||||
let preserved_messages = session.get_history(&topic_id).unwrap().clone();
|
||||
@ -1272,7 +1212,6 @@ mod tests {
|
||||
max_tool_iterations: 1,
|
||||
llm_timeout_secs: 30,
|
||||
memory_maintenance_timeout_secs: 600,
|
||||
max_retries: 3,
|
||||
tool_result_max_chars: 100_000,
|
||||
context_tool_result_trim_chars: 100_000,
|
||||
max_images_in_context: 1,
|
||||
@ -1296,15 +1235,7 @@ mod tests {
|
||||
.unwrap();
|
||||
|
||||
let outbound = session_manager
|
||||
.handle_message(
|
||||
"test-channel",
|
||||
"user-1",
|
||||
"chat-1",
|
||||
"hello",
|
||||
Vec::new(),
|
||||
None,
|
||||
None,
|
||||
)
|
||||
.handle_message("test-channel", "user-1", "chat-1", "hello", Vec::new(), None, None)
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
@ -1329,7 +1260,6 @@ mod tests {
|
||||
max_tool_iterations: 1,
|
||||
llm_timeout_secs: 30,
|
||||
memory_maintenance_timeout_secs: 600,
|
||||
max_retries: 3,
|
||||
tool_result_max_chars: 100_000,
|
||||
context_tool_result_trim_chars: 100_000,
|
||||
max_images_in_context: 1,
|
||||
@ -1408,7 +1338,6 @@ mod tests {
|
||||
max_tool_iterations: 1,
|
||||
llm_timeout_secs: 30,
|
||||
memory_maintenance_timeout_secs: 600,
|
||||
max_retries: 3,
|
||||
tool_result_max_chars: 100_000,
|
||||
context_tool_result_trim_chars: 100_000,
|
||||
max_images_in_context: 1,
|
||||
@ -1505,7 +1434,6 @@ mod tests {
|
||||
max_tool_iterations: 1,
|
||||
llm_timeout_secs: 30,
|
||||
memory_maintenance_timeout_secs: 600,
|
||||
max_retries: 3,
|
||||
tool_result_max_chars: 100_000,
|
||||
context_tool_result_trim_chars: 100_000,
|
||||
max_images_in_context: 1,
|
||||
@ -1596,7 +1524,6 @@ mod tests {
|
||||
max_tool_iterations: 1,
|
||||
llm_timeout_secs: 1,
|
||||
memory_maintenance_timeout_secs: 600,
|
||||
max_retries: 3,
|
||||
tool_result_max_chars: 100_000,
|
||||
context_tool_result_trim_chars: 100_000,
|
||||
max_images_in_context: 1,
|
||||
@ -1686,7 +1613,6 @@ mod tests {
|
||||
max_tool_iterations: 1,
|
||||
llm_timeout_secs: 30,
|
||||
memory_maintenance_timeout_secs: 600,
|
||||
max_retries: 3,
|
||||
tool_result_max_chars: 100_000,
|
||||
context_tool_result_trim_chars: 100_000,
|
||||
max_images_in_context: 1,
|
||||
@ -1758,7 +1684,6 @@ mod tests {
|
||||
max_tool_iterations: 1,
|
||||
llm_timeout_secs: 30,
|
||||
memory_maintenance_timeout_secs: 600,
|
||||
max_retries: 3,
|
||||
tool_result_max_chars: 100_000,
|
||||
context_tool_result_trim_chars: 100_000,
|
||||
max_images_in_context: 1,
|
||||
@ -1808,8 +1733,7 @@ mod tests {
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_run_memory_maintenance_for_all_scopes_scans_all_scopes_even_without_recent_updates()
|
||||
{
|
||||
async fn test_run_memory_maintenance_for_all_scopes_scans_all_scopes_even_without_recent_updates() {
|
||||
let mock_response_content = serde_json::to_string(&json!({
|
||||
"user_facts": ["用户在做AI产品"],
|
||||
"preferences": [],
|
||||
@ -1840,7 +1764,6 @@ mod tests {
|
||||
max_tool_iterations: 1,
|
||||
llm_timeout_secs: 30,
|
||||
memory_maintenance_timeout_secs: 600,
|
||||
max_retries: 3,
|
||||
tool_result_max_chars: 100_000,
|
||||
context_tool_result_trim_chars: 100_000,
|
||||
max_images_in_context: 1,
|
||||
@ -1908,7 +1831,6 @@ mod tests {
|
||||
max_tool_iterations: 1,
|
||||
llm_timeout_secs: 1,
|
||||
memory_maintenance_timeout_secs: 600,
|
||||
max_retries: 3,
|
||||
tool_result_max_chars: 100_000,
|
||||
context_tool_result_trim_chars: 100_000,
|
||||
max_images_in_context: 1,
|
||||
@ -2061,18 +1983,11 @@ mod tests {
|
||||
|
||||
let all_memories = store.list_memories_for_scope("user", scope_key).unwrap();
|
||||
// 过滤掉 _meta 记录
|
||||
let user_memories: Vec<_> = all_memories
|
||||
.iter()
|
||||
.filter(|m| m.namespace != "_meta")
|
||||
.collect();
|
||||
let user_memories: Vec<_> = all_memories.iter().filter(|m| m.namespace != "_meta").collect();
|
||||
// 合并 2 条为 1 条，删除 1 条，7 - 2 + 1 = 5 条
|
||||
assert_eq!(user_memories.len(), 5);
|
||||
// 验证合并后的记忆存在
|
||||
assert!(
|
||||
user_memories
|
||||
.iter()
|
||||
.any(|m| m.namespace == "user" && m.memory_key == "work")
|
||||
);
|
||||
assert!(user_memories.iter().any(|m| m.namespace == "user" && m.memory_key == "work"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
@ -2095,19 +2010,22 @@ mod tests {
|
||||
let store = Arc::new(SessionStore::in_memory().unwrap());
|
||||
let bus = MessageBus::new(4);
|
||||
let emitter =
|
||||
BusToolCallEmitter::new(bus.clone(), "test-channel", "chat-1", HashMap::new(), store);
|
||||
BusToolCallEmitter::new(
|
||||
bus.clone(),
|
||||
"test-channel",
|
||||
"chat-1",
|
||||
HashMap::new(),
|
||||
store,
|
||||
);
|
||||
|
||||
emitter
|
||||
.handle(ChatMessage::tool("call-1", "calculator", "2"))
|
||||
.await;
|
||||
|
||||
let msg = tokio::time::timeout(
|
||||
std::time::Duration::from_millis(500),
|
||||
bus.consume_outbound(),
|
||||
)
|
||||
.await
|
||||
.expect("timeout waiting for outbound message")
|
||||
.expect("bus outbound closed");
|
||||
let msg = tokio::time::timeout(std::time::Duration::from_millis(500), bus.consume_outbound())
|
||||
.await
|
||||
.expect("timeout waiting for outbound message")
|
||||
.expect("bus outbound closed");
|
||||
assert_eq!(msg.event_kind, OutboundEventKind::ToolResult);
|
||||
}
|
||||
|
||||
@ -2124,9 +2042,9 @@ mod tests {
|
||||
store.clone(),
|
||||
store.clone(),
|
||||
Arc::new(NoopSessionMessageSender),
|
||||
HashSet::new(),
|
||||
HashSet::new(),
|
||||
"Asia/Shanghai".to_string(),
|
||||
HashSet::new(),
|
||||
HashSet::new(),
|
||||
Default::default(),
|
||||
)
|
||||
.build(),
|
||||
@ -2165,9 +2083,9 @@ mod tests {
|
||||
store.clone(),
|
||||
store.clone(),
|
||||
Arc::new(NoopSessionMessageSender),
|
||||
HashSet::new(),
|
||||
HashSet::new(),
|
||||
"Asia/Shanghai".to_string(),
|
||||
HashSet::new(),
|
||||
HashSet::new(),
|
||||
Default::default(),
|
||||
)
|
||||
.build(),
|
||||
@ -2193,11 +2111,7 @@ mod tests {
|
||||
|
||||
for turn in 0..100 {
|
||||
session
|
||||
.append_persisted_message(
|
||||
"chat-1",
|
||||
Some(&topic_id),
|
||||
ChatMessage::user(format!("user-{turn}")),
|
||||
)
|
||||
.append_persisted_message("chat-1", Some(&topic_id), ChatMessage::user(format!("user-{turn}")))
|
||||
.unwrap();
|
||||
}
|
||||
|
||||
@ -2246,9 +2160,9 @@ mod tests {
|
||||
store.clone(),
|
||||
store.clone(),
|
||||
Arc::new(NoopSessionMessageSender),
|
||||
HashSet::new(),
|
||||
HashSet::new(),
|
||||
"Asia/Shanghai".to_string(),
|
||||
HashSet::new(),
|
||||
HashSet::new(),
|
||||
Default::default(),
|
||||
)
|
||||
.build(),
|
||||
@ -2274,11 +2188,7 @@ mod tests {
|
||||
|
||||
for turn in 0..100 {
|
||||
session
|
||||
.append_persisted_message(
|
||||
"chat-1",
|
||||
Some(&topic_id),
|
||||
ChatMessage::user(format!("user-{turn}")),
|
||||
)
|
||||
.append_persisted_message("chat-1", Some(&topic_id), ChatMessage::user(format!("user-{turn}")))
|
||||
.unwrap();
|
||||
}
|
||||
|
||||
|
||||
@ -7,11 +7,6 @@ use crate::storage::{
|
||||
ConversationRepository, SessionRecord, SkillEventRepository, persistent_session_id,
|
||||
};
|
||||
|
||||
/// 内存中缓存的 topic 历史上限。
|
||||
/// 超过此值时，驱逐非活跃 topic（不在 chat_topic_ids 当前引用中的 topic）。
|
||||
/// 活跃 topic 永不被驱逐，避免影响正在进行的对话。
|
||||
const MAX_CACHED_TOPICS: usize = 32;
|
||||
|
||||
fn preview_text(content: &str, max_chars: usize) -> String {
|
||||
let mut preview = content.chars().take(max_chars).collect::<String>();
|
||||
if content.chars().count() > max_chars {
|
||||
@ -24,7 +19,6 @@ pub(crate) struct SessionHistory {
|
||||
channel_name: String,
|
||||
/// 按 topic_id 键化的内存历史缓存。
|
||||
/// 不同 topic 的历史独立存储，互不干扰，支持多话题并发执行。
|
||||
/// 超过 `MAX_CACHED_TOPICS` 时自动驱逐非活跃 topic。
|
||||
topic_histories: HashMap<String, Vec<ChatMessage>>,
|
||||
/// UI 状态：每个 chat 当前活跃的 topic（按 chat_id 键）。
|
||||
chat_topic_ids: HashMap<String, String>,
|
||||
@ -40,56 +34,6 @@ pub(crate) struct SessionHistory {
|
||||
}
|
||||
|
||||
impl SessionHistory {
|
||||
/// 当缓存 topic 数超过 `MAX_CACHED_TOPICS` 时，驱逐非活跃 topic。
|
||||
///
|
||||
/// 活跃判定（任一满足即活跃，不驱逐）：
|
||||
/// 1. 在 `chat_topic_ids` 的 values 中（UI 当前引用的 topic）
|
||||
/// 2. 在 `compression_in_flight` 中（正在压缩的 topic）
|
||||
/// 3. `topic_serial_lock` 被持有（有活跃 agent 任务正在处理该 topic）
|
||||
///
|
||||
/// 第 3 项防止驱逐正在 agent 处理中的 topic：agent 处理使用 `original_topic_id`
|
||||
/// 而非 UI 状态 `chat_topic_ids`，用户切换 topic 后原 topic 不在 UI 集合中，
|
||||
/// 但 agent 仍在处理（持有 serial lock），此时不应驱逐。
|
||||
fn evict_inactive_if_needed(&mut self) {
|
||||
if self.topic_histories.len() <= MAX_CACHED_TOPICS {
|
||||
return;
|
||||
}
|
||||
|
||||
// 收集当前活跃 topic 集合
|
||||
let active: HashSet<&str> = self
|
||||
.chat_topic_ids
|
||||
.values()
|
||||
.map(|s| s.as_str())
|
||||
.collect();
|
||||
|
||||
// 找一个非活跃 topic 驱逐
|
||||
let to_evict = self.topic_histories.keys().find(|tid| {
|
||||
if active.contains(tid.as_str()) || self.compression_in_flight.contains(*tid) {
|
||||
return false;
|
||||
}
|
||||
// 检查是否有活跃 agent 任务（serial lock 被持有）
|
||||
// try_lock 成功 = 锁空闲 = 无活跃任务 = 可驱逐
|
||||
// try_lock 失败 = 锁被持有 = 有活跃任务 = 不驱逐
|
||||
if let Some(lock) = self.topic_serial_locks.get(*tid) {
|
||||
if lock.try_lock().is_err() {
|
||||
return false;
|
||||
}
|
||||
}
|
||||
true
|
||||
});
|
||||
|
||||
if let Some(tid) = to_evict.cloned() {
|
||||
let msg_count = self.topic_histories.get(&tid).map(|h| h.len()).unwrap_or(0);
|
||||
self.topic_histories.remove(&tid);
|
||||
tracing::info!(
|
||||
topic_id = %tid,
|
||||
evicted_messages = msg_count,
|
||||
remaining_topics = self.topic_histories.len(),
|
||||
"Evicted inactive topic history to respect MAX_CACHED_TOPICS"
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) fn new(
|
||||
channel_name: impl Into<String>,
|
||||
conversations: Arc<dyn ConversationRepository>,
|
||||
@ -159,7 +103,6 @@ impl SessionHistory {
|
||||
}
|
||||
|
||||
self.topic_histories.insert(tid.to_string(), history);
|
||||
self.evict_inactive_if_needed();
|
||||
Ok(())
|
||||
}
|
||||
|
||||
@ -172,9 +115,7 @@ impl SessionHistory {
|
||||
}
|
||||
|
||||
pub(crate) fn get_or_create_history(&mut self, topic_id: &str) -> &mut Vec<ChatMessage> {
|
||||
self.topic_histories
|
||||
.entry(topic_id.to_string())
|
||||
.or_default()
|
||||
self.topic_histories.entry(topic_id.to_string()).or_default()
|
||||
}
|
||||
|
||||
pub(crate) fn get_history(&self, topic_id: &str) -> Option<&Vec<ChatMessage>> {
|
||||
@ -183,7 +124,6 @@ impl SessionHistory {
|
||||
|
||||
pub(crate) fn set_history(&mut self, topic_id: &str, history: Vec<ChatMessage>) {
|
||||
self.topic_histories.insert(topic_id.to_string(), history);
|
||||
self.evict_inactive_if_needed();
|
||||
}
|
||||
|
||||
/// 设置指定 chat 的当前 topic（UI 状态）
|
||||
@ -208,10 +148,6 @@ impl SessionHistory {
|
||||
pub(crate) fn remove_history(&mut self, topic_id: &str) {
|
||||
self.topic_histories.remove(topic_id);
|
||||
self.compression_in_flight.remove(topic_id);
|
||||
// 清理 serial lock，防止 topic_serial_locks 无限增长
|
||||
// （仅在无活跃任务时安全移除；有活跃任务时 lock 被 Arc clone 持有，
|
||||
// 移除 HashMap entry 不影响正在使用 lock 的任务）
|
||||
self.topic_serial_locks.remove(topic_id);
|
||||
}
|
||||
|
||||
/// 清空指定 chat/topic 的内存历史和 DB 消息。
|
||||
@ -223,7 +159,6 @@ impl SessionHistory {
|
||||
) -> Result<(), AgentError> {
|
||||
if let Some(tid) = topic_id {
|
||||
if let Some(history) = self.topic_histories.get_mut(tid) {
|
||||
#[cfg(debug_assertions)]
|
||||
let len = history.len();
|
||||
history.clear();
|
||||
#[cfg(debug_assertions)]
|
||||
@ -321,7 +256,6 @@ impl SessionHistory {
|
||||
/// 清空所有内存历史（主要用于测试全局重置）。
|
||||
/// 不遍历清 DB，生产环境如需清 DB 应由调用方显式调用。
|
||||
pub(crate) fn clear_all_history(&mut self) -> Result<(), AgentError> {
|
||||
#[cfg(debug_assertions)]
|
||||
let total: usize = self.topic_histories.values().map(|h| h.len()).sum();
|
||||
self.topic_histories.clear();
|
||||
self.compression_in_flight.clear();
|
||||
@ -342,7 +276,6 @@ impl SessionHistory {
|
||||
.load_messages_for_topic(topic_id, Some(&sid))
|
||||
.map_err(|err| AgentError::Other(format!("session history reload error: {}", err)))?;
|
||||
self.topic_histories.insert(topic_id.to_string(), history);
|
||||
self.evict_inactive_if_needed();
|
||||
Ok(())
|
||||
}
|
||||
|
||||
|
||||
@ -52,12 +52,9 @@ impl SessionLifecycleService {
|
||||
channel_name: &str,
|
||||
chat_id: &str,
|
||||
) -> Result<Arc<Mutex<Session>>, AgentError> {
|
||||
self.session_pool
|
||||
.ensure_session_for_chat_id(channel_name, chat_id)
|
||||
.await?;
|
||||
self.session_pool.ensure_session_for_chat_id(channel_name, chat_id).await?;
|
||||
self.touch(channel_name).await;
|
||||
self.session_pool
|
||||
.get_for_chat_id(channel_name, chat_id)
|
||||
self.session_pool.get_for_chat_id(channel_name, chat_id)
|
||||
.await
|
||||
.ok_or_else(|| AgentError::Other("Session not found".to_string()))
|
||||
}
|
||||
|
||||
@ -3,7 +3,7 @@ use std::sync::Arc;
|
||||
|
||||
use async_trait::async_trait;
|
||||
|
||||
use crate::bus::{BusError, MessageBus, OutboundMessage};
|
||||
use crate::bus::{MessageBus, OutboundMessage};
|
||||
use crate::tools::{SessionMessageSender, SessionSendOutcome, SessionSendRequest, ToolContext};
|
||||
|
||||
pub(crate) struct BusSessionMessageSender {
|
||||
@ -55,28 +55,15 @@ impl SessionMessageSender for BusSessionMessageSender {
|
||||
if attachment_count > 0 {
|
||||
outbound.media = request.attachments.clone();
|
||||
}
|
||||
match self.bus.publish_outbound(outbound).await {
|
||||
Ok(()) => {
|
||||
published_messages += 1;
|
||||
tracing::info!(
|
||||
channel = %channel_name,
|
||||
chat_id = %chat_id,
|
||||
content_len = content_len,
|
||||
attachment_count = attachment_count,
|
||||
"Published session text message to outbound bus"
|
||||
);
|
||||
}
|
||||
Err(BusError::Dropped) => {
|
||||
tracing::warn!(
|
||||
channel = %channel_name,
|
||||
chat_id = %chat_id,
|
||||
"Outbound bus full, dropping session text message"
|
||||
);
|
||||
}
|
||||
Err(BusError::Closed) => {
|
||||
return Err(anyhow::anyhow!("Outbound bus closed"));
|
||||
}
|
||||
}
|
||||
self.bus.publish_outbound(outbound).await?;
|
||||
published_messages += 1;
|
||||
tracing::info!(
|
||||
channel = %channel_name,
|
||||
chat_id = %chat_id,
|
||||
content_len = content_len,
|
||||
attachment_count = attachment_count,
|
||||
"Published session text message to outbound bus"
|
||||
);
|
||||
} else {
|
||||
for attachment in request.attachments {
|
||||
let media_path = attachment.path.clone();
|
||||
@ -90,28 +77,15 @@ impl SessionMessageSender for BusSessionMessageSender {
|
||||
metadata.clone(),
|
||||
);
|
||||
outbound.media = vec![attachment];
|
||||
match self.bus.publish_outbound(outbound).await {
|
||||
Ok(()) => {
|
||||
published_messages += 1;
|
||||
tracing::info!(
|
||||
channel = %channel_name,
|
||||
chat_id = %chat_id,
|
||||
media_type = %media_type,
|
||||
media_path = %media_path,
|
||||
"Published session attachment to outbound bus"
|
||||
);
|
||||
}
|
||||
Err(BusError::Dropped) => {
|
||||
tracing::warn!(
|
||||
channel = %channel_name,
|
||||
chat_id = %chat_id,
|
||||
"Outbound bus full, dropping session attachment"
|
||||
);
|
||||
}
|
||||
Err(BusError::Closed) => {
|
||||
return Err(anyhow::anyhow!("Outbound bus closed"));
|
||||
}
|
||||
}
|
||||
self.bus.publish_outbound(outbound).await?;
|
||||
published_messages += 1;
|
||||
tracing::info!(
|
||||
channel = %channel_name,
|
||||
chat_id = %chat_id,
|
||||
media_type = %media_type,
|
||||
media_path = %media_path,
|
||||
"Published session attachment to outbound bus"
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
@ -148,7 +122,7 @@ mod tests {
|
||||
// 使用临时目录确保跨平台兼容
|
||||
attachments: vec![MediaItem::new(
|
||||
&std::env::temp_dir().join("demo.png").display().to_string(),
|
||||
"image",
|
||||
"image"
|
||||
)],
|
||||
},
|
||||
)
|
||||
@ -169,4 +143,4 @@ mod tests {
|
||||
assert_eq!(msg.media.len(), 1);
|
||||
assert_eq!(msg.media[0].media_type, "image");
|
||||
}
|
||||
}
|
||||
}
|
||||
@ -49,10 +49,7 @@ impl SessionPool {
|
||||
}
|
||||
|
||||
/// 确保定时任务专用 Session 存在
|
||||
pub(crate) async fn ensure_scheduler_session(
|
||||
&self,
|
||||
channel_name: &str,
|
||||
) -> Result<(), AgentError> {
|
||||
pub(crate) async fn ensure_scheduler_session(&self, channel_name: &str) -> Result<(), AgentError> {
|
||||
self.ensure_session_internal(channel_name, true).await
|
||||
}
|
||||
|
||||
@ -62,11 +59,7 @@ impl SessionPool {
|
||||
/// session 创建（含配置加载、agent 工厂构造），再次持锁插入并处理竞态。
|
||||
/// 避免跨 `session_factory.create().await` 持有全局锁导致所有 channel 的
|
||||
/// session 访问串行化。
|
||||
async fn ensure_session_internal(
|
||||
&self,
|
||||
channel_name: &str,
|
||||
is_scheduler: bool,
|
||||
) -> Result<(), AgentError> {
|
||||
async fn ensure_session_internal(&self, channel_name: &str, is_scheduler: bool) -> Result<(), AgentError> {
|
||||
// Fast path: 已存在直接返回（短暂持锁）
|
||||
{
|
||||
let inner = self.inner.lock().await;
|
||||
@ -116,26 +109,14 @@ impl SessionPool {
|
||||
}
|
||||
|
||||
/// 获取定时任务专用 Session
|
||||
pub(crate) async fn get_scheduler_session(
|
||||
&self,
|
||||
channel_name: &str,
|
||||
) -> Option<Arc<Mutex<Session>>> {
|
||||
self.inner
|
||||
.lock()
|
||||
.await
|
||||
.scheduler_sessions
|
||||
.get(channel_name)
|
||||
.cloned()
|
||||
pub(crate) async fn get_scheduler_session(&self, channel_name: &str) -> Option<Arc<Mutex<Session>>> {
|
||||
self.inner.lock().await.scheduler_sessions.get(channel_name).cloned()
|
||||
}
|
||||
|
||||
/// 根据 chat_id 自动选择 Session
|
||||
/// - scheduler/ 开头：返回定时任务专用 Session
|
||||
/// - 其他：返回主 Session
|
||||
pub(crate) async fn get_for_chat_id(
|
||||
&self,
|
||||
channel_name: &str,
|
||||
chat_id: &str,
|
||||
) -> Option<Arc<Mutex<Session>>> {
|
||||
pub(crate) async fn get_for_chat_id(&self, channel_name: &str, chat_id: &str) -> Option<Arc<Mutex<Session>>> {
|
||||
if is_scheduler_chat_id(chat_id) {
|
||||
self.get_scheduler_session(channel_name).await
|
||||
} else {
|
||||
@ -144,11 +125,7 @@ impl SessionPool {
|
||||
}
|
||||
|
||||
/// 确保 Session 存在（根据 chat_id 自动选择）
|
||||
pub(crate) async fn ensure_session_for_chat_id(
|
||||
&self,
|
||||
channel_name: &str,
|
||||
chat_id: &str,
|
||||
) -> Result<(), AgentError> {
|
||||
pub(crate) async fn ensure_session_for_chat_id(&self, channel_name: &str, chat_id: &str) -> Result<(), AgentError> {
|
||||
if is_scheduler_chat_id(chat_id) {
|
||||
self.ensure_scheduler_session(channel_name).await
|
||||
} else {
|
||||
|
||||
@ -1,6 +1,6 @@
|
||||
use axum::{
|
||||
body::Body,
|
||||
http::{Response, StatusCode, Uri, header},
|
||||
http::{header, Response, StatusCode, Uri},
|
||||
};
|
||||
use rust_embed::RustEmbed;
|
||||
|
||||
@ -16,7 +16,11 @@ pub async fn static_handler(uri: Uri) -> Response<Body> {
|
||||
let path = uri.path().trim_start_matches('/');
|
||||
|
||||
// 处理根路径，返回 index.html
|
||||
let path = if path.is_empty() { "index.html" } else { path };
|
||||
let path = if path.is_empty() {
|
||||
"index.html"
|
||||
} else {
|
||||
path
|
||||
};
|
||||
|
||||
match StaticAssets::get(path) {
|
||||
Some(content) => {
|
||||
@ -50,4 +54,4 @@ pub async fn static_handler(uri: Uri) -> Response<Body> {
|
||||
.unwrap()
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
@ -6,14 +6,13 @@ use tokio::sync::RwLock;
|
||||
use crate::config::TaskConfig;
|
||||
use crate::mcp::McpClientManager;
|
||||
use crate::skills::SkillRuntime;
|
||||
use crate::storage::{
|
||||
MemoryRepository, SchedulerJobRepository, SkillEventRepository, TodoRepository,
|
||||
};
|
||||
use crate::storage::{MemoryRepository, SchedulerJobRepository, SkillEventRepository, TodoRepository};
|
||||
use crate::tools::todo_write::TodoItem;
|
||||
use crate::tools::{
|
||||
BashTool, CalculatorTool, FileEditTool, FileReadTool, FileWriteTool, HttpRequestTool,
|
||||
MemoryManageTool, MemorySearchTool, SchedulerManageTool, SessionMessageSender, SessionSendTool,
|
||||
ShellSessionManager, SkillActivateTool, SkillManageTool, SubAgentRuntime, TaskTool, TimeTool,
|
||||
BashTool, CalculatorTool, FileEditTool, FileReadTool, FileWriteTool,
|
||||
HttpRequestTool, MemoryManageTool, MemorySearchTool,
|
||||
SchedulerManageTool, SessionMessageSender, SessionSendTool, ShellSessionManager,
|
||||
SkillActivateTool, SkillManageTool, SubAgentRuntime, TaskTool, TimeTool,
|
||||
TodoReadTool, TodoWriteTool, ToolRegistry, WebFetchTool,
|
||||
};
|
||||
|
||||
@ -73,12 +72,18 @@ impl ToolRegistryFactory {
|
||||
self
|
||||
}
|
||||
|
||||
pub(crate) fn with_subagent_runtime(mut self, runtime: Arc<dyn SubAgentRuntime>) -> Self {
|
||||
pub(crate) fn with_subagent_runtime(
|
||||
mut self,
|
||||
runtime: Arc<dyn SubAgentRuntime>,
|
||||
) -> Self {
|
||||
self.subagent_runtime = Some(runtime);
|
||||
self
|
||||
}
|
||||
|
||||
pub(crate) fn with_mcp_manager(mut self, manager: Arc<McpClientManager>) -> Self {
|
||||
pub(crate) fn with_mcp_manager(
|
||||
mut self,
|
||||
manager: Arc<McpClientManager>,
|
||||
) -> Self {
|
||||
self.mcp_manager = Some(manager);
|
||||
self
|
||||
}
|
||||
@ -113,14 +118,8 @@ impl ToolRegistryFactory {
|
||||
}
|
||||
if self.is_enabled("todo_write") {
|
||||
if let Some(ref state) = self.todo_state {
|
||||
registry.register(TodoWriteTool::new(
|
||||
state.clone(),
|
||||
self.todo_repository.clone(),
|
||||
));
|
||||
registry.register(TodoReadTool::new(
|
||||
state.clone(),
|
||||
self.todo_repository.clone(),
|
||||
));
|
||||
registry.register(TodoWriteTool::new(state.clone(), self.todo_repository.clone()));
|
||||
registry.register(TodoReadTool::new(state.clone(), self.todo_repository.clone()));
|
||||
}
|
||||
}
|
||||
if self.is_enabled("session_send") {
|
||||
@ -227,14 +226,8 @@ impl ToolRegistryFactory {
|
||||
// Todo 追踪工具
|
||||
if self.is_enabled("todo_write") {
|
||||
if let Some(ref state) = self.todo_state {
|
||||
registry.register(TodoWriteTool::new(
|
||||
state.clone(),
|
||||
self.todo_repository.clone(),
|
||||
));
|
||||
registry.register(TodoReadTool::new(
|
||||
state.clone(),
|
||||
self.todo_repository.clone(),
|
||||
));
|
||||
registry.register(TodoWriteTool::new(state.clone(), self.todo_repository.clone()));
|
||||
registry.register(TodoReadTool::new(state.clone(), self.todo_repository.clone()));
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@ -10,16 +10,16 @@ use crate::command::handlers::get_current::GetCurrentSessionCommandHandler;
|
||||
use crate::command::handlers::help::HelpCommandHandler;
|
||||
use crate::command::handlers::list_channels::ListChannelsCommandHandler;
|
||||
use crate::command::handlers::list_memories::ListMemoriesCommandHandler;
|
||||
use crate::command::handlers::list_skills::ListSkillsCommandHandler;
|
||||
use crate::command::handlers::list_scheduler_jobs::ListSchedulerJobsCommandHandler;
|
||||
use crate::command::handlers::list_todos::ListTodosCommandHandler;
|
||||
use crate::command::handlers::memory_crud::MemoryCrudCommandHandler;
|
||||
use crate::command::handlers::list_sessions::ListSessionsCommandHandler;
|
||||
use crate::command::handlers::list_sessions_by_channel::ListSessionsByChannelCommandHandler;
|
||||
use crate::command::handlers::list_skills::ListSkillsCommandHandler;
|
||||
use crate::command::handlers::list_todos::ListTodosCommandHandler;
|
||||
use crate::command::handlers::list_topics::ListTopicsCommandHandler;
|
||||
use crate::command::handlers::load_chat_messages::LoadChatMessagesCommandHandler;
|
||||
use crate::command::handlers::load_task_messages::LoadTaskMessagesCommandHandler;
|
||||
use crate::command::handlers::load_topic::LoadTopicCommandHandler;
|
||||
use crate::command::handlers::memory_crud::MemoryCrudCommandHandler;
|
||||
use crate::command::handlers::rename_topic::RenameTopicCommandHandler;
|
||||
use crate::command::handlers::save_session::SaveSessionCommandHandler;
|
||||
use crate::command::handlers::save_topic::SaveTopicCommandHandler;
|
||||
@ -27,28 +27,25 @@ use crate::command::handlers::session::SessionCommandHandler;
|
||||
use crate::command::handlers::stop_execution::StopExecutionCommandHandler;
|
||||
use crate::command::handlers::switch_topic::SwitchTopicCommandHandler;
|
||||
use crate::gateway::agent_factory::build_system_prompt_provider;
|
||||
use crate::protocol::{MediaSummary, WsInbound, WsOutbound, parse_inbound, serialize_outbound};
|
||||
use crate::protocol::{WsInbound, WsOutbound, MediaSummary, parse_inbound, serialize_outbound};
|
||||
use crate::storage::persistent_session_id;
|
||||
use crate::tools::task::repository::TaskRepository;
|
||||
use crate::utils::current_timestamp;
|
||||
use axum::extract::Query;
|
||||
use crate::tools::task::types::TaskSessionState;
|
||||
use axum::extract::State;
|
||||
use axum::extract::ws::{Message as WsMessage, WebSocket, WebSocketUpgrade};
|
||||
use axum::http::StatusCode;
|
||||
use axum::response::{IntoResponse, Response};
|
||||
use axum::response::Response;
|
||||
use base64::{Engine as _, engine::general_purpose::STANDARD};
|
||||
use futures_util::{SinkExt, StreamExt};
|
||||
use std::collections::HashMap;
|
||||
use std::path::PathBuf;
|
||||
use std::sync::Arc;
|
||||
use tokio::sync::mpsc;
|
||||
use tokio_util::sync::CancellationToken;
|
||||
|
||||
const WS_CHANNEL_NAME: &str = "websocket";
|
||||
|
||||
/// Default media directory for WebSocket uploads
|
||||
fn default_ws_media_dir() -> PathBuf {
|
||||
let home = crate::platform::picobot_home_dir();
|
||||
let home = dirs::home_dir().unwrap_or_else(|| PathBuf::from("."));
|
||||
home.join(".picobot").join("media").join("ws")
|
||||
}
|
||||
|
||||
@ -71,9 +68,7 @@ fn build_media_filename(media_type: &str, file_name: Option<&str>) -> String {
|
||||
|
||||
/// Process attachments with base64 content: save to local file and return MediaItem with correct path
|
||||
/// Keeps content_base64 for frontend display/download
|
||||
fn process_attachments_with_base64(
|
||||
attachments: Vec<MediaSummary>,
|
||||
) -> Result<Vec<MediaItem>, AgentError> {
|
||||
fn process_attachments_with_base64(attachments: Vec<MediaSummary>) -> Result<Vec<MediaItem>, AgentError> {
|
||||
if attachments.is_empty() {
|
||||
return Ok(Vec::new());
|
||||
}
|
||||
@ -87,16 +82,15 @@ fn process_attachments_with_base64(
|
||||
.map(|att| {
|
||||
// If content_base64 exists, save to file and update path
|
||||
if let Some(base64_content) = &att.content_base64 {
|
||||
let decoded = STANDARD.decode(base64_content).map_err(|error| {
|
||||
AgentError::Other(format!("Failed to decode base64: {}", error))
|
||||
})?;
|
||||
let decoded = STANDARD
|
||||
.decode(base64_content)
|
||||
.map_err(|error| AgentError::Other(format!("Failed to decode base64: {}", error)))?;
|
||||
|
||||
let filename = build_media_filename(&att.media_type, att.file_name.as_deref());
|
||||
let file_path = media_dir.join(&filename);
|
||||
|
||||
std::fs::write(&file_path, decoded).map_err(|error| {
|
||||
AgentError::Other(format!("Failed to write media file: {}", error))
|
||||
})?;
|
||||
std::fs::write(&file_path, decoded)
|
||||
.map_err(|error| AgentError::Other(format!("Failed to write media file: {}", error)))?;
|
||||
|
||||
tracing::info!(
|
||||
filename = %filename,
|
||||
@ -129,29 +123,7 @@ fn process_attachments_with_base64(
|
||||
.collect()
|
||||
}
|
||||
|
||||
#[derive(serde::Deserialize)]
|
||||
pub struct WsAuthQuery {
|
||||
/// 可选的认证 token（浏览器原生 WebSocket 不支持自定义 header，通过 query param 传递）
|
||||
pub token: Option<String>,
|
||||
}
|
||||
|
||||
pub async fn ws_handler(
|
||||
ws: WebSocketUpgrade,
|
||||
State(state): State<Arc<GatewayState>>,
|
||||
Query(query): Query<WsAuthQuery>,
|
||||
auth_cfg: Option<axum::Extension<crate::gateway::auth::AuthConfig>>,
|
||||
) -> Response {
|
||||
// 若启用了认证（auth_cfg 存在且 token 已配置），校验 query param 中的 token
|
||||
if let Some(axum::Extension(cfg)) = auth_cfg {
|
||||
if let Some(ref expected) = cfg.token {
|
||||
let provided = query.token.as_deref();
|
||||
if !crate::gateway::auth::token_matches(provided, &Some(expected.clone())) {
|
||||
tracing::warn!("WebSocket connection rejected: missing or invalid token");
|
||||
return (StatusCode::UNAUTHORIZED, "missing or invalid token").into_response();
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub async fn ws_handler(ws: WebSocketUpgrade, State(state): State<Arc<GatewayState>>) -> Response {
|
||||
ws.on_upgrade(|socket| async {
|
||||
handle_socket(socket, state).await;
|
||||
})
|
||||
@ -164,8 +136,10 @@ async fn handle_socket(ws: WebSocket, state: Arc<GatewayState>) {
|
||||
let store = state.session_manager.store();
|
||||
|
||||
// 1. 查询 websocket 和 cli 两个通道的 Sessions（兼容旧版本 cli 通道创建的会话）
|
||||
let mut websocket_sessions = store.list_sessions("websocket", false).unwrap_or_default();
|
||||
let cli_channel_sessions = store.list_sessions("cli", false).unwrap_or_default();
|
||||
let mut websocket_sessions = store.list_sessions("websocket", false)
|
||||
.unwrap_or_default();
|
||||
let cli_channel_sessions = store.list_sessions("cli", false)
|
||||
.unwrap_or_default();
|
||||
websocket_sessions.extend(cli_channel_sessions);
|
||||
websocket_sessions.sort_by_key(|s| -(s.last_active_at));
|
||||
|
||||
@ -206,7 +180,9 @@ async fn handle_socket(ws: WebSocket, state: Arc<GatewayState>) {
|
||||
|
||||
// 连接建立后立即发送通道列表（合并 websocket + ChannelManager 动态通道）
|
||||
let channels = state.channel_manager.build_channel_list().await;
|
||||
let _ = sender.send(WsOutbound::ChannelList { channels }).await;
|
||||
let _ = sender
|
||||
.send(WsOutbound::ChannelList { channels })
|
||||
.await;
|
||||
|
||||
// 3. 发送合并后的 Session 列表（已在上面合并了 websocket + cli 通道）
|
||||
// 如果刚创建了新会话，确保它也在列表中
|
||||
@ -246,34 +222,13 @@ async fn handle_socket(ws: WebSocket, state: Arc<GatewayState>) {
|
||||
|
||||
let mut receiver = receiver;
|
||||
let session_id_for_sender = runtime_session_id.clone();
|
||||
// 每连接独立的关闭信号：writer 超时/错误时 cancel，通知主 loop 退出
|
||||
// 不能用共享的 shutdown_token——那是 CliChannel 级别的，cancel 会关闭所有连接
|
||||
let writer_closed = CancellationToken::new();
|
||||
let writer_closed_clone = writer_closed.clone();
|
||||
tokio::spawn(async move {
|
||||
while let Some(msg) = receiver.recv().await {
|
||||
if let Ok(text) = serialize_outbound(&msg) {
|
||||
let send_result = tokio::time::timeout(
|
||||
std::time::Duration::from_secs(30),
|
||||
ws_sender.send(WsMessage::Text(text.into())),
|
||||
)
|
||||
.await;
|
||||
match send_result {
|
||||
Ok(Ok(())) => {}
|
||||
Ok(Err(_)) => {
|
||||
#[cfg(debug_assertions)]
|
||||
tracing::debug!(session_id = %session_id_for_sender, "WebSocket send error");
|
||||
writer_closed_clone.cancel();
|
||||
break;
|
||||
}
|
||||
Err(_) => {
|
||||
tracing::warn!(
|
||||
session_id = %session_id_for_sender,
|
||||
"WebSocket send timed out after 30s, closing connection"
|
||||
);
|
||||
writer_closed_clone.cancel();
|
||||
break;
|
||||
}
|
||||
if ws_sender.send(WsMessage::Text(text.into())).await.is_err() {
|
||||
#[cfg(debug_assertions)]
|
||||
tracing::debug!(session_id = %session_id_for_sender, "WebSocket send error");
|
||||
break;
|
||||
}
|
||||
}
|
||||
}
|
||||
@ -281,16 +236,11 @@ async fn handle_socket(ws: WebSocket, state: Arc<GatewayState>) {
|
||||
|
||||
loop {
|
||||
tokio::select! {
|
||||
// 监听全局 shutdown 信号（来自 CliChannel::stop()）
|
||||
// 监听 shutdown 信号（来自 CliChannel::stop()）
|
||||
_ = shutdown_token.cancelled() => {
|
||||
tracing::info!(session_id = %current_session_id, "WebSocket shutdown signal received, closing connection");
|
||||
break;
|
||||
}
|
||||
// 监听 writer 退出信号（writer 超时或错误，仅关闭当前连接）
|
||||
_ = writer_closed.cancelled() => {
|
||||
tracing::info!(session_id = %current_session_id, "WebSocket writer closed, shutting down connection");
|
||||
break;
|
||||
}
|
||||
// 监听 WebSocket 消息
|
||||
msg = ws_receiver.next() => {
|
||||
let Some(msg) = msg else {
|
||||
@ -354,6 +304,7 @@ async fn handle_socket(ws: WebSocket, state: Arc<GatewayState>) {
|
||||
tracing::info!(session_id = %runtime_session_id, current_session_id = %current_session_id, "CLI session ended");
|
||||
}
|
||||
|
||||
|
||||
async fn handle_inbound(
|
||||
state: &Arc<GatewayState>,
|
||||
sender: &mpsc::Sender<WsOutbound>,
|
||||
@ -443,11 +394,7 @@ async fn handle_inbound(
|
||||
let store = state.session_manager.store();
|
||||
let skills = state.session_manager.skills();
|
||||
let skills_for_handler = skills.clone();
|
||||
let provider_config = state
|
||||
.config
|
||||
.read()
|
||||
.await
|
||||
.get_provider_config("default")
|
||||
let provider_config = state.config.read().await.get_provider_config("default")
|
||||
.map_err(|e| AgentError::Other(e.to_string()))?;
|
||||
let prompt_repository = state.session_manager.store().clone();
|
||||
|
||||
@ -470,13 +417,9 @@ async fn handle_inbound(
|
||||
// 注册 list_sessions 处理器
|
||||
router.register(Box::new(ListSessionsCommandHandler::new(store.clone())));
|
||||
// 注册 list_sessions_by_channel 处理器
|
||||
router.register(Box::new(ListSessionsByChannelCommandHandler::new(
|
||||
store.clone(),
|
||||
)));
|
||||
router.register(Box::new(ListSessionsByChannelCommandHandler::new(store.clone())));
|
||||
// 注册 list_channels 处理器
|
||||
router.register(Box::new(ListChannelsCommandHandler::new(Arc::new(
|
||||
state.channel_manager.clone(),
|
||||
))));
|
||||
router.register(Box::new(ListChannelsCommandHandler::new(Arc::new(state.channel_manager.clone()))));
|
||||
// 注册 list_topics 处理器
|
||||
router.register(Box::new(ListTopicsCommandHandler::new(store.clone())));
|
||||
// 注册 switch_topic 处理器
|
||||
@ -517,9 +460,7 @@ async fn handle_inbound(
|
||||
let metadata = router.metadata_arc();
|
||||
router.register(Box::new(HelpCommandHandler::new(metadata)));
|
||||
// 注册 list_scheduler_jobs 处理器
|
||||
router.register(Box::new(ListSchedulerJobsCommandHandler::new(
|
||||
store.clone(),
|
||||
)));
|
||||
router.register(Box::new(ListSchedulerJobsCommandHandler::new(store.clone())));
|
||||
// 注册 list_memories 处理器
|
||||
router.register(Box::new(ListMemoriesCommandHandler::new(store.clone())));
|
||||
// 注册 list_skills 处理器
|
||||
@ -583,118 +524,66 @@ async fn handle_inbound(
|
||||
*current_topic_id = Some(topic_id.clone());
|
||||
|
||||
// 加载并发送该话题的历史消息
|
||||
if let Err(e) = send_topic_history(
|
||||
&store,
|
||||
current_session_id,
|
||||
topic_id,
|
||||
sender,
|
||||
&state.task_repository,
|
||||
)
|
||||
.await
|
||||
{
|
||||
if let Err(e) = send_topic_history(&store, current_session_id, topic_id, sender, &state.task_repository).await {
|
||||
tracing::warn!(error = %e, topic_id = %topic_id, "Failed to send topic history");
|
||||
}
|
||||
}
|
||||
// 加载子智能体任务消息
|
||||
if let Some(task_session_id) = response.metadata.get("task_session_id") {
|
||||
// 提前提取 task_id，用于给历史消息打标记
|
||||
let task_id = response
|
||||
.metadata
|
||||
.get("task_id")
|
||||
.cloned()
|
||||
.unwrap_or_default();
|
||||
if let Err(e) = send_task_messages(
|
||||
&store,
|
||||
task_session_id,
|
||||
sender,
|
||||
Some(task_id.clone()),
|
||||
Some(&state.task_repository),
|
||||
)
|
||||
.await
|
||||
{
|
||||
let task_id = response.metadata.get("task_id").cloned().unwrap_or_default();
|
||||
if let Err(e) = send_task_messages(&store, task_session_id, sender, Some(task_id.clone()), Some(&state.task_repository)).await {
|
||||
tracing::warn!(error = %e, task_session_id = %task_session_id, "Failed to send task messages");
|
||||
}
|
||||
|
||||
// 发送 TaskMessagesLoaded 元数据
|
||||
let description = response
|
||||
.metadata
|
||||
.get("task_description")
|
||||
.cloned()
|
||||
.unwrap_or_default();
|
||||
let subagent_type = response
|
||||
.metadata
|
||||
.get("task_subagent_type")
|
||||
.cloned()
|
||||
.unwrap_or_default();
|
||||
let status = response
|
||||
.metadata
|
||||
.get("task_status")
|
||||
.cloned()
|
||||
.unwrap_or_default();
|
||||
let description = response.metadata.get("task_description").cloned().unwrap_or_default();
|
||||
let subagent_type = response.metadata.get("task_subagent_type").cloned().unwrap_or_default();
|
||||
let status = response.metadata.get("task_status").cloned().unwrap_or_default();
|
||||
let summary = response.metadata.get("task_summary").cloned();
|
||||
let token_stats = response.metadata.get("task_token_stats").and_then(|json| {
|
||||
serde_json::from_str::<crate::protocol::TopicTokenStats>(json).ok()
|
||||
});
|
||||
|
||||
let _ = sender
|
||||
.send(WsOutbound::TaskMessagesLoaded {
|
||||
task_id,
|
||||
description,
|
||||
subagent_type,
|
||||
status,
|
||||
summary,
|
||||
token_stats,
|
||||
})
|
||||
.await;
|
||||
let _ = sender.send(WsOutbound::TaskMessagesLoaded {
|
||||
task_id,
|
||||
description,
|
||||
subagent_type,
|
||||
status,
|
||||
summary,
|
||||
}).await;
|
||||
}
|
||||
|
||||
// 处理定时任务列表
|
||||
if let Some(jobs_json) = response.metadata.get("scheduler_jobs") {
|
||||
if let Ok(jobs) =
|
||||
serde_json::from_str::<Vec<crate::protocol::SchedulerJobSummary>>(jobs_json)
|
||||
{
|
||||
if let Ok(jobs) = serde_json::from_str::<Vec<crate::protocol::SchedulerJobSummary>>(jobs_json) {
|
||||
let _ = sender.send(WsOutbound::SchedulerJobList { jobs }).await;
|
||||
}
|
||||
}
|
||||
|
||||
// 处理技能列表
|
||||
if let Some(skills_json) = response.metadata.get("skills") {
|
||||
if let Ok(skills) =
|
||||
serde_json::from_str::<Vec<crate::protocol::SkillSummary>>(skills_json)
|
||||
{
|
||||
if let Ok(skills) = serde_json::from_str::<Vec<crate::protocol::SkillSummary>>(skills_json) {
|
||||
let _ = sender.send(WsOutbound::SkillList { skills }).await;
|
||||
}
|
||||
}
|
||||
|
||||
// 处理 Todo 列表
|
||||
if let Some(todos_json) = response.metadata.get("todos") {
|
||||
if let Ok(todos) =
|
||||
serde_json::from_str::<Vec<crate::protocol::TodoItemSummary>>(todos_json)
|
||||
{
|
||||
let scope_key = response
|
||||
.metadata
|
||||
.get("todos_scope_key")
|
||||
.cloned()
|
||||
.unwrap_or_default();
|
||||
tracing::debug!(todo_count = todos.len(), %scope_key, "list_todos command response");
|
||||
if let Ok(todos) = serde_json::from_str::<Vec<crate::protocol::TodoItemSummary>>(todos_json) {
|
||||
let scope_key = response.metadata.get("todos_scope_key").cloned().unwrap_or_default();
|
||||
tracing::info!(todo_count = todos.len(), %scope_key, "list_todos command response");
|
||||
let _ = sender.send(WsOutbound::TodoList { todos, scope_key }).await;
|
||||
}
|
||||
}
|
||||
|
||||
// 处理记忆列表
|
||||
if let Some(memories_json) = response.metadata.get("memories") {
|
||||
if let Ok(memories) =
|
||||
serde_json::from_str::<Vec<crate::protocol::MemorySummary>>(memories_json)
|
||||
{
|
||||
if let Ok(memories) = serde_json::from_str::<Vec<crate::protocol::MemorySummary>>(memories_json) {
|
||||
let _ = sender.send(WsOutbound::MemoryList { memories }).await;
|
||||
}
|
||||
}
|
||||
|
||||
// 记忆 CRUD 后自动刷新列表
|
||||
if response.metadata.get("memory_updated").map(|v| v.as_str()) == Some("true") {
|
||||
if let Ok(records) =
|
||||
store.list_memories_for_scope("user", crate::storage::GLOBAL_SCOPE_KEY)
|
||||
{
|
||||
if let Ok(records) = store.list_memories_for_scope("user", crate::storage::GLOBAL_SCOPE_KEY) {
|
||||
let memories: Vec<crate::protocol::MemorySummary> = records
|
||||
.into_iter()
|
||||
.filter(|m| m.namespace != "_meta")
|
||||
@ -713,17 +602,15 @@ async fn handle_inbound(
|
||||
|
||||
// 处理加载聊天消息请求
|
||||
if let Some(load_chat_id) = response.metadata.get("load_chat_id") {
|
||||
let load_chat_channel = response
|
||||
.metadata
|
||||
.get("load_chat_channel")
|
||||
let load_chat_channel = response.metadata.get("load_chat_channel")
|
||||
.cloned()
|
||||
.unwrap_or_default();
|
||||
// session_id = "{channel}:{chat_id}" (cli channel 例外)
|
||||
let session_id =
|
||||
crate::storage::persistent_session_id(&load_chat_channel, load_chat_id);
|
||||
if let Err(e) =
|
||||
send_task_messages(&store, &session_id, sender, None, None).await
|
||||
{
|
||||
let session_id = crate::storage::persistent_session_id(
|
||||
&load_chat_channel,
|
||||
load_chat_id,
|
||||
);
|
||||
if let Err(e) = send_task_messages(&store, &session_id, sender, None, None).await {
|
||||
tracing::warn!(
|
||||
error = %e,
|
||||
channel = %load_chat_channel,
|
||||
@ -736,22 +623,12 @@ async fn handle_inbound(
|
||||
|
||||
if current_topic_id.is_none() {
|
||||
if let Some(topics_json) = response.metadata.get("topics") {
|
||||
match serde_json::from_str::<Vec<crate::protocol::TopicSummary>>(
|
||||
topics_json,
|
||||
) {
|
||||
match serde_json::from_str::<Vec<crate::protocol::TopicSummary>>(topics_json) {
|
||||
Ok(topics) => {
|
||||
if let Some(first_topic) = topics.first() {
|
||||
let topic_id = first_topic.topic_id.clone();
|
||||
*current_topic_id = Some(topic_id.clone());
|
||||
if let Err(e) = send_topic_history(
|
||||
&store,
|
||||
current_session_id,
|
||||
&topic_id,
|
||||
sender,
|
||||
&state.task_repository,
|
||||
)
|
||||
.await
|
||||
{
|
||||
if let Err(e) = send_topic_history(&store, current_session_id, &topic_id, sender, &state.task_repository).await {
|
||||
tracing::warn!(error = %e, topic_id = %topic_id, "Failed to send initial topic history");
|
||||
}
|
||||
}
|
||||
@ -785,6 +662,13 @@ async fn handle_inbound(
|
||||
}
|
||||
}
|
||||
|
||||
fn current_timestamp() -> i64 {
|
||||
std::time::SystemTime::now()
|
||||
.duration_since(std::time::UNIX_EPOCH)
|
||||
.unwrap()
|
||||
.as_millis() as i64
|
||||
}
|
||||
|
||||
fn resolve_ws_sender_id(sender_id: Option<&str>, runtime_session_id: &str) -> String {
|
||||
sender_id
|
||||
.map(str::trim)
|
||||
@ -802,21 +686,10 @@ async fn send_topic_history(
|
||||
task_repository: &Arc<dyn TaskRepository>,
|
||||
) -> Result<(), Box<dyn std::error::Error>> {
|
||||
// 加载话题消息，按 session_id 过滤，避免混入子智能体消息
|
||||
let messages = store.load_messages_for_topic_full(topic_id, Some(session_id))?;
|
||||
let messages = store.load_messages_for_topic(topic_id, Some(session_id))?;
|
||||
|
||||
tracing::info!(topic_id = %topic_id, message_count = messages.len(), "Sending topic history");
|
||||
|
||||
// 收集已有 tool_result 的 tool_call_id 集合，用于判断任务是否已有结果
|
||||
let mut tool_call_ids_with_results: std::collections::HashSet<String> =
|
||||
std::collections::HashSet::new();
|
||||
for msg in &messages {
|
||||
if msg.role == "tool" {
|
||||
if let Some(ref tcid) = msg.tool_call_id {
|
||||
tool_call_ids_with_results.insert(tcid.clone());
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// 将消息转换为 WsOutbound 并发送
|
||||
for msg in messages {
|
||||
for outbound in chat_message_to_ws_outbound(&msg) {
|
||||
@ -824,9 +697,9 @@ async fn send_topic_history(
|
||||
}
|
||||
}
|
||||
|
||||
// 查询该话题下所有子智能体任务，补发 TaskStarted 事件
|
||||
// 查询该话题下所有运行中的子智能体任务，补发 TaskStarted 事件
|
||||
// 解决页面刷新后 navigateToTaskId 丢失的问题
|
||||
let tasks = match task_repository.list_tasks_for_topic(topic_id).await {
|
||||
let running_tasks = match task_repository.list_tasks_for_topic(topic_id).await {
|
||||
Ok(tasks) => tasks,
|
||||
Err(e) => {
|
||||
tracing::warn!(error = %e, topic_id = %topic_id, "Failed to list tasks for topic");
|
||||
@ -834,40 +707,28 @@ async fn send_topic_history(
|
||||
}
|
||||
};
|
||||
|
||||
for task in tasks {
|
||||
// 判断是否需要补发 TaskStarted：
|
||||
// - 如果该任务的 tool_call_id 已有对应的 tool_result，前端会显示结果，不需要补发
|
||||
// - 否则（Running 状态或已完成但结果未进入历史），补发 TaskStarted 以便前端显示"查看实时进度"
|
||||
let has_tool_result = task
|
||||
.tool_call_id
|
||||
.as_ref()
|
||||
.map(|tcid| tool_call_ids_with_results.contains(tcid))
|
||||
.unwrap_or(false);
|
||||
if has_tool_result {
|
||||
continue;
|
||||
for task in running_tasks {
|
||||
if task.state == TaskSessionState::Running {
|
||||
// 判断是否为孙智能体：parent_session_id 以 "sub:" 开头表示父会话是子智能体
|
||||
let parent_task_id = extract_parent_task_id(&task);
|
||||
|
||||
tracing::info!(
|
||||
task_id = %task.id,
|
||||
description = %task.description,
|
||||
parent_task_id = ?parent_task_id,
|
||||
"Re-sending TaskStarted for running task after topic history load"
|
||||
);
|
||||
let _ = sender
|
||||
.send(WsOutbound::TaskStarted {
|
||||
task_id: task.id.clone(),
|
||||
description: task.description.clone(),
|
||||
subagent_type: task.subagent_type.clone(),
|
||||
topic_id: Some(topic_id.to_string()),
|
||||
parent_task_id,
|
||||
tool_call_id: None,
|
||||
})
|
||||
.await;
|
||||
}
|
||||
|
||||
// 判断是否为孙智能体：parent_session_id 以 "sub:" 开头表示父会话是子智能体
|
||||
let parent_task_id = extract_parent_task_id(&task);
|
||||
|
||||
tracing::info!(
|
||||
task_id = %task.id,
|
||||
description = %task.description,
|
||||
parent_task_id = ?parent_task_id,
|
||||
tool_call_id = ?task.tool_call_id,
|
||||
state = ?task.state,
|
||||
"Re-sending TaskStarted for task without tool_result after topic history load"
|
||||
);
|
||||
let _ = sender
|
||||
.send(WsOutbound::TaskStarted {
|
||||
task_id: task.id.clone(),
|
||||
description: task.description.clone(),
|
||||
subagent_type: task.subagent_type.clone(),
|
||||
topic_id: Some(topic_id.to_string()),
|
||||
parent_task_id,
|
||||
tool_call_id: task.tool_call_id.clone(),
|
||||
})
|
||||
.await;
|
||||
}
|
||||
|
||||
Ok(())
|
||||
@ -885,17 +746,6 @@ async fn send_task_messages(
|
||||
|
||||
tracing::info!(session_id = %session_id, message_count = messages.len(), "Sending task messages");
|
||||
|
||||
// 收集已有 tool_result 的 tool_call_id 集合，用于判断子任务是否已有结果
|
||||
let mut tool_call_ids_with_results: std::collections::HashSet<String> =
|
||||
std::collections::HashSet::new();
|
||||
for msg in &messages {
|
||||
if msg.role == "tool" {
|
||||
if let Some(ref tcid) = msg.tool_call_id {
|
||||
tool_call_ids_with_results.insert(tcid.clone());
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
for msg in messages {
|
||||
let mut outbounds = chat_message_to_ws_outbound(&msg);
|
||||
if let Some(ref task_id) = subagent_task_id {
|
||||
@ -914,33 +764,23 @@ async fn send_task_messages(
|
||||
match repo.list_tasks_for_session(session_id).await {
|
||||
Ok(child_tasks) => {
|
||||
for child in child_tasks {
|
||||
// 如果该子任务的 tool_call_id 已有对应的 tool_result，前端会显示结果，不需要补发
|
||||
let has_tool_result = child
|
||||
.tool_call_id
|
||||
.as_ref()
|
||||
.map(|tcid| tool_call_ids_with_results.contains(tcid))
|
||||
.unwrap_or(false);
|
||||
if has_tool_result {
|
||||
continue;
|
||||
if child.state == TaskSessionState::Running {
|
||||
tracing::info!(
|
||||
child_task_id = %child.id,
|
||||
parent_task_id = %parent_task_id,
|
||||
"Re-sending TaskStarted for child task after sub-agent view re-enter"
|
||||
);
|
||||
let _ = sender
|
||||
.send(WsOutbound::TaskStarted {
|
||||
task_id: child.id.clone(),
|
||||
description: child.description.clone(),
|
||||
subagent_type: child.subagent_type.clone(),
|
||||
topic_id: child.parent_topic_id.clone(),
|
||||
parent_task_id: Some(parent_task_id.clone()),
|
||||
tool_call_id: None,
|
||||
})
|
||||
.await;
|
||||
}
|
||||
|
||||
tracing::info!(
|
||||
child_task_id = %child.id,
|
||||
parent_task_id = %parent_task_id,
|
||||
tool_call_id = ?child.tool_call_id,
|
||||
state = ?child.state,
|
||||
"Re-sending TaskStarted for child task without tool_result after sub-agent view re-enter"
|
||||
);
|
||||
let _ = sender
|
||||
.send(WsOutbound::TaskStarted {
|
||||
task_id: child.id.clone(),
|
||||
description: child.description.clone(),
|
||||
subagent_type: child.subagent_type.clone(),
|
||||
topic_id: child.parent_topic_id.clone(),
|
||||
parent_task_id: Some(parent_task_id.clone()),
|
||||
tool_call_id: child.tool_call_id.clone(),
|
||||
})
|
||||
.await;
|
||||
}
|
||||
}
|
||||
Err(e) => {
|
||||
@ -1025,15 +865,10 @@ fn chat_message_to_ws_outbound(msg: &crate::bus::ChatMessage) -> Vec<WsOutbound>
|
||||
let media_type = mime_type
|
||||
.as_ref()
|
||||
.map(|m| {
|
||||
if m.starts_with("image/") {
|
||||
"image"
|
||||
} else if m.starts_with("audio/") {
|
||||
"audio"
|
||||
} else if m.starts_with("video/") {
|
||||
"video"
|
||||
} else {
|
||||
"file"
|
||||
}
|
||||
if m.starts_with("image/") { "image" }
|
||||
else if m.starts_with("audio/") { "audio" }
|
||||
else if m.starts_with("video/") { "video" }
|
||||
else { "file" }
|
||||
})
|
||||
.unwrap_or("file");
|
||||
|
||||
@ -1057,8 +892,7 @@ fn chat_message_to_ws_outbound(msg: &crate::bus::ChatMessage) -> Vec<WsOutbound>
|
||||
"assistant" => {
|
||||
if let Some(tool_calls) = &msg.tool_calls {
|
||||
let mut outbound = Vec::new();
|
||||
let has_content_or_reasoning =
|
||||
!msg.content.trim().is_empty() || msg.reasoning_content.is_some();
|
||||
let has_content_or_reasoning = !msg.content.trim().is_empty() || msg.reasoning_content.is_some();
|
||||
if has_content_or_reasoning {
|
||||
outbound.push(WsOutbound::AssistantResponse {
|
||||
id: msg.id.clone(),
|
||||
@ -1073,11 +907,7 @@ fn chat_message_to_ws_outbound(msg: &crate::bus::ChatMessage) -> Vec<WsOutbound>
|
||||
});
|
||||
}
|
||||
// AssistantResponse 已携带 reasoning 时，ToolCall 不再重复
|
||||
let tc_reasoning = if has_content_or_reasoning {
|
||||
None
|
||||
} else {
|
||||
msg.reasoning_content.clone()
|
||||
};
|
||||
let tc_reasoning = if has_content_or_reasoning { None } else { msg.reasoning_content.clone() };
|
||||
for tool_call in tool_calls {
|
||||
outbound.push(WsOutbound::ToolCall {
|
||||
id: tool_call.id.clone(),
|
||||
@ -1110,16 +940,10 @@ fn chat_message_to_ws_outbound(msg: &crate::bus::ChatMessage) -> Vec<WsOutbound>
|
||||
}
|
||||
}
|
||||
"tool" => {
|
||||
let tool_state = msg
|
||||
.tool_state
|
||||
.as_ref()
|
||||
.unwrap_or(&ToolMessageState::Completed);
|
||||
let tool_state = msg.tool_state.as_ref().unwrap_or(&ToolMessageState::Completed);
|
||||
match tool_state {
|
||||
ToolMessageState::Completed => vec![WsOutbound::ToolResult {
|
||||
id: msg
|
||||
.tool_call_id
|
||||
.clone()
|
||||
.unwrap_or_else(|| uuid::Uuid::new_v4().to_string()),
|
||||
id: msg.tool_call_id.clone().unwrap_or_else(|| uuid::Uuid::new_v4().to_string()),
|
||||
tool_call_id: msg.tool_call_id.clone().unwrap_or_default(),
|
||||
tool_name: msg.tool_name.clone().unwrap_or_default(),
|
||||
content: msg.content.clone(),
|
||||
@ -1130,10 +954,7 @@ fn chat_message_to_ws_outbound(msg: &crate::bus::ChatMessage) -> Vec<WsOutbound>
|
||||
timestamp: Some(msg.timestamp / 1000),
|
||||
}],
|
||||
ToolMessageState::PendingUserAction => vec![WsOutbound::ToolPending {
|
||||
id: msg
|
||||
.tool_call_id
|
||||
.clone()
|
||||
.unwrap_or_else(|| uuid::Uuid::new_v4().to_string()),
|
||||
id: msg.tool_call_id.clone().unwrap_or_else(|| uuid::Uuid::new_v4().to_string()),
|
||||
tool_call_id: msg.tool_call_id.clone().unwrap_or_default(),
|
||||
tool_name: msg.tool_name.clone().unwrap_or_default(),
|
||||
content: msg.content.clone(),
|
||||
@ -1162,7 +983,7 @@ fn chat_message_to_ws_outbound(msg: &crate::bus::ChatMessage) -> Vec<WsOutbound>
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::{build_media_filename, process_attachments_with_base64, resolve_ws_sender_id};
|
||||
use super::{resolve_ws_sender_id, build_media_filename, process_attachments_with_base64};
|
||||
use crate::protocol::MediaSummary;
|
||||
use base64::{Engine as _, engine::general_purpose::STANDARD};
|
||||
|
||||
|
||||
@ -20,6 +20,5 @@ pub mod scheduler;
|
||||
pub mod skills;
|
||||
pub mod storage;
|
||||
pub mod text;
|
||||
pub mod tools;
|
||||
pub mod topic_description;
|
||||
pub mod utils;
|
||||
pub mod tools;
|
||||
|
||||
@ -3,7 +3,7 @@ use chrono_tz::Tz;
|
||||
use std::path::PathBuf;
|
||||
use tracing_appender::rolling::{RollingFileAppender, Rotation};
|
||||
use tracing_subscriber::{
|
||||
fmt, fmt::time::FormatTime, layer::SubscriberExt, util::SubscriberInitExt, EnvFilter,
|
||||
EnvFilter, fmt, fmt::time::FormatTime, layer::SubscriberExt, util::SubscriberInitExt,
|
||||
};
|
||||
|
||||
#[derive(Clone, Copy, Debug)]
|
||||
@ -28,13 +28,13 @@ impl FormatTime for ConfiguredTimestamp {
|
||||
|
||||
/// Get the default log directory path: ~/.picobot/logs
|
||||
pub fn get_default_log_dir() -> PathBuf {
|
||||
let home = crate::platform::picobot_home_dir();
|
||||
let home = dirs::home_dir().unwrap_or_else(|| PathBuf::from("."));
|
||||
home.join(".picobot").join("logs")
|
||||
}
|
||||
|
||||
/// Get the default config file path: ~/.picobot/config.json
|
||||
pub fn get_default_config_path() -> PathBuf {
|
||||
let home = crate::platform::picobot_home_dir();
|
||||
let home = dirs::home_dir().unwrap_or_else(|| PathBuf::from("."));
|
||||
home.join(".picobot").join("config.json")
|
||||
}
|
||||
|
||||
@ -45,9 +45,7 @@ pub fn init_logging(timezone: Tz) {
|
||||
static INIT: Once = Once::new();
|
||||
|
||||
let mut initialized = false;
|
||||
INIT.call_once(|| {
|
||||
initialized = true;
|
||||
});
|
||||
INIT.call_once(|| { initialized = true; });
|
||||
if !initialized {
|
||||
// Already initialized (e.g. after gateway restart), skip
|
||||
return;
|
||||
|
||||
17
src/main.rs
17
src/main.rs
@ -40,14 +40,11 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
|
||||
if std::env::args().len() <= 1 {
|
||||
cmd.print_help()?;
|
||||
println!();
|
||||
return Ok(());
|
||||
return Ok(())
|
||||
}
|
||||
|
||||
match Command::parse() {
|
||||
Command::Init {
|
||||
force,
|
||||
skip_channels,
|
||||
} => {
|
||||
Command::Init { force, skip_channels } => {
|
||||
let mut wizard = picobot::cli::InitWizard::new();
|
||||
wizard.run(force, skip_channels).await?;
|
||||
}
|
||||
@ -59,12 +56,12 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
|
||||
picobot::client::run(&url).await?;
|
||||
}
|
||||
Command::Gateway { host, port } => {
|
||||
let mut should_restart = true;
|
||||
while should_restart {
|
||||
should_restart = picobot::gateway::run(host.clone(), port).await?;
|
||||
if should_restart {
|
||||
tracing::info!("Gateway restarting...");
|
||||
loop {
|
||||
let should_restart = picobot::gateway::run(host.clone(), port).await?;
|
||||
if !should_restart {
|
||||
break;
|
||||
}
|
||||
tracing::info!("Gateway restarting...");
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@ -7,20 +7,17 @@
|
||||
//! - Dynamically registers MCP tools via the Tool trait adapter
|
||||
|
||||
use std::collections::HashMap;
|
||||
use std::sync::Arc;
|
||||
use parking_lot::Mutex;
|
||||
use std::sync::{Arc, Mutex};
|
||||
use tokio::sync::RwLock;
|
||||
|
||||
use http::{HeaderName, HeaderValue};
|
||||
use rmcp::{
|
||||
RoleClient, ServiceExt,
|
||||
model::{CallToolRequestParams, CallToolResult, ServerInfo, Tool},
|
||||
RoleClient, ServiceExt,
|
||||
service::RunningService,
|
||||
transport::TokioChildProcess,
|
||||
transport::streamable_http_client::{
|
||||
StreamableHttpClientTransport, StreamableHttpClientTransportConfig,
|
||||
},
|
||||
transport::streamable_http_client::{StreamableHttpClientTransport, StreamableHttpClientTransportConfig},
|
||||
};
|
||||
use http::{HeaderName, HeaderValue};
|
||||
use tokio::io::{AsyncBufReadExt, BufReader};
|
||||
use tokio::process::Command;
|
||||
|
||||
@ -32,7 +29,7 @@ use std::process::Stdio;
|
||||
/// Resolve ${ENV_VAR} placeholders in a value string
|
||||
fn resolve_env_placeholders_in_value(value: &str) -> String {
|
||||
let re = regex::Regex::new(r"\$\{([A-Z_][A-Z0-9_]*)\}").expect("invalid regex");
|
||||
re.replace_all(value, |caps: ®ex::Captures<'_>| {
|
||||
re.replace_all(value, |caps: ®ex::Captures| {
|
||||
let var_name = &caps[1];
|
||||
env::var(var_name).unwrap_or_else(|_| caps[0].to_string())
|
||||
})
|
||||
@ -67,11 +64,7 @@ fn resolve_command_path(command: &str) -> Option<PathBuf> {
|
||||
}
|
||||
|
||||
// If it has a path separator (relative path), don't search PATH
|
||||
if path
|
||||
.parent()
|
||||
.map(|p| !p.as_os_str().is_empty())
|
||||
.unwrap_or(false)
|
||||
{
|
||||
if path.parent().map(|p| !p.as_os_str().is_empty()).unwrap_or(false) {
|
||||
return None;
|
||||
}
|
||||
|
||||
@ -215,10 +208,7 @@ impl McpClientManager {
|
||||
"Failed to connect to MCP server after all retries"
|
||||
);
|
||||
// Record error for status reporting
|
||||
self.connection_errors
|
||||
.write()
|
||||
.await
|
||||
.insert(key.clone(), e.to_string());
|
||||
self.connection_errors.write().await.insert(key.clone(), e.to_string());
|
||||
failed += 1;
|
||||
} else {
|
||||
// Clear any previous error on successful connection
|
||||
@ -248,23 +238,18 @@ impl McpClientManager {
|
||||
}
|
||||
|
||||
/// Connect to a single MCP server
|
||||
pub async fn connect_server(
|
||||
&self,
|
||||
key: &str,
|
||||
config: &McpServerConfig,
|
||||
) -> anyhow::Result<McpServerInfo> {
|
||||
pub async fn connect_server(&self, key: &str, config: &McpServerConfig) -> anyhow::Result<McpServerInfo> {
|
||||
let effective_name = config.effective_name(key);
|
||||
tracing::info!(key = %key, name = %effective_name, transport_type = %config.transport_type, "Connecting to MCP server");
|
||||
|
||||
let transport = config.transport().map_err(|e| anyhow::anyhow!("{}", e))?;
|
||||
let client = match transport {
|
||||
McpTransportConfig::Stdio {
|
||||
command,
|
||||
args,
|
||||
env,
|
||||
cwd,
|
||||
} => self.connect_stdio(key, &command, &args, &env, &cwd).await?,
|
||||
McpTransportConfig::Http { url, headers } => self.connect_http(&url, &headers).await?,
|
||||
McpTransportConfig::Stdio { command, args, env, cwd } => {
|
||||
self.connect_stdio(key, &command, &args, &env, &cwd).await?
|
||||
}
|
||||
McpTransportConfig::Http { url, headers } => {
|
||||
self.connect_http(&url, &headers).await?
|
||||
}
|
||||
};
|
||||
|
||||
// Get server info (returns Option<Arc<ServerInfo>> in rmcp 1.8+)
|
||||
@ -314,10 +299,7 @@ impl McpClientManager {
|
||||
if resolved_command.is_none() {
|
||||
let path = std::path::Path::new(command);
|
||||
let is_absolute = path.is_absolute();
|
||||
let has_separator = path
|
||||
.parent()
|
||||
.map(|p| !p.as_os_str().is_empty())
|
||||
.unwrap_or(false);
|
||||
let has_separator = path.parent().map(|p| !p.as_os_str().is_empty()).unwrap_or(false);
|
||||
if !is_absolute && !has_separator && path.extension().is_none() {
|
||||
// Bare name not found on Windows
|
||||
let path_env = std::env::var("PATH").unwrap_or_default();
|
||||
@ -326,8 +308,7 @@ impl McpClientManager {
|
||||
Current PATH: {}. \
|
||||
Suggestion: use the full absolute path to the executable, \
|
||||
or ensure the tool is installed and its directory is in PATH.",
|
||||
command,
|
||||
path_env
|
||||
command, path_env
|
||||
));
|
||||
}
|
||||
}
|
||||
@ -403,8 +384,7 @@ impl McpClientManager {
|
||||
);
|
||||
}
|
||||
// Also collect into the shared buffer (cap at 50 lines)
|
||||
{
|
||||
let mut buf = stderr_lines_for_task.lock();
|
||||
if let Ok(mut buf) = stderr_lines_for_task.lock() {
|
||||
if buf.len() < 50 {
|
||||
buf.push(line);
|
||||
}
|
||||
@ -416,14 +396,17 @@ impl McpClientManager {
|
||||
// Use default client handler (empty tuple)
|
||||
let client = ().serve(transport).await.map_err(|e| {
|
||||
// Include stderr summary in error if available
|
||||
let stderr_summary = {
|
||||
let buf = stderr_lines.lock();
|
||||
if buf.is_empty() {
|
||||
String::new()
|
||||
} else {
|
||||
format!("\nstderr:\n {}", buf.join("\n "))
|
||||
}
|
||||
};
|
||||
let stderr_summary = stderr_lines
|
||||
.lock()
|
||||
.ok()
|
||||
.map(|buf| {
|
||||
if buf.is_empty() {
|
||||
String::new()
|
||||
} else {
|
||||
format!("\nstderr:\n {}", buf.join("\n "))
|
||||
}
|
||||
})
|
||||
.unwrap_or_default();
|
||||
anyhow::anyhow!(
|
||||
"Failed to establish MCP stdio connection '{}': {}{}",
|
||||
effective_command.display(),
|
||||
@ -433,8 +416,7 @@ impl McpClientManager {
|
||||
})?;
|
||||
|
||||
// Track that we have a stdio (child process) connection
|
||||
self.stdio_client_count
|
||||
.fetch_add(1, std::sync::atomic::Ordering::SeqCst);
|
||||
self.stdio_client_count.fetch_add(1, std::sync::atomic::Ordering::SeqCst);
|
||||
|
||||
Ok(client)
|
||||
}
|
||||
@ -464,21 +446,25 @@ impl McpClientManager {
|
||||
.iter()
|
||||
.filter_map(|(key, value)| {
|
||||
// Try to parse header name and value
|
||||
HeaderName::try_from(key.clone()).ok().and_then(|name| {
|
||||
HeaderValue::try_from(value.clone())
|
||||
.ok()
|
||||
.map(|val| (name, val))
|
||||
})
|
||||
HeaderName::try_from(key.clone())
|
||||
.ok()
|
||||
.and_then(|name| {
|
||||
HeaderValue::try_from(value.clone())
|
||||
.ok()
|
||||
.map(|val| (name, val))
|
||||
})
|
||||
})
|
||||
.collect();
|
||||
|
||||
// Create transport config with custom headers
|
||||
let config =
|
||||
StreamableHttpClientTransportConfig::with_uri(url).custom_headers(custom_headers);
|
||||
let config = StreamableHttpClientTransportConfig::with_uri(url)
|
||||
.custom_headers(custom_headers);
|
||||
|
||||
// Create transport using reqwest client (default)
|
||||
let transport =
|
||||
StreamableHttpClientTransport::with_client(reqwest::Client::default(), config);
|
||||
let transport = StreamableHttpClientTransport::with_client(
|
||||
reqwest::Client::default(),
|
||||
config,
|
||||
);
|
||||
|
||||
// Connect
|
||||
let client = ().serve(transport).await?;
|
||||
@ -511,9 +497,7 @@ impl McpClientManager {
|
||||
info_map
|
||||
.values()
|
||||
.flat_map(|info| {
|
||||
info.tools
|
||||
.iter()
|
||||
.map(|tool| (info.key.clone(), tool.clone()))
|
||||
info.tools.iter().map(|tool| (info.key.clone(), tool.clone()))
|
||||
})
|
||||
.collect()
|
||||
}
|
||||
@ -577,9 +561,7 @@ impl McpClientManager {
|
||||
/// gateway restart where old MCP processes may still be running when
|
||||
/// new ones start.
|
||||
pub async fn shutdown_all(&self) -> anyhow::Result<()> {
|
||||
let stdio_count = self
|
||||
.stdio_client_count
|
||||
.load(std::sync::atomic::Ordering::SeqCst);
|
||||
let stdio_count = self.stdio_client_count.load(std::sync::atomic::Ordering::SeqCst);
|
||||
|
||||
// Drop all clients (triggers cancellation + graceful shutdown in rmcp)
|
||||
self.disconnect_all().await?;
|
||||
@ -596,8 +578,7 @@ impl McpClientManager {
|
||||
);
|
||||
tokio::time::sleep(std::time::Duration::from_secs(6)).await;
|
||||
tracing::info!("MCP child process cleanup wait complete");
|
||||
self.stdio_client_count
|
||||
.store(0, std::sync::atomic::Ordering::SeqCst);
|
||||
self.stdio_client_count.store(0, std::sync::atomic::Ordering::SeqCst);
|
||||
}
|
||||
|
||||
Ok(())
|
||||
@ -785,10 +766,7 @@ impl McpInitializer {
|
||||
///
|
||||
/// This should be called after the gateway is ready to accept tools.
|
||||
/// Waits for connections to complete before registering tools.
|
||||
pub async fn register_tools(
|
||||
&mut self,
|
||||
registry: &mut crate::tools::ToolRegistry,
|
||||
) -> anyhow::Result<()> {
|
||||
pub async fn register_tools(&mut self, registry: &mut crate::tools::ToolRegistry) -> anyhow::Result<()> {
|
||||
if let Some(manager) = self.manager.clone() {
|
||||
// Wait for connections to complete first
|
||||
self.wait_for_connections().await?;
|
||||
@ -811,15 +789,9 @@ mod tests {
|
||||
// On all platforms, this should return None (either because it's not found,
|
||||
// or because non-Windows always returns None)
|
||||
#[cfg(windows)]
|
||||
assert!(
|
||||
result.is_none(),
|
||||
"Expected None for nonexistent command on Windows"
|
||||
);
|
||||
assert!(result.is_none(), "Expected None for nonexistent command on Windows");
|
||||
#[cfg(not(windows))]
|
||||
assert!(
|
||||
result.is_none(),
|
||||
"Expected None on non-Windows (always returns None)"
|
||||
);
|
||||
assert!(result.is_none(), "Expected None on non-Windows (always returns None)");
|
||||
}
|
||||
|
||||
#[test]
|
||||
@ -834,10 +806,7 @@ mod tests {
|
||||
#[cfg(windows)]
|
||||
assert!(result.is_some(), "Expected to find {} on Windows", cmd);
|
||||
#[cfg(not(windows))]
|
||||
assert!(
|
||||
result.is_none(),
|
||||
"Expected None on non-Windows for absolute path"
|
||||
);
|
||||
assert!(result.is_none(), "Expected None on non-Windows for absolute path");
|
||||
}
|
||||
|
||||
#[test]
|
||||
@ -851,10 +820,7 @@ mod tests {
|
||||
fn test_resolve_command_path_finds_known_windows_binary() {
|
||||
// cmd.exe should always be in C:\Windows\System32 which is in PATH
|
||||
let result = resolve_command_path("cmd");
|
||||
assert!(
|
||||
result.is_some(),
|
||||
"Expected to find cmd.exe via PATH on Windows"
|
||||
);
|
||||
assert!(result.is_some(), "Expected to find cmd.exe via PATH on Windows");
|
||||
if let Some(p) = result {
|
||||
assert!(
|
||||
p.to_string_lossy().to_lowercase().ends_with("cmd.exe"),
|
||||
@ -863,4 +829,4 @@ mod tests {
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
@ -203,11 +203,7 @@ mod tests {
|
||||
let config = McpServerConfig::stdio(
|
||||
"filesystem",
|
||||
"npx",
|
||||
vec![
|
||||
"-y".to_string(),
|
||||
"@modelcontextprotocol/server-filesystem".to_string(),
|
||||
"/tmp".to_string(),
|
||||
],
|
||||
vec!["-y".to_string(), "@modelcontextprotocol/server-filesystem".to_string(), "/tmp".to_string()],
|
||||
);
|
||||
|
||||
assert_eq!(config.name, Some("filesystem".to_string()));
|
||||
@ -271,14 +267,11 @@ mod tests {
|
||||
match transport {
|
||||
McpTransportConfig::Stdio { command, args, .. } => {
|
||||
assert_eq!(command, "npx");
|
||||
assert_eq!(
|
||||
args,
|
||||
vec![
|
||||
"-y",
|
||||
"@modelcontextprotocol/server-filesystem",
|
||||
"/home/user"
|
||||
]
|
||||
);
|
||||
assert_eq!(args, vec![
|
||||
"-y",
|
||||
"@modelcontextprotocol/server-filesystem",
|
||||
"/home/user"
|
||||
]);
|
||||
}
|
||||
_ => panic!("Expected stdio transport"),
|
||||
}
|
||||
@ -286,18 +279,12 @@ mod tests {
|
||||
// Check WebSearch server (streamableHttp)
|
||||
let websearch = config.mcp_servers.get("WebSearch").unwrap();
|
||||
assert_eq!(websearch.transport_type, "streamableHttp");
|
||||
assert_eq!(
|
||||
websearch.name,
|
||||
Some("AliyunBailianMCP_WebSearch".to_string())
|
||||
);
|
||||
assert_eq!(websearch.name, Some("AliyunBailianMCP_WebSearch".to_string()));
|
||||
assert!(websearch.is_active);
|
||||
let transport = websearch.transport().unwrap();
|
||||
match transport {
|
||||
McpTransportConfig::Http { url, headers } => {
|
||||
assert_eq!(
|
||||
url,
|
||||
"https://dashscope.aliyuncs.com/api/v1/mcps/WebSearch/mcp"
|
||||
);
|
||||
assert_eq!(url, "https://dashscope.aliyuncs.com/api/v1/mcps/WebSearch/mcp");
|
||||
assert_eq!(
|
||||
headers.get("Authorization"),
|
||||
Some(&"Bearer ${DASHSCOPE_API_KEY}".to_string())
|
||||
@ -398,31 +385,17 @@ mod tests {
|
||||
#[test]
|
||||
fn test_http_type_alias() {
|
||||
// Both "http" and "streamableHttp" should work
|
||||
let json_http =
|
||||
r#"{"mcpServers": {"test": {"type": "http", "baseUrl": "http://localhost"}}}"#;
|
||||
let json_http = r#"{"mcpServers": {"test": {"type": "http", "baseUrl": "http://localhost"}}}"#;
|
||||
let json_streamable = r#"{"mcpServers": {"test": {"type": "streamableHttp", "baseUrl": "http://localhost"}}}"#;
|
||||
|
||||
let config_http: McpConfig = serde_json::from_str(json_http).unwrap();
|
||||
let config_streamable: McpConfig = serde_json::from_str(json_streamable).unwrap();
|
||||
|
||||
let transport_http = config_http
|
||||
.mcp_servers
|
||||
.get("test")
|
||||
.unwrap()
|
||||
.transport()
|
||||
.unwrap();
|
||||
let transport_streamable = config_streamable
|
||||
.mcp_servers
|
||||
.get("test")
|
||||
.unwrap()
|
||||
.transport()
|
||||
.unwrap();
|
||||
let transport_http = config_http.mcp_servers.get("test").unwrap().transport().unwrap();
|
||||
let transport_streamable = config_streamable.mcp_servers.get("test").unwrap().transport().unwrap();
|
||||
|
||||
assert!(matches!(transport_http, McpTransportConfig::Http { .. }));
|
||||
assert!(matches!(
|
||||
transport_streamable,
|
||||
McpTransportConfig::Http { .. }
|
||||
));
|
||||
assert!(matches!(transport_streamable, McpTransportConfig::Http { .. }));
|
||||
}
|
||||
|
||||
#[test]
|
||||
@ -447,4 +420,4 @@ mod tests {
|
||||
let server = config.mcp_servers.get("test").unwrap();
|
||||
assert!(server.cwd.is_none());
|
||||
}
|
||||
}
|
||||
}
|
||||
@ -11,12 +11,10 @@
|
||||
//!
|
||||
//! MCP is completely optional and disabled by default.
|
||||
|
||||
pub mod client;
|
||||
pub mod config;
|
||||
pub mod client;
|
||||
pub mod tool_adapter;
|
||||
|
||||
pub use client::{
|
||||
McpClient, McpClientManager, McpInitializer, McpServerInfo, McpServerStatus, McpStatusResponse,
|
||||
};
|
||||
pub use config::{McpConfig, McpServerConfig, McpTransportConfig};
|
||||
pub use tool_adapter::{McpToolWrapper, register_mcp_tools};
|
||||
pub use client::{McpClientManager, McpClient, McpServerInfo, McpInitializer, McpServerStatus, McpStatusResponse};
|
||||
pub use tool_adapter::{McpToolWrapper, register_mcp_tools};
|
||||
@ -8,22 +8,6 @@ use rmcp::model::Tool;
|
||||
use crate::mcp::client::McpClientManager;
|
||||
use crate::tools::traits::{Tool as PicoBotTool, ToolResult};
|
||||
|
||||
/// Sanitize a tool name to comply with OpenAI's function name pattern `^[a-zA-Z0-9_-]+$`.
|
||||
/// Any character outside [a-zA-Z0-9_-] (e.g. '.', ':', '/') is replaced with '_'.
|
||||
/// This is applied to the LLM-facing name only; `McpToolWrapper` retains the original
|
||||
/// `server_key` and `tool_name` for routing tool calls to the correct MCP server.
|
||||
fn sanitize_tool_name(name: &str) -> String {
|
||||
name.chars()
|
||||
.map(|c| {
|
||||
if c.is_ascii_alphanumeric() || c == '_' || c == '-' {
|
||||
c
|
||||
} else {
|
||||
'_'
|
||||
}
|
||||
})
|
||||
.collect()
|
||||
}
|
||||
|
||||
/// Wrapper that adapts an MCP tool to PicoBot's Tool trait
|
||||
#[derive(Clone)]
|
||||
pub struct McpToolWrapper {
|
||||
@ -41,20 +25,13 @@ pub struct McpToolWrapper {
|
||||
|
||||
impl McpToolWrapper {
|
||||
/// Create a new tool wrapper
|
||||
pub fn new(manager: Arc<McpClientManager>, server_key: String, tool_info: Tool) -> Self {
|
||||
pub fn new(
|
||||
manager: Arc<McpClientManager>,
|
||||
server_key: String,
|
||||
tool_info: Tool,
|
||||
) -> Self {
|
||||
let tool_name = tool_info.name.clone().into_owned();
|
||||
let raw_name = format!("mcp_{}_{}", server_key, tool_name);
|
||||
let full_name = sanitize_tool_name(&raw_name);
|
||||
if full_name != raw_name {
|
||||
tracing::warn!(
|
||||
original = %raw_name,
|
||||
sanitized = %full_name,
|
||||
server_key = %server_key,
|
||||
tool_name = %tool_name,
|
||||
"MCP tool name contained characters invalid for OpenAI function name pattern \
|
||||
(^[a-zA-Z0-9_-]+$); sanitized to comply"
|
||||
);
|
||||
}
|
||||
let full_name = format!("mcp_{}_{}", server_key, tool_name);
|
||||
Self {
|
||||
manager,
|
||||
server_key,
|
||||
@ -151,7 +128,11 @@ pub async fn register_mcp_tools(
|
||||
let all_tools = manager.all_tools().await;
|
||||
|
||||
for (server_key, tool_info) in all_tools {
|
||||
let wrapper = McpToolWrapper::new(manager.clone(), server_key.clone(), tool_info);
|
||||
let wrapper = McpToolWrapper::new(
|
||||
manager.clone(),
|
||||
server_key.clone(),
|
||||
tool_info,
|
||||
);
|
||||
|
||||
tracing::info!(
|
||||
name = %wrapper.name(),
|
||||
@ -172,7 +153,10 @@ mod tests {
|
||||
|
||||
#[test]
|
||||
fn test_extract_text_content_from_text() {
|
||||
let result = CallToolResult::success(vec![Content::text("Hello"), Content::text("World")]);
|
||||
let result = CallToolResult::success(vec![
|
||||
Content::text("Hello"),
|
||||
Content::text("World"),
|
||||
]);
|
||||
|
||||
let text = extract_text_content(&result);
|
||||
assert_eq!(text, "Hello\nWorld");
|
||||
@ -191,11 +175,10 @@ mod tests {
|
||||
fn test_mcp_tool_wrapper_name() {
|
||||
let manager = Arc::new(McpClientManager::new());
|
||||
// Create a minimal tool info using rmcp's Tool constructor
|
||||
let schema: serde_json::Map<String, serde_json::Value> =
|
||||
serde_json::json!({"type": "object"})
|
||||
.as_object()
|
||||
.unwrap()
|
||||
.clone();
|
||||
let schema: serde_json::Map<String, serde_json::Value> = serde_json::json!({"type": "object"})
|
||||
.as_object()
|
||||
.unwrap()
|
||||
.clone();
|
||||
let tool_info = Tool::new("echo", "Echo tool", schema);
|
||||
|
||||
let wrapper = McpToolWrapper::new(manager, "filesystem".to_string(), tool_info);
|
||||
@ -203,45 +186,4 @@ mod tests {
|
||||
assert_eq!(wrapper.original_name(), "echo");
|
||||
assert_eq!(wrapper.server_key(), "filesystem");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_mcp_tool_wrapper_name_sanitizes_invalid_chars() {
|
||||
// OpenAI requires function names to match ^[a-zA-Z0-9_-]+$.
|
||||
// server_key and tool_name from MCP servers may contain '.', ':', '/', etc.
|
||||
let manager = Arc::new(McpClientManager::new());
|
||||
let schema: serde_json::Map<String, serde_json::Value> =
|
||||
serde_json::json!({"type": "object"})
|
||||
.as_object()
|
||||
.unwrap()
|
||||
.clone();
|
||||
let tool_info = Tool::new("tools.list:read", "Namespaced tool", schema);
|
||||
|
||||
let wrapper = McpToolWrapper::new(manager, "github.api".to_string(), tool_info);
|
||||
// mcp_github.api_tools.list:read → mcp_github_api_tools_list_read
|
||||
assert_eq!(wrapper.name(), "mcp_github_api_tools_list_read");
|
||||
// Original identifiers preserved for routing
|
||||
assert_eq!(wrapper.original_name(), "tools.list:read");
|
||||
assert_eq!(wrapper.server_key(), "github.api");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_sanitize_tool_name_matches_openai_pattern() {
|
||||
let re = regex::Regex::new(r"^[a-zA-Z0-9_-]+$").unwrap();
|
||||
for input in [
|
||||
"mcp_filesystem_echo",
|
||||
"mcp_github.api_tools.list:read",
|
||||
"mcp_a/b@c d",
|
||||
"mcp_中文_tool",
|
||||
] {
|
||||
let sanitized = sanitize_tool_name(input);
|
||||
assert!(
|
||||
re.is_match(&sanitized),
|
||||
"sanitized name {:?} (from {:?}) does not match OpenAI pattern",
|
||||
sanitized,
|
||||
input
|
||||
);
|
||||
}
|
||||
// Empty stays empty
|
||||
assert_eq!(sanitize_tool_name(""), "");
|
||||
}
|
||||
}
|
||||
}
|
||||
@ -300,8 +300,7 @@ mod tests {
|
||||
fn test_truncate_args_utf8_boundary() {
|
||||
// Test that truncation respects UTF-8 character boundaries
|
||||
// Each Chinese character is 3 bytes in UTF-8
|
||||
let long_args =
|
||||
serde_json::json!({"key": "测试测试测试测试测试测试测试测试测试测试测试测试测试测试"});
|
||||
let long_args = serde_json::json!({"key": "测试测试测试测试测试测试测试测试测试测试测试测试测试测试"});
|
||||
let truncated = truncate_args(&long_args, 50);
|
||||
assert!(truncated.ends_with("...truncated"));
|
||||
// Verify the truncated string is valid UTF-8 (no panic occurred)
|
||||
|
||||
@ -93,47 +93,22 @@ impl ShellInfo {
|
||||
}
|
||||
|
||||
/// Dangerous command patterns for safety guards.
|
||||
///
|
||||
/// Returns patterns filtered by the current platform. Platform-specific
|
||||
/// rules (e.g. `format` on Windows, `rm` on Unix) are only injected on
|
||||
/// their target platform to avoid false positives.
|
||||
pub fn dangerous_command_patterns() -> Vec<String> {
|
||||
dangerous_command_patterns_for_platform(Platform::current())
|
||||
}
|
||||
|
||||
/// Platform-specific dangerous command patterns.
|
||||
///
|
||||
/// Exposed primarily for testing. Callers should prefer
|
||||
/// [`dangerous_command_patterns`] which auto-detects the platform.
|
||||
pub fn dangerous_command_patterns_for_platform(platform: Platform) -> Vec<String> {
|
||||
let mut patterns: Vec<String> = Vec::new();
|
||||
|
||||
// Cross-platform: fork bomb
|
||||
patterns.push(r":\(\)\s*\{.*\};\s*:".to_string());
|
||||
|
||||
match platform {
|
||||
Platform::Unix => {
|
||||
// Unix dangerous commands
|
||||
patterns.push(r"\brm\s+-[rf]{1,2}\b".to_string());
|
||||
patterns.push(r"\bchmod\s+-[Rr]".to_string());
|
||||
patterns.push(r"\bchown\s+-[Rr]".to_string());
|
||||
}
|
||||
Platform::Windows => {
|
||||
// Windows cmd dangerous commands.
|
||||
// `format` requires a drive letter (`[a-z]:`) somewhere after it,
|
||||
// so legitimate uses like `dart format lib/` or
|
||||
// `pytest --format json` (no drive letter) are not matched.
|
||||
patterns.push(r"\bformat\s+.*[a-z]:".to_string());
|
||||
patterns.push(r"\bdel\s+/[fq]\b".to_string());
|
||||
patterns.push(r"\brmdir\s+/s\b".to_string());
|
||||
// PowerShell dangerous commands. Patterns are lowercase because
|
||||
// `guard_command` lowercases the command string before matching.
|
||||
patterns.push(r"\bremove-item\s+.*-recurse".to_string());
|
||||
patterns.push(r"\bremove-item\s+.*-force".to_string());
|
||||
}
|
||||
}
|
||||
|
||||
patterns
|
||||
vec![
|
||||
// Unix dangerous commands
|
||||
r"\brm\s+-[rf]{1,2}\b".to_string(),
|
||||
r"\bchmod\s+-[Rr]".to_string(),
|
||||
r"\bchown\s+-[Rr]".to_string(),
|
||||
// Windows dangerous commands
|
||||
r"\bdel\s+/[fq]\b".to_string(),
|
||||
r"\brmdir\s+/s\b".to_string(),
|
||||
r"\bformat\s+".to_string(),
|
||||
// PowerShell dangerous commands
|
||||
r"\bRemove-Item\s+.*-Recurse".to_string(),
|
||||
r"\bRemove-Item\s+.*-Force".to_string(),
|
||||
// Fork bomb (cross-platform)
|
||||
r":\(\)\s*\{.*\};\s*:".to_string(),
|
||||
]
|
||||
}
|
||||
|
||||
/// Check whether a child process is blocked waiting for stdin input.
|
||||
@ -151,7 +126,11 @@ pub fn is_process_waiting_on_stdin(pid: u32) -> Option<bool> {
|
||||
if wchan.is_empty() {
|
||||
return None;
|
||||
}
|
||||
Some(wchan.contains("tty_read") || wchan.contains("n_tty_read") || wchan == "pipe_wait")
|
||||
Some(
|
||||
wchan.contains("tty_read")
|
||||
|| wchan.contains("n_tty_read")
|
||||
|| wchan == "pipe_wait",
|
||||
)
|
||||
}
|
||||
#[cfg(target_os = "macos")]
|
||||
{
|
||||
@ -345,14 +324,6 @@ pub fn home_dir() -> Option<PathBuf> {
|
||||
.or_else(|| dirs::home_dir())
|
||||
}
|
||||
|
||||
/// 返回 PicoBot 主目录，若无法确定则回退到当前目录 `"."`。
|
||||
///
|
||||
/// 包装 [`home_dir`] 并内置 fallback，消除各模块重复的
|
||||
/// `dirs::home_dir().unwrap_or_else(|| PathBuf::from("."))` 模式。
|
||||
pub fn picobot_home_dir() -> PathBuf {
|
||||
home_dir().unwrap_or_else(|| PathBuf::from("."))
|
||||
}
|
||||
|
||||
/// Atomically rename a file, handling platform differences.
|
||||
///
|
||||
/// On Windows, `fs::rename` fails if the destination exists, so we need to
|
||||
@ -464,64 +435,9 @@ mod tests {
|
||||
fn test_dangerous_patterns() {
|
||||
let patterns = dangerous_command_patterns();
|
||||
assert!(!patterns.is_empty());
|
||||
// Cross-platform fork bomb rule is always present
|
||||
assert!(patterns.iter().any(|p| p.contains(r":\(\)")));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_dangerous_patterns_unix() {
|
||||
let patterns = dangerous_command_patterns_for_platform(Platform::Unix);
|
||||
// Unix-specific rules
|
||||
// Should contain patterns for both platforms
|
||||
assert!(patterns.iter().any(|p| p.contains("rm")));
|
||||
assert!(patterns.iter().any(|p| p.contains("chmod")));
|
||||
assert!(patterns.iter().any(|p| p.contains("chown")));
|
||||
// Windows-specific rules must NOT be present on Unix
|
||||
assert!(!patterns.iter().any(|p| p.contains("format")));
|
||||
assert!(!patterns.iter().any(|p| p.contains("del")));
|
||||
assert!(!patterns.iter().any(|p| p.contains("remove-item")));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_dangerous_patterns_windows() {
|
||||
let patterns = dangerous_command_patterns_for_platform(Platform::Windows);
|
||||
// Windows-specific rules
|
||||
assert!(patterns.iter().any(|p| p.contains("del")));
|
||||
assert!(patterns.iter().any(|p| p.contains("format")));
|
||||
assert!(patterns.iter().any(|p| p.contains("remove-item")));
|
||||
// Unix-specific rules must NOT be present on Windows.
|
||||
// Use r"\brm\s" to match rm-as-command without matching rmdir.
|
||||
assert!(!patterns.iter().any(|p| p.contains(r"\brm\s")));
|
||||
assert!(!patterns.iter().any(|p| p.contains("chmod")));
|
||||
assert!(!patterns.iter().any(|p| p.contains("chown")));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_format_pattern_precision() {
|
||||
let patterns = dangerous_command_patterns_for_platform(Platform::Windows);
|
||||
let format_pat = patterns
|
||||
.iter()
|
||||
.find(|p| p.contains("format"))
|
||||
.expect("format pattern should exist on Windows");
|
||||
let re = regex::Regex::new(format_pat).unwrap();
|
||||
|
||||
// Helper: guard_command lowercases before matching, so tests must too.
|
||||
let m = |cmd: &str| re.is_match(&cmd.to_lowercase());
|
||||
|
||||
// Truly dangerous commands — should match
|
||||
assert!(m("format c:"));
|
||||
assert!(m("format /q c:"));
|
||||
assert!(m("format d: /fs:ntfs"));
|
||||
assert!(m("echo ok | format c:"));
|
||||
assert!(m("echo ok; format c:"));
|
||||
// Sub-shell invocation must also be caught
|
||||
assert!(m(r#"cmd /c "format C:""#));
|
||||
|
||||
// Legitimate commands containing literal "format " — should NOT match
|
||||
assert!(!m("dart format lib/"));
|
||||
assert!(!m("buf format -w"));
|
||||
assert!(!m("pytest --format json"));
|
||||
assert!(!m(r#"echo "please format the disk""#));
|
||||
assert!(!m(r#"git log --pretty=format:"%h""#));
|
||||
}
|
||||
|
||||
#[test]
|
||||
@ -529,4 +445,4 @@ mod tests {
|
||||
assert_eq!(xml_escape("a & b"), "a & b");
|
||||
assert_eq!(xml_escape("<tag>"), "<tag>");
|
||||
}
|
||||
}
|
||||
}
|
||||
@ -42,22 +42,6 @@ pub struct TopicSummary {
|
||||
pub last_active_at: i64,
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub description: Option<String>,
|
||||
/// Token 用量统计（与 command::handlers::TopicSummary 对应）。
|
||||
/// 老消息或未触发 LLM 调用的 topic 为 None。
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub token_stats: Option<TopicTokenStats>,
|
||||
}
|
||||
|
||||
/// Topic 维度的 token 统计（与 command::handlers::TopicTokenStats 对应）。
|
||||
#[derive(Debug, Clone, Default, Serialize, Deserialize)]
|
||||
pub struct TopicTokenStats {
|
||||
pub prompt_tokens: u64,
|
||||
pub completion_tokens: u64,
|
||||
pub total_tokens: u64,
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub last_prompt_tokens: Option<u32>,
|
||||
#[serde(default)]
|
||||
pub context_window_tokens: u32,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
@ -256,7 +240,9 @@ pub enum WsOutbound {
|
||||
channel_name: Option<String>,
|
||||
},
|
||||
#[serde(rename = "channel_list")]
|
||||
ChannelList { channels: Vec<Channel> },
|
||||
ChannelList {
|
||||
channels: Vec<Channel>,
|
||||
},
|
||||
#[serde(rename = "topic_list")]
|
||||
TopicList {
|
||||
topics: Vec<TopicSummary>,
|
||||
@ -276,10 +262,7 @@ pub enum WsOutbound {
|
||||
message_count: i64,
|
||||
},
|
||||
#[serde(rename = "session_saved")]
|
||||
SessionSaved {
|
||||
session_id: String,
|
||||
filepath: String,
|
||||
},
|
||||
SessionSaved { session_id: String, filepath: String },
|
||||
#[serde(rename = "task_messages_loaded")]
|
||||
TaskMessagesLoaded {
|
||||
task_id: String,
|
||||
@ -288,17 +271,19 @@ pub enum WsOutbound {
|
||||
status: String,
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
summary: Option<String>,
|
||||
/// 子代理 session 的 token 用量统计（cost 累计 + context 瞬时）。
|
||||
/// 无 assistant 消息时为 None。前端用于在子代理视图下显示 token 统计。
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
token_stats: Option<TopicTokenStats>,
|
||||
},
|
||||
#[serde(rename = "scheduler_job_list")]
|
||||
SchedulerJobList { jobs: Vec<SchedulerJobSummary> },
|
||||
SchedulerJobList {
|
||||
jobs: Vec<SchedulerJobSummary>,
|
||||
},
|
||||
#[serde(rename = "memory_list")]
|
||||
MemoryList { memories: Vec<MemorySummary> },
|
||||
MemoryList {
|
||||
memories: Vec<MemorySummary>,
|
||||
},
|
||||
#[serde(rename = "skill_list")]
|
||||
SkillList { skills: Vec<SkillSummary> },
|
||||
SkillList {
|
||||
skills: Vec<SkillSummary>,
|
||||
},
|
||||
#[serde(rename = "execution_cancelled")]
|
||||
ExecutionCancelled { message: String },
|
||||
#[serde(rename = "stream_delta")]
|
||||
|
||||
@ -15,8 +15,7 @@ pub(crate) fn ws_outbound_from_chat_message(message: &ChatMessage) -> Vec<WsOutb
|
||||
"assistant" => {
|
||||
if let Some(tool_calls) = &message.tool_calls {
|
||||
let mut outbound = Vec::new();
|
||||
let has_content_or_reasoning =
|
||||
!message.content.trim().is_empty() || message.reasoning_content.is_some();
|
||||
let has_content_or_reasoning = !message.content.trim().is_empty() || message.reasoning_content.is_some();
|
||||
if has_content_or_reasoning {
|
||||
outbound.push(WsOutbound::AssistantResponse {
|
||||
id: message.id.clone(),
|
||||
@ -32,11 +31,7 @@ pub(crate) fn ws_outbound_from_chat_message(message: &ChatMessage) -> Vec<WsOutb
|
||||
}
|
||||
|
||||
// AssistantResponse 已携带 reasoning 时，ToolCall 不再重复
|
||||
let tc_reasoning = if has_content_or_reasoning {
|
||||
None
|
||||
} else {
|
||||
message.reasoning_content.clone()
|
||||
};
|
||||
let tc_reasoning = if has_content_or_reasoning { None } else { message.reasoning_content.clone() };
|
||||
outbound.extend(tool_calls.iter().map(|tool_call| WsOutbound::ToolCall {
|
||||
id: tool_call.id.clone(),
|
||||
tool_call_id: tool_call.id.clone(),
|
||||
@ -71,10 +66,7 @@ pub(crate) fn ws_outbound_from_chat_message(message: &ChatMessage) -> Vec<WsOutb
|
||||
.unwrap_or(&ToolMessageState::Completed)
|
||||
{
|
||||
ToolMessageState::Completed => vec![WsOutbound::ToolResult {
|
||||
id: message
|
||||
.tool_call_id
|
||||
.clone()
|
||||
.unwrap_or_else(|| uuid::Uuid::new_v4().to_string()),
|
||||
id: message.tool_call_id.clone().unwrap_or_else(|| uuid::Uuid::new_v4().to_string()),
|
||||
tool_call_id: message.tool_call_id.clone().unwrap_or_default(),
|
||||
tool_name: message.tool_name.clone().unwrap_or_default(),
|
||||
content: message.content.clone(),
|
||||
@ -85,10 +77,7 @@ pub(crate) fn ws_outbound_from_chat_message(message: &ChatMessage) -> Vec<WsOutb
|
||||
timestamp: None,
|
||||
}],
|
||||
ToolMessageState::PendingUserAction => vec![WsOutbound::ToolPending {
|
||||
id: message
|
||||
.tool_call_id
|
||||
.clone()
|
||||
.unwrap_or_else(|| uuid::Uuid::new_v4().to_string()),
|
||||
id: message.tool_call_id.clone().unwrap_or_else(|| uuid::Uuid::new_v4().to_string()),
|
||||
tool_call_id: message.tool_call_id.clone().unwrap_or_default(),
|
||||
tool_name: message.tool_name.clone().unwrap_or_default(),
|
||||
content: message.content.clone(),
|
||||
@ -118,10 +107,7 @@ pub(crate) fn ws_outbound_from_outbound_message(message: &OutboundMessage) -> Ve
|
||||
})
|
||||
.collect();
|
||||
vec![WsOutbound::AssistantResponse {
|
||||
id: message
|
||||
.message_id
|
||||
.clone()
|
||||
.unwrap_or_else(|| uuid::Uuid::new_v4().to_string()),
|
||||
id: message.message_id.clone().unwrap_or_else(|| uuid::Uuid::new_v4().to_string()),
|
||||
content: message.content.clone(),
|
||||
role: message.role.clone(),
|
||||
attachments,
|
||||
@ -190,16 +176,8 @@ pub(crate) fn ws_outbound_from_outbound_message(message: &OutboundMessage) -> Ve
|
||||
}],
|
||||
OutboundEventKind::TaskStarted => vec![WsOutbound::TaskStarted {
|
||||
task_id: message.metadata.get("task_id").cloned().unwrap_or_default(),
|
||||
description: message
|
||||
.metadata
|
||||
.get("task_description")
|
||||
.cloned()
|
||||
.unwrap_or_default(),
|
||||
subagent_type: message
|
||||
.metadata
|
||||
.get("task_subagent_type")
|
||||
.cloned()
|
||||
.unwrap_or_default(),
|
||||
description: message.metadata.get("task_description").cloned().unwrap_or_default(),
|
||||
subagent_type: message.metadata.get("task_subagent_type").cloned().unwrap_or_default(),
|
||||
topic_id: message.metadata.get("topic_id").cloned(),
|
||||
parent_task_id: message.metadata.get("parent_task_id").cloned(),
|
||||
tool_call_id: message.metadata.get("tool_call_id").cloned(),
|
||||
|
||||
@ -2,16 +2,26 @@ use async_trait::async_trait;
|
||||
use reqwest::Client;
|
||||
use serde::{Deserialize, Serialize};
|
||||
use std::collections::HashMap;
|
||||
use std::sync::OnceLock;
|
||||
use std::time::Duration;
|
||||
|
||||
use super::traits::Usage;
|
||||
use super::{ChatCompletionRequest, ChatCompletionResponse, LLMProvider, Tool, ToolCall};
|
||||
use crate::domain::messages::ContentBlock;
|
||||
use crate::utils::format_error_chain;
|
||||
|
||||
const INTERNAL_MODEL_EXTRA_KEYS: &[&str] = &["supported_content_types"];
|
||||
|
||||
fn format_error_chain(error: &(dyn std::error::Error + 'static)) -> String {
|
||||
let mut details = vec![error.to_string()];
|
||||
let mut current = error.source();
|
||||
|
||||
while let Some(source) = current {
|
||||
details.push(source.to_string());
|
||||
current = source.source();
|
||||
}
|
||||
|
||||
details.join("\ncaused by: ")
|
||||
}
|
||||
|
||||
fn serialize_content_blocks<S>(
|
||||
blocks: &[serde_json::Value],
|
||||
serializer: S,
|
||||
@ -31,9 +41,7 @@ fn convert_content_blocks(
|
||||
) -> Vec<serde_json::Value> {
|
||||
// 检查是否有图片且模型不支持
|
||||
if !supports_images {
|
||||
let has_images = blocks
|
||||
.iter()
|
||||
.any(|b| matches!(b, ContentBlock::ImageUrl { .. }));
|
||||
let has_images = blocks.iter().any(|b| matches!(b, ContentBlock::ImageUrl { .. }));
|
||||
|
||||
if has_images {
|
||||
let image_count = blocks
|
||||
@ -71,8 +79,10 @@ fn convert_content_blocks(
|
||||
|
||||
// 添加通知文本块
|
||||
if !notices.is_empty() {
|
||||
let notice_text =
|
||||
format!("[系统提示] 以下图片未能成功入模：\n{}", notices.join("\n"));
|
||||
let notice_text = format!(
|
||||
"[系统提示] 以下图片未能成功入模：\n{}",
|
||||
notices.join("\n")
|
||||
);
|
||||
converted_blocks.push(serde_json::json!({ "type": "text", "text": notice_text }));
|
||||
}
|
||||
|
||||
@ -94,10 +104,10 @@ fn convert_content_blocks(
|
||||
|
||||
fn convert_image_url_to_anthropic(url: &str) -> serde_json::Value {
|
||||
// data:image/png;base64,... -> Anthropic image block
|
||||
static RE: OnceLock<regex::Regex> = OnceLock::new();
|
||||
let re =
|
||||
RE.get_or_init(|| regex::Regex::new(r"data:(image/\w+);base64,(.+)").expect("valid regex"));
|
||||
if let Some(caps) = re.captures(url) {
|
||||
if let Some(caps) = regex::Regex::new(r"data:(image/\w+);base64,(.+)")
|
||||
.ok()
|
||||
.and_then(|re| re.captures(url))
|
||||
{
|
||||
let media_type = caps.get(1).map(|m| m.as_str()).unwrap_or("image/png");
|
||||
let data = caps.get(2).map(|d| d.as_str()).unwrap_or("");
|
||||
return serde_json::json!({
|
||||
@ -125,7 +135,6 @@ pub struct AnthropicProvider {
|
||||
api_key: String,
|
||||
base_url: String,
|
||||
extra_headers: HashMap<String, String>,
|
||||
#[cfg_attr(not(debug_assertions), allow(dead_code))]
|
||||
llm_timeout_secs: u64,
|
||||
model_id: String,
|
||||
temperature: Option<f32>,
|
||||
@ -173,7 +182,9 @@ impl AnthropicProvider {
|
||||
self.model_extra
|
||||
.get("supported_content_types")
|
||||
.and_then(|value| value.as_array())
|
||||
.map(|types| types.iter().any(|t| t.as_str() == Some(content_type)))
|
||||
.map(|types| {
|
||||
types.iter().any(|t| t.as_str() == Some(content_type))
|
||||
})
|
||||
.unwrap_or(true)
|
||||
}
|
||||
|
||||
@ -391,234 +402,3 @@ impl LLMProvider for AnthropicProvider {
|
||||
&self.model_id
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use crate::domain::messages::ContentBlock;
|
||||
use std::collections::HashMap;
|
||||
|
||||
/// 构造一个最小 Provider 用于测试配置驱动的方法
|
||||
fn make_provider(model_extra: HashMap<String, serde_json::Value>) -> AnthropicProvider {
|
||||
AnthropicProvider::new(
|
||||
"test".to_string(),
|
||||
"key".to_string(),
|
||||
"https://api.test".to_string(),
|
||||
HashMap::new(),
|
||||
30,
|
||||
"claude-test".to_string(),
|
||||
None,
|
||||
None,
|
||||
model_extra,
|
||||
)
|
||||
}
|
||||
|
||||
// ---- convert_image_url_to_anthropic ----
|
||||
|
||||
#[test]
|
||||
fn test_convert_data_url_extracts_media_type_and_base64() {
|
||||
let url = "data:image/png;base64,iVBORw0KGgo=";
|
||||
let v = convert_image_url_to_anthropic(url);
|
||||
assert_eq!(v["type"], "image");
|
||||
assert_eq!(v["source"]["type"], "base64");
|
||||
assert_eq!(v["source"]["media_type"], "image/png");
|
||||
assert_eq!(v["source"]["data"], "iVBORw0KGgo=");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_convert_data_url_jpeg() {
|
||||
let url = "data:image/jpeg;base64,/9j/4AAQ";
|
||||
let v = convert_image_url_to_anthropic(url);
|
||||
assert_eq!(v["source"]["media_type"], "image/jpeg");
|
||||
assert_eq!(v["source"]["data"], "/9j/4AAQ");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_convert_regular_url_uses_url_source() {
|
||||
let url = "https://example.com/img.png";
|
||||
let v = convert_image_url_to_anthropic(url);
|
||||
assert_eq!(v["type"], "image");
|
||||
assert_eq!(v["source"]["type"], "url");
|
||||
assert_eq!(v["source"]["url"], url);
|
||||
}
|
||||
|
||||
// ---- convert_content_blocks: 图片不支持时的过滤 ----
|
||||
|
||||
#[test]
|
||||
fn test_convert_blocks_filters_images_when_unsupported() {
|
||||
let blocks = vec![
|
||||
ContentBlock::text("hello"),
|
||||
ContentBlock::image_url("data:image/png;base64,abc"),
|
||||
ContentBlock::image_url("data:image/png;base64,def"),
|
||||
];
|
||||
let result = convert_content_blocks(false, "test", "claude-test", &blocks, 0);
|
||||
// 文本块保留，图片块被替换为通知
|
||||
assert_eq!(result.len(), 2);
|
||||
assert_eq!(result[0]["type"], "text");
|
||||
assert_eq!(result[0]["text"], "hello");
|
||||
// 第二个是合并的图片通知
|
||||
assert_eq!(result[1]["type"], "text");
|
||||
let notice = result[1]["text"].as_str().unwrap();
|
||||
assert!(notice.contains("第 1 张图片"));
|
||||
assert!(notice.contains("第 2 张图片"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_convert_blocks_keeps_images_when_supported() {
|
||||
let blocks = vec![
|
||||
ContentBlock::text("hi"),
|
||||
ContentBlock::image_url("data:image/png;base64,abc"),
|
||||
];
|
||||
let result = convert_content_blocks(true, "test", "claude-test", &blocks, 0);
|
||||
assert_eq!(result.len(), 2);
|
||||
assert_eq!(result[0]["type"], "text");
|
||||
assert_eq!(result[1]["type"], "image");
|
||||
assert_eq!(result[1]["source"]["data"], "abc");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_convert_blocks_text_only_passthrough() {
|
||||
let blocks = vec![ContentBlock::text("just text")];
|
||||
let result = convert_content_blocks(false, "test", "claude-test", &blocks, 0);
|
||||
assert_eq!(result.len(), 1);
|
||||
assert_eq!(result[0]["type"], "text");
|
||||
}
|
||||
|
||||
// ---- request_model_extra: 内部字段过滤 ----
|
||||
|
||||
#[test]
|
||||
fn test_request_model_extra_filters_internal_keys() {
|
||||
let mut extra = HashMap::new();
|
||||
extra.insert(
|
||||
"supported_content_types".to_string(),
|
||||
serde_json::json!(["text"]),
|
||||
);
|
||||
extra.insert("top_p".to_string(), serde_json::json!(0.9));
|
||||
let provider = make_provider(extra);
|
||||
let filtered = provider.request_model_extra();
|
||||
// 内部字段被过滤
|
||||
assert!(!filtered.contains_key("supported_content_types"));
|
||||
// 业务字段保留
|
||||
assert_eq!(filtered.get("top_p").and_then(|v| v.as_f64()), Some(0.9));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_request_model_extra_empty_when_only_internal() {
|
||||
let mut extra = HashMap::new();
|
||||
extra.insert(
|
||||
"supported_content_types".to_string(),
|
||||
serde_json::json!(["text", "image"]),
|
||||
);
|
||||
let provider = make_provider(extra);
|
||||
assert!(provider.request_model_extra().is_empty());
|
||||
}
|
||||
|
||||
// ---- supports_images: 配置驱动 ----
|
||||
|
||||
#[test]
|
||||
fn test_supports_images_default_true() {
|
||||
let provider = make_provider(HashMap::new());
|
||||
assert!(provider.supports_images());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_supports_images_disabled_via_config() {
|
||||
let mut extra = HashMap::new();
|
||||
extra.insert(
|
||||
"supported_content_types".to_string(),
|
||||
serde_json::json!(["text"]),
|
||||
);
|
||||
let provider = make_provider(extra);
|
||||
assert!(!provider.supports_images());
|
||||
}
|
||||
|
||||
// ---- AnthropicResponse 反序列化 ----
|
||||
|
||||
#[test]
|
||||
fn test_deserialize_response_with_text_and_tool_use() {
|
||||
let json = r#"{
|
||||
"id": "msg_001",
|
||||
"model": "claude-3-sonnet",
|
||||
"content": [
|
||||
{"type": "text", "text": "I'll use a tool"},
|
||||
{"type": "tool_use", "id": "call_1", "name": "bash", "input": {"cmd": "ls"}}
|
||||
],
|
||||
"usage": {"input_tokens": 10, "output_tokens": 20}
|
||||
}"#;
|
||||
let resp: AnthropicResponse = serde_json::from_str(json).unwrap();
|
||||
assert_eq!(resp.id, "msg_001");
|
||||
assert_eq!(resp.content.len(), 2);
|
||||
match &resp.content[0] {
|
||||
AnthropicContent::Text { text } => assert_eq!(text, "I'll use a tool"),
|
||||
_ => panic!("expected Text"),
|
||||
}
|
||||
match &resp.content[1] {
|
||||
AnthropicContent::ToolUse { id, name, input } => {
|
||||
assert_eq!(id, "call_1");
|
||||
assert_eq!(name, "bash");
|
||||
assert_eq!(input["cmd"], "ls");
|
||||
}
|
||||
_ => panic!("expected ToolUse"),
|
||||
}
|
||||
assert_eq!(resp.usage.input_tokens, 10);
|
||||
assert_eq!(resp.usage.output_tokens, 20);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_deserialize_response_thinking_variant() {
|
||||
let json = r#"{
|
||||
"id": "msg_002",
|
||||
"model": "claude-3",
|
||||
"content": [
|
||||
{"type": "thinking", "thinking": "internal reasoning"}
|
||||
],
|
||||
"usage": {"input_tokens": 5, "output_tokens": 5}
|
||||
}"#;
|
||||
let resp: AnthropicResponse = serde_json::from_str(json).unwrap();
|
||||
assert_eq!(resp.content.len(), 1);
|
||||
match &resp.content[0] {
|
||||
AnthropicContent::Thinking { thinking } => {
|
||||
assert_eq!(thinking, "internal reasoning");
|
||||
}
|
||||
_ => panic!("expected Thinking"),
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_deserialize_response_empty_content() {
|
||||
let json = r#"{
|
||||
"id": "msg_003",
|
||||
"model": "claude-3",
|
||||
"content": [],
|
||||
"usage": {"input_tokens": 1, "output_tokens": 1}
|
||||
}"#;
|
||||
let resp: AnthropicResponse = serde_json::from_str(json).unwrap();
|
||||
assert!(resp.content.is_empty());
|
||||
}
|
||||
|
||||
// ---- format_error_chain ----
|
||||
|
||||
#[test]
|
||||
fn test_format_error_chain_single() {
|
||||
let err = std::io::Error::new(std::io::ErrorKind::Other, "single error");
|
||||
let chain = format_error_chain(&err);
|
||||
assert_eq!(chain, "single error");
|
||||
}
|
||||
|
||||
/// 用 thiserror 构造真正的嵌套 source 链，验证 "caused by" 拼接
|
||||
#[derive(Debug, thiserror::Error)]
|
||||
enum OuterError {
|
||||
#[error("outer wrapper")]
|
||||
Wrapped(#[source] std::io::Error),
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_format_error_chain_nested() {
|
||||
let inner = std::io::Error::new(std::io::ErrorKind::Other, "root cause");
|
||||
let outer = OuterError::Wrapped(inner);
|
||||
let chain = format_error_chain(&outer);
|
||||
assert!(chain.contains("outer wrapper"));
|
||||
assert!(chain.contains("caused by"));
|
||||
assert!(chain.contains("root cause"));
|
||||
}
|
||||
}
|
||||
|
||||
@ -10,13 +10,8 @@ use std::time::Duration;
|
||||
use super::traits::{StreamCallback, StreamDelta, Usage};
|
||||
use super::{ChatCompletionRequest, ChatCompletionResponse, LLMProvider, ToolCall};
|
||||
use crate::domain::messages::ContentBlock;
|
||||
use crate::utils::format_error_chain;
|
||||
|
||||
const INTERNAL_MODEL_EXTRA_KEYS: &[&str] = &[
|
||||
"tool_call_arguments_json",
|
||||
"mock_response_content",
|
||||
"supported_content_types",
|
||||
];
|
||||
const INTERNAL_MODEL_EXTRA_KEYS: &[&str] = &["tool_call_arguments_json", "mock_response_content", "supported_content_types"];
|
||||
|
||||
/// 流式响应中的工具调用增量
|
||||
#[derive(Debug, Default)]
|
||||
@ -33,8 +28,6 @@ struct StreamingAccumulator {
|
||||
reasoning_content: Option<String>,
|
||||
tool_calls: BTreeMap<usize, StreamingToolCall>,
|
||||
response_id: String,
|
||||
/// 流式末帧返回的 usage（需要 stream_options.include_usage=true）
|
||||
usage: Option<OpenAIUsage>,
|
||||
}
|
||||
|
||||
impl StreamingAccumulator {
|
||||
@ -56,17 +49,8 @@ impl StreamingAccumulator {
|
||||
}
|
||||
|
||||
/// 添加工具调用增量
|
||||
fn add_tool_call(
|
||||
&mut self,
|
||||
index: usize,
|
||||
id: Option<&str>,
|
||||
name: Option<&str>,
|
||||
arguments: Option<&str>,
|
||||
) {
|
||||
let entry = self
|
||||
.tool_calls
|
||||
.entry(index)
|
||||
.or_insert_with(StreamingToolCall::default);
|
||||
fn add_tool_call(&mut self, index: usize, id: Option<&str>, name: Option<&str>, arguments: Option<&str>) {
|
||||
let entry = self.tool_calls.entry(index).or_insert_with(StreamingToolCall::default);
|
||||
|
||||
// 只在 id 非空时才更新，防止流式响应中后续 chunk 的空 id 覆盖之前的值
|
||||
if let Some(id) = id {
|
||||
@ -92,18 +76,9 @@ impl StreamingAccumulator {
|
||||
}
|
||||
}
|
||||
|
||||
/// 设置 usage（来自流式末帧的 usage 字段）
|
||||
/// 跳过 total_tokens=0 的占位帧，避免覆盖真实值。
|
||||
fn set_usage(&mut self, usage: OpenAIUsage) {
|
||||
if usage.total_tokens > 0 {
|
||||
self.usage = Some(usage);
|
||||
}
|
||||
}
|
||||
|
||||
/// 构建最终的 ChatCompletionResponse
|
||||
fn build_response(self, model: String) -> ChatCompletionResponse {
|
||||
let tool_calls: Vec<ToolCall> = self
|
||||
.tool_calls
|
||||
let tool_calls: Vec<ToolCall> = self.tool_calls
|
||||
.into_iter()
|
||||
.filter(|(_, call)| !call.id.is_empty() && !call.name.is_empty())
|
||||
.map(|(_, call)| {
|
||||
@ -127,23 +102,27 @@ impl StreamingAccumulator {
|
||||
content: self.content,
|
||||
reasoning_content: self.reasoning_content,
|
||||
tool_calls,
|
||||
usage: self
|
||||
.usage
|
||||
.clone()
|
||||
.map(|u| Usage {
|
||||
prompt_tokens: u.prompt_tokens,
|
||||
completion_tokens: u.completion_tokens,
|
||||
total_tokens: u.total_tokens,
|
||||
})
|
||||
.unwrap_or(Usage {
|
||||
prompt_tokens: 0,
|
||||
completion_tokens: 0,
|
||||
total_tokens: 0,
|
||||
}),
|
||||
usage: Usage {
|
||||
prompt_tokens: 0,
|
||||
completion_tokens: 0,
|
||||
total_tokens: 0,
|
||||
},
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn format_error_chain(error: &(dyn std::error::Error + 'static)) -> String {
|
||||
let mut details = vec![error.to_string()];
|
||||
let mut current = error.source();
|
||||
|
||||
while let Some(source) = current {
|
||||
details.push(source.to_string());
|
||||
current = source.source();
|
||||
}
|
||||
|
||||
details.join("\ncaused by: ")
|
||||
}
|
||||
|
||||
fn format_transport_error_context(
|
||||
provider_name: &str,
|
||||
model_id: &str,
|
||||
@ -170,13 +149,10 @@ fn convert_content_blocks(
|
||||
) -> Value {
|
||||
// 检查是否有图片且模型不支持
|
||||
if !supports_images {
|
||||
let has_images = blocks
|
||||
.iter()
|
||||
.any(|b| matches!(b, ContentBlock::ImageUrl { .. }));
|
||||
let has_images = blocks.iter().any(|b| matches!(b, ContentBlock::ImageUrl { .. }));
|
||||
|
||||
if has_images {
|
||||
let image_count = blocks
|
||||
.iter()
|
||||
let image_count = blocks.iter()
|
||||
.filter(|b| matches!(b, ContentBlock::ImageUrl { .. }))
|
||||
.count();
|
||||
|
||||
@ -210,8 +186,10 @@ fn convert_content_blocks(
|
||||
|
||||
// 添加通知文本块
|
||||
if !notices.is_empty() {
|
||||
let notice_text =
|
||||
format!("[系统提示] 以下图片未能成功入模：\n{}", notices.join("\n"));
|
||||
let notice_text = format!(
|
||||
"[系统提示] 以下图片未能成功入模：\n{}",
|
||||
notices.join("\n")
|
||||
);
|
||||
converted_blocks.push(json!({ "type": "text", "text": notice_text }));
|
||||
}
|
||||
|
||||
@ -325,7 +303,9 @@ impl OpenAIProvider {
|
||||
self.model_extra
|
||||
.get("supported_content_types")
|
||||
.and_then(|value| value.as_array())
|
||||
.map(|types| types.iter().any(|t| t.as_str() == Some(content_type)))
|
||||
.map(|types| {
|
||||
types.iter().any(|t| t.as_str() == Some(content_type))
|
||||
})
|
||||
.unwrap_or(true)
|
||||
}
|
||||
|
||||
@ -359,9 +339,7 @@ impl OpenAIProvider {
|
||||
Value::String(raw)
|
||||
} else {
|
||||
// Invalid JSON string - wrap it as a proper JSON string
|
||||
Value::String(
|
||||
serde_json::to_string(&raw).unwrap_or_else(|_| "null".to_string()),
|
||||
)
|
||||
Value::String(serde_json::to_string(&raw).unwrap_or_else(|_| "null".to_string()))
|
||||
}
|
||||
}
|
||||
value => Value::String(
|
||||
@ -392,8 +370,6 @@ impl OpenAIProvider {
|
||||
let mut body = self.build_request_body(request);
|
||||
// 启用流式输出
|
||||
body["stream"] = json!(true);
|
||||
// 请求在流式末帧返回 usage（DeepSeek/OpenAI 兼容协议）
|
||||
body["stream_options"] = json!({ "include_usage": true });
|
||||
|
||||
let mut req_builder = self
|
||||
.client
|
||||
@ -438,7 +414,7 @@ impl OpenAIProvider {
|
||||
// 读取 SSE 流
|
||||
let mut stream = resp.bytes_stream();
|
||||
let mut buffer = String::new();
|
||||
let mut raw_body = String::new(); // 完整原始响应，用于非 SSE JSON 回退
|
||||
let mut raw_body = String::new(); // 完整原始响应，用于非 SSE JSON 回退
|
||||
let mut done_received = false;
|
||||
|
||||
while let Some(chunk_result) = stream.next().await {
|
||||
@ -459,8 +435,7 @@ impl OpenAIProvider {
|
||||
}
|
||||
|
||||
// SSE 格式: data: {...} 或 data:{...}（某些 API 如 139 云没有空格）
|
||||
let data_opt = line_trimmed
|
||||
.strip_prefix("data: ")
|
||||
let data_opt = line_trimmed.strip_prefix("data: ")
|
||||
.or_else(|| line_trimmed.strip_prefix("data:"));
|
||||
|
||||
if let Some(data) = data_opt {
|
||||
@ -478,26 +453,13 @@ impl OpenAIProvider {
|
||||
accumulator.set_response_id(id.to_string());
|
||||
}
|
||||
|
||||
// 提取流式末帧的 usage（stream_options.include_usage=true 时返回）
|
||||
if let Some(usage_val) = json.get("usage") {
|
||||
if !usage_val.is_null() {
|
||||
if let Ok(u) =
|
||||
serde_json::from_value::<OpenAIUsage>(usage_val.clone())
|
||||
{
|
||||
accumulator.set_usage(u);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// 提取 choices
|
||||
if let Some(choices) = json.get("choices").and_then(|c| c.as_array()) {
|
||||
for choice in choices {
|
||||
// 尝试从 delta 提取（标准 OpenAI 流式格式）
|
||||
if let Some(delta) = choice.get("delta") {
|
||||
// 提取内容增量
|
||||
if let Some(content) =
|
||||
delta.get("content").and_then(|c| c.as_str())
|
||||
{
|
||||
if let Some(content) = delta.get("content").and_then(|c| c.as_str()) {
|
||||
accumulator.add_content(content);
|
||||
if let Some(cb) = &stream_callback {
|
||||
cb(StreamDelta {
|
||||
@ -508,9 +470,7 @@ impl OpenAIProvider {
|
||||
}
|
||||
|
||||
// 提取推理内容增量
|
||||
if let Some(reasoning) =
|
||||
delta.get("reasoning_content").and_then(|r| r.as_str())
|
||||
{
|
||||
if let Some(reasoning) = delta.get("reasoning_content").and_then(|r| r.as_str()) {
|
||||
accumulator.add_reasoning_content(reasoning);
|
||||
if let Some(cb) = &stream_callback {
|
||||
cb(StreamDelta {
|
||||
@ -521,43 +481,28 @@ impl OpenAIProvider {
|
||||
}
|
||||
|
||||
// 提取工具调用增量
|
||||
if let Some(tool_calls) =
|
||||
delta.get("tool_calls").and_then(|t| t.as_array())
|
||||
{
|
||||
if let Some(tool_calls) = delta.get("tool_calls").and_then(|t| t.as_array()) {
|
||||
for tool_call in tool_calls {
|
||||
let index = tool_call
|
||||
.get("index")
|
||||
.and_then(|i| i.as_u64())
|
||||
.unwrap_or(0)
|
||||
as usize;
|
||||
let index = tool_call.get("index").and_then(|i| i.as_u64()).unwrap_or(0) as usize;
|
||||
|
||||
let id =
|
||||
tool_call.get("id").and_then(|v| v.as_str());
|
||||
let name = tool_call
|
||||
.get("function")
|
||||
let id = tool_call.get("id").and_then(|v| v.as_str());
|
||||
let name = tool_call.get("function")
|
||||
.and_then(|f| f.get("name"))
|
||||
.and_then(|n| n.as_str());
|
||||
let arguments = tool_call
|
||||
.get("function")
|
||||
let arguments = tool_call.get("function")
|
||||
.and_then(|f| f.get("arguments"))
|
||||
.and_then(|a| a.as_str());
|
||||
|
||||
accumulator
|
||||
.add_tool_call(index, id, name, arguments);
|
||||
accumulator.add_tool_call(index, id, name, arguments);
|
||||
}
|
||||
}
|
||||
}
|
||||
// 尝试从 message 提取（某些非标准 API 格式）
|
||||
else if let Some(message) = choice.get("message") {
|
||||
if let Some(content) =
|
||||
message.get("content").and_then(|c| c.as_str())
|
||||
{
|
||||
if let Some(content) = message.get("content").and_then(|c| c.as_str()) {
|
||||
accumulator.add_content(content);
|
||||
}
|
||||
if let Some(reasoning) = message
|
||||
.get("reasoning_content")
|
||||
.and_then(|r| r.as_str())
|
||||
{
|
||||
if let Some(reasoning) = message.get("reasoning_content").and_then(|r| r.as_str()) {
|
||||
accumulator.add_reasoning_content(reasoning);
|
||||
}
|
||||
}
|
||||
@ -588,8 +533,7 @@ impl OpenAIProvider {
|
||||
}
|
||||
|
||||
// 同样支持 data: {...} 和 data:{...} 两种格式
|
||||
let data_opt = line_trimmed
|
||||
.strip_prefix("data: ")
|
||||
let data_opt = line_trimmed.strip_prefix("data: ")
|
||||
.or_else(|| line_trimmed.strip_prefix("data:"));
|
||||
|
||||
if let Some(data) = data_opt {
|
||||
@ -602,22 +546,11 @@ impl OpenAIProvider {
|
||||
accumulator.set_response_id(id.to_string());
|
||||
}
|
||||
|
||||
// 提取流式末帧的 usage（与主循环一致）
|
||||
if let Some(usage_val) = json.get("usage") {
|
||||
if !usage_val.is_null() {
|
||||
if let Ok(u) = serde_json::from_value::<OpenAIUsage>(usage_val.clone())
|
||||
{
|
||||
accumulator.set_usage(u);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if let Some(choices) = json.get("choices").and_then(|c| c.as_array()) {
|
||||
for choice in choices {
|
||||
// 尝试从 delta 提取
|
||||
if let Some(delta) = choice.get("delta") {
|
||||
if let Some(content) = delta.get("content").and_then(|c| c.as_str())
|
||||
{
|
||||
if let Some(content) = delta.get("content").and_then(|c| c.as_str()) {
|
||||
accumulator.add_content(content);
|
||||
if let Some(cb) = &stream_callback {
|
||||
cb(StreamDelta {
|
||||
@ -626,9 +559,7 @@ impl OpenAIProvider {
|
||||
});
|
||||
}
|
||||
}
|
||||
if let Some(reasoning) =
|
||||
delta.get("reasoning_content").and_then(|r| r.as_str())
|
||||
{
|
||||
if let Some(reasoning) = delta.get("reasoning_content").and_then(|r| r.as_str()) {
|
||||
accumulator.add_reasoning_content(reasoning);
|
||||
if let Some(cb) = &stream_callback {
|
||||
cb(StreamDelta {
|
||||
@ -637,22 +568,14 @@ impl OpenAIProvider {
|
||||
});
|
||||
}
|
||||
}
|
||||
if let Some(tool_calls) =
|
||||
delta.get("tool_calls").and_then(|t| t.as_array())
|
||||
{
|
||||
if let Some(tool_calls) = delta.get("tool_calls").and_then(|t| t.as_array()) {
|
||||
for tool_call in tool_calls {
|
||||
let index = tool_call
|
||||
.get("index")
|
||||
.and_then(|i| i.as_u64())
|
||||
.unwrap_or(0)
|
||||
as usize;
|
||||
let index = tool_call.get("index").and_then(|i| i.as_u64()).unwrap_or(0) as usize;
|
||||
let id = tool_call.get("id").and_then(|v| v.as_str());
|
||||
let name = tool_call
|
||||
.get("function")
|
||||
let name = tool_call.get("function")
|
||||
.and_then(|f| f.get("name"))
|
||||
.and_then(|n| n.as_str());
|
||||
let arguments = tool_call
|
||||
.get("function")
|
||||
let arguments = tool_call.get("function")
|
||||
.and_then(|f| f.get("arguments"))
|
||||
.and_then(|a| a.as_str());
|
||||
accumulator.add_tool_call(index, id, name, arguments);
|
||||
@ -661,14 +584,10 @@ impl OpenAIProvider {
|
||||
}
|
||||
// 尝试从 message 提取（某些非标准 API 格式）
|
||||
else if let Some(message) = choice.get("message") {
|
||||
if let Some(content) =
|
||||
message.get("content").and_then(|c| c.as_str())
|
||||
{
|
||||
if let Some(content) = message.get("content").and_then(|c| c.as_str()) {
|
||||
accumulator.add_content(content);
|
||||
}
|
||||
if let Some(reasoning) =
|
||||
message.get("reasoning_content").and_then(|r| r.as_str())
|
||||
{
|
||||
if let Some(reasoning) = message.get("reasoning_content").and_then(|r| r.as_str()) {
|
||||
accumulator.add_reasoning_content(reasoning);
|
||||
}
|
||||
}
|
||||
@ -684,8 +603,7 @@ impl OpenAIProvider {
|
||||
// 服务器可能返回的是非 SSE 格式的纯 JSON，尝试直接反序列化整个响应体
|
||||
if response.content.is_empty() && response.tool_calls.is_empty() {
|
||||
if let Ok(openai_resp) = serde_json::from_str::<OpenAIResponse>(&raw_body) {
|
||||
let fallback_content = openai_resp
|
||||
.choices
|
||||
let fallback_content = openai_resp.choices
|
||||
.first()
|
||||
.and_then(|c| c.message.content.as_deref())
|
||||
.unwrap_or("")
|
||||
@ -696,37 +614,24 @@ impl OpenAIProvider {
|
||||
"Streaming accumulator empty, falling back to non-SSE JSON parsing"
|
||||
);
|
||||
response.content = fallback_content;
|
||||
response.reasoning_content = openai_resp
|
||||
.choices
|
||||
response.reasoning_content = openai_resp.choices
|
||||
.first()
|
||||
.and_then(|c| c.message.reasoning_content.clone());
|
||||
response.tool_calls = openai_resp
|
||||
.choices
|
||||
response.tool_calls = openai_resp.choices
|
||||
.first()
|
||||
.map(|c| {
|
||||
c.message
|
||||
.tool_calls
|
||||
.iter()
|
||||
.map(|tc| ToolCall {
|
||||
id: tc.id.clone(),
|
||||
name: tc.function.name.clone(),
|
||||
arguments: match &tc.function.arguments {
|
||||
OAIFunctionArguments::Json(args) => args.clone(),
|
||||
OAIFunctionArguments::String(args) => {
|
||||
serde_json::from_str(args)
|
||||
.unwrap_or(serde_json::Value::Null)
|
||||
}
|
||||
},
|
||||
})
|
||||
.collect()
|
||||
c.message.tool_calls.iter().map(|tc| ToolCall {
|
||||
id: tc.id.clone(),
|
||||
name: tc.function.name.clone(),
|
||||
arguments: match &tc.function.arguments {
|
||||
OAIFunctionArguments::Json(args) => args.clone(),
|
||||
OAIFunctionArguments::String(args) => {
|
||||
serde_json::from_str(args).unwrap_or(serde_json::Value::Null)
|
||||
}
|
||||
},
|
||||
}).collect()
|
||||
})
|
||||
.unwrap_or_default();
|
||||
// 回退场景下也从非流式响应提取 usage
|
||||
response.usage = Usage {
|
||||
prompt_tokens: openai_resp.usage.prompt_tokens,
|
||||
completion_tokens: openai_resp.usage.completion_tokens,
|
||||
total_tokens: openai_resp.usage.total_tokens,
|
||||
};
|
||||
}
|
||||
}
|
||||
}
|
||||
@ -751,11 +656,9 @@ impl OpenAIProvider {
|
||||
// result that precedes its parent assistant (e.g. after compaction
|
||||
// boundary splits), leading to API 400 errors:
|
||||
// "insufficient tool messages following tool_calls message".
|
||||
let mut resolved_tool_ids: std::collections::HashSet<&str> =
|
||||
std::collections::HashSet::new();
|
||||
let mut resolved_tool_ids: std::collections::HashSet<&str> = std::collections::HashSet::new();
|
||||
let mut with_parent: std::collections::HashSet<&str> = std::collections::HashSet::new();
|
||||
let mut skip_assistant_indices: std::collections::HashSet<usize> =
|
||||
std::collections::HashSet::new();
|
||||
let mut skip_assistant_indices: std::collections::HashSet<usize> = std::collections::HashSet::new();
|
||||
|
||||
for (i, m) in request.messages.iter().enumerate().rev() {
|
||||
if m.role == "tool" {
|
||||
@ -767,9 +670,8 @@ impl OpenAIProvider {
|
||||
if m.role == "assistant" {
|
||||
if let Some(ref calls) = m.tool_calls {
|
||||
if !calls.is_empty() {
|
||||
let all_resolved = calls
|
||||
.iter()
|
||||
.all(|tc| resolved_tool_ids.contains(tc.id.as_str()));
|
||||
let all_resolved =
|
||||
calls.iter().all(|tc| resolved_tool_ids.contains(tc.id.as_str()));
|
||||
if all_resolved {
|
||||
for tc in calls {
|
||||
with_parent.insert(tc.id.as_str());
|
||||
@ -793,8 +695,7 @@ impl OpenAIProvider {
|
||||
// ^ reverse scan sees tool(A) after assistant → "resolved"
|
||||
// but API requires tool(A) to be IMMEDIATELY after assistant
|
||||
{
|
||||
let mut pending_tool_ids: std::collections::HashSet<&str> =
|
||||
std::collections::HashSet::new();
|
||||
let mut pending_tool_ids: std::collections::HashSet<&str> = std::collections::HashSet::new();
|
||||
let mut pending_assistant_idx: Option<usize> = None;
|
||||
|
||||
for (i, m) in request.messages.iter().enumerate() {
|
||||
@ -991,38 +892,30 @@ impl OpenAIProvider {
|
||||
/// avoid flooding logs on every request — see callers in `chat` and
|
||||
/// `chat_streaming_internal`.
|
||||
fn format_message_sequence(body: &Value) -> Vec<String> {
|
||||
body["messages"]
|
||||
.as_array()
|
||||
.map(|msgs| {
|
||||
msgs.iter()
|
||||
.enumerate()
|
||||
.map(|(i, m)| {
|
||||
let role = m.get("role").and_then(|r| r.as_str()).unwrap_or("?");
|
||||
match role {
|
||||
"assistant" => {
|
||||
let tc_count = m
|
||||
.get("tool_calls")
|
||||
.and_then(|t| t.as_array())
|
||||
.map(|a| a.len())
|
||||
.unwrap_or(0);
|
||||
if tc_count > 0 {
|
||||
format!("[{}] assistant(tool_calls={})", i, tc_count)
|
||||
} else {
|
||||
format!("[{}] assistant", i)
|
||||
}
|
||||
}
|
||||
"tool" => {
|
||||
let tcid = m
|
||||
.get("tool_call_id")
|
||||
.and_then(|t| t.as_str())
|
||||
.unwrap_or("??");
|
||||
format!("[{}] tool(id={})", i, tcid)
|
||||
}
|
||||
_ => format!("[{}] {}", i, role),
|
||||
body["messages"].as_array()
|
||||
.map(|msgs| msgs.iter().enumerate().map(|(i, m)| {
|
||||
let role = m.get("role").and_then(|r| r.as_str()).unwrap_or("?");
|
||||
match role {
|
||||
"assistant" => {
|
||||
let tc_count = m.get("tool_calls")
|
||||
.and_then(|t| t.as_array())
|
||||
.map(|a| a.len())
|
||||
.unwrap_or(0);
|
||||
if tc_count > 0 {
|
||||
format!("[{}] assistant(tool_calls={})", i, tc_count)
|
||||
} else {
|
||||
format!("[{}] assistant", i)
|
||||
}
|
||||
})
|
||||
.collect()
|
||||
})
|
||||
}
|
||||
"tool" => {
|
||||
let tcid = m.get("tool_call_id")
|
||||
.and_then(|t| t.as_str())
|
||||
.unwrap_or("??");
|
||||
format!("[{}] tool(id={})", i, tcid)
|
||||
}
|
||||
_ => format!("[{}] {}", i, role),
|
||||
}
|
||||
}).collect())
|
||||
.unwrap_or_default()
|
||||
}
|
||||
|
||||
@ -1069,7 +962,7 @@ struct OAIFunction {
|
||||
arguments: OAIFunctionArguments,
|
||||
}
|
||||
|
||||
#[derive(Deserialize, Default, Clone, Debug)]
|
||||
#[derive(Deserialize, Default)]
|
||||
struct OpenAIUsage {
|
||||
#[serde(default)]
|
||||
prompt_tokens: u32,
|
||||
@ -1248,10 +1141,7 @@ impl LLMProvider for OpenAIProvider {
|
||||
callback: StreamCallback,
|
||||
) -> Result<ChatCompletionResponse, Box<dyn std::error::Error + Send + Sync>> {
|
||||
if self.is_streaming_enabled() {
|
||||
match self
|
||||
.chat_streaming_internal(&request, Some(&callback))
|
||||
.await
|
||||
{
|
||||
match self.chat_streaming_internal(&request, Some(&callback)).await {
|
||||
Ok(response) => return Ok(response),
|
||||
Err(e) => {
|
||||
tracing::debug!(
|
||||
@ -1581,12 +1471,7 @@ mod tests {
|
||||
let mut accumulator = StreamingAccumulator::new();
|
||||
|
||||
// 第一个 chunk：包含完整的 id 和 name
|
||||
accumulator.add_tool_call(
|
||||
0,
|
||||
Some("call_abc123"),
|
||||
Some("memory_search"),
|
||||
Some("{\"action\":\""),
|
||||
);
|
||||
accumulator.add_tool_call(0, Some("call_abc123"), Some("memory_search"), Some("{\"action\":\""));
|
||||
// 第二个 chunk：只有参数增量
|
||||
accumulator.add_tool_call(0, None, None, Some("list"));
|
||||
// 第三个 chunk：参数继续
|
||||
@ -1602,10 +1487,7 @@ mod tests {
|
||||
assert_eq!(response.tool_calls.len(), 1);
|
||||
assert_eq!(response.tool_calls[0].id, "call_abc123");
|
||||
assert_eq!(response.tool_calls[0].name, "memory_search");
|
||||
assert_eq!(
|
||||
response.tool_calls[0].arguments,
|
||||
json!({"action":"list", "limit": 20})
|
||||
);
|
||||
assert_eq!(response.tool_calls[0].arguments, json!({"action":"list", "limit": 20}));
|
||||
}
|
||||
|
||||
#[test]
|
||||
@ -1613,12 +1495,7 @@ mod tests {
|
||||
let mut accumulator = StreamingAccumulator::new();
|
||||
|
||||
// 第一个工具调用
|
||||
accumulator.add_tool_call(
|
||||
0,
|
||||
Some("call_1"),
|
||||
Some("calculator"),
|
||||
Some("{\"expr\": \"1+1\"}"),
|
||||
);
|
||||
accumulator.add_tool_call(0, Some("call_1"), Some("calculator"), Some("{\"expr\": \"1+1\"}"));
|
||||
// 第二个工具调用（id 和 name 只在第一个 chunk 出现）
|
||||
accumulator.add_tool_call(1, Some("call_2"), Some("get_time"), Some("{}"));
|
||||
|
||||
@ -1723,10 +1600,7 @@ mod tests {
|
||||
"supported_content_types".to_string(),
|
||||
Value::Array(vec![Value::String("text".to_string())]),
|
||||
),
|
||||
(
|
||||
"custom_param".to_string(),
|
||||
Value::String("value".to_string()),
|
||||
),
|
||||
("custom_param".to_string(), Value::String("value".to_string())),
|
||||
]),
|
||||
);
|
||||
|
||||
@ -1863,8 +1737,7 @@ mod tests {
|
||||
let messages = body["messages"].as_array().unwrap();
|
||||
|
||||
// Assistant should keep tool_calls (valid immediate sequence)
|
||||
let tool_calls = messages[0]
|
||||
.get("tool_calls")
|
||||
let tool_calls = messages[0].get("tool_calls")
|
||||
.and_then(|t| t.as_array())
|
||||
.expect("tool_calls should be preserved when immediately followed");
|
||||
assert_eq!(tool_calls.len(), 1);
|
||||
|
||||
@ -1,6 +1,6 @@
|
||||
use crate::config::LLMProviderConfig;
|
||||
use crate::domain::messages::{ContentBlock, ToolCall};
|
||||
use crate::domain::tools::Tool;
|
||||
use crate::config::LLMProviderConfig;
|
||||
use async_trait::async_trait;
|
||||
use serde::{Deserialize, Serialize};
|
||||
use std::collections::HashMap;
|
||||
|
||||
@ -59,9 +59,7 @@ pub trait AgentTaskExecutor: Send + Sync {
|
||||
pub trait MaintenanceExecutor: Send + Sync {
|
||||
async fn cleanup_expired_sessions(&self) -> usize;
|
||||
|
||||
async fn run_memory_maintenance_for_all_scopes(
|
||||
&self,
|
||||
) -> anyhow::Result<Vec<MaintenanceRunSummary>>;
|
||||
async fn run_memory_maintenance_for_all_scopes(&self) -> anyhow::Result<Vec<MaintenanceRunSummary>>;
|
||||
}
|
||||
|
||||
pub struct Scheduler {
|
||||
@ -298,11 +296,7 @@ impl Scheduler {
|
||||
match job.kind {
|
||||
SchedulerJobKind::OutboundMessage => {
|
||||
let message = build_outbound_message(job)?;
|
||||
// publish_outbound 失败（bus 满或关闭）不视为 job 失败：
|
||||
// 通知丢弃是预期的背压行为，标记 job 失败会触发 misfire 重试风暴
|
||||
if let Err(e) = self.bus.publish_outbound(message).await {
|
||||
tracing::warn!(error = %e, job_id = %job.id, "Dropping outbound for scheduler job");
|
||||
}
|
||||
self.bus.publish_outbound(message).await?;
|
||||
}
|
||||
SchedulerJobKind::InternalEvent => {
|
||||
execute_internal_event(self.maintenance_executor.as_ref(), job).await?;
|
||||
@ -315,9 +309,7 @@ impl Scheduler {
|
||||
)
|
||||
.await?;
|
||||
for message in outbound_messages {
|
||||
if let Err(e) = self.bus.publish_outbound(message).await {
|
||||
tracing::warn!(error = %e, job_id = %job.id, "Dropping outbound for scheduler agent task");
|
||||
}
|
||||
self.bus.publish_outbound(message).await?;
|
||||
}
|
||||
}
|
||||
SchedulerJobKind::SilentAgentTask => {
|
||||
@ -413,8 +405,7 @@ impl Scheduler {
|
||||
"silent_agent_task".to_string(),
|
||||
);
|
||||
|
||||
if let Err(e) = self
|
||||
.bus
|
||||
self.bus
|
||||
.publish_outbound(OutboundMessage::error_notification(
|
||||
channel,
|
||||
chat_id,
|
||||
@ -428,10 +419,7 @@ impl Scheduler {
|
||||
metadata,
|
||||
))
|
||||
.await
|
||||
{
|
||||
tracing::warn!(error = %e, job_id = %job.id, "Dropping silent agent task failure notification");
|
||||
}
|
||||
Ok(())
|
||||
.map_err(|error| anyhow::anyhow!(error.to_string()))
|
||||
}
|
||||
}
|
||||
|
||||
@ -464,15 +452,11 @@ fn scheduler_job_definition_matches(
|
||||
existing: &SchedulerJobRecord,
|
||||
) -> bool {
|
||||
let input_schedule = serde_json::from_value::<SchedulerSchedule>(input.schedule.clone()).ok();
|
||||
let existing_schedule = deserialize_schedule(
|
||||
&existing.schedule,
|
||||
existing.interval_secs,
|
||||
existing.startup_delay_secs,
|
||||
)
|
||||
.ok();
|
||||
let existing_schedule =
|
||||
deserialize_schedule(&existing.schedule, existing.interval_secs, existing.startup_delay_secs)
|
||||
.ok();
|
||||
let input_target = serde_json::from_value::<SchedulerJobTarget>(input.target.clone()).ok();
|
||||
let existing_target =
|
||||
serde_json::from_value::<SchedulerJobTarget>(existing.target.clone()).ok();
|
||||
let existing_target = serde_json::from_value::<SchedulerJobTarget>(existing.target.clone()).ok();
|
||||
let targets_match = match (input_target, existing_target) {
|
||||
(Some(input_target), Some(existing_target)) => {
|
||||
input_target.channel == existing_target.channel
|
||||
@ -829,10 +813,7 @@ fn convert_weekday_field(expression: &str) -> String {
|
||||
let weekday_field = parts[5];
|
||||
let converted = convert_cron_weekday(weekday_field);
|
||||
|
||||
format!(
|
||||
"{} {} {} {} {} {}",
|
||||
parts[0], parts[1], parts[2], parts[3], parts[4], converted
|
||||
)
|
||||
format!("{} {} {} {} {} {}", parts[0], parts[1], parts[2], parts[3], parts[4], converted)
|
||||
}
|
||||
|
||||
/// 转换星期表达式中的数字
|
||||
@ -843,10 +824,9 @@ fn convert_cron_weekday(field: &str) -> String {
|
||||
|
||||
// 处理列表（逗号分隔）
|
||||
let items: Vec<&str> = field.split(',').collect();
|
||||
let converted_items: Vec<String> = items
|
||||
.iter()
|
||||
.map(|item| convert_weekday_item(item.trim()))
|
||||
.collect();
|
||||
let converted_items: Vec<String> = items.iter().map(|item| {
|
||||
convert_weekday_item(item.trim())
|
||||
}).collect();
|
||||
|
||||
converted_items.join(",")
|
||||
}
|
||||
@ -885,14 +865,14 @@ fn convert_weekday_range_or_value(item: &str) -> String {
|
||||
/// 转换单个星期数字
|
||||
fn convert_single_weekday(day: &str) -> String {
|
||||
match day {
|
||||
"0" | "7" => "1".to_string(), // 周日 -> 1
|
||||
"1" => "2".to_string(), // 周一 -> 2
|
||||
"2" => "3".to_string(), // 周二 -> 3
|
||||
"3" => "4".to_string(), // 周三 -> 4
|
||||
"4" => "5".to_string(), // 周四 -> 5
|
||||
"5" => "6".to_string(), // 周五 -> 6
|
||||
"6" => "7".to_string(), // 周六 -> 7
|
||||
_ => day.to_string(), // 其他（如字母）保持不变
|
||||
"0" | "7" => "1".to_string(), // 周日 -> 1
|
||||
"1" => "2".to_string(), // 周一 -> 2
|
||||
"2" => "3".to_string(), // 周二 -> 3
|
||||
"3" => "4".to_string(), // 周三 -> 4
|
||||
"4" => "5".to_string(), // 周四 -> 5
|
||||
"5" => "6".to_string(), // 周五 -> 6
|
||||
"6" => "7".to_string(), // 周六 -> 7
|
||||
_ => day.to_string(), // 其他（如字母）保持不变
|
||||
}
|
||||
}
|
||||
|
||||
@ -949,9 +929,7 @@ async fn execute_internal_event(
|
||||
Ok(())
|
||||
}
|
||||
"memory_maintenance" => {
|
||||
let results = maintenance_executor
|
||||
.run_memory_maintenance_for_all_scopes()
|
||||
.await?;
|
||||
let results = maintenance_executor.run_memory_maintenance_for_all_scopes().await?;
|
||||
for result in &results {
|
||||
tracing::info!(
|
||||
job_id = %job.id,
|
||||
@ -1306,10 +1284,10 @@ impl TryFrom<serde_json::Value> for SchedulerJobTarget {
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use chrono::{Datelike, Timelike};
|
||||
use crate::bus::MessageBus;
|
||||
use crate::config::BUILTIN_MEMORY_MAINTENANCE_JOB_ID;
|
||||
use crate::storage::{SchedulerJobUpsert, SessionStore};
|
||||
use chrono::{Datelike, Timelike};
|
||||
|
||||
#[derive(Clone)]
|
||||
struct TestAgentTaskExecutor;
|
||||
@ -1347,9 +1325,7 @@ mod tests {
|
||||
0
|
||||
}
|
||||
|
||||
async fn run_memory_maintenance_for_all_scopes(
|
||||
&self,
|
||||
) -> anyhow::Result<Vec<MaintenanceRunSummary>> {
|
||||
async fn run_memory_maintenance_for_all_scopes(&self) -> anyhow::Result<Vec<MaintenanceRunSummary>> {
|
||||
Ok(Vec::new())
|
||||
}
|
||||
}
|
||||
@ -1616,19 +1592,17 @@ mod tests {
|
||||
|
||||
let probe_runtime = RuntimeJob::from_config(
|
||||
&config_job,
|
||||
Utc.timestamp_millis_opt(1_700_000_000_000)
|
||||
.single()
|
||||
.unwrap(),
|
||||
Utc.timestamp_millis_opt(1_700_000_000_000).single().unwrap(),
|
||||
SchedulerMisfirePolicy::Skip,
|
||||
chrono_tz::Asia::Shanghai,
|
||||
)
|
||||
.unwrap();
|
||||
let probe_existing = store.get_scheduler_job("agent.heartbeat").unwrap().unwrap();
|
||||
let probe_existing = store
|
||||
.get_scheduler_job("agent.heartbeat")
|
||||
.unwrap()
|
||||
.unwrap();
|
||||
let probe_upsert = probe_runtime.to_upsert();
|
||||
assert!(scheduler_job_definition_matches(
|
||||
&probe_upsert,
|
||||
&probe_existing
|
||||
));
|
||||
assert!(scheduler_job_definition_matches(&probe_upsert, &probe_existing));
|
||||
|
||||
let (agent_task_executor, maintenance_service) = test_scheduler_services();
|
||||
let scheduler = Scheduler::new(
|
||||
@ -1648,7 +1622,10 @@ mod tests {
|
||||
|
||||
scheduler.sync_config_jobs().unwrap();
|
||||
|
||||
let saved = store.get_scheduler_job("agent.heartbeat").unwrap().unwrap();
|
||||
let saved = store
|
||||
.get_scheduler_job("agent.heartbeat")
|
||||
.unwrap()
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(saved.next_fire_at, Some(persisted_next_fire_at));
|
||||
assert_eq!(saved.run_count, 3);
|
||||
@ -1746,6 +1723,7 @@ mod tests {
|
||||
);
|
||||
}
|
||||
|
||||
|
||||
#[test]
|
||||
fn debug_cron_weekday_definitions() {
|
||||
// 重大发现：cron crate 的星期定义是反常规的！
|
||||
@ -1764,16 +1742,9 @@ mod tests {
|
||||
];
|
||||
|
||||
// 从周六（2026-04-25）开始测试
|
||||
let saturday = Utc
|
||||
.with_ymd_and_hms(2026, 4, 25, 10, 0, 0)
|
||||
.single()
|
||||
.unwrap();
|
||||
let saturday = Utc.with_ymd_and_hms(2026, 4, 25, 10, 0, 0).single().unwrap();
|
||||
let shanghai_saturday = saturday.with_timezone(&chrono_tz::Asia::Shanghai);
|
||||
println!(
|
||||
"\n=== 从周六 {} ({:?}) 开始测试 ===",
|
||||
shanghai_saturday,
|
||||
shanghai_saturday.weekday()
|
||||
);
|
||||
println!("\n=== 从周六 {} ({:?}) 开始测试 ===", shanghai_saturday, shanghai_saturday.weekday());
|
||||
|
||||
for (expr, desc) in &test_cases {
|
||||
let schedule = parse_scheduler_cron(expr).unwrap();
|
||||
@ -1786,49 +1757,21 @@ mod tests {
|
||||
let schedule_workday = parse_scheduler_cron("0 9 * * 1-5").unwrap();
|
||||
|
||||
let sat_next = schedule_workday.after(&shanghai_saturday).next().unwrap();
|
||||
println!(
|
||||
"周六 -> 1-5 下次执行: {} (星期: {:?})",
|
||||
sat_next,
|
||||
sat_next.weekday()
|
||||
);
|
||||
assert_eq!(
|
||||
sat_next.weekday(),
|
||||
chrono::Weekday::Mon,
|
||||
"1-5 应该从周六跳到周一"
|
||||
);
|
||||
println!("周六 -> 1-5 下次执行: {} (星期: {:?})", sat_next, sat_next.weekday());
|
||||
assert_eq!(sat_next.weekday(), chrono::Weekday::Mon, "1-5 应该从周六跳到周一");
|
||||
|
||||
// 从周日开始
|
||||
let sunday = Utc
|
||||
.with_ymd_and_hms(2026, 4, 26, 10, 0, 0)
|
||||
.single()
|
||||
.unwrap();
|
||||
let sunday = Utc.with_ymd_and_hms(2026, 4, 26, 10, 0, 0).single().unwrap();
|
||||
let shanghai_sunday = sunday.with_timezone(&chrono_tz::Asia::Shanghai);
|
||||
let sun_next = schedule_workday.after(&shanghai_sunday).next().unwrap();
|
||||
println!(
|
||||
"周日 -> 1-5 下次执行: {} (星期: {:?})",
|
||||
sun_next,
|
||||
sun_next.weekday()
|
||||
);
|
||||
assert_eq!(
|
||||
sun_next.weekday(),
|
||||
chrono::Weekday::Mon,
|
||||
"1-5 应该从周日跳到周一"
|
||||
);
|
||||
println!("周日 -> 1-5 下次执行: {} (星期: {:?})", sun_next, sun_next.weekday());
|
||||
assert_eq!(sun_next.weekday(), chrono::Weekday::Mon, "1-5 应该从周日跳到周一");
|
||||
|
||||
// 从周一早上7点开始
|
||||
let shanghai_monday = chrono_tz::Asia::Shanghai
|
||||
.with_ymd_and_hms(2026, 4, 27, 7, 0, 0)
|
||||
.single()
|
||||
.unwrap();
|
||||
println!(
|
||||
"周一早上7点 -> 1-5 下次执行: {} (星期: {:?})",
|
||||
let shanghai_monday = chrono_tz::Asia::Shanghai.with_ymd_and_hms(2026, 4, 27, 7, 0, 0).single().unwrap();
|
||||
println!("周一早上7点 -> 1-5 下次执行: {} (星期: {:?})",
|
||||
schedule_workday.after(&shanghai_monday).next().unwrap(),
|
||||
schedule_workday
|
||||
.after(&shanghai_monday)
|
||||
.next()
|
||||
.unwrap()
|
||||
.weekday()
|
||||
);
|
||||
schedule_workday.after(&shanghai_monday).next().unwrap().weekday());
|
||||
}
|
||||
|
||||
/// 测试标准 cron 星期转换功能
|
||||
@ -1841,103 +1784,62 @@ mod tests {
|
||||
#[test]
|
||||
fn standard_cron_weekday_conversion() {
|
||||
// 测试：标准 cron 的 1-5 应该表示周一到周五
|
||||
let saturday = Utc
|
||||
.with_ymd_and_hms(2026, 4, 25, 10, 0, 0)
|
||||
.single()
|
||||
.unwrap();
|
||||
let saturday = Utc.with_ymd_and_hms(2026, 4, 25, 10, 0, 0).single().unwrap();
|
||||
let shanghai_saturday = saturday.with_timezone(&chrono_tz::Asia::Shanghai);
|
||||
|
||||
// 现在使用标准 cron：1-5 表示周一到周五
|
||||
let schedule_std = parse_scheduler_cron("0 9 * * 1-5").unwrap();
|
||||
|
||||
let sat_next = schedule_std.after(&shanghai_saturday).next().unwrap();
|
||||
println!(
|
||||
"周六 -> 标准 cron 1-5 下次执行: {} (星期: {:?})",
|
||||
sat_next,
|
||||
sat_next.weekday()
|
||||
);
|
||||
assert_eq!(
|
||||
sat_next.weekday(),
|
||||
chrono::Weekday::Mon,
|
||||
"标准 cron 1-5 应该从周六跳到周一"
|
||||
);
|
||||
println!("周六 -> 标准 cron 1-5 下次执行: {} (星期: {:?})", sat_next, sat_next.weekday());
|
||||
assert_eq!(sat_next.weekday(), chrono::Weekday::Mon, "标准 cron 1-5 应该从周六跳到周一");
|
||||
|
||||
// 从周日开始
|
||||
let sunday = Utc
|
||||
.with_ymd_and_hms(2026, 4, 26, 10, 0, 0)
|
||||
.single()
|
||||
.unwrap();
|
||||
let sunday = Utc.with_ymd_and_hms(2026, 4, 26, 10, 0, 0).single().unwrap();
|
||||
let shanghai_sunday = sunday.with_timezone(&chrono_tz::Asia::Shanghai);
|
||||
let sun_next = schedule_std.after(&shanghai_sunday).next().unwrap();
|
||||
println!(
|
||||
"周日 -> 标准 cron 1-5 下次执行: {} (星期: {:?})",
|
||||
sun_next,
|
||||
sun_next.weekday()
|
||||
);
|
||||
assert_eq!(
|
||||
sun_next.weekday(),
|
||||
chrono::Weekday::Mon,
|
||||
"标准 cron 1-5 应该从周日跳到周一"
|
||||
);
|
||||
println!("周日 -> 标准 cron 1-5 下次执行: {} (星期: {:?})", sun_next, sun_next.weekday());
|
||||
assert_eq!(sun_next.weekday(), chrono::Weekday::Mon, "标准 cron 1-5 应该从周日跳到周一");
|
||||
|
||||
// 从周一开始（上海时间周一早上7点）
|
||||
let shanghai_monday = chrono_tz::Asia::Shanghai
|
||||
.with_ymd_and_hms(2026, 4, 27, 7, 0, 0)
|
||||
.single()
|
||||
.unwrap();
|
||||
let shanghai_monday = chrono_tz::Asia::Shanghai.with_ymd_and_hms(2026, 4, 27, 7, 0, 0).single().unwrap();
|
||||
let mon_next = schedule_std.after(&shanghai_monday).next().unwrap();
|
||||
println!(
|
||||
"周一早上 -> 标准 cron 1-5 下次执行: {} (星期: {:?})",
|
||||
mon_next,
|
||||
mon_next.weekday()
|
||||
);
|
||||
assert_eq!(
|
||||
mon_next.weekday(),
|
||||
chrono::Weekday::Mon,
|
||||
"标准 cron 1-5 应该当天执行"
|
||||
);
|
||||
println!("周一早上 -> 标准 cron 1-5 下次执行: {} (星期: {:?})", mon_next, mon_next.weekday());
|
||||
assert_eq!(mon_next.weekday(), chrono::Weekday::Mon, "标准 cron 1-5 应该当天执行");
|
||||
assert_eq!(mon_next.hour(), 9, "应该是上海时间9点");
|
||||
|
||||
// 从周五开始（应该下周周一）
|
||||
let friday = Utc.with_ymd_and_hms(2026, 5, 1, 10, 0, 0).single().unwrap(); // 周五
|
||||
let shanghai_friday = friday.with_timezone(&chrono_tz::Asia::Shanghai);
|
||||
let fri_next = schedule_std.after(&shanghai_friday).next().unwrap();
|
||||
println!(
|
||||
"周五 -> 标准 cron 1-5 下次执行: {} (星期: {:?})",
|
||||
fri_next,
|
||||
fri_next.weekday()
|
||||
);
|
||||
assert_eq!(
|
||||
fri_next.weekday(),
|
||||
chrono::Weekday::Mon,
|
||||
"标准 cron 1-5 应该从周五跳到下周一"
|
||||
);
|
||||
println!("周五 -> 标准 cron 1-5 下次执行: {} (星期: {:?})", fri_next, fri_next.weekday());
|
||||
assert_eq!(fri_next.weekday(), chrono::Weekday::Mon, "标准 cron 1-5 应该从周五跳到下周一");
|
||||
}
|
||||
|
||||
/// 测试转换辅助函数
|
||||
#[test]
|
||||
fn test_weekday_conversion_helper() {
|
||||
// 测试单个值
|
||||
assert_eq!(convert_single_weekday("0"), "1"); // 周日
|
||||
assert_eq!(convert_single_weekday("1"), "2"); // 周一
|
||||
assert_eq!(convert_single_weekday("5"), "6"); // 周五
|
||||
assert_eq!(convert_single_weekday("6"), "7"); // 周六
|
||||
assert_eq!(convert_single_weekday("7"), "1"); // 周日（标准 cron 兼容写法）
|
||||
assert_eq!(convert_single_weekday("0"), "1"); // 周日
|
||||
assert_eq!(convert_single_weekday("1"), "2"); // 周一
|
||||
assert_eq!(convert_single_weekday("5"), "6"); // 周五
|
||||
assert_eq!(convert_single_weekday("6"), "7"); // 周六
|
||||
assert_eq!(convert_single_weekday("7"), "1"); // 周日（标准 cron 兼容写法）
|
||||
|
||||
// 测试范围
|
||||
assert_eq!(convert_weekday_range_or_value("1-5"), "2-6"); // 周一到周五
|
||||
assert_eq!(convert_weekday_range_or_value("0-6"), "1-7"); // 周日到周六
|
||||
assert_eq!(convert_weekday_range_or_value("0-7"), "1-1"); // 周日（循环）
|
||||
assert_eq!(convert_weekday_range_or_value("1-5"), "2-6"); // 周一到周五
|
||||
assert_eq!(convert_weekday_range_or_value("0-6"), "1-7"); // 周日到周六
|
||||
assert_eq!(convert_weekday_range_or_value("0-7"), "1-1"); // 周日（循环）
|
||||
|
||||
// 测试列表
|
||||
assert_eq!(convert_cron_weekday("1,3,5"), "2,4,6"); // 周一、三、五
|
||||
assert_eq!(convert_cron_weekday("0,6"), "1,7"); // 周日和周六
|
||||
assert_eq!(convert_cron_weekday("1,3,5"), "2,4,6"); // 周一、三、五
|
||||
assert_eq!(convert_cron_weekday("0,6"), "1,7"); // 周日和周六
|
||||
|
||||
// 测试步长
|
||||
assert_eq!(convert_weekday_item("*/2"), "*/2"); // 步长保持不变
|
||||
assert_eq!(convert_weekday_item("*/2"), "*/2"); // 步长保持不变
|
||||
|
||||
// 测试混合
|
||||
assert_eq!(convert_cron_weekday("1-5,7"), "2-6,1"); // 周一到周五 + 周日
|
||||
assert_eq!(convert_cron_weekday("1-5,7"), "2-6,1"); // 周一到周五 + 周日
|
||||
|
||||
// 测试特殊字符
|
||||
assert_eq!(convert_cron_weekday("*"), "*");
|
||||
|
||||
@ -1,22 +1,17 @@
|
||||
use crate::platform::{
|
||||
atomic_rename, home_dir as platform_home_dir, path_to_uri, xml_escape as platform_xml_escape,
|
||||
};
|
||||
use crate::platform::{atomic_rename, home_dir as platform_home_dir, path_to_uri, xml_escape as platform_xml_escape};
|
||||
use serde::{Deserialize, Serialize};
|
||||
use serde_json::json;
|
||||
use std::collections::{HashMap, HashSet};
|
||||
use std::fs;
|
||||
use std::path::{Path, PathBuf};
|
||||
use std::sync::Arc;
|
||||
use parking_lot::RwLock;
|
||||
use std::sync::{Arc, RwLock};
|
||||
|
||||
#[cfg(test)]
|
||||
static SKILL_TEST_ENV_LOCK: std::sync::Mutex<()> = std::sync::Mutex::new(());
|
||||
|
||||
#[cfg(test)]
|
||||
pub(crate) fn acquire_skill_test_env_lock() -> std::sync::MutexGuard<'static, ()> {
|
||||
SKILL_TEST_ENV_LOCK
|
||||
.lock()
|
||||
.unwrap_or_else(|err| err.into_inner())
|
||||
SKILL_TEST_ENV_LOCK.lock().unwrap_or_else(|err| err.into_inner())
|
||||
}
|
||||
|
||||
use crate::config::SkillsConfig;
|
||||
@ -137,7 +132,7 @@ impl SkillRuntime {
|
||||
|
||||
pub fn reload(&self) -> Result<SkillCatalog, String> {
|
||||
let catalog = SkillCatalog::discover(&self.config);
|
||||
let mut guard = self.catalog.write();
|
||||
let mut guard = self.catalog.write().expect("skills rwlock poisoned");
|
||||
*guard = catalog.clone();
|
||||
Ok(catalog)
|
||||
}
|
||||
@ -145,16 +140,18 @@ impl SkillRuntime {
|
||||
pub fn is_empty(&self) -> bool {
|
||||
self.catalog
|
||||
.read()
|
||||
.expect("skills rwlock poisoned")
|
||||
.is_empty()
|
||||
}
|
||||
|
||||
pub fn len(&self) -> usize {
|
||||
self.catalog.read().len()
|
||||
self.catalog.read().expect("skills rwlock poisoned").len()
|
||||
}
|
||||
|
||||
pub fn system_index_prompt(&self) -> Option<String> {
|
||||
self.catalog
|
||||
.read()
|
||||
.expect("skills rwlock poisoned")
|
||||
.system_index_prompt()
|
||||
}
|
||||
|
||||
@ -166,36 +163,42 @@ impl SkillRuntime {
|
||||
) -> Option<String> {
|
||||
self.catalog
|
||||
.read()
|
||||
.expect("skills rwlock poisoned")
|
||||
.system_index_prompt_filtered(allowed, denied)
|
||||
}
|
||||
|
||||
pub fn discovery_event_payload(&self) -> serde_json::Value {
|
||||
self.catalog
|
||||
.read()
|
||||
.expect("skills rwlock poisoned")
|
||||
.discovery_event_payload()
|
||||
}
|
||||
|
||||
pub fn offered_event_payload(&self) -> serde_json::Value {
|
||||
self.catalog
|
||||
.read()
|
||||
.expect("skills rwlock poisoned")
|
||||
.offered_event_payload()
|
||||
}
|
||||
|
||||
pub fn activation_payload(&self, name: &str) -> Result<String, String> {
|
||||
self.catalog
|
||||
.read()
|
||||
.expect("skills rwlock poisoned")
|
||||
.activation_payload(name)
|
||||
}
|
||||
|
||||
pub fn activation_event_payload(&self, name: &str) -> Result<serde_json::Value, String> {
|
||||
self.catalog
|
||||
.read()
|
||||
.expect("skills rwlock poisoned")
|
||||
.activation_event_payload(name)
|
||||
}
|
||||
|
||||
pub fn list_skills(&self) -> Vec<Skill> {
|
||||
self.catalog
|
||||
.read()
|
||||
.expect("skills rwlock poisoned")
|
||||
.skills
|
||||
.clone()
|
||||
}
|
||||
@ -206,28 +209,22 @@ impl SkillRuntime {
|
||||
let catalog = SkillCatalog::discover_without_state(&self.config, &cwd);
|
||||
let disable_state = load_skill_disable_state(&cwd);
|
||||
|
||||
catalog
|
||||
.skills
|
||||
.iter()
|
||||
.map(|skill| {
|
||||
let disabled_scopes = disable_state.disabled_scopes_for(&skill.name);
|
||||
SkillWithStatus {
|
||||
name: skill.name.clone(),
|
||||
description: skill.description.clone(),
|
||||
source: skill.source.as_str().to_string(),
|
||||
path: skill.path.display().to_string(),
|
||||
disabled_in_scopes: disabled_scopes
|
||||
.iter()
|
||||
.map(|s| s.as_str().to_string())
|
||||
.collect(),
|
||||
}
|
||||
})
|
||||
.collect()
|
||||
catalog.skills.iter().map(|skill| {
|
||||
let disabled_scopes = disable_state.disabled_scopes_for(&skill.name);
|
||||
SkillWithStatus {
|
||||
name: skill.name.clone(),
|
||||
description: skill.description.clone(),
|
||||
source: skill.source.as_str().to_string(),
|
||||
path: skill.path.display().to_string(),
|
||||
disabled_in_scopes: disabled_scopes.iter().map(|s| s.as_str().to_string()).collect(),
|
||||
}
|
||||
}).collect()
|
||||
}
|
||||
|
||||
pub fn get_skill(&self, name: &str) -> Option<Skill> {
|
||||
self.catalog
|
||||
.read()
|
||||
.expect("skills rwlock poisoned")
|
||||
.find_skill(name)
|
||||
.cloned()
|
||||
}
|
||||
@ -324,8 +321,8 @@ impl SkillRuntime {
|
||||
|
||||
pub fn has_skill_definition(&self, name: &str) -> Result<bool, String> {
|
||||
validate_skill_name(name)?;
|
||||
let cwd =
|
||||
std::env::current_dir().map_err(|err| format!("failed to get current dir: {}", err))?;
|
||||
let cwd = std::env::current_dir()
|
||||
.map_err(|err| format!("failed to get current dir: {}", err))?;
|
||||
Ok(SkillCatalog::discover_without_state(&self.config, &cwd)
|
||||
.find_skill(name)
|
||||
.is_some())
|
||||
@ -361,8 +358,8 @@ impl SkillRuntime {
|
||||
let _ = self.reload()?;
|
||||
}
|
||||
|
||||
let cwd =
|
||||
std::env::current_dir().map_err(|err| format!("failed to get current dir: {}", err))?;
|
||||
let cwd = std::env::current_dir()
|
||||
.map_err(|err| format!("failed to get current dir: {}", err))?;
|
||||
let effective_state = load_skill_disable_state(&cwd);
|
||||
let disabled_in_scopes = effective_state.disabled_scopes_for(name);
|
||||
|
||||
@ -759,9 +756,8 @@ fn skill_file_path(scope: SkillScope, name: &str) -> Result<PathBuf, String> {
|
||||
|
||||
fn skill_state_path(scope: SkillScope) -> Result<PathBuf, String> {
|
||||
match scope {
|
||||
SkillScope::User => {
|
||||
user_skill_state_path().ok_or_else(|| "failed to resolve home directory".to_string())
|
||||
}
|
||||
SkillScope::User => user_skill_state_path()
|
||||
.ok_or_else(|| "failed to resolve home directory".to_string()),
|
||||
SkillScope::Project => {
|
||||
let cwd = std::env::current_dir()
|
||||
.map_err(|err| format!("failed to get current dir: {}", err))?;
|
||||
@ -970,7 +966,10 @@ impl SystemPromptProvider for SkillPromptProvider {
|
||||
// 读取所选专家的技能策略；无专家或无策略时走全局索引（主智能体默认）
|
||||
let content = match context.session_id.as_deref() {
|
||||
Some(sid) => {
|
||||
let policy = self.experts.selected_expert_for(sid).map(|e| e.capability);
|
||||
let policy = self
|
||||
.experts
|
||||
.selected_expert_for(sid)
|
||||
.map(|e| e.capability);
|
||||
match policy {
|
||||
Some(p) if p.has_skill_policy() => self.skills.system_index_prompt_filtered(
|
||||
p.allowed_skills.as_deref(),
|
||||
@ -1059,11 +1058,7 @@ mod tests {
|
||||
let skill_dir = dir.path().join("demo");
|
||||
fs::create_dir_all(&skill_dir).unwrap();
|
||||
let skill_md = skill_dir.join("SKILL.md");
|
||||
fs::write(
|
||||
&skill_md,
|
||||
"---\r\ndescription: demo skill\r\n---\r\nStep A\r\nStep B",
|
||||
)
|
||||
.unwrap();
|
||||
fs::write(&skill_md, "---\r\ndescription: demo skill\r\n---\r\nStep A\r\nStep B").unwrap();
|
||||
|
||||
let skill = parse_skill_file(&skill_md, SkillSource::Project).unwrap();
|
||||
assert_eq!(skill.name, "demo");
|
||||
@ -1133,10 +1128,7 @@ mod tests {
|
||||
|
||||
// 验证 location 包含正确的 file:// URI 格式
|
||||
let expected_uri = path_to_uri(&skill_path);
|
||||
assert!(prompt.contains(&format!(
|
||||
"<location>{}</location>",
|
||||
platform_xml_escape(&expected_uri)
|
||||
)));
|
||||
assert!(prompt.contains(&format!("<location>{}</location>", platform_xml_escape(&expected_uri))));
|
||||
assert!(prompt.contains("</available_skills>"));
|
||||
}
|
||||
|
||||
@ -1396,17 +1388,13 @@ mod tests {
|
||||
max_listed_skills: 32,
|
||||
});
|
||||
|
||||
let disabled = runtime
|
||||
.disable_skill(SkillScope::Project, "demo", true)
|
||||
.unwrap();
|
||||
let disabled = runtime.disable_skill(SkillScope::Project, "demo", true).unwrap();
|
||||
assert!(disabled.changed);
|
||||
assert_eq!(disabled.disabled_in_scopes, vec![SkillScope::Project]);
|
||||
assert!(!disabled.available);
|
||||
assert!(runtime.get_skill("demo").is_none());
|
||||
|
||||
let enabled = runtime
|
||||
.enable_skill(SkillScope::Project, "demo", true)
|
||||
.unwrap();
|
||||
let enabled = runtime.enable_skill(SkillScope::Project, "demo", true).unwrap();
|
||||
assert!(enabled.changed);
|
||||
assert!(enabled.disabled_in_scopes.is_empty());
|
||||
assert!(enabled.available);
|
||||
@ -1439,22 +1427,16 @@ mod tests {
|
||||
max_listed_skills: 32,
|
||||
});
|
||||
|
||||
let user_disabled = runtime
|
||||
.disable_skill(SkillScope::User, "demo", true)
|
||||
.unwrap();
|
||||
let user_disabled = runtime.disable_skill(SkillScope::User, "demo", true).unwrap();
|
||||
assert_eq!(user_disabled.disabled_in_scopes, vec![SkillScope::User]);
|
||||
assert!(runtime.get_skill("demo").is_none());
|
||||
|
||||
let project_enabled = runtime
|
||||
.enable_skill(SkillScope::Project, "demo", true)
|
||||
.unwrap();
|
||||
let project_enabled = runtime.enable_skill(SkillScope::Project, "demo", true).unwrap();
|
||||
assert!(!project_enabled.available);
|
||||
assert_eq!(project_enabled.disabled_in_scopes, vec![SkillScope::User]);
|
||||
assert!(runtime.get_skill("demo").is_none());
|
||||
|
||||
let user_enabled = runtime
|
||||
.enable_skill(SkillScope::User, "demo", true)
|
||||
.unwrap();
|
||||
let user_enabled = runtime.enable_skill(SkillScope::User, "demo", true).unwrap();
|
||||
assert!(user_enabled.available);
|
||||
assert!(user_enabled.disabled_in_scopes.is_empty());
|
||||
assert!(runtime.get_skill("demo").is_some());
|
||||
@ -1524,9 +1506,7 @@ mod tests {
|
||||
});
|
||||
|
||||
assert_eq!(catalog.len(), 1);
|
||||
let payload = catalog
|
||||
.activation_event_payload("demo-user-openclaw")
|
||||
.unwrap();
|
||||
let payload = catalog.activation_event_payload("demo-user-openclaw").unwrap();
|
||||
assert_eq!(payload["source"], "user_openclaw");
|
||||
}
|
||||
|
||||
@ -1585,9 +1565,7 @@ mod tests {
|
||||
);
|
||||
|
||||
// After enabling, list_skills_with_status should report no disabled scopes
|
||||
runtime
|
||||
.enable_skill(SkillScope::Project, "demo", true)
|
||||
.unwrap();
|
||||
runtime.enable_skill(SkillScope::Project, "demo", true).unwrap();
|
||||
let skills_after = runtime.list_skills_with_status();
|
||||
assert_eq!(skills_after.len(), 1);
|
||||
assert!(skills_after[0].disabled_in_scopes.is_empty());
|
||||
|
||||
@ -51,39 +51,6 @@ pub(super) fn ensure_messages_schema(conn: &Connection) -> Result<(), StorageErr
|
||||
)?;
|
||||
}
|
||||
|
||||
// Token usage 字段（仅 assistant 消息有值，来自 LLM 响应）
|
||||
if !has_column(conn, "messages", "prompt_tokens")? {
|
||||
add_column_if_missing(
|
||||
conn,
|
||||
"ALTER TABLE messages ADD COLUMN prompt_tokens INTEGER",
|
||||
)?;
|
||||
}
|
||||
if !has_column(conn, "messages", "completion_tokens")? {
|
||||
add_column_if_missing(
|
||||
conn,
|
||||
"ALTER TABLE messages ADD COLUMN completion_tokens INTEGER",
|
||||
)?;
|
||||
}
|
||||
if !has_column(conn, "messages", "total_tokens")? {
|
||||
add_column_if_missing(conn, "ALTER TABLE messages ADD COLUMN total_tokens INTEGER")?;
|
||||
}
|
||||
if !has_column(conn, "messages", "context_window_tokens")? {
|
||||
add_column_if_missing(
|
||||
conn,
|
||||
"ALTER TABLE messages ADD COLUMN context_window_tokens INTEGER",
|
||||
)?;
|
||||
}
|
||||
|
||||
// is_compacted: 1 表示该消息是被压缩消费掉的原始消息（前端可见、LLM 不可见）。
|
||||
// 压缩摘要消息 is_compacted=0（LLM 可见），通过 system_context='history_compaction*'
|
||||
// 在前端查询中被排除。保留的原消息（system_guards / 最新 user）is_compacted=0，不重复。
|
||||
if !has_column(conn, "messages", "is_compacted")? {
|
||||
add_column_if_missing(
|
||||
conn,
|
||||
"ALTER TABLE messages ADD COLUMN is_compacted INTEGER NOT NULL DEFAULT 0",
|
||||
)?;
|
||||
}
|
||||
|
||||
// 创建 topic_id 索引（如果不存在）
|
||||
conn.execute(
|
||||
"CREATE INDEX IF NOT EXISTS idx_messages_topic_seq ON messages(topic_id, seq) WHERE topic_id IS NOT NULL",
|
||||
@ -145,18 +112,6 @@ pub(super) fn ensure_scheduler_schema(conn: &Connection) -> Result<(), StorageEr
|
||||
}
|
||||
|
||||
pub(super) fn ensure_memory_scope_key_migration(conn: &Connection) -> Result<(), StorageError> {
|
||||
// 用 PRAGMA user_version 追踪迁移是否已完成，避免每次启动都执行全表 DELETE + UPDATE。
|
||||
// user_version 是 SQLite 内置的 32 位整数，持久化在数据库文件头中。
|
||||
// 版本 0：未迁移；版本 1：memory_scope_key 迁移已完成。
|
||||
const MEMORY_SCOPE_KEY_MIGRATION_VERSION: i64 = 1;
|
||||
|
||||
let current_version: i64 = conn.query_row("PRAGMA user_version", [], |row| row.get(0))?;
|
||||
|
||||
if current_version >= MEMORY_SCOPE_KEY_MIGRATION_VERSION {
|
||||
// 已迁移过，跳过
|
||||
return Ok(());
|
||||
}
|
||||
|
||||
// 步骤1：去重。多条记录 scope_key 不同，改为 "default" 后会违反唯一约束。
|
||||
// 对每个 (scope_kind, namespace, memory_key) 组合保留 updated_at 最新的一条。
|
||||
conn.execute(
|
||||
@ -181,16 +136,6 @@ pub(super) fn ensure_memory_scope_key_migration(conn: &Connection) -> Result<(),
|
||||
"UPDATE memories SET scope_key = 'default' WHERE scope_key != 'default'",
|
||||
[],
|
||||
)?;
|
||||
|
||||
// 步骤3：记录迁移版本，后续启动直接跳过
|
||||
conn.execute(
|
||||
&format!(
|
||||
"PRAGMA user_version = {}",
|
||||
MEMORY_SCOPE_KEY_MIGRATION_VERSION
|
||||
),
|
||||
[],
|
||||
)?;
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
|
||||
@ -1,10 +1,6 @@
|
||||
#[cfg(not(test))]
|
||||
use std::path::{Path, PathBuf};
|
||||
|
||||
use std::collections::HashMap;
|
||||
|
||||
use crate::utils::current_timestamp;
|
||||
|
||||
use r2d2::Pool;
|
||||
use r2d2_sqlite::SqliteConnectionManager;
|
||||
use rusqlite::{Connection, OptionalExtension, TransactionBehavior, params};
|
||||
@ -28,10 +24,10 @@ pub use ports::{
|
||||
SkillEventRepository, TodoRepository,
|
||||
};
|
||||
pub use records::{
|
||||
allowed_namespace_names, get_namespace_description, is_valid_namespace,
|
||||
ALLOWED_MEMORY_NAMESPACES, GLOBAL_SCOPE_KEY, MemoryRecord, MemoryUpsert, SchedulerJobRecord,
|
||||
SchedulerJobState, SchedulerJobStatus, SchedulerJobUpsert, SessionRecord, SessionTokenStats,
|
||||
SkillEventRecord, TodoRecord, TopicRecord, allowed_namespace_names, get_namespace_description,
|
||||
is_valid_namespace,
|
||||
SchedulerJobState, SchedulerJobStatus, SchedulerJobUpsert, SessionRecord, SkillEventRecord,
|
||||
TodoRecord, TopicRecord,
|
||||
};
|
||||
|
||||
#[derive(Clone)]
|
||||
@ -106,11 +102,6 @@ impl SessionStore {
|
||||
tool_call_id TEXT,
|
||||
tool_name TEXT,
|
||||
tool_calls_json TEXT,
|
||||
tool_duration_ms INTEGER,
|
||||
prompt_tokens INTEGER,
|
||||
completion_tokens INTEGER,
|
||||
total_tokens INTEGER,
|
||||
context_window_tokens INTEGER,
|
||||
created_at INTEGER NOT NULL,
|
||||
FOREIGN KEY(session_id) REFERENCES sessions(id) ON DELETE CASCADE,
|
||||
FOREIGN KEY(topic_id) REFERENCES topics(id) ON DELETE SET NULL,
|
||||
@ -237,11 +228,14 @@ impl SessionStore {
|
||||
|
||||
drop(conn);
|
||||
|
||||
let manager = SqliteConnectionManager::file(db_uri).with_init(|c| {
|
||||
c.busy_timeout(std::time::Duration::from_secs(30))?;
|
||||
Ok(())
|
||||
});
|
||||
let pool = Pool::builder().max_size(8).build(manager)?;
|
||||
let manager = SqliteConnectionManager::file(db_uri)
|
||||
.with_init(|c| {
|
||||
c.busy_timeout(std::time::Duration::from_secs(30))?;
|
||||
Ok(())
|
||||
});
|
||||
let pool = Pool::builder()
|
||||
.max_size(8)
|
||||
.build(manager)?;
|
||||
|
||||
Ok(Self { pool })
|
||||
}
|
||||
@ -251,7 +245,8 @@ impl SessionStore {
|
||||
// Use a temp file so the database survives across pool connections.
|
||||
// Temp dir is cleaned by the OS eventually; tests that need cleanup
|
||||
// can call std::fs::remove_file on the path.
|
||||
let path = std::env::temp_dir().join(format!("picobot_test_{}.db", uuid::Uuid::new_v4()));
|
||||
let path = std::env::temp_dir()
|
||||
.join(format!("picobot_test_{}.db", uuid::Uuid::new_v4()));
|
||||
let conn = Connection::open(&path)?;
|
||||
let path_str = path.to_string_lossy().to_string();
|
||||
// ignore unused mut warning for manager in tests
|
||||
@ -309,12 +304,7 @@ impl SessionStore {
|
||||
chat_id: &str,
|
||||
) -> Result<SessionRecord, StorageError> {
|
||||
let session_id = persistent_session_id(channel_name, chat_id);
|
||||
self.ensure_session(
|
||||
&session_id,
|
||||
channel_name,
|
||||
chat_id,
|
||||
&format!("{}:{}", channel_name, chat_id),
|
||||
)
|
||||
self.ensure_session(&session_id, channel_name, chat_id, &format!("{}:{}", channel_name, chat_id))
|
||||
}
|
||||
|
||||
/// 确保指定 session_id 的会话存在（如果不存在则创建）
|
||||
@ -522,11 +512,7 @@ impl SessionStore {
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub fn update_topic_description(
|
||||
&self,
|
||||
topic_id: &str,
|
||||
description: &str,
|
||||
) -> Result<(), StorageError> {
|
||||
pub fn update_topic_description(&self, topic_id: &str, description: &str) -> Result<(), StorageError> {
|
||||
let now = current_timestamp();
|
||||
let conn = self.pool.get()?;
|
||||
conn.execute(
|
||||
@ -608,8 +594,8 @@ impl SessionStore {
|
||||
"
|
||||
INSERT INTO messages (
|
||||
id, session_id, topic_id, seq, role, content,
|
||||
system_context, reasoning_content, media_refs_json, tool_call_id, tool_name, tool_calls_json, tool_duration_ms, prompt_tokens, completion_tokens, total_tokens, context_window_tokens, created_at
|
||||
) VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8, ?9, ?10, ?11, ?12, ?13, ?14, ?15, ?16, ?17, ?18)
|
||||
system_context, reasoning_content, media_refs_json, tool_call_id, tool_name, tool_calls_json, tool_duration_ms, created_at
|
||||
) VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8, ?9, ?10, ?11, ?12, ?13, ?14)
|
||||
",
|
||||
params![
|
||||
message.id,
|
||||
@ -625,10 +611,6 @@ impl SessionStore {
|
||||
message.tool_name,
|
||||
tool_calls_json,
|
||||
message.tool_duration_ms.map(|v| v as i64),
|
||||
message.usage.as_ref().map(|u| u.prompt_tokens as i64),
|
||||
message.usage.as_ref().map(|u| u.completion_tokens as i64),
|
||||
message.usage.as_ref().map(|u| u.total_tokens as i64),
|
||||
message.usage.as_ref().and_then(|u| u.context_window_tokens.map(|v| v as i64)),
|
||||
message.timestamp,
|
||||
],
|
||||
)?;
|
||||
@ -690,8 +672,8 @@ impl SessionStore {
|
||||
INSERT INTO messages (
|
||||
id, session_id, topic_id, seq, role, content,
|
||||
system_context, reasoning_content, media_refs_json,
|
||||
tool_call_id, tool_name, tool_calls_json, tool_duration_ms, prompt_tokens, completion_tokens, total_tokens, context_window_tokens, created_at
|
||||
) VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8, ?9, ?10, ?11, ?12, ?13, ?14, ?15, ?16, ?17, ?18)
|
||||
tool_call_id, tool_name, tool_calls_json, tool_duration_ms, created_at
|
||||
) VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8, ?9, ?10, ?11, ?12, ?13, ?14)
|
||||
",
|
||||
params![
|
||||
message.id,
|
||||
@ -707,10 +689,6 @@ impl SessionStore {
|
||||
message.tool_name,
|
||||
tool_calls_json,
|
||||
message.tool_duration_ms.map(|v| v as i64),
|
||||
message.usage.as_ref().map(|u| u.prompt_tokens as i64),
|
||||
message.usage.as_ref().map(|u| u.completion_tokens as i64),
|
||||
message.usage.as_ref().map(|u| u.total_tokens as i64),
|
||||
message.usage.as_ref().and_then(|u| u.context_window_tokens.map(|v| v as i64)),
|
||||
message.timestamp,
|
||||
],
|
||||
)?;
|
||||
@ -832,7 +810,12 @@ impl SessionStore {
|
||||
archived_at = NULL
|
||||
WHERE id = ?1 AND deleted_at IS NULL
|
||||
",
|
||||
params![session_id, inserted_count, active_user_turn_count, now,],
|
||||
params![
|
||||
session_id,
|
||||
inserted_count,
|
||||
active_user_turn_count,
|
||||
now,
|
||||
],
|
||||
)?;
|
||||
|
||||
tx.commit()?;
|
||||
@ -960,105 +943,6 @@ impl SessionStore {
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// 压缩该 topic 的历史：保留原始消息（标记 is_compacted=1，前端可见、LLM 不可见），
|
||||
/// 并插入压缩摘要消息（is_compacted=0，LLM 可见，前端通过 system_context 过滤排除）。
|
||||
///
|
||||
/// 与 `replace_topic_history` 的区别：不删除原消息，仅打标记，从而让前端仍能展示
|
||||
/// 完整原始对话，同时 LLM 只看到压缩后的精简历史。
|
||||
///
|
||||
/// `new_messages` 是 `compress_two_segment` 的输出，包含：
|
||||
/// - 保留原样的消息（system_guards / 最新 user，保留原 ID）
|
||||
/// - 压缩摘要消息（system_context = history_compaction_*，新 ID）
|
||||
pub fn compact_topic_history(
|
||||
&self,
|
||||
session_id: &str,
|
||||
topic_id: &str,
|
||||
new_messages: &[ChatMessage],
|
||||
) -> Result<(), StorageError> {
|
||||
let mut conn = self.pool.get()?;
|
||||
let tx = conn.transaction_with_behavior(TransactionBehavior::Immediate)?;
|
||||
let now = current_timestamp();
|
||||
|
||||
// 分离摘要消息与保留消息（保留消息携带原 ID，摘要消息是新构造的）
|
||||
let (summaries, preserved): (Vec<&ChatMessage>, Vec<&ChatMessage>) =
|
||||
new_messages.iter().partition(|m| {
|
||||
m.system_context
|
||||
.as_deref()
|
||||
.map_or(false, |sc| sc.starts_with("history_compaction"))
|
||||
});
|
||||
|
||||
// 先删除该 topic 下已有的旧压缩摘要（system_context LIKE 'history_compaction%'）。
|
||||
// 旧摘要已被新摘要替代，保留它们只会累积垃圾行（前端和 LLM 都看不到，但占存储）。
|
||||
tx.execute(
|
||||
"DELETE FROM messages \
|
||||
WHERE topic_id = ?1 AND session_id = ?2 \
|
||||
AND system_context LIKE 'history_compaction%'",
|
||||
params![topic_id, session_id],
|
||||
)?;
|
||||
|
||||
// 将该 topic 中未被保留的原消息标记为 is_compacted=1（仅更新尚未标记的行，避免重复写）。
|
||||
// 保留消息（system_guards / 最新 user）保持 is_compacted=0，不重复插入。
|
||||
let preserved_ids: Vec<String> = preserved.iter().map(|m| m.id.clone()).collect();
|
||||
if preserved_ids.is_empty() {
|
||||
tx.execute(
|
||||
"UPDATE messages SET is_compacted = 1 \
|
||||
WHERE topic_id = ?1 AND session_id = ?2 AND is_compacted = 0",
|
||||
params![topic_id, session_id],
|
||||
)?;
|
||||
} else {
|
||||
let placeholders = (0..preserved_ids.len())
|
||||
.map(|_| "?")
|
||||
.collect::<Vec<_>>()
|
||||
.join(",");
|
||||
let sql = format!(
|
||||
"UPDATE messages SET is_compacted = 1 \
|
||||
WHERE topic_id = ? AND session_id = ? AND is_compacted = 0 \
|
||||
AND id NOT IN ({})",
|
||||
placeholders
|
||||
);
|
||||
let mut params_vec: Vec<String> = vec![topic_id.to_string(), session_id.to_string()];
|
||||
params_vec.extend(preserved_ids.iter().cloned());
|
||||
tx.execute(&sql, rusqlite::params_from_iter(params_vec))?;
|
||||
}
|
||||
|
||||
// 插入压缩摘要消息（is_compacted=0，由列默认值保证）
|
||||
let start_seq: i64 = tx.query_row(
|
||||
"SELECT COALESCE(MAX(seq), 0) + 1 FROM messages WHERE session_id = ?1",
|
||||
params![session_id],
|
||||
|row| row.get(0),
|
||||
)?;
|
||||
for (i, message) in summaries.iter().enumerate() {
|
||||
let seq = start_seq + i as i64;
|
||||
insert_message_with_topic_seq(&tx, session_id, topic_id, seq, message)?;
|
||||
}
|
||||
|
||||
// 更新 topic / session 计数（基于该 topic 全部消息，含被压缩的原始消息）
|
||||
let topic_count: i64 = tx.query_row(
|
||||
"SELECT COUNT(*) FROM messages WHERE session_id = ?1 AND topic_id = ?2",
|
||||
params![session_id, topic_id],
|
||||
|row| row.get(0),
|
||||
)?;
|
||||
tx.execute(
|
||||
"UPDATE topics SET message_count = ?2, last_active_at = ?3, updated_at = ?3 WHERE id = ?1",
|
||||
params![topic_id, topic_count, now],
|
||||
)?;
|
||||
let (total_count, user_turn_count): (i64, i64) = tx.query_row(
|
||||
"SELECT COUNT(*), COALESCE(SUM(CASE WHEN role = 'user' THEN 1 ELSE 0 END), 0) \
|
||||
FROM messages WHERE session_id = ?1",
|
||||
params![session_id],
|
||||
|row| Ok((row.get(0)?, row.get(1)?)),
|
||||
)?;
|
||||
tx.execute(
|
||||
"UPDATE sessions SET message_count = ?2, user_turn_count = ?3, \
|
||||
updated_at = ?4, last_active_at = ?4, archived_at = NULL \
|
||||
WHERE id = ?1 AND deleted_at IS NULL",
|
||||
params![session_id, total_count, user_turn_count, now],
|
||||
)?;
|
||||
|
||||
tx.commit()?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub fn mark_agent_prompt_reinjected(&self, session_id: &str) -> Result<(), StorageError> {
|
||||
let now = current_timestamp();
|
||||
let conn = self.pool.get()?;
|
||||
@ -1657,8 +1541,6 @@ impl SessionStore {
|
||||
load_messages_after(&conn, session_id, 0)
|
||||
}
|
||||
|
||||
/// LLM 视角：只返回 is_compacted = 0 的消息（压缩摘要 + 未被压缩的新消息）。
|
||||
/// 被压缩消费掉的原始消息（is_compacted = 1）对 LLM 不可见，以节省 context。
|
||||
pub fn load_messages_for_topic(
|
||||
&self,
|
||||
topic_id: &str,
|
||||
@ -1669,53 +1551,9 @@ impl SessionStore {
|
||||
if let Some(sid) = session_id {
|
||||
let mut stmt = conn.prepare(
|
||||
"
|
||||
SELECT id, role, content, system_context, reasoning_content, media_refs_json, created_at, tool_call_id, tool_name, tool_calls_json, tool_duration_ms, prompt_tokens, completion_tokens, total_tokens, context_window_tokens
|
||||
FROM messages
|
||||
WHERE topic_id = ?1 AND session_id = ?2 AND is_compacted = 0
|
||||
ORDER BY seq ASC
|
||||
",
|
||||
)?;
|
||||
let rows = stmt.query_map(params![topic_id, sid], map_chat_message_row)?;
|
||||
let mut messages = Vec::new();
|
||||
for row in rows {
|
||||
messages.push(row?);
|
||||
}
|
||||
Ok(messages)
|
||||
} else {
|
||||
let mut stmt = conn.prepare(
|
||||
"
|
||||
SELECT id, role, content, system_context, reasoning_content, media_refs_json, created_at, tool_call_id, tool_name, tool_calls_json, tool_duration_ms, prompt_tokens, completion_tokens, total_tokens, context_window_tokens
|
||||
FROM messages
|
||||
WHERE topic_id = ?1 AND is_compacted = 0
|
||||
ORDER BY seq ASC
|
||||
",
|
||||
)?;
|
||||
let rows = stmt.query_map(params![topic_id], map_chat_message_row)?;
|
||||
let mut messages = Vec::new();
|
||||
for row in rows {
|
||||
messages.push(row?);
|
||||
}
|
||||
Ok(messages)
|
||||
}
|
||||
}
|
||||
|
||||
/// UI 视角：返回原始消息（含被压缩消费的 is_compacted=1 消息）+ 未压缩新消息，
|
||||
/// 排除压缩摘要消息（system_context LIKE 'history_compaction%'）。
|
||||
/// 用于前端历史展示、/current、/save-topic、topic 描述生成等场景。
|
||||
pub fn load_messages_for_topic_full(
|
||||
&self,
|
||||
topic_id: &str,
|
||||
session_id: Option<&str>,
|
||||
) -> Result<Vec<ChatMessage>, StorageError> {
|
||||
let conn = self.pool.get()?;
|
||||
|
||||
if let Some(sid) = session_id {
|
||||
let mut stmt = conn.prepare(
|
||||
"
|
||||
SELECT id, role, content, system_context, reasoning_content, media_refs_json, created_at, tool_call_id, tool_name, tool_calls_json, tool_duration_ms, prompt_tokens, completion_tokens, total_tokens, context_window_tokens
|
||||
SELECT id, role, content, system_context, reasoning_content, media_refs_json, created_at, tool_call_id, tool_name, tool_calls_json, tool_duration_ms
|
||||
FROM messages
|
||||
WHERE topic_id = ?1 AND session_id = ?2
|
||||
AND (system_context IS NULL OR system_context NOT LIKE 'history_compaction%')
|
||||
ORDER BY seq ASC
|
||||
",
|
||||
)?;
|
||||
@ -1728,10 +1566,9 @@ impl SessionStore {
|
||||
} else {
|
||||
let mut stmt = conn.prepare(
|
||||
"
|
||||
SELECT id, role, content, system_context, reasoning_content, media_refs_json, created_at, tool_call_id, tool_name, tool_calls_json, tool_duration_ms, prompt_tokens, completion_tokens, total_tokens, context_window_tokens
|
||||
SELECT id, role, content, system_context, reasoning_content, media_refs_json, created_at, tool_call_id, tool_name, tool_calls_json, tool_duration_ms
|
||||
FROM messages
|
||||
WHERE topic_id = ?1
|
||||
AND (system_context IS NULL OR system_context NOT LIKE 'history_compaction%')
|
||||
ORDER BY seq ASC
|
||||
",
|
||||
)?;
|
||||
@ -1744,19 +1581,9 @@ impl SessionStore {
|
||||
}
|
||||
}
|
||||
|
||||
/// 获取指定话题的消息数量。
|
||||
///
|
||||
/// 使用 `SELECT COUNT(*)` 在数据库侧计数，避免将所有消息
|
||||
/// （含 content、tool_calls_json 等大字段反序列化）加载到内存。
|
||||
/// 查询命中 `idx_messages_topic_seq(topic_id, seq)` 索引。
|
||||
/// 获取指定话题的消息数量（动态计算，确保准确）
|
||||
pub fn get_topic_message_count(&self, topic_id: &str) -> Result<usize, StorageError> {
|
||||
let conn = self.pool.get()?;
|
||||
let count: i64 = conn.query_row(
|
||||
"SELECT COUNT(*) FROM messages WHERE topic_id = ?1",
|
||||
params![topic_id],
|
||||
|row| row.get(0),
|
||||
)?;
|
||||
Ok(count as usize)
|
||||
self.load_messages_for_topic(topic_id, None).map(|msgs| msgs.len())
|
||||
}
|
||||
|
||||
pub fn load_all_messages(&self, session_id: &str) -> Result<Vec<ChatMessage>, StorageError> {
|
||||
@ -1778,184 +1605,6 @@ impl SessionStore {
|
||||
.map_err(StorageError::from)
|
||||
}
|
||||
|
||||
/// 批量查询多个 topic 的 token 消耗统计（cost 累计 + context 瞬时）。
|
||||
///
|
||||
/// 按 `topic_id` 聚合而非 `session_id`：一个 session 可包含多个 topic，
|
||||
/// 若按 session_id 聚合会导致同 session 下的所有 topic 显示相同的总和。
|
||||
///
|
||||
/// 子代理隔离：子代理消息持久化时 session_id='sub:...'，topic_id=父 topic_id
|
||||
/// （见 task::runtime PersistingEmittedMessageHandler 构造），因此不能仅靠
|
||||
/// topic_id 隔离。此处用 `session_id NOT LIKE 'sub:%'` 显式排除子代理消息，
|
||||
/// 与项目约定一致（session 列表同样过滤 'sub:%'）。子代理 token 不计入父 topic，
|
||||
/// 保持"子代理分别计算"语义。
|
||||
pub fn batch_topic_token_stats(
|
||||
&self,
|
||||
topic_ids: &[&str],
|
||||
) -> Result<HashMap<String, SessionTokenStats>, StorageError> {
|
||||
if topic_ids.is_empty() {
|
||||
return Ok(HashMap::new());
|
||||
}
|
||||
let conn = self.pool.get()?;
|
||||
|
||||
let placeholders = (0..topic_ids.len())
|
||||
.map(|i| format!("?{}", i + 1))
|
||||
.collect::<Vec<_>>()
|
||||
.join(", ");
|
||||
// topic_id IN (...) 自动排除 NULL topic_id 的旧消息；
|
||||
// session_id NOT LIKE 'sub:%' 排除子代理消息（其 topic_id=父 topic_id）
|
||||
let sum_sql = format!(
|
||||
"SELECT topic_id, \
|
||||
COALESCE(SUM(prompt_tokens), 0) AS sum_prompt, \
|
||||
COALESCE(SUM(completion_tokens), 0) AS sum_completion, \
|
||||
COALESCE(SUM(total_tokens), 0) AS sum_total \
|
||||
FROM messages \
|
||||
WHERE topic_id IN ({placeholders}) AND role = 'assistant' \
|
||||
AND session_id NOT LIKE 'sub:%' \
|
||||
GROUP BY topic_id"
|
||||
);
|
||||
|
||||
let mut stmt = conn.prepare(&sum_sql)?;
|
||||
let params: Vec<&dyn rusqlite::ToSql> = topic_ids
|
||||
.iter()
|
||||
.map(|s| s as &dyn rusqlite::ToSql)
|
||||
.collect();
|
||||
let sum_rows = stmt.query_map(params.as_slice(), |row| {
|
||||
Ok((
|
||||
row.get::<_, String>(0)?,
|
||||
SessionTokenStats {
|
||||
prompt_tokens: row.get::<_, i64>(1)? as u64,
|
||||
completion_tokens: row.get::<_, i64>(2)? as u64,
|
||||
total_tokens: row.get::<_, i64>(3)? as u64,
|
||||
last_prompt_tokens: None,
|
||||
context_window_tokens: None,
|
||||
},
|
||||
))
|
||||
})?;
|
||||
|
||||
let mut stats: HashMap<String, SessionTokenStats> = HashMap::new();
|
||||
for row in sum_rows {
|
||||
let (tid, s) = row?;
|
||||
stats.insert(tid, s);
|
||||
}
|
||||
|
||||
// 查找每个 topic 中最新的**有 usage 数据的** assistant 消息，
|
||||
// 读取其 prompt_tokens 和 context_window_tokens。
|
||||
// 过滤 prompt_tokens IS NOT NULL 确保跳过 error/cancel 消息（usage 为 NULL）；
|
||||
// session_id NOT LIKE 'sub:%' 排除子代理消息，避免取到子代理的 context_window。
|
||||
//
|
||||
// 注意：seq 是 session 级递增（见 append_message_with_topic），主 session 与
|
||||
// 子代理 session 各自独立计数，可能存在相同 seq。外层 WHERE 必须再次过滤
|
||||
// session_id NOT LIKE 'sub:%'，否则 JOIN 会同时匹配主消息和子代理消息，
|
||||
// 导致重复行并使 stats.entry(tid) 被覆盖，结果不确定。
|
||||
let last_sql = format!(
|
||||
"SELECT m.topic_id, m.prompt_tokens, m.context_window_tokens \
|
||||
FROM messages m \
|
||||
INNER JOIN ( \
|
||||
SELECT topic_id, MAX(seq) AS max_seq \
|
||||
FROM messages \
|
||||
WHERE topic_id IN ({placeholders}) AND role = 'assistant' \
|
||||
AND prompt_tokens IS NOT NULL \
|
||||
AND session_id NOT LIKE 'sub:%' \
|
||||
GROUP BY topic_id \
|
||||
) latest ON m.topic_id = latest.topic_id AND m.seq = latest.max_seq \
|
||||
WHERE m.session_id NOT LIKE 'sub:%'"
|
||||
);
|
||||
let mut stmt2 = conn.prepare(&last_sql)?;
|
||||
let params2: Vec<&dyn rusqlite::ToSql> = topic_ids
|
||||
.iter()
|
||||
.map(|s| s as &dyn rusqlite::ToSql)
|
||||
.collect();
|
||||
let last_rows = stmt2.query_map(params2.as_slice(), |row| {
|
||||
Ok((
|
||||
row.get::<_, String>(0)?,
|
||||
row.get::<_, Option<i64>>(1)?,
|
||||
row.get::<_, Option<i64>>(2)?,
|
||||
))
|
||||
})?;
|
||||
for row in last_rows {
|
||||
let (tid, last_prompt, last_ctx_window) = row?;
|
||||
let entry = stats.entry(tid).or_insert(SessionTokenStats {
|
||||
prompt_tokens: 0,
|
||||
completion_tokens: 0,
|
||||
total_tokens: 0,
|
||||
last_prompt_tokens: None,
|
||||
context_window_tokens: None,
|
||||
});
|
||||
entry.last_prompt_tokens = last_prompt.map(|v| v as u32);
|
||||
entry.context_window_tokens = last_ctx_window.map(|v| v as u32);
|
||||
}
|
||||
|
||||
Ok(stats)
|
||||
}
|
||||
|
||||
/// 查询单个 session 的 token 消耗统计（cost 累计 + context 瞬时）。
|
||||
///
|
||||
/// 按 `session_id` 精确匹配查询，**不过滤** `sub:%`——专门用于子代理 session
|
||||
/// （session_id = `sub:...`）的 token 统计。子代理没有 topic，一个 session 即
|
||||
/// 一个完整执行单元，因此按 session 维度聚合而非 topic 维度。
|
||||
///
|
||||
/// 与 `batch_topic_token_stats` 共享同一套 SQL 模式（SUM + last），差异仅在
|
||||
/// WHERE 条件：单 session 精确匹配，无 `NOT LIKE 'sub:%'` 过滤，无 topic 维度。
|
||||
pub fn get_session_token_stats(
|
||||
&self,
|
||||
session_id: &str,
|
||||
) -> Result<Option<SessionTokenStats>, StorageError> {
|
||||
let conn = self.pool.get()?;
|
||||
|
||||
// 1. SUM 查询：累计 prompt/completion/total
|
||||
let sum_sql = "SELECT \
|
||||
COALESCE(SUM(prompt_tokens), 0), \
|
||||
COALESCE(SUM(completion_tokens), 0), \
|
||||
COALESCE(SUM(total_tokens), 0) \
|
||||
FROM messages \
|
||||
WHERE session_id = ?1 AND role = 'assistant'";
|
||||
let mut stmt = conn.prepare(sum_sql)?;
|
||||
let sum_row = stmt.query_row(params![session_id], |row| {
|
||||
Ok((
|
||||
row.get::<_, i64>(0)? as u64,
|
||||
row.get::<_, i64>(1)? as u64,
|
||||
row.get::<_, i64>(2)? as u64,
|
||||
))
|
||||
})?;
|
||||
let (prompt_tokens, completion_tokens, total_tokens) = sum_row;
|
||||
|
||||
// 无 assistant 消息时直接返回 None
|
||||
if total_tokens == 0 && prompt_tokens == 0 && completion_tokens == 0 {
|
||||
// 需要二次确认是否真的没有 assistant 消息（usage 全 0 也可能是合法的）
|
||||
let count_sql =
|
||||
"SELECT COUNT(*) FROM messages WHERE session_id = ?1 AND role = 'assistant'";
|
||||
let count: i64 = conn.query_row(count_sql, params![session_id], |row| row.get(0))?;
|
||||
if count == 0 {
|
||||
return Ok(None);
|
||||
}
|
||||
}
|
||||
|
||||
// 2. last 查询：最新有 usage 的 assistant 消息的 prompt_tokens + context_window_tokens
|
||||
let last_sql = "SELECT prompt_tokens, context_window_tokens \
|
||||
FROM messages \
|
||||
WHERE session_id = ?1 AND role = 'assistant' AND prompt_tokens IS NOT NULL \
|
||||
ORDER BY seq DESC LIMIT 1";
|
||||
let mut stmt2 = conn.prepare(last_sql)?;
|
||||
let last_row = stmt2
|
||||
.query_row(params![session_id], |row| {
|
||||
Ok((row.get::<_, Option<i64>>(0)?, row.get::<_, Option<i64>>(1)?))
|
||||
})
|
||||
.optional()?;
|
||||
|
||||
let (last_prompt_tokens, context_window_tokens) = match last_row {
|
||||
Some((lp, lcw)) => (lp.map(|v| v as u32), lcw.map(|v| v as u32)),
|
||||
None => (None, None),
|
||||
};
|
||||
|
||||
Ok(Some(SessionTokenStats {
|
||||
prompt_tokens,
|
||||
completion_tokens,
|
||||
total_tokens,
|
||||
last_prompt_tokens,
|
||||
context_window_tokens,
|
||||
}))
|
||||
}
|
||||
|
||||
pub fn replace_todos(
|
||||
&self,
|
||||
scope_key: &str,
|
||||
@ -1970,7 +1619,10 @@ impl SessionStore {
|
||||
let now = current_timestamp();
|
||||
|
||||
// Delete existing todos for this scope_key
|
||||
tx.execute("DELETE FROM todos WHERE scope_key = ?1", params![scope_key])?;
|
||||
tx.execute(
|
||||
"DELETE FROM todos WHERE scope_key = ?1",
|
||||
params![scope_key],
|
||||
)?;
|
||||
|
||||
// Insert new todos
|
||||
for item in items {
|
||||
@ -2017,7 +1669,7 @@ impl SessionStore {
|
||||
for row in rows {
|
||||
result.push(row?);
|
||||
}
|
||||
drop(stmt); // 释放 stmt 借用，才能 commit
|
||||
drop(stmt); // 释放 stmt 借用，才能 commit
|
||||
tx.commit()?;
|
||||
Ok(result)
|
||||
}
|
||||
@ -2071,7 +1723,7 @@ pub fn persistent_session_id(channel_name: &str, chat_id: &str) -> String {
|
||||
|
||||
#[cfg(not(test))]
|
||||
fn default_session_db_path() -> Result<PathBuf, std::io::Error> {
|
||||
let home = crate::platform::picobot_home_dir();
|
||||
let home = dirs::home_dir().unwrap_or_else(|| PathBuf::from("."));
|
||||
Ok(home.join(".picobot").join("storage").join("sessions.db"))
|
||||
}
|
||||
|
||||
@ -2091,8 +1743,8 @@ fn insert_message_with_seq(
|
||||
"
|
||||
INSERT INTO messages (
|
||||
id, session_id, seq, role, content,
|
||||
system_context, reasoning_content, media_refs_json, tool_call_id, tool_name, tool_calls_json, tool_duration_ms, prompt_tokens, completion_tokens, total_tokens, context_window_tokens, created_at
|
||||
) VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8, ?9, ?10, ?11, ?12, ?13, ?14, ?15, ?16, ?17)
|
||||
system_context, reasoning_content, media_refs_json, tool_call_id, tool_name, tool_calls_json, tool_duration_ms, created_at
|
||||
) VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8, ?9, ?10, ?11, ?12, ?13)
|
||||
",
|
||||
params![
|
||||
message.id,
|
||||
@ -2107,10 +1759,6 @@ fn insert_message_with_seq(
|
||||
message.tool_name,
|
||||
tool_calls_json,
|
||||
message.tool_duration_ms.map(|v| v as i64),
|
||||
message.usage.as_ref().map(|u| u.prompt_tokens as i64),
|
||||
message.usage.as_ref().map(|u| u.completion_tokens as i64),
|
||||
message.usage.as_ref().map(|u| u.total_tokens as i64),
|
||||
message.usage.as_ref().and_then(|u| u.context_window_tokens.map(|v| v as i64)),
|
||||
message.timestamp,
|
||||
],
|
||||
)?;
|
||||
@ -2139,8 +1787,8 @@ fn insert_message_with_topic_seq(
|
||||
"
|
||||
INSERT INTO messages (
|
||||
id, session_id, topic_id, seq, role, content,
|
||||
system_context, reasoning_content, media_refs_json, tool_call_id, tool_name, tool_calls_json, tool_duration_ms, prompt_tokens, completion_tokens, total_tokens, context_window_tokens, created_at
|
||||
) VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8, ?9, ?10, ?11, ?12, ?13, ?14, ?15, ?16, ?17, ?18)
|
||||
system_context, reasoning_content, media_refs_json, tool_call_id, tool_name, tool_calls_json, tool_duration_ms, created_at
|
||||
) VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8, ?9, ?10, ?11, ?12, ?13, ?14)
|
||||
",
|
||||
params![
|
||||
message.id,
|
||||
@ -2156,10 +1804,6 @@ fn insert_message_with_topic_seq(
|
||||
message.tool_name,
|
||||
tool_calls_json,
|
||||
message.tool_duration_ms.map(|v| v as i64),
|
||||
message.usage.as_ref().map(|u| u.prompt_tokens as i64),
|
||||
message.usage.as_ref().map(|u| u.completion_tokens as i64),
|
||||
message.usage.as_ref().map(|u| u.total_tokens as i64),
|
||||
message.usage.as_ref().and_then(|u| u.context_window_tokens.map(|v| v as i64)),
|
||||
message.timestamp,
|
||||
],
|
||||
)?;
|
||||
@ -2180,8 +1824,6 @@ fn clone_message_for_compaction(message: &ChatMessage, timestamp: i64) -> ChatMe
|
||||
tool_state: message.tool_state.clone(),
|
||||
tool_duration_ms: message.tool_duration_ms,
|
||||
tool_calls: message.tool_calls.clone(),
|
||||
// 压缩克隆不保留 usage：压缩产生的是合成消息，不代表真实 LLM 调用
|
||||
usage: None,
|
||||
}
|
||||
}
|
||||
|
||||
@ -2193,7 +1835,7 @@ fn load_messages_between(
|
||||
) -> Result<Vec<ChatMessage>, StorageError> {
|
||||
let mut stmt = conn.prepare(
|
||||
"
|
||||
SELECT id, role, content, system_context, reasoning_content, media_refs_json, created_at, tool_call_id, tool_name, tool_calls_json, tool_duration_ms, prompt_tokens, completion_tokens, total_tokens, context_window_tokens
|
||||
SELECT id, role, content, system_context, reasoning_content, media_refs_json, created_at, tool_call_id, tool_name, tool_calls_json, tool_duration_ms
|
||||
FROM messages
|
||||
WHERE session_id = ?1 AND seq > ?2 AND seq <= ?3
|
||||
ORDER BY seq ASC
|
||||
@ -2239,7 +1881,6 @@ fn load_messages_between(
|
||||
tool_state: None,
|
||||
tool_duration_ms: row.get::<_, Option<i64>>(10)?.map(|v| v as u64),
|
||||
tool_calls,
|
||||
usage: map_usage_row(row, 11, 12, 13, 14)?,
|
||||
})
|
||||
},
|
||||
)?;
|
||||
@ -2258,7 +1899,7 @@ fn load_messages_after(
|
||||
) -> Result<Vec<ChatMessage>, StorageError> {
|
||||
let mut stmt = conn.prepare(
|
||||
"
|
||||
SELECT id, role, content, system_context, reasoning_content, media_refs_json, created_at, tool_call_id, tool_name, tool_calls_json, tool_duration_ms, prompt_tokens, completion_tokens, total_tokens, context_window_tokens
|
||||
SELECT id, role, content, system_context, reasoning_content, media_refs_json, created_at, tool_call_id, tool_name, tool_calls_json, tool_duration_ms
|
||||
FROM messages
|
||||
WHERE session_id = ?1 AND seq > ?2
|
||||
ORDER BY seq ASC
|
||||
@ -2301,7 +1942,6 @@ fn load_messages_after(
|
||||
tool_state: None,
|
||||
tool_duration_ms: row.get::<_, Option<i64>>(10)?.map(|v| v as u64),
|
||||
tool_calls,
|
||||
usage: map_usage_row(row, 11, 12, 13, 14)?,
|
||||
})
|
||||
})?;
|
||||
|
||||
@ -2310,6 +1950,13 @@ fn load_messages_after(
|
||||
messages.push(row?);
|
||||
}
|
||||
Ok(messages)
|
||||
}
|
||||
|
||||
fn current_timestamp() -> i64 {
|
||||
std::time::SystemTime::now()
|
||||
.duration_since(std::time::UNIX_EPOCH)
|
||||
.expect("system clock before unix epoch")
|
||||
.as_millis() as i64
|
||||
}
|
||||
|
||||
fn quote_fts_query(query: &str) -> String {
|
||||
|
||||
@ -28,14 +28,6 @@ pub trait ConversationRepository: Send + Sync + 'static {
|
||||
session_id: Option<&str>,
|
||||
) -> Result<Vec<ChatMessage>, StorageError>;
|
||||
|
||||
/// UI 视角：返回原始消息（含被压缩消费的 is_compacted=1 消息）+ 未压缩新消息，
|
||||
/// 排除压缩摘要消息（system_context LIKE 'history_compaction%'）。
|
||||
fn load_messages_for_topic_full(
|
||||
&self,
|
||||
topic_id: &str,
|
||||
session_id: Option<&str>,
|
||||
) -> Result<Vec<ChatMessage>, StorageError>;
|
||||
|
||||
fn append_message(&self, session_id: &str, message: &ChatMessage) -> Result<(), StorageError>;
|
||||
|
||||
fn append_message_with_topic(
|
||||
@ -82,15 +74,6 @@ pub trait ConversationRepository: Send + Sync + 'static {
|
||||
topic_id: &str,
|
||||
messages: &[ChatMessage],
|
||||
) -> Result<(), StorageError>;
|
||||
|
||||
/// 压缩该 topic 的历史：保留原始消息（标记 is_compacted=1）+ 插入压缩摘要。
|
||||
/// 不删除原消息，让前端仍能展示完整原始对话，LLM 只看压缩后的精简历史。
|
||||
fn compact_topic_history(
|
||||
&self,
|
||||
session_id: &str,
|
||||
topic_id: &str,
|
||||
new_messages: &[ChatMessage],
|
||||
) -> Result<(), StorageError>;
|
||||
}
|
||||
|
||||
pub trait PromptInjectionRepository: Send + Sync + 'static {
|
||||
@ -289,23 +272,6 @@ impl ConversationRepository for super::SessionStore {
|
||||
) -> Result<(), StorageError> {
|
||||
super::SessionStore::replace_topic_history(self, session_id, topic_id, messages)
|
||||
}
|
||||
|
||||
fn load_messages_for_topic_full(
|
||||
&self,
|
||||
topic_id: &str,
|
||||
session_id: Option<&str>,
|
||||
) -> Result<Vec<ChatMessage>, StorageError> {
|
||||
super::SessionStore::load_messages_for_topic_full(self, topic_id, session_id)
|
||||
}
|
||||
|
||||
fn compact_topic_history(
|
||||
&self,
|
||||
session_id: &str,
|
||||
topic_id: &str,
|
||||
new_messages: &[ChatMessage],
|
||||
) -> Result<(), StorageError> {
|
||||
super::SessionStore::compact_topic_history(self, session_id, topic_id, new_messages)
|
||||
}
|
||||
}
|
||||
|
||||
impl PromptInjectionRepository for super::SessionStore {
|
||||
|
||||
@ -8,38 +8,18 @@ pub const GLOBAL_SCOPE_KEY: &str = "default";
|
||||
/// 每个命名空间代表一类记忆内容，用于分类管理和检索。
|
||||
/// 禁止使用未在此列表中的 namespace 创建记忆。
|
||||
pub const ALLOWED_MEMORY_NAMESPACES: &[(&str, &str)] = &[
|
||||
(
|
||||
"user",
|
||||
"用户记忆：存储用户长期偏好、身份背景和历史协作信息，实现跨会话的个性化服务与持续协作",
|
||||
),
|
||||
(
|
||||
"semantic",
|
||||
"语义记忆：存储结构化或非结构化知识内容，支持知识检索、问答增强和长期知识积累",
|
||||
),
|
||||
(
|
||||
"episodic",
|
||||
"情景记忆：记录历史对话、任务执行过程及关键事件，支持经验回溯、案例复用和行为追踪",
|
||||
),
|
||||
(
|
||||
"skill",
|
||||
"技能记忆：存储技能定义、工作流、工具调用策略及最佳实践，支持能力复用与自动化执行",
|
||||
),
|
||||
(
|
||||
"environment",
|
||||
"环境记忆：存储外部系统状态、运行环境配置和实时资源信息，为智能决策提供环境感知能力",
|
||||
),
|
||||
(
|
||||
"reflection",
|
||||
"反思记忆：沉淀任务执行过程中的成功经验、失败原因和优化建议，支持智能体持续学习与自我改进",
|
||||
),
|
||||
("user", "用户记忆：存储用户长期偏好、身份背景和历史协作信息，实现跨会话的个性化服务与持续协作"),
|
||||
("semantic", "语义记忆：存储结构化或非结构化知识内容，支持知识检索、问答增强和长期知识积累"),
|
||||
("episodic", "情景记忆：记录历史对话、任务执行过程及关键事件，支持经验回溯、案例复用和行为追踪"),
|
||||
("skill", "技能记忆：存储技能定义、工作流、工具调用策略及最佳实践，支持能力复用与自动化执行"),
|
||||
("environment", "环境记忆：存储外部系统状态、运行环境配置和实时资源信息，为智能决策提供环境感知能力"),
|
||||
("reflection", "反思记忆：沉淀任务执行过程中的成功经验、失败原因和优化建议，支持智能体持续学习与自我改进"),
|
||||
("other", "其他记忆：不属于以上分类的其他记忆内容"),
|
||||
];
|
||||
|
||||
/// 验证 namespace 是否在允许列表中
|
||||
pub fn is_valid_namespace(namespace: &str) -> bool {
|
||||
ALLOWED_MEMORY_NAMESPACES
|
||||
.iter()
|
||||
.any(|(name, _)| *name == namespace)
|
||||
ALLOWED_MEMORY_NAMESPACES.iter().any(|(name, _)| *name == namespace)
|
||||
}
|
||||
|
||||
/// 获取 namespace 的中文描述
|
||||
@ -52,10 +32,7 @@ pub fn get_namespace_description(namespace: &str) -> Option<&'static str> {
|
||||
|
||||
/// 获取所有允许的 namespace 名称列表（用于 JSON schema enum）
|
||||
pub fn allowed_namespace_names() -> Vec<&'static str> {
|
||||
ALLOWED_MEMORY_NAMESPACES
|
||||
.iter()
|
||||
.map(|(name, _)| *name)
|
||||
.collect()
|
||||
ALLOWED_MEMORY_NAMESPACES.iter().map(|(name, _)| *name).collect()
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
@ -111,16 +88,6 @@ pub struct TopicRecord {
|
||||
pub message_count: i64,
|
||||
}
|
||||
|
||||
/// 单个 session 的 token 用量统计（聚合结果）。
|
||||
#[derive(Debug, Clone, Default, Serialize, Deserialize)]
|
||||
pub struct SessionTokenStats {
|
||||
pub prompt_tokens: u64,
|
||||
pub completion_tokens: u64,
|
||||
pub total_tokens: u64,
|
||||
pub last_prompt_tokens: Option<u32>,
|
||||
pub context_window_tokens: Option<u32>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct MemoryRecord {
|
||||
pub id: String,
|
||||
|
||||
Some files were not shown because too many files have changed in this diff Show More
Loading…
x
Reference in New Issue
Block a user