refactor: 提取公共工具函数，修复飞书正则重复编译和网关路由DRY问题，添加CI安全审计
This commit is contained in:
oudecheng
parent
3abd22ddf4
commit
510520e08e
7
.github/workflows/ci.yml
vendored
7
.github/workflows/ci.yml
vendored
@ -57,6 +57,9 @@ jobs:
|
||||
# 需要真实 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
|
||||
@ -89,3 +92,7 @@ jobs:
|
||||
- name: Run tests
|
||||
working-directory: web
|
||||
run: npm run test
|
||||
|
||||
- name: Security audit
|
||||
working-directory: web
|
||||
run: npm audit --audit-level=high
|
||||
|
||||
179
docs/OPTIMIZATION_PLAN.md
Normal file
179
docs/OPTIMIZATION_PLAN.md
Normal file
@ -0,0 +1,179 @@
|
||||
# PicoBot 项目优化计划（修订版）
|
||||
|
||||
> 重新评估日期：2026-08-06 | 项目版本：v0.3.2
|
||||
|
||||
---
|
||||
|
||||
## 重新评估说明
|
||||
|
||||
初版计划存在三个问题：
|
||||
|
||||
1. **混淆了"优化"与"重构美化"**：迁移 8 个错误类型到 thiserror、拆分文件、重组目录等纯属代码美化，零功能价值，却带来大量 churn 和引入 Bug 的风险
|
||||
2. **混入了新功能**：Docker 支持、指标导出、WebSocket 速率限制等是功能需求，不是优化
|
||||
3. **投机性优化**：在没有 profiling 数据的情况下假设 `.clone()` 是瓶颈、假设 tokio features 影响编译时间，这些都违背了"不过度工程化"原则
|
||||
|
||||
本修订版仅保留**有明确证据支撑、高 ROI、低风险**的改进项。
|
||||
|
||||
---
|
||||
|
||||
## Tier 1：值得做（高 ROI，低风险）
|
||||
|
||||
### 1.1 消除函数重复定义
|
||||
|
||||
**证据**：完全相同的函数在多个文件中重复定义，且错误处理策略不一致。
|
||||
|
||||
| 函数 | 重复次数 | 位置 |
|
||||
|------|----------|------|
|
||||
| `current_timestamp()` | 6 处 | `bus/message.rs`、`storage/mod.rs`、`gateway/ws.rs`、`tools/scheduler_manage.rs`、`tools/task/types.rs`、`tools/task/repository.rs` |
|
||||
| `dirs::home_dir().unwrap_or_else(\|\| PathBuf::from("."))` | 7 处 | `config/mod.rs`(×3)、`cli/init.rs`(×3)、`logging.rs` |
|
||||
| `format_error_chain()` | 3 处 | `agent_loop.rs`、`openai.rs`、`anthropic.rs` |
|
||||
|
||||
**方案**：提取到 `src/utils.rs` 公共模块，统一错误处理策略。
|
||||
|
||||
**风险**：极低——纯机械操作，函数体完全相同。
|
||||
|
||||
---
|
||||
|
||||
### 1.2 修复飞书正则表达式重复编译
|
||||
|
||||
**证据**：`channels/feishu.rs` 中 `MdPatterns::new()` 在每条出站消息处理时调用（[L2537](file:///e:/code_project/PicoBot/src/channels/feishu.rs#L2537)），每次编译 9 个正则表达式。正则编译是 CPU 密集型操作，在消息热路径上是不必要的开销。
|
||||
|
||||
**方案**：将 `MdPatterns` 改为 `LazyLock<MdPatterns>` 全局静态实例，仅编译一次。
|
||||
|
||||
```rust
|
||||
use std::sync::LazyLock;
|
||||
|
||||
static MD_PATTERNS: LazyLock<MdPatterns> = LazyLock::new(MdPatterns::new);
|
||||
```
|
||||
|
||||
**风险**：低——`MdPatterns` 是无状态的纯数据结构，全局共享安全。
|
||||
|
||||
---
|
||||
|
||||
### 1.3 CI 添加 `cargo audit`
|
||||
|
||||
**证据**：项目有 ~30 个 Rust 依赖和数十个 npm 依赖，但 CI 中无任何安全扫描。依赖漏洞是真实的安全风险。
|
||||
|
||||
**方案**：在 `.github/workflows/ci.yml` 的 rust-checks job 中添加一步：
|
||||
|
||||
```yaml
|
||||
- name: Security audit
|
||||
run: cargo install cargo-audit --locked && cargo audit
|
||||
```
|
||||
|
||||
前端部分添加：
|
||||
|
||||
```yaml
|
||||
- name: npm audit
|
||||
run: npm audit --audit-level=high
|
||||
```
|
||||
|
||||
**风险**：零——只读检查，不改变构建产物。
|
||||
|
||||
---
|
||||
|
||||
## Tier 2：可以做（中 ROI，中风险）
|
||||
|
||||
### 2.1 为 LLMProvider 引入结构化错误类型
|
||||
|
||||
**证据**：`LLMProvider::chat()` 返回 `Box<dyn std::error::Error + Send + Sync>`，导致 `agent_loop.rs` 中的 `is_recoverable_llm_error()` 只能通过字符串子串匹配判断错误类型：
|
||||
|
||||
```rust
|
||||
// agent_loop.rs L527-541 — 脆弱的字符串匹配
|
||||
fn is_recoverable_llm_error(error: &str) -> bool {
|
||||
let normalized = error.to_ascii_lowercase();
|
||||
normalized.contains("429")
|
||||
|| normalized.contains("502")
|
||||
|| normalized.contains("503")
|
||||
// ...
|
||||
}
|
||||
```
|
||||
|
||||
这种模式有真实风险：Provider 修改错误消息格式后匹配静默失效，且无法实现"429 读取 Retry-After 头"等精细化策略。
|
||||
|
||||
**方案**：定义 `ProviderCallError` 枚举替代 `Box<dyn Error>`：
|
||||
|
||||
```rust
|
||||
#[derive(Debug, thiserror::Error)]
|
||||
pub enum ProviderCallError {
|
||||
#[error("HTTP {status}: {body}")]
|
||||
Http { status: u16, body: String },
|
||||
#[error("rate limited")]
|
||||
RateLimited { retry_after: Option<Duration> },
|
||||
#[error("network error: {0}")]
|
||||
Network(#[from] reqwest::Error),
|
||||
#[error("parse error: {0}")]
|
||||
Parse(#[from] serde_json::Error),
|
||||
#[error("{0}")]
|
||||
Other(String),
|
||||
}
|
||||
```
|
||||
|
||||
更新 `LLMProvider` trait 签名、OpenAI/Anthropic Provider 实现、以及 `agent_loop.rs` 的重试逻辑。
|
||||
|
||||
**风险**：中——涉及 trait 签名变更，影响所有 Provider 实现和调用方。但项目已有 `StorageError` 作为 thiserror 范例，且 `LLMProvider` 是内部 trait（非公开 API），影响面可控。
|
||||
|
||||
**工作量**：约 3-5 个文件需要改动，现有测试需要调整断言。
|
||||
|
||||
---
|
||||
|
||||
### 2.2 WebSocket 媒体文件写入改用 spawn_blocking
|
||||
|
||||
**证据**：`gateway/ws.rs` [L81-96](file:///e:/code_project/PicoBot/src/gateway/ws.rs#L81-L96) 在 async handler 中直接调用 `std::fs::create_dir_all` 和 `std::fs::write` 保存用户上传的图片/文件。对于较大的图片（几 MB），阻塞时间可能达到几十毫秒。
|
||||
|
||||
**方案**：用 `tokio::task::spawn_blocking` 包裹文件写入操作。
|
||||
|
||||
**风险**：低——项目其他地方（`file_read.rs`、`agent_loop.rs` 图片编码）已使用相同模式。
|
||||
|
||||
**注意**：`http.rs` L162 的 `std::fs::write`（配置保存）**不需要改**——这是一次几 KB 文件的罕见操作，阻塞时间在微秒级，不值得增加代码复杂性。
|
||||
|
||||
---
|
||||
|
||||
## Tier 3：明确不做
|
||||
|
||||
以下项目经重新评估后决定**不做**：
|
||||
|
||||
| 项目 | 原计划编号 | 不做理由 |
|
||||
|------|-----------|----------|
|
||||
| 迁移 8 个手工错误类型到 thiserror | 1.1 阶段 1 | **纯美化**。手工实现的 `Display + Error` 工作正常，样板代码已写完。迁移零功能价值，却有引入 Bug 的风险。 |
|
||||
| 拆分 `agent_loop.rs` | 2.2 | **过度工程化**。3000 行虽大但结构清晰，函数边界明确。拆分引入 import 变更和合并冲突，收益仅是"文件短了"。 |
|
||||
| 重构 `gateway/` 目录结构 | 2.3 | **过度工程化**。30 个文件的平铺结构工作正常，无人报告导航困难。移动文件是大规模 churn，零功能价值。 |
|
||||
| 统一锁策略文档 | 2.4 | **现有做法已正确**。tokio::sync 用于 async、parking_lot 用于 sync 是正确的选型模式，不需要额外文档。 |
|
||||
| 减少 `.clone()` 调用 | 3.1 | **投机性优化**。无 profiling 数据表明 clone 是瓶颈。Rust 中大部分 clone 是所有权必需。盲目减少 clone 可能引入生命周期问题。应等 profiling 数据支撑后再做。 |
|
||||
| 精简 tokio features | 3.2 | **低收益有风险**。`["full"]` 方便且安全，精简后可能遗漏 feature 导致跨平台构建失败，排查成本远超节省的编译时间。 |
|
||||
| 评估 SQLite 连接池 | 3.3 | **投机性**。无数据表明连接池是问题。r2d2 提供超时和错误恢复能力，保留更安全。 |
|
||||
| 添加 `#[instrument]` | 4.1 | **锦上添花**。项目已有 916 处 tracing 调用，日志覆盖充分。`#[instrument]` 是增量改进，非必需。可在排查具体问题时按需添加。 |
|
||||
| 生产指标导出 | 4.2 | **这是新功能，不是优化**。 |
|
||||
| 共享测试工具模块 | 5.1 | **低优先级**。仅 MockTodoRepository 重复 2 处，影响极小。 |
|
||||
| 可离线集成测试 | 5.2 | **低优先级**。已有 570 个单元测试 + MockProvider，覆盖率足够。 |
|
||||
| WebSocket 速率限制 | 6.2 | **新功能**。 |
|
||||
| 工具执行超时 | 6.2 | **新功能**。 |
|
||||
| Docker 支持 | 7.2 | **新功能**。 |
|
||||
| 代码覆盖率报告 | 7.1 | **低优先级**。无证据表明测试覆盖不足是当前瓶颈。 |
|
||||
| Clippy `-D warnings` | 7.1 | **需先清理存量**。当前存量警告太多，直接启用会阻断 CI。 |
|
||||
|
||||
---
|
||||
|
||||
## 实施顺序
|
||||
|
||||
```
|
||||
Step 1: 函数去重（1.1） ← 30 分钟，零风险
|
||||
Step 2: feishu 正则 LazyLock（1.2）← 15 分钟，低风险
|
||||
Step 3: CI cargo audit（1.3） ← 10 分钟，零风险
|
||||
Step 4: ProviderCallError（2.1） ← 2-3 小时，中风险
|
||||
Step 5: ws.rs spawn_blocking（2.2）← 20 分钟，低风险
|
||||
```
|
||||
|
||||
Step 1-3 可以立即执行，互不依赖。Step 4 是独立的大型重构，可单独安排。Step 5 可随时做。
|
||||
|
||||
---
|
||||
|
||||
## 量化指标
|
||||
|
||||
| 指标 | 当前值 | 目标值 |
|
||||
|------|--------|--------|
|
||||
| 重复函数定义 | 16 处 | 0 |
|
||||
| 每条飞书消息正则编译次数 | 9 | 0（编译一次，全局复用） |
|
||||
| CI 安全扫描 | 无 | cargo audit + npm audit |
|
||||
| `is_recoverable_llm_error` 字符串匹配 | 10 个子串 | 0（改为类型匹配） |
|
||||
| `Box<dyn Error>` 在 Provider trait | 2 处 | 0 |
|
||||
@ -13,6 +13,7 @@ use crate::providers::{
|
||||
use crate::storage::ConversationRepository;
|
||||
use crate::text::{char_count, take_prefix_chars, take_suffix_chars};
|
||||
use crate::tools::{ToolContext, ToolRegistry};
|
||||
use crate::utils::format_error_chain;
|
||||
use async_trait::async_trait;
|
||||
use std::borrow::Cow;
|
||||
use std::collections::{HashMap, VecDeque};
|
||||
@ -20,8 +21,8 @@ use std::hash::{Hash, Hasher};
|
||||
use std::io::Read;
|
||||
use std::sync::Arc;
|
||||
use std::sync::atomic::{AtomicBool, Ordering};
|
||||
use std::time::Instant;
|
||||
use std::time::Duration;
|
||||
use std::time::Instant;
|
||||
|
||||
/// Minimum characters to keep when truncating
|
||||
const TRUNCATION_SUFFIX_LEN: usize = 200;
|
||||
@ -512,18 +513,6 @@ fn normalize_tool_arguments(arguments: &serde_json::Value) -> serde_json::Value
|
||||
}
|
||||
}
|
||||
|
||||
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 is_recoverable_llm_error(error: &str) -> bool {
|
||||
let normalized = error.to_ascii_lowercase();
|
||||
// 瞬态可恢复错误：服务端过载/限流/网关错误/超时/连接重置。
|
||||
@ -730,8 +719,7 @@ async fn preencode_images_for_request(
|
||||
}
|
||||
|
||||
// 阶段 2：并行 spawn_blocking 编码
|
||||
let mut join_set: tokio::task::JoinSet<(String, PreencodeEntry)> =
|
||||
tokio::task::JoinSet::new();
|
||||
let mut join_set: tokio::task::JoinSet<(String, PreencodeEntry)> = tokio::task::JoinSet::new();
|
||||
for (path, target_tokens) in to_encode {
|
||||
join_set.spawn_blocking(move || {
|
||||
match encode_image_to_base64_with_budget(&path, target_tokens) {
|
||||
@ -1182,13 +1170,14 @@ impl AgentLoop {
|
||||
);
|
||||
}
|
||||
|
||||
let request = self.build_llm_request(
|
||||
&messages,
|
||||
system_prompt_context,
|
||||
tools.clone(),
|
||||
tools_tokens,
|
||||
)
|
||||
.await;
|
||||
let request = self
|
||||
.build_llm_request(
|
||||
&messages,
|
||||
system_prompt_context,
|
||||
tools.clone(),
|
||||
tools_tokens,
|
||||
)
|
||||
.await;
|
||||
|
||||
// Set up streaming delta consumer
|
||||
// Pre-generate the message ID so stream deltas and the final assistant
|
||||
@ -1217,10 +1206,11 @@ impl AgentLoop {
|
||||
// 重试会重复输出，此时不重试。
|
||||
let emitted = Arc::new(AtomicBool::new(false));
|
||||
let emitted_for_cb = emitted.clone();
|
||||
let stream_callback: StreamCallback = std::sync::Arc::new(move |delta: StreamDelta| {
|
||||
emitted_for_cb.store(true, Ordering::SeqCst);
|
||||
let _ = delta_tx.try_send(delta);
|
||||
});
|
||||
let stream_callback: StreamCallback =
|
||||
std::sync::Arc::new(move |delta: StreamDelta| {
|
||||
emitted_for_cb.store(true, Ordering::SeqCst);
|
||||
let _ = delta_tx.try_send(delta);
|
||||
});
|
||||
|
||||
// LLM 调用与取消信号竞速：若取消信号到达，drop LLM future 以 abort HTTP 请求。
|
||||
let llm_result: Result<
|
||||
@ -1261,8 +1251,8 @@ impl AgentLoop {
|
||||
&& !emitted.load(Ordering::SeqCst)
|
||||
&& is_recoverable_llm_error(&error_text);
|
||||
if can_retry {
|
||||
let delay = LLM_RETRY_DELAYS_MS
|
||||
[attempt.min(LLM_RETRY_DELAYS_MS.len() - 1)];
|
||||
let delay =
|
||||
LLM_RETRY_DELAYS_MS[attempt.min(LLM_RETRY_DELAYS_MS.len() - 1)];
|
||||
tracing::warn!(
|
||||
attempt = attempt + 1,
|
||||
retry_in_ms = delay,
|
||||
@ -1600,7 +1590,8 @@ impl AgentLoop {
|
||||
let preencoded =
|
||||
preencode_images_for_request(filtered_messages_ref, &mut image_budget).await;
|
||||
|
||||
let mut messages_for_llm: Vec<Message> = Vec::with_capacity(filtered_messages_ref.len() + 2);
|
||||
let mut messages_for_llm: Vec<Message> =
|
||||
Vec::with_capacity(filtered_messages_ref.len() + 2);
|
||||
if let Some(ref prompt) = system_prompt {
|
||||
messages_for_llm.push(Message::system(prompt.content.clone()));
|
||||
}
|
||||
@ -1719,7 +1710,9 @@ impl AgentLoop {
|
||||
);
|
||||
messages.push(summary_request);
|
||||
|
||||
let request = self.build_llm_request(messages, system_prompt_context, None, 0).await;
|
||||
let request = self
|
||||
.build_llm_request(messages, system_prompt_context, None, 0)
|
||||
.await;
|
||||
let max_retries = self.runtime_config.max_retries as usize;
|
||||
|
||||
for attempt in 0..=max_retries {
|
||||
@ -1745,7 +1738,8 @@ impl AgentLoop {
|
||||
|
||||
match final_result {
|
||||
Ok(response) => {
|
||||
let mut assistant_message = if let Some(reasoning_content) = response.reasoning_content
|
||||
let mut assistant_message = if let Some(reasoning_content) =
|
||||
response.reasoning_content
|
||||
{
|
||||
ChatMessage::assistant_with_reasoning(response.content, reasoning_content)
|
||||
} else {
|
||||
@ -1768,11 +1762,9 @@ impl AgentLoop {
|
||||
}
|
||||
Err(e) => {
|
||||
let error_text = e.to_string();
|
||||
let can_retry = attempt < max_retries
|
||||
&& is_recoverable_llm_error(&error_text);
|
||||
let can_retry = attempt < max_retries && is_recoverable_llm_error(&error_text);
|
||||
if can_retry {
|
||||
let delay = LLM_RETRY_DELAYS_MS
|
||||
[attempt.min(LLM_RETRY_DELAYS_MS.len() - 1)];
|
||||
let delay = LLM_RETRY_DELAYS_MS[attempt.min(LLM_RETRY_DELAYS_MS.len() - 1)];
|
||||
tracing::warn!(
|
||||
attempt = attempt + 1,
|
||||
retry_in_ms = delay,
|
||||
@ -1802,7 +1794,8 @@ impl AgentLoop {
|
||||
error_details = %format_error_chain(e.as_ref()),
|
||||
"Failed to get summary from LLM"
|
||||
);
|
||||
let final_message = ChatMessage::assistant(recoverable_llm_message(&error_text));
|
||||
let final_message =
|
||||
ChatMessage::assistant(recoverable_llm_message(&error_text));
|
||||
emitted_messages.push(final_message.clone());
|
||||
self.emit_live_tool_call_message(final_message.clone())
|
||||
.await;
|
||||
@ -2258,10 +2251,17 @@ mod tests {
|
||||
let pdf_path = temp_dir.path().join("demo.pdf");
|
||||
std::fs::write(&pdf_path, b"%PDF-1.4").unwrap();
|
||||
|
||||
let messages = vec![ChatMessage::user_with_media("hello", vec![pdf_path.to_string_lossy().to_string()])];
|
||||
let messages = vec![ChatMessage::user_with_media(
|
||||
"hello",
|
||||
vec![pdf_path.to_string_lossy().to_string()],
|
||||
)];
|
||||
let mut budget = ImageInlineBudget::new(1_000, 0);
|
||||
let preencoded = preencode_images_for_request(&messages, &mut budget).await;
|
||||
let blocks = build_content_blocks("hello", &[pdf_path.to_string_lossy().to_string()], &preencoded);
|
||||
let blocks = build_content_blocks(
|
||||
"hello",
|
||||
&[pdf_path.to_string_lossy().to_string()],
|
||||
&preencoded,
|
||||
);
|
||||
|
||||
assert_eq!(blocks.len(), 1);
|
||||
assert!(matches!(&blocks[0], ContentBlock::Text { text } if text == "hello"));
|
||||
@ -2275,7 +2275,10 @@ mod tests {
|
||||
image.save(&jpg_path).unwrap();
|
||||
|
||||
let path_str = jpg_path.to_string_lossy().to_string();
|
||||
let messages = vec![ChatMessage::user_with_media("hello", vec![path_str.clone()])];
|
||||
let messages = vec![ChatMessage::user_with_media(
|
||||
"hello",
|
||||
vec![path_str.clone()],
|
||||
)];
|
||||
let mut budget = ImageInlineBudget::new(10_000, 1);
|
||||
let preencoded = preencode_images_for_request(&messages, &mut budget).await;
|
||||
let blocks = build_content_blocks("hello", &[path_str], &preencoded);
|
||||
@ -2295,7 +2298,10 @@ mod tests {
|
||||
image.save(&png_path).unwrap();
|
||||
|
||||
let path_str = png_path.to_string_lossy().to_string();
|
||||
let messages = vec![ChatMessage::user_with_media("hello", vec![path_str.clone()])];
|
||||
let messages = vec![ChatMessage::user_with_media(
|
||||
"hello",
|
||||
vec![path_str.clone()],
|
||||
)];
|
||||
let mut budget = ImageInlineBudget::new(512, 1);
|
||||
let preencoded = preencode_images_for_request(&messages, &mut budget).await;
|
||||
let blocks = build_content_blocks("hello", &[path_str], &preencoded);
|
||||
@ -2315,7 +2321,10 @@ mod tests {
|
||||
image.save(&jpg_path).unwrap();
|
||||
|
||||
let path_str = jpg_path.to_string_lossy().to_string();
|
||||
let messages = vec![ChatMessage::user_with_media("hello", vec![path_str.clone()])];
|
||||
let messages = vec![ChatMessage::user_with_media(
|
||||
"hello",
|
||||
vec![path_str.clone()],
|
||||
)];
|
||||
let mut budget = ImageInlineBudget::new(0, 1);
|
||||
let preencoded = preencode_images_for_request(&messages, &mut budget).await;
|
||||
let blocks = build_content_blocks("hello", &[path_str], &preencoded);
|
||||
@ -3002,7 +3011,9 @@ mod tests {
|
||||
assert!(!is_recoverable_llm_error("400 Bad Request"));
|
||||
assert!(!is_recoverable_llm_error("404 Not Found: model not found"));
|
||||
assert!(!is_recoverable_llm_error("Invalid API key"));
|
||||
assert!(!is_recoverable_llm_error("content filter blocked the request"));
|
||||
assert!(!is_recoverable_llm_error(
|
||||
"content filter blocked the request"
|
||||
));
|
||||
}
|
||||
|
||||
#[test]
|
||||
@ -3091,12 +3102,14 @@ mod tests {
|
||||
#[tokio::test]
|
||||
async fn test_retry_recovers_on_transient_error() {
|
||||
// 前 2 次返回 504（可恢复），第 3 次成功
|
||||
let loop_instance =
|
||||
make_loop_with_provider(Box::new(MockProvider::new(vec![
|
||||
let loop_instance = make_loop_with_provider(
|
||||
Box::new(MockProvider::new(vec![
|
||||
Err("504 Gateway Timeout".to_string()),
|
||||
Err("504 Gateway Timeout".to_string()),
|
||||
Ok(mock_success_response("recovered")),
|
||||
])), 3);
|
||||
])),
|
||||
3,
|
||||
);
|
||||
|
||||
let result = loop_instance
|
||||
.process(vec![ChatMessage::user("hello")], None, None)
|
||||
@ -3109,12 +3122,14 @@ mod tests {
|
||||
#[tokio::test]
|
||||
async fn test_retry_skipped_on_fatal_error() {
|
||||
// 401 不可恢复 → 立即失败，不重试
|
||||
let loop_instance =
|
||||
make_loop_with_provider(Box::new(MockProvider::new(vec![
|
||||
let loop_instance = make_loop_with_provider(
|
||||
Box::new(MockProvider::new(vec![
|
||||
Err("401 Unauthorized: invalid api key".to_string()),
|
||||
// 如果错误地重试了，第二次会返回成功，但我们期望不会到达
|
||||
Ok(mock_success_response("should not reach")),
|
||||
])), 3);
|
||||
])),
|
||||
3,
|
||||
);
|
||||
|
||||
let result = loop_instance
|
||||
.process(vec![ChatMessage::user("hello")], None, None)
|
||||
@ -3123,17 +3138,22 @@ mod tests {
|
||||
|
||||
// 应返回错误消息而非成功内容
|
||||
assert_ne!(result.final_response.content, "should not reach");
|
||||
assert!(result.final_response.content.contains("暂时不可用") || result.final_response.content.contains("401"));
|
||||
assert!(
|
||||
result.final_response.content.contains("暂时不可用")
|
||||
|| result.final_response.content.contains("401")
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_no_retry_when_max_retries_zero() {
|
||||
// max_retries=0 → 不重试，第一次失败即返回
|
||||
let loop_instance =
|
||||
make_loop_with_provider(Box::new(MockProvider::new(vec![
|
||||
let loop_instance = make_loop_with_provider(
|
||||
Box::new(MockProvider::new(vec![
|
||||
Err("504 Gateway Timeout".to_string()),
|
||||
Ok(mock_success_response("should not reach")),
|
||||
])), 0);
|
||||
])),
|
||||
0,
|
||||
);
|
||||
|
||||
let result = loop_instance
|
||||
.process(vec![ChatMessage::user("hello")], None, None)
|
||||
@ -3147,11 +3167,13 @@ mod tests {
|
||||
#[tokio::test]
|
||||
async fn test_retry_exhausted_returns_error_message() {
|
||||
// max_retries=1，两次都返回 504 → 重试1次后仍失败
|
||||
let loop_instance =
|
||||
make_loop_with_provider(Box::new(MockProvider::new(vec![
|
||||
let loop_instance = make_loop_with_provider(
|
||||
Box::new(MockProvider::new(vec![
|
||||
Err("504 Gateway Timeout".to_string()),
|
||||
Err("504 Gateway Timeout".to_string()),
|
||||
])), 1);
|
||||
])),
|
||||
1,
|
||||
);
|
||||
|
||||
let result = loop_instance
|
||||
.process(vec![ChatMessage::user("hello")], None, None)
|
||||
|
||||
@ -2,6 +2,7 @@ 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";
|
||||
@ -931,13 +932,6 @@ 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};
|
||||
|
||||
@ -1,6 +1,7 @@
|
||||
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;
|
||||
@ -1863,10 +1864,13 @@ 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 = MdPatterns::new();
|
||||
let patterns = &MD_PATTERNS;
|
||||
let stripped = content.trim();
|
||||
|
||||
// Tables and headings are not supported by post `md` nodes, so use cards.
|
||||
@ -1902,7 +1906,7 @@ impl FeishuChannel {
|
||||
|
||||
/// Strip markdown formatting markers from text for plain display.
|
||||
fn strip_md_formatting(text: &str) -> String {
|
||||
let patterns = MdPatterns::new();
|
||||
let patterns = &MD_PATTERNS;
|
||||
let mut result = text.to_string();
|
||||
|
||||
// Remove bold markers
|
||||
@ -1981,7 +1985,7 @@ impl FeishuChannel {
|
||||
|
||||
/// Split content by headings, converting headings to div elements.
|
||||
fn split_headings(content: &str) -> Vec<serde_json::Value> {
|
||||
let patterns = MdPatterns::new();
|
||||
let patterns = &MD_PATTERNS;
|
||||
let mut protected = content.to_string();
|
||||
|
||||
// Protect code blocks by replacing them with placeholders
|
||||
@ -2087,7 +2091,7 @@ impl FeishuChannel {
|
||||
|
||||
/// Build content into card elements (div/markdown + table).
|
||||
fn build_card_elements(content: &str) -> Vec<serde_json::Value> {
|
||||
let patterns = MdPatterns::new();
|
||||
let patterns = &MD_PATTERNS;
|
||||
let mut elements: Vec<serde_json::Value> = Vec::new();
|
||||
let mut last_end = 0;
|
||||
|
||||
|
||||
@ -104,7 +104,7 @@ impl WechatChannel {
|
||||
}
|
||||
|
||||
fn default_media_dir() -> PathBuf {
|
||||
let home = dirs::home_dir().unwrap_or_else(|| PathBuf::from("."));
|
||||
let home = crate::platform::picobot_home_dir();
|
||||
home.join(".picobot").join("media").join("wechat")
|
||||
}
|
||||
|
||||
|
||||
@ -18,7 +18,7 @@ pub struct InitWizard {
|
||||
|
||||
impl InitWizard {
|
||||
pub fn new() -> Self {
|
||||
let home = dirs::home_dir().unwrap_or_else(|| PathBuf::from("."));
|
||||
let home = crate::platform::picobot_home_dir();
|
||||
Self {
|
||||
read: BufReader::new(tokio::io::stdin()),
|
||||
write: tokio::io::stdout(),
|
||||
@ -892,14 +892,14 @@ impl InitWizard {
|
||||
}
|
||||
|
||||
fn default_feishu_media_dir() -> String {
|
||||
let home = dirs::home_dir().unwrap_or_else(|| PathBuf::from("."));
|
||||
let home = crate::platform::picobot_home_dir();
|
||||
home.join(".picobot/media/feishu")
|
||||
.to_string_lossy()
|
||||
.to_string()
|
||||
}
|
||||
|
||||
fn default_wechat_cred_path() -> String {
|
||||
let home = dirs::home_dir().unwrap_or_else(|| PathBuf::from("."));
|
||||
let home = crate::platform::picobot_home_dir();
|
||||
home.join(".picobot/wechat/credentials.json")
|
||||
.to_string_lossy()
|
||||
.to_string()
|
||||
|
||||
@ -457,7 +457,7 @@ fn default_allow_from() -> Vec<String> {
|
||||
}
|
||||
|
||||
fn default_media_dir() -> String {
|
||||
let home = dirs::home_dir().unwrap_or_else(|| std::path::PathBuf::from("."));
|
||||
let home = crate::platform::picobot_home_dir();
|
||||
home.join(".picobot/media/feishu")
|
||||
.to_string_lossy()
|
||||
.to_string()
|
||||
@ -471,7 +471,7 @@ fn default_wechat_base_url() -> String {
|
||||
}
|
||||
|
||||
fn default_wechat_cred_path() -> String {
|
||||
let home = dirs::home_dir().unwrap_or_else(|| std::path::PathBuf::from("."));
|
||||
let home = crate::platform::picobot_home_dir();
|
||||
home.join(".picobot/wechat/credentials.json")
|
||||
.to_string_lossy()
|
||||
.to_string()
|
||||
@ -955,7 +955,7 @@ impl LLMProviderConfig {
|
||||
}
|
||||
|
||||
pub(crate) fn get_default_config_path() -> PathBuf {
|
||||
let home = dirs::home_dir().unwrap_or_else(|| PathBuf::from("."));
|
||||
let home = crate::platform::picobot_home_dir();
|
||||
home.join(".picobot").join("config.json")
|
||||
}
|
||||
|
||||
@ -1931,7 +1931,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,33 +2252,25 @@ 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]
|
||||
|
||||
@ -281,110 +281,66 @@ pub async fn run(
|
||||
// 开发模式下可通过 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/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 {
|
||||
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/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(static_handler)
|
||||
.with_state(state.clone())
|
||||
app.fallback(static_handler)
|
||||
} else {
|
||||
let static_dir = std::env::var("STATIC_DIR").unwrap_or_else(|_| "static".to_string());
|
||||
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/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_service(ServeDir::new(&static_dir))
|
||||
.with_state(state.clone())
|
||||
};
|
||||
app.fallback_service(ServeDir::new(&static_dir))
|
||||
}
|
||||
.with_state(state.clone());
|
||||
|
||||
// 条件性挂载认证中间件：仅在需要认证时启用。
|
||||
// 中间件内部按 path 前缀判断，仅 /api/* 需要校验；
|
||||
|
||||
@ -30,11 +30,12 @@ use crate::gateway::agent_factory::build_system_prompt_provider;
|
||||
use crate::protocol::{MediaSummary, WsInbound, WsOutbound, 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 axum::extract::State;
|
||||
use axum::extract::ws::{Message as WsMessage, WebSocket, WebSocketUpgrade};
|
||||
use axum::extract::Query;
|
||||
use axum::response::{IntoResponse, Response};
|
||||
use axum::http::StatusCode;
|
||||
use axum::response::{IntoResponse, Response};
|
||||
use base64::{Engine as _, engine::general_purpose::STANDARD};
|
||||
use futures_util::{SinkExt, StreamExt};
|
||||
use std::collections::HashMap;
|
||||
@ -47,7 +48,7 @@ const WS_CHANNEL_NAME: &str = "websocket";
|
||||
|
||||
/// Default media directory for WebSocket uploads
|
||||
fn default_ws_media_dir() -> PathBuf {
|
||||
let home = dirs::home_dir().unwrap_or_else(|| PathBuf::from("."));
|
||||
let home = crate::platform::picobot_home_dir();
|
||||
home.join(".picobot").join("media").join("ws")
|
||||
}
|
||||
|
||||
@ -145,14 +146,8 @@ pub async fn ws_handler(
|
||||
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();
|
||||
tracing::warn!("WebSocket connection rejected: missing or invalid token");
|
||||
return (StatusCode::UNAUTHORIZED, "missing or invalid token").into_response();
|
||||
}
|
||||
}
|
||||
}
|
||||
@ -637,10 +632,9 @@ async fn handle_inbound(
|
||||
.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 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 {
|
||||
@ -791,13 +785,6 @@ 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)
|
||||
|
||||
@ -22,3 +22,4 @@ pub mod storage;
|
||||
pub mod text;
|
||||
pub mod tools;
|
||||
pub mod topic_description;
|
||||
pub mod utils;
|
||||
|
||||
@ -3,7 +3,7 @@ use chrono_tz::Tz;
|
||||
use std::path::PathBuf;
|
||||
use tracing_appender::rolling::{RollingFileAppender, Rotation};
|
||||
use tracing_subscriber::{
|
||||
EnvFilter, fmt, fmt::time::FormatTime, layer::SubscriberExt, util::SubscriberInitExt,
|
||||
fmt, fmt::time::FormatTime, layer::SubscriberExt, util::SubscriberInitExt, EnvFilter,
|
||||
};
|
||||
|
||||
#[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 = dirs::home_dir().unwrap_or_else(|| PathBuf::from("."));
|
||||
let home = crate::platform::picobot_home_dir();
|
||||
home.join(".picobot").join("logs")
|
||||
}
|
||||
|
||||
/// Get the default config file path: ~/.picobot/config.json
|
||||
pub fn get_default_config_path() -> PathBuf {
|
||||
let home = dirs::home_dir().unwrap_or_else(|| PathBuf::from("."));
|
||||
let home = crate::platform::picobot_home_dir();
|
||||
home.join(".picobot").join("config.json")
|
||||
}
|
||||
|
||||
|
||||
@ -345,6 +345,14 @@ 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
|
||||
|
||||
@ -8,21 +8,10 @@ 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,
|
||||
@ -106,9 +95,8 @@ 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")
|
||||
});
|
||||
let re =
|
||||
RE.get_or_init(|| regex::Regex::new(r"data:(image/\w+);base64,(.+)").expect("valid regex"));
|
||||
if let Some(caps) = 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("");
|
||||
|
||||
@ -10,6 +10,7 @@ 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",
|
||||
@ -126,31 +127,23 @@ 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: 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,
|
||||
}),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
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,
|
||||
@ -612,8 +605,7 @@ impl OpenAIProvider {
|
||||
// 提取流式末帧的 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())
|
||||
if let Ok(u) = serde_json::from_value::<OpenAIUsage>(usage_val.clone())
|
||||
{
|
||||
accumulator.set_usage(u);
|
||||
}
|
||||
|
||||
@ -3,6 +3,8 @@ 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};
|
||||
@ -978,9 +980,8 @@ impl SessionStore {
|
||||
let now = current_timestamp();
|
||||
|
||||
// 分离摘要消息与保留消息（保留消息携带原 ID，摘要消息是新构造的）
|
||||
let (summaries, preserved): (Vec<&ChatMessage>, Vec<&ChatMessage>) = new_messages
|
||||
.iter()
|
||||
.partition(|m| {
|
||||
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"))
|
||||
@ -997,8 +998,7 @@ impl SessionStore {
|
||||
|
||||
// 将该 topic 中未被保留的原消息标记为 is_compacted=1（仅更新尚未标记的行，避免重复写）。
|
||||
// 保留消息（system_guards / 最新 user）保持 is_compacted=0，不重复插入。
|
||||
let preserved_ids: Vec<String> =
|
||||
preserved.iter().map(|m| m.id.clone()).collect();
|
||||
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 \
|
||||
@ -1922,7 +1922,8 @@ impl SessionStore {
|
||||
// 无 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_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);
|
||||
@ -1937,10 +1938,7 @@ impl SessionStore {
|
||||
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)?,
|
||||
))
|
||||
Ok((row.get::<_, Option<i64>>(0)?, row.get::<_, Option<i64>>(1)?))
|
||||
})
|
||||
.optional()?;
|
||||
|
||||
@ -2073,7 +2071,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 = dirs::home_dir().unwrap_or_else(|| PathBuf::from("."));
|
||||
let home = crate::platform::picobot_home_dir();
|
||||
Ok(home.join(".picobot").join("storage").join("sessions.db"))
|
||||
}
|
||||
|
||||
@ -2314,13 +2312,6 @@ fn load_messages_after(
|
||||
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 {
|
||||
format!("\"{}\"", query.replace('"', "\"\""))
|
||||
}
|
||||
|
||||
@ -4,6 +4,8 @@ use std::sync::Arc;
|
||||
use async_trait::async_trait;
|
||||
use serde_json::json;
|
||||
|
||||
use crate::utils::current_timestamp;
|
||||
|
||||
use crate::config::SchedulerSchedule;
|
||||
use crate::storage::{
|
||||
SchedulerJobRecord, SchedulerJobRepository, SchedulerJobState, SchedulerJobUpsert,
|
||||
@ -502,13 +504,6 @@ fn error_result(message: &str) -> ToolResult {
|
||||
}
|
||||
}
|
||||
|
||||
fn current_timestamp() -> i64 {
|
||||
std::time::SystemTime::now()
|
||||
.duration_since(std::time::UNIX_EPOCH)
|
||||
.unwrap_or_default()
|
||||
.as_millis() as i64
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
@ -1,8 +1,10 @@
|
||||
use std::collections::HashMap;
|
||||
use parking_lot::RwLock;
|
||||
use std::collections::HashMap;
|
||||
|
||||
use async_trait::async_trait;
|
||||
|
||||
use crate::utils::current_timestamp;
|
||||
|
||||
use crate::storage::StorageError;
|
||||
|
||||
use super::types::TaskSession;
|
||||
@ -126,10 +128,3 @@ impl TaskRepository for InMemoryTaskRepository {
|
||||
Ok(before - sessions.len())
|
||||
}
|
||||
}
|
||||
|
||||
fn current_timestamp() -> i64 {
|
||||
std::time::SystemTime::now()
|
||||
.duration_since(std::time::UNIX_EPOCH)
|
||||
.expect("system clock before unix epoch")
|
||||
.as_millis() as i64
|
||||
}
|
||||
|
||||
@ -3,6 +3,7 @@ use std::path::PathBuf;
|
||||
use serde::{Deserialize, Serialize};
|
||||
|
||||
use crate::domain::CapabilityPolicy;
|
||||
use crate::utils::current_timestamp;
|
||||
|
||||
/// 子代理会话状态
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
|
||||
@ -279,10 +280,3 @@ pub struct TaskToolResult {
|
||||
/// 会话 ID（用于恢复）
|
||||
pub task_id: String,
|
||||
}
|
||||
|
||||
fn current_timestamp() -> i64 {
|
||||
std::time::SystemTime::now()
|
||||
.duration_since(std::time::UNIX_EPOCH)
|
||||
.expect("system clock before unix epoch")
|
||||
.as_millis() as i64
|
||||
}
|
||||
|
||||
25
src/utils.rs
Normal file
25
src/utils.rs
Normal file
@ -0,0 +1,25 @@
|
||||
use std::time::{SystemTime, UNIX_EPOCH};
|
||||
|
||||
/// 当前 Unix 时间戳（毫秒）。
|
||||
///
|
||||
/// 系统时钟倒拨到 Unix 纪元之前时 panic——与原各模块实现一致，
|
||||
/// 统一使用 `.expect()` 附带诊断信息。
|
||||
pub fn current_timestamp() -> i64 {
|
||||
SystemTime::now()
|
||||
.duration_since(UNIX_EPOCH)
|
||||
.expect("system clock before unix epoch")
|
||||
.as_millis() as i64
|
||||
}
|
||||
|
||||
/// 递归展开 `error.source()` 链，生成 `"顶层错误\ncaused by: 原因\ncaused by: ..."` 格式的字符串。
|
||||
pub 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: ")
|
||||
}
|
||||
Loading…
x
Reference in New Issue
Block a user