refactor: 迁移 parking_lot 锁并优化阻塞 IO 与内存管理
## 锁迁移:std::sync → parking_lot 消除锁中毒(poison)导致的级联崩溃风险。parking_lot 锁不会中毒, 且性能更优。迁移覆盖全部生产代码: - experts/mod.rs: 4 RwLock + 13 expect - skills/mod.rs: 1 RwLock + 11 expect - tools/registry.rs: 1 RwLock + 2 expect - gateway/model_selection.rs: 1 RwLock + 2 expect - tools/task/repository.rs: 1 RwLock + 4 unwrap - tools/task/runtime.rs: 2 RwLock + 17 expect - gateway/session.rs + task/runtime.rs: stream_message_id Mutex - gateway/processor.rs: description_generation_in_flight Mutex - command/handler.rs + help.rs: metadata Mutex(公开 API) - mcp/client.rs: stderr_lines Mutex 测试代码中的 std::sync::Mutex(串行化锁 + TestObserver)有意保留, 已通过 unwrap_or_else(|err| err.into_inner()) 做中毒恢复。 ## P1: 阻塞 IO 迁移到 spawn_blocking 将 3 处阻塞 async worker 的操作迁移到 blocking 线程池: - file_read.rs: read_to_string + 行处理 + base64 编码整体包入 spawn_blocking - agent_loop.rs: 新增 preencode_images_for_request 两阶段预编码 (顺序分配预算 → 并行 spawn_blocking 编码),build_llm_request 改为 async - wechat.rs: media_to_send_content 改为 async,std::fs::read 用 spawn_blocking 包裹 ## P2: session_history topic_histories 内存上限 新增 MAX_CACHED_TOPICS=32 上限和 evict_inactive_if_needed 方法。 超限时驱逐非活跃 topic(不在 chat_topic_ids、不在 compression_in_flight、 serial_lock 未被持有)。活跃 topic 永不误驱逐。 remove_history 同步清理 topic_serial_locks,防止无限增长。 ## P3: 减少 panic 面 agent_loop.rs retry 循环的 response.expect(...) 改为 ok_or_else(...)? 返回 AgentError::Other,逻辑 bug 不再导致整个 agent 崩溃。 ## 对抗性审查修复 - preencode_images_for_request: 用 seen HashSet 去重,防止同 path 重复 编码导致 HashMap entry 覆盖(NoBudget 覆盖 Encoded 等) - evict_inactive_if_needed: 检查 topic_serial_lock.try_lock(),防止驱逐 正在 agent 处理中的 topic(original_topic_id 不在 chat_topic_ids 但 agent 仍持锁) - remove_history: 清理 topic_serial_locks ## 验证 - cargo check: 通过(仅既有 lifetime 警告) - cargo test: 559 passed / 3 failed(均为环境/sandbox 权限问题,与本次改动无关)
This commit is contained in:
parent
fa420b713f
commit
bf8c227634
1
Cargo.lock
generated
1
Cargo.lock
generated
@ -1656,6 +1656,7 @@ dependencies = [
|
|||||||
"libc",
|
"libc",
|
||||||
"meval",
|
"meval",
|
||||||
"mime_guess",
|
"mime_guess",
|
||||||
|
"parking_lot",
|
||||||
"prost",
|
"prost",
|
||||||
"r2d2",
|
"r2d2",
|
||||||
"r2d2_sqlite",
|
"r2d2_sqlite",
|
||||||
|
|||||||
@ -56,6 +56,7 @@ r2d2 = "0.8"
|
|||||||
r2d2_sqlite = "0.34"
|
r2d2_sqlite = "0.34"
|
||||||
rustls = { version = "0.23", features = ["ring"] }
|
rustls = { version = "0.23", features = ["ring"] }
|
||||||
subtle = "2.6"
|
subtle = "2.6"
|
||||||
|
parking_lot = "0.12"
|
||||||
wechatbot = { path = "vendor/wechatbot" }
|
wechatbot = { path = "vendor/wechatbot" }
|
||||||
encoding_rs = "0.8"
|
encoding_rs = "0.8"
|
||||||
libc = "0.2"
|
libc = "0.2"
|
||||||
|
|||||||
@ -15,7 +15,7 @@ use crate::text::{char_count, take_prefix_chars, take_suffix_chars};
|
|||||||
use crate::tools::{ToolContext, ToolRegistry};
|
use crate::tools::{ToolContext, ToolRegistry};
|
||||||
use async_trait::async_trait;
|
use async_trait::async_trait;
|
||||||
use std::borrow::Cow;
|
use std::borrow::Cow;
|
||||||
use std::collections::VecDeque;
|
use std::collections::{HashMap, VecDeque};
|
||||||
use std::hash::{Hash, Hasher};
|
use std::hash::{Hash, Hasher};
|
||||||
use std::io::Read;
|
use std::io::Read;
|
||||||
use std::sync::Arc;
|
use std::sync::Arc;
|
||||||
@ -43,19 +43,21 @@ const JPEG_QUALITY_STEPS: &[u8] = &[82, 72, 60, 48, 36];
|
|||||||
const MIN_COMPRESSED_IMAGE_SIDE: u32 = 64;
|
const MIN_COMPRESSED_IMAGE_SIDE: u32 = 64;
|
||||||
const IMAGE_INPUT_NOTICE_PREFIX: &str = "[系统提示] 以下图片未能成功入模:";
|
const IMAGE_INPUT_NOTICE_PREFIX: &str = "[系统提示] 以下图片未能成功入模:";
|
||||||
|
|
||||||
/// Build content blocks from text and media paths
|
/// Build content blocks from text and media paths.
|
||||||
|
/// `preencoded` 包含已通过 spawn_blocking 预编码的图片结果(path → PreencodeEntry)。
|
||||||
|
/// 预算分配已在 `preencode_images_for_request` 中完成,此处只查表组装。
|
||||||
fn build_content_blocks(
|
fn build_content_blocks(
|
||||||
text: &str,
|
text: &str,
|
||||||
media_paths: &[String],
|
media_paths: &[String],
|
||||||
budget: &mut ImageInlineBudget,
|
preencoded: &HashMap<String, PreencodeEntry>,
|
||||||
) -> Vec<ContentBlock> {
|
) -> Vec<ContentBlock> {
|
||||||
build_content_blocks_with_image_budget(text, media_paths, budget)
|
build_content_blocks_with_image_budget(text, media_paths, preencoded)
|
||||||
}
|
}
|
||||||
|
|
||||||
fn build_content_blocks_with_image_budget(
|
fn build_content_blocks_with_image_budget(
|
||||||
text: &str,
|
text: &str,
|
||||||
media_paths: &[String],
|
media_paths: &[String],
|
||||||
budget: &mut ImageInlineBudget,
|
preencoded: &HashMap<String, PreencodeEntry>,
|
||||||
) -> Vec<ContentBlock> {
|
) -> Vec<ContentBlock> {
|
||||||
let mut blocks = Vec::new();
|
let mut blocks = Vec::new();
|
||||||
let mut skipped_image_notices = Vec::new();
|
let mut skipped_image_notices = Vec::new();
|
||||||
@ -72,27 +74,33 @@ fn build_content_blocks_with_image_budget(
|
|||||||
continue;
|
continue;
|
||||||
}
|
}
|
||||||
|
|
||||||
let Some(target_tokens) = budget.take_next_image_tokens() else {
|
match preencoded.get(path) {
|
||||||
tracing::warn!(media_path = %path, "Skipping image media ref because no LLM context budget remains");
|
Some(PreencodeEntry::Encoded(mime_type, base64_data)) => {
|
||||||
skipped_image_notices.push(format!(
|
|
||||||
"- {}:模型上下文预算不足,当前轮无法读取这张图片,请直接告知用户图片未成功入模。",
|
|
||||||
display_media_name(path)
|
|
||||||
));
|
|
||||||
continue;
|
|
||||||
};
|
|
||||||
|
|
||||||
match encode_image_to_base64_with_budget(path, target_tokens) {
|
|
||||||
Ok((mime_type, base64_data)) => {
|
|
||||||
let url = format!("data:{};base64,{}", mime_type, base64_data);
|
let url = format!("data:{};base64,{}", mime_type, base64_data);
|
||||||
blocks.push(ContentBlock::image_url(url));
|
blocks.push(ContentBlock::image_url(url));
|
||||||
}
|
}
|
||||||
Err(err) => {
|
Some(PreencodeEntry::NoBudget) => {
|
||||||
tracing::warn!(media_path = %path, target_tokens = target_tokens, error = %err, "Skipping image media ref after compression failed");
|
tracing::warn!(media_path = %path, "Skipping image media ref because no LLM context budget remains");
|
||||||
|
skipped_image_notices.push(format!(
|
||||||
|
"- {}:模型上下文预算不足,当前轮无法读取这张图片,请直接告知用户图片未成功入模。",
|
||||||
|
display_media_name(path)
|
||||||
|
));
|
||||||
|
}
|
||||||
|
Some(PreencodeEntry::Failed) => {
|
||||||
|
tracing::warn!(media_path = %path, "Skipping image media ref after compression failed");
|
||||||
|
skipped_image_notices.push(format!(
|
||||||
|
"- {}:图片压缩或编码失败,当前轮无法读取这张图片,请直接告知用户图片未成功入模。",
|
||||||
|
display_media_name(path)
|
||||||
|
));
|
||||||
|
}
|
||||||
|
None => {
|
||||||
|
// 不在 preencoded map 中:可能是 preencode 时被 filter_images_by_age_and_count
|
||||||
|
// 过滤掉的旧图片,或 preencode 逻辑遗漏。按编码失败处理。
|
||||||
|
tracing::warn!(media_path = %path, "Image media ref not found in preencoded map");
|
||||||
skipped_image_notices.push(format!(
|
skipped_image_notices.push(format!(
|
||||||
"- {}:图片压缩或编码失败,当前轮无法读取这张图片,请直接告知用户图片未成功入模。",
|
"- {}:图片压缩或编码失败,当前轮无法读取这张图片,请直接告知用户图片未成功入模。",
|
||||||
display_media_name(path)
|
display_media_name(path)
|
||||||
));
|
));
|
||||||
continue;
|
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
@ -655,11 +663,14 @@ fn canonicalise_json(value: &serde_json::Value) -> serde_json::Value {
|
|||||||
}
|
}
|
||||||
|
|
||||||
/// Convert ChatMessage to LLM Message format
|
/// Convert ChatMessage to LLM Message format
|
||||||
fn chat_message_to_llm_message(m: &ChatMessage, image_budget: &mut ImageInlineBudget) -> Message {
|
fn chat_message_to_llm_message(
|
||||||
|
m: &ChatMessage,
|
||||||
|
preencoded: &HashMap<String, PreencodeEntry>,
|
||||||
|
) -> Message {
|
||||||
let content = if m.media_refs.is_empty() {
|
let content = if m.media_refs.is_empty() {
|
||||||
vec![ContentBlock::text(&m.content)]
|
vec![ContentBlock::text(&m.content)]
|
||||||
} else {
|
} else {
|
||||||
build_content_blocks(&m.content, &m.media_refs, image_budget)
|
build_content_blocks(&m.content, &m.media_refs, preencoded)
|
||||||
};
|
};
|
||||||
|
|
||||||
Message {
|
Message {
|
||||||
@ -672,6 +683,97 @@ fn chat_message_to_llm_message(m: &ChatMessage, image_budget: &mut ImageInlineBu
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
/// 预编码请求中所有图片:顺序分配预算后并行 spawn_blocking 编码。
|
||||||
|
///
|
||||||
|
/// 两阶段设计:
|
||||||
|
/// 1. 顺序扫描所有消息的 media_refs,对每张支持编码的图片调用 `budget.take_next_image_tokens()`
|
||||||
|
/// 收集 `(path, target_tokens)` 列表。预算分配必须顺序以保证确定性。
|
||||||
|
/// 预算耗尽后剩余图片不编码(返回 `PreencodeStatus::NoBudget`)。
|
||||||
|
/// 2. 对收集到的列表并行 `spawn_blocking` 执行 `encode_image_to_base64_with_budget`
|
||||||
|
/// (IO + JPEG 压缩,CPU 密集),避免阻塞 tokio worker。
|
||||||
|
///
|
||||||
|
/// 返回 `HashMap<path, PreencodeEntry>`,`build_content_blocks_with_image_budget`
|
||||||
|
/// 据此组装结果或生成跳过提示,不再调用 `take_next_image_tokens`(budget 已在此消费完)。
|
||||||
|
async fn preencode_images_for_request(
|
||||||
|
messages: &[ChatMessage],
|
||||||
|
budget: &mut ImageInlineBudget,
|
||||||
|
) -> HashMap<String, PreencodeEntry> {
|
||||||
|
// 阶段 1:顺序分配预算,收集待编码列表 + 记录无预算的图片
|
||||||
|
let mut to_encode: Vec<(String, usize)> = Vec::new();
|
||||||
|
let mut result: HashMap<String, PreencodeEntry> = HashMap::new();
|
||||||
|
// 同一 path 跨多条消息重复出现时,只处理第一次(只扣减一次预算、只编码一次)。
|
||||||
|
// 后续出现直接查表复用,避免 HashMap entry 覆盖导致的状态丢失
|
||||||
|
// (如 NoBudget 覆盖 Encoded,或第二次编码失败覆盖第一次成功)。
|
||||||
|
let mut seen: std::collections::HashSet<String> = std::collections::HashSet::new();
|
||||||
|
|
||||||
|
for m in messages {
|
||||||
|
for path in &m.media_refs {
|
||||||
|
if supported_image_mime_type(path).is_none() {
|
||||||
|
continue; // 非图片,跳过(build_content_blocks 也会跳过)
|
||||||
|
}
|
||||||
|
if !seen.insert(path.clone()) {
|
||||||
|
continue; // 同一 path 已处理,跳过(build_content_blocks 会查表复用)
|
||||||
|
}
|
||||||
|
match budget.take_next_image_tokens() {
|
||||||
|
Some(target_tokens) => {
|
||||||
|
to_encode.push((path.clone(), target_tokens));
|
||||||
|
}
|
||||||
|
None => {
|
||||||
|
result.insert(path.clone(), PreencodeEntry::NoBudget);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
if to_encode.is_empty() {
|
||||||
|
return result;
|
||||||
|
}
|
||||||
|
|
||||||
|
// 阶段 2:并行 spawn_blocking 编码
|
||||||
|
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) {
|
||||||
|
Ok((mime, base64)) => (path, PreencodeEntry::Encoded(mime, base64)),
|
||||||
|
Err(err) => {
|
||||||
|
tracing::warn!(
|
||||||
|
media_path = %path,
|
||||||
|
target_tokens = target_tokens,
|
||||||
|
error = %err,
|
||||||
|
"Image preencode failed"
|
||||||
|
);
|
||||||
|
(path, PreencodeEntry::Failed)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
});
|
||||||
|
}
|
||||||
|
|
||||||
|
while let Some(join_result) = join_set.join_next().await {
|
||||||
|
match join_result {
|
||||||
|
Ok((path, entry)) => {
|
||||||
|
result.insert(path, entry);
|
||||||
|
}
|
||||||
|
Err(e) => {
|
||||||
|
tracing::warn!(error = %e, "Image preencode task panicked");
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
result
|
||||||
|
}
|
||||||
|
|
||||||
|
/// 单张图片的预编码结果
|
||||||
|
#[derive(Debug, Clone)]
|
||||||
|
enum PreencodeEntry {
|
||||||
|
/// 成功编码:(mime, base64)
|
||||||
|
Encoded(String, String),
|
||||||
|
/// 预算不足
|
||||||
|
NoBudget,
|
||||||
|
/// 编码失败
|
||||||
|
Failed,
|
||||||
|
}
|
||||||
|
|
||||||
/// AgentLoop - Stateless agent that processes messages with tool calling support.
|
/// AgentLoop - Stateless agent that processes messages with tool calling support.
|
||||||
/// History is managed externally by SessionManager.
|
/// History is managed externally by SessionManager.
|
||||||
pub struct AgentLoop {
|
pub struct AgentLoop {
|
||||||
@ -1085,7 +1187,8 @@ impl AgentLoop {
|
|||||||
system_prompt_context,
|
system_prompt_context,
|
||||||
tools.clone(),
|
tools.clone(),
|
||||||
tools_tokens,
|
tools_tokens,
|
||||||
);
|
)
|
||||||
|
.await;
|
||||||
|
|
||||||
// Set up streaming delta consumer
|
// Set up streaming delta consumer
|
||||||
// Pre-generate the message ID so stream deltas and the final assistant
|
// Pre-generate the message ID so stream deltas and the final assistant
|
||||||
@ -1205,7 +1308,11 @@ impl AgentLoop {
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
let response = response.expect("retry loop must set response or return");
|
let response = response.ok_or_else(|| {
|
||||||
|
AgentError::Other(
|
||||||
|
"retry loop exited without setting response or returning".to_string(),
|
||||||
|
)
|
||||||
|
})?;
|
||||||
|
|
||||||
// Signal stream end if handler exists
|
// Signal stream end if handler exists
|
||||||
let had_streaming = self.emitted_message_handler.is_some();
|
let had_streaming = self.emitted_message_handler.is_some();
|
||||||
@ -1460,7 +1567,7 @@ impl AgentLoop {
|
|||||||
/// 优化:token 估算直接基于 ChatMessage 累加字段长度(轻量),
|
/// 优化:token 估算直接基于 ChatMessage 累加字段长度(轻量),
|
||||||
/// 不再构造中间 text_only_messages: Vec<Message>(避免 N 条消息 × 5 字段双克隆)。
|
/// 不再构造中间 text_only_messages: Vec<Message>(避免 N 条消息 × 5 字段双克隆)。
|
||||||
/// 工具 token 估算在循环外预算后传入(tools_tokens)。
|
/// 工具 token 估算在循环外预算后传入(tools_tokens)。
|
||||||
fn build_llm_request(
|
async fn build_llm_request(
|
||||||
&self,
|
&self,
|
||||||
messages: &[ChatMessage],
|
messages: &[ChatMessage],
|
||||||
system_prompt_context: Option<&SystemPromptContext>,
|
system_prompt_context: Option<&SystemPromptContext>,
|
||||||
@ -1490,6 +1597,11 @@ impl AgentLoop {
|
|||||||
image_token_budget_for_request(&self.runtime_config, text_tokens, tools_tokens);
|
image_token_budget_for_request(&self.runtime_config, text_tokens, tools_tokens);
|
||||||
let mut image_budget = ImageInlineBudget::new(image_tokens, image_count);
|
let mut image_budget = ImageInlineBudget::new(image_tokens, image_count);
|
||||||
|
|
||||||
|
// 两阶段图片编码:先顺序分配预算,收集 (path, target_tokens) 列表;
|
||||||
|
// 再并行 spawn_blocking 编码,避免阻塞 async worker。
|
||||||
|
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 {
|
if let Some(ref prompt) = system_prompt {
|
||||||
messages_for_llm.push(Message::system(prompt.content.clone()));
|
messages_for_llm.push(Message::system(prompt.content.clone()));
|
||||||
@ -1497,7 +1609,7 @@ impl AgentLoop {
|
|||||||
messages_for_llm.extend(
|
messages_for_llm.extend(
|
||||||
filtered_messages_ref
|
filtered_messages_ref
|
||||||
.iter()
|
.iter()
|
||||||
.map(|message| chat_message_to_llm_message(message, &mut image_budget)),
|
.map(|message| chat_message_to_llm_message(message, &preencoded)),
|
||||||
);
|
);
|
||||||
|
|
||||||
ChatCompletionRequest {
|
ChatCompletionRequest {
|
||||||
@ -1609,7 +1721,7 @@ impl AgentLoop {
|
|||||||
);
|
);
|
||||||
messages.push(summary_request);
|
messages.push(summary_request);
|
||||||
|
|
||||||
let request = self.build_llm_request(messages, system_prompt_context, None, 0);
|
let request = self.build_llm_request(messages, system_prompt_context, None, 0).await;
|
||||||
let max_retries = self.runtime_config.max_retries as usize;
|
let max_retries = self.runtime_config.max_retries as usize;
|
||||||
|
|
||||||
for attempt in 0..=max_retries {
|
for attempt in 0..=max_retries {
|
||||||
@ -2055,8 +2167,8 @@ mod tests {
|
|||||||
}],
|
}],
|
||||||
);
|
);
|
||||||
|
|
||||||
let mut image_budget = ImageInlineBudget::new(0, 0);
|
let preencoded = HashMap::new();
|
||||||
let provider_message = chat_message_to_llm_message(&chat_message, &mut image_budget);
|
let provider_message = chat_message_to_llm_message(&chat_message, &preencoded);
|
||||||
|
|
||||||
assert_eq!(provider_message.role, "assistant");
|
assert_eq!(provider_message.role, "assistant");
|
||||||
assert_eq!(provider_message.tool_calls.as_ref().unwrap().len(), 1);
|
assert_eq!(provider_message.tool_calls.as_ref().unwrap().len(), 1);
|
||||||
@ -2075,8 +2187,8 @@ mod tests {
|
|||||||
let chat_message =
|
let chat_message =
|
||||||
ChatMessage::assistant_with_reasoning("final answer", "hidden chain of thought");
|
ChatMessage::assistant_with_reasoning("final answer", "hidden chain of thought");
|
||||||
|
|
||||||
let mut image_budget = ImageInlineBudget::new(0, 0);
|
let preencoded = HashMap::new();
|
||||||
let provider_message = chat_message_to_llm_message(&chat_message, &mut image_budget);
|
let provider_message = chat_message_to_llm_message(&chat_message, &preencoded);
|
||||||
|
|
||||||
assert_eq!(provider_message.role, "assistant");
|
assert_eq!(provider_message.role, "assistant");
|
||||||
assert_eq!(
|
assert_eq!(
|
||||||
@ -2122,36 +2234,33 @@ mod tests {
|
|||||||
);
|
);
|
||||||
}
|
}
|
||||||
|
|
||||||
#[test]
|
#[tokio::test]
|
||||||
fn test_build_content_blocks_skips_non_image_media_refs() {
|
async fn test_build_content_blocks_skips_non_image_media_refs() {
|
||||||
let temp_dir = tempdir().unwrap();
|
let temp_dir = tempdir().unwrap();
|
||||||
let pdf_path = temp_dir.path().join("demo.pdf");
|
let pdf_path = temp_dir.path().join("demo.pdf");
|
||||||
std::fs::write(&pdf_path, b"%PDF-1.4").unwrap();
|
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 mut budget = ImageInlineBudget::new(1_000, 0);
|
let mut budget = ImageInlineBudget::new(1_000, 0);
|
||||||
let blocks = build_content_blocks(
|
let preencoded = preencode_images_for_request(&messages, &mut budget).await;
|
||||||
"hello",
|
let blocks = build_content_blocks("hello", &[pdf_path.to_string_lossy().to_string()], &preencoded);
|
||||||
&[pdf_path.to_string_lossy().to_string()],
|
|
||||||
&mut budget,
|
|
||||||
);
|
|
||||||
|
|
||||||
assert_eq!(blocks.len(), 1);
|
assert_eq!(blocks.len(), 1);
|
||||||
assert!(matches!(&blocks[0], ContentBlock::Text { text } if text == "hello"));
|
assert!(matches!(&blocks[0], ContentBlock::Text { text } if text == "hello"));
|
||||||
}
|
}
|
||||||
|
|
||||||
#[test]
|
#[tokio::test]
|
||||||
fn test_build_content_blocks_keeps_supported_images() {
|
async fn test_build_content_blocks_keeps_supported_images() {
|
||||||
let temp_dir = tempdir().unwrap();
|
let temp_dir = tempdir().unwrap();
|
||||||
let jpg_path = temp_dir.path().join("demo.jpg");
|
let jpg_path = temp_dir.path().join("demo.jpg");
|
||||||
let image = image::DynamicImage::new_rgb8(8, 8);
|
let image = image::DynamicImage::new_rgb8(8, 8);
|
||||||
image.save(&jpg_path).unwrap();
|
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 mut budget = ImageInlineBudget::new(10_000, 1);
|
let mut budget = ImageInlineBudget::new(10_000, 1);
|
||||||
let blocks = build_content_blocks(
|
let preencoded = preencode_images_for_request(&messages, &mut budget).await;
|
||||||
"hello",
|
let blocks = build_content_blocks("hello", &[path_str], &preencoded);
|
||||||
&[jpg_path.to_string_lossy().to_string()],
|
|
||||||
&mut budget,
|
|
||||||
);
|
|
||||||
|
|
||||||
assert_eq!(blocks.len(), 2);
|
assert_eq!(blocks.len(), 2);
|
||||||
assert!(matches!(&blocks[0], ContentBlock::Text { text } if text == "hello"));
|
assert!(matches!(&blocks[0], ContentBlock::Text { text } if text == "hello"));
|
||||||
@ -2160,19 +2269,18 @@ mod tests {
|
|||||||
);
|
);
|
||||||
}
|
}
|
||||||
|
|
||||||
#[test]
|
#[tokio::test]
|
||||||
fn test_build_content_blocks_compresses_images_to_budget() {
|
async fn test_build_content_blocks_compresses_images_to_budget() {
|
||||||
let temp_dir = tempdir().unwrap();
|
let temp_dir = tempdir().unwrap();
|
||||||
let png_path = temp_dir.path().join("large.png");
|
let png_path = temp_dir.path().join("large.png");
|
||||||
let image = image::DynamicImage::new_rgb8(512, 512);
|
let image = image::DynamicImage::new_rgb8(512, 512);
|
||||||
image.save(&png_path).unwrap();
|
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 mut budget = ImageInlineBudget::new(512, 1);
|
let mut budget = ImageInlineBudget::new(512, 1);
|
||||||
let blocks = build_content_blocks(
|
let preencoded = preencode_images_for_request(&messages, &mut budget).await;
|
||||||
"hello",
|
let blocks = build_content_blocks("hello", &[path_str], &preencoded);
|
||||||
&[png_path.to_string_lossy().to_string()],
|
|
||||||
&mut budget,
|
|
||||||
);
|
|
||||||
|
|
||||||
assert_eq!(blocks.len(), 2);
|
assert_eq!(blocks.len(), 2);
|
||||||
assert!(matches!(&blocks[0], ContentBlock::Text { text } if text == "hello"));
|
assert!(matches!(&blocks[0], ContentBlock::Text { text } if text == "hello"));
|
||||||
@ -2181,19 +2289,18 @@ mod tests {
|
|||||||
);
|
);
|
||||||
}
|
}
|
||||||
|
|
||||||
#[test]
|
#[tokio::test]
|
||||||
fn test_build_content_blocks_adds_user_visible_notice_when_image_cannot_be_sent() {
|
async fn test_build_content_blocks_adds_user_visible_notice_when_image_cannot_be_sent() {
|
||||||
let temp_dir = tempdir().unwrap();
|
let temp_dir = tempdir().unwrap();
|
||||||
let jpg_path = temp_dir.path().join("demo.jpg");
|
let jpg_path = temp_dir.path().join("demo.jpg");
|
||||||
let image = image::DynamicImage::new_rgb8(8, 8);
|
let image = image::DynamicImage::new_rgb8(8, 8);
|
||||||
image.save(&jpg_path).unwrap();
|
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 mut budget = ImageInlineBudget::new(0, 1);
|
let mut budget = ImageInlineBudget::new(0, 1);
|
||||||
let blocks = build_content_blocks(
|
let preencoded = preencode_images_for_request(&messages, &mut budget).await;
|
||||||
"hello",
|
let blocks = build_content_blocks("hello", &[path_str], &preencoded);
|
||||||
&[jpg_path.to_string_lossy().to_string()],
|
|
||||||
&mut budget,
|
|
||||||
);
|
|
||||||
|
|
||||||
assert_eq!(blocks.len(), 2);
|
assert_eq!(blocks.len(), 2);
|
||||||
assert!(matches!(&blocks[0], ContentBlock::Text { text } if text == "hello"));
|
assert!(matches!(&blocks[0], ContentBlock::Text { text } if text == "hello"));
|
||||||
|
|||||||
@ -61,16 +61,23 @@ impl WechatChannel {
|
|||||||
.any(|pattern| pattern == "*" || pattern == sender_id)
|
.any(|pattern| pattern == "*" || pattern == sender_id)
|
||||||
}
|
}
|
||||||
|
|
||||||
fn media_to_send_content(
|
async fn media_to_send_content(
|
||||||
media: &MediaItem,
|
media: &MediaItem,
|
||||||
caption: Option<String>,
|
caption: Option<String>,
|
||||||
) -> Result<SendContent, ChannelError> {
|
) -> Result<SendContent, ChannelError> {
|
||||||
let data = std::fs::read(&media.path).map_err(|error| {
|
// 媒体文件读取是阻塞 IO,放到 blocking 线程池避免阻塞 async worker。
|
||||||
ChannelError::SendError(format!(
|
let path = media.path.clone();
|
||||||
"WeChat media read failed for '{}': {}",
|
let data = tokio::task::spawn_blocking(move || std::fs::read(&path))
|
||||||
media.path, error
|
.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
|
||||||
|
))
|
||||||
|
})?;
|
||||||
|
|
||||||
if data.is_empty() {
|
if data.is_empty() {
|
||||||
return Err(ChannelError::SendError(format!(
|
return Err(ChannelError::SendError(format!(
|
||||||
@ -355,7 +362,7 @@ impl Channel for WechatChannel {
|
|||||||
} else {
|
} else {
|
||||||
None
|
None
|
||||||
};
|
};
|
||||||
let content = Self::media_to_send_content(media, caption)?;
|
let content = Self::media_to_send_content(media, caption).await?;
|
||||||
self.bot
|
self.bot
|
||||||
.send_media(&msg.chat_id, content)
|
.send_media(&msg.chat_id, content)
|
||||||
.await
|
.await
|
||||||
@ -403,21 +410,21 @@ mod tests {
|
|||||||
assert!(filename.ends_with(".silk"));
|
assert!(filename.ends_with(".silk"));
|
||||||
}
|
}
|
||||||
|
|
||||||
#[test]
|
#[tokio::test]
|
||||||
fn media_to_send_content_maps_image() {
|
async fn media_to_send_content_maps_image() {
|
||||||
let file = NamedTempFile::new().unwrap();
|
let file = NamedTempFile::new().unwrap();
|
||||||
std::fs::write(file.path(), b"demo-image").unwrap();
|
std::fs::write(file.path(), b"demo-image").unwrap();
|
||||||
let image_path = file.path().with_extension("png");
|
let image_path = file.path().with_extension("png");
|
||||||
std::fs::rename(file.path(), &image_path).unwrap();
|
std::fs::rename(file.path(), &image_path).unwrap();
|
||||||
|
|
||||||
let media = MediaItem::new(image_path.to_string_lossy().to_string(), "image");
|
let media = MediaItem::new(image_path.to_string_lossy().to_string(), "image");
|
||||||
let content = WechatChannel::media_to_send_content(&media, None).unwrap();
|
let content = WechatChannel::media_to_send_content(&media, None).await.unwrap();
|
||||||
|
|
||||||
assert!(matches!(content, SendContent::Image { .. }));
|
assert!(matches!(content, SendContent::Image { .. }));
|
||||||
}
|
}
|
||||||
|
|
||||||
#[test]
|
#[tokio::test]
|
||||||
fn media_to_send_content_maps_generic_file() {
|
async fn media_to_send_content_maps_generic_file() {
|
||||||
let file = NamedTempFile::new().unwrap();
|
let file = NamedTempFile::new().unwrap();
|
||||||
std::fs::write(file.path(), b"hello").unwrap();
|
std::fs::write(file.path(), b"hello").unwrap();
|
||||||
let doc_path = file.path().with_extension("md");
|
let doc_path = file.path().with_extension("md");
|
||||||
@ -425,7 +432,7 @@ mod tests {
|
|||||||
|
|
||||||
let media = MediaItem::new(doc_path.to_string_lossy().to_string(), "file");
|
let media = MediaItem::new(doc_path.to_string_lossy().to_string(), "file");
|
||||||
let content =
|
let content =
|
||||||
WechatChannel::media_to_send_content(&media, Some("note".to_string())).unwrap();
|
WechatChannel::media_to_send_content(&media, Some("note".to_string())).await.unwrap();
|
||||||
|
|
||||||
match content {
|
match content {
|
||||||
SendContent::File {
|
SendContent::File {
|
||||||
|
|||||||
@ -78,7 +78,7 @@ pub trait InChatCommandHandler: Send + Sync {
|
|||||||
/// 负责将命令分发到合适的处理器
|
/// 负责将命令分发到合适的处理器
|
||||||
pub struct CommandRouter {
|
pub struct CommandRouter {
|
||||||
handlers: Vec<Box<dyn CommandHandler>>,
|
handlers: Vec<Box<dyn CommandHandler>>,
|
||||||
metadata: Arc<std::sync::Mutex<Vec<CommandMetadata>>>,
|
metadata: Arc<parking_lot::Mutex<Vec<CommandMetadata>>>,
|
||||||
}
|
}
|
||||||
|
|
||||||
impl CommandRouter {
|
impl CommandRouter {
|
||||||
@ -86,7 +86,7 @@ impl CommandRouter {
|
|||||||
pub fn new() -> Self {
|
pub fn new() -> Self {
|
||||||
Self {
|
Self {
|
||||||
handlers: Vec::new(),
|
handlers: Vec::new(),
|
||||||
metadata: Arc::new(std::sync::Mutex::new(Vec::new())),
|
metadata: Arc::new(parking_lot::Mutex::new(Vec::new())),
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
@ -96,13 +96,13 @@ impl CommandRouter {
|
|||||||
/// * `handler` - 要注册的处理器
|
/// * `handler` - 要注册的处理器
|
||||||
pub fn register(&mut self, handler: Box<dyn CommandHandler>) {
|
pub fn register(&mut self, handler: Box<dyn CommandHandler>) {
|
||||||
if let Some(meta) = handler.metadata() {
|
if let Some(meta) = handler.metadata() {
|
||||||
self.metadata.lock().unwrap().push(meta);
|
self.metadata.lock().push(meta);
|
||||||
}
|
}
|
||||||
self.handlers.push(handler);
|
self.handlers.push(handler);
|
||||||
}
|
}
|
||||||
|
|
||||||
/// 获取已注册命令的元数据列表(用于 Help 命令)
|
/// 获取已注册命令的元数据列表(用于 Help 命令)
|
||||||
pub fn metadata_arc(&self) -> Arc<std::sync::Mutex<Vec<CommandMetadata>>> {
|
pub fn metadata_arc(&self) -> Arc<parking_lot::Mutex<Vec<CommandMetadata>>> {
|
||||||
self.metadata.clone()
|
self.metadata.clone()
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|||||||
@ -3,7 +3,8 @@ use crate::command::context::CommandContext;
|
|||||||
use crate::command::handler::{CommandHandler, CommandMetadata};
|
use crate::command::handler::{CommandHandler, CommandMetadata};
|
||||||
use crate::command::response::{CommandError, CommandResponse, MessageKind};
|
use crate::command::response::{CommandError, CommandResponse, MessageKind};
|
||||||
use async_trait::async_trait;
|
use async_trait::async_trait;
|
||||||
use std::sync::{Arc, Mutex};
|
use parking_lot::Mutex;
|
||||||
|
use std::sync::Arc;
|
||||||
|
|
||||||
/// Help 命令处理器
|
/// Help 命令处理器
|
||||||
///
|
///
|
||||||
@ -41,7 +42,7 @@ impl CommandHandler for HelpCommandHandler {
|
|||||||
_cmd: Command,
|
_cmd: Command,
|
||||||
ctx: CommandContext,
|
ctx: CommandContext,
|
||||||
) -> Result<CommandResponse, CommandError> {
|
) -> Result<CommandResponse, CommandError> {
|
||||||
let metadata = self.metadata.lock().unwrap();
|
let metadata = self.metadata.lock();
|
||||||
let help_text = format_help(&metadata);
|
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))
|
||||||
|
|||||||
@ -5,7 +5,8 @@ use serde::{Deserialize, Serialize};
|
|||||||
use std::collections::{HashMap, HashSet};
|
use std::collections::{HashMap, HashSet};
|
||||||
use std::fs;
|
use std::fs;
|
||||||
use std::path::{Path, PathBuf};
|
use std::path::{Path, PathBuf};
|
||||||
use std::sync::{Arc, RwLock};
|
use std::sync::Arc;
|
||||||
|
use parking_lot::RwLock;
|
||||||
|
|
||||||
#[cfg(test)]
|
#[cfg(test)]
|
||||||
static EXPERT_TEST_ENV_LOCK: std::sync::Mutex<()> = std::sync::Mutex::new(());
|
static EXPERT_TEST_ENV_LOCK: std::sync::Mutex<()> = std::sync::Mutex::new(());
|
||||||
@ -296,7 +297,6 @@ impl ExpertRuntime {
|
|||||||
let config = self
|
let config = self
|
||||||
.config
|
.config
|
||||||
.read()
|
.read()
|
||||||
.expect("experts config rwlock poisoned")
|
|
||||||
.clone();
|
.clone();
|
||||||
let catalog = ExpertCatalog::discover_with_state(
|
let catalog = ExpertCatalog::discover_with_state(
|
||||||
&config,
|
&config,
|
||||||
@ -305,8 +305,7 @@ impl ExpertRuntime {
|
|||||||
);
|
);
|
||||||
let mut guard = self
|
let mut guard = self
|
||||||
.catalog
|
.catalog
|
||||||
.write()
|
.write();
|
||||||
.expect("experts catalog rwlock poisoned");
|
|
||||||
*guard = catalog.clone();
|
*guard = catalog.clone();
|
||||||
Ok(catalog)
|
Ok(catalog)
|
||||||
}
|
}
|
||||||
@ -315,7 +314,7 @@ impl ExpertRuntime {
|
|||||||
/// 用于前端保存配置后即时生效,无需重启网关。
|
/// 用于前端保存配置后即时生效,无需重启网关。
|
||||||
pub fn update_config(&self, new_config: ExpertsConfig) -> Result<(), String> {
|
pub fn update_config(&self, new_config: ExpertsConfig) -> Result<(), String> {
|
||||||
{
|
{
|
||||||
let mut guard = self.config.write().expect("experts config rwlock poisoned");
|
let mut guard = self.config.write();
|
||||||
*guard = new_config;
|
*guard = new_config;
|
||||||
}
|
}
|
||||||
self.reload()?;
|
self.reload()?;
|
||||||
@ -326,7 +325,6 @@ impl ExpertRuntime {
|
|||||||
pub fn list_experts(&self) -> Vec<Expert> {
|
pub fn list_experts(&self) -> Vec<Expert> {
|
||||||
self.catalog
|
self.catalog
|
||||||
.read()
|
.read()
|
||||||
.expect("experts catalog rwlock poisoned")
|
|
||||||
.experts
|
.experts
|
||||||
.clone()
|
.clone()
|
||||||
}
|
}
|
||||||
@ -336,7 +334,6 @@ impl ExpertRuntime {
|
|||||||
let config = self
|
let config = self
|
||||||
.config
|
.config
|
||||||
.read()
|
.read()
|
||||||
.expect("experts config rwlock poisoned")
|
|
||||||
.clone();
|
.clone();
|
||||||
let catalog = ExpertCatalog::discover_without_state(&config, &self.cwd);
|
let catalog = ExpertCatalog::discover_without_state(&config, &self.cwd);
|
||||||
let disable_state = load_expert_disable_state(&self.cwd);
|
let disable_state = load_expert_disable_state(&self.cwd);
|
||||||
@ -366,7 +363,6 @@ impl ExpertRuntime {
|
|||||||
pub fn get_expert(&self, name: &str) -> Option<Expert> {
|
pub fn get_expert(&self, name: &str) -> Option<Expert> {
|
||||||
self.catalog
|
self.catalog
|
||||||
.read()
|
.read()
|
||||||
.expect("experts catalog rwlock poisoned")
|
|
||||||
.find_expert(name)
|
.find_expert(name)
|
||||||
.cloned()
|
.cloned()
|
||||||
}
|
}
|
||||||
@ -481,7 +477,6 @@ impl ExpertRuntime {
|
|||||||
let config = self
|
let config = self
|
||||||
.config
|
.config
|
||||||
.read()
|
.read()
|
||||||
.expect("experts config rwlock poisoned")
|
|
||||||
.clone();
|
.clone();
|
||||||
let catalog = ExpertCatalog::discover_without_state(&config, &self.cwd);
|
let catalog = ExpertCatalog::discover_without_state(&config, &self.cwd);
|
||||||
Ok(catalog.find_expert(name).is_some())
|
Ok(catalog.find_expert(name).is_some())
|
||||||
@ -516,8 +511,7 @@ impl ExpertRuntime {
|
|||||||
{
|
{
|
||||||
let mut state = self
|
let mut state = self
|
||||||
.disable_state
|
.disable_state
|
||||||
.write()
|
.write();
|
||||||
.expect("experts disable_state rwlock poisoned");
|
|
||||||
match scope {
|
match scope {
|
||||||
ExpertScope::User => {
|
ExpertScope::User => {
|
||||||
if enabled {
|
if enabled {
|
||||||
@ -541,8 +535,7 @@ impl ExpertRuntime {
|
|||||||
|
|
||||||
let state = self
|
let state = self
|
||||||
.disable_state
|
.disable_state
|
||||||
.read()
|
.read();
|
||||||
.expect("experts disable_state rwlock poisoned");
|
|
||||||
let disabled_in_scopes = state.disabled_scopes_for(name);
|
let disabled_in_scopes = state.disabled_scopes_for(name);
|
||||||
|
|
||||||
Ok(ExpertAvailabilityChange {
|
Ok(ExpertAvailabilityChange {
|
||||||
@ -567,8 +560,7 @@ impl ExpertRuntime {
|
|||||||
{
|
{
|
||||||
let mut sessions = self
|
let mut sessions = self
|
||||||
.session_experts
|
.session_experts
|
||||||
.write()
|
.write();
|
||||||
.expect("experts session_experts rwlock poisoned");
|
|
||||||
sessions.insert(session_id.to_string(), expert_name.to_string());
|
sessions.insert(session_id.to_string(), expert_name.to_string());
|
||||||
}
|
}
|
||||||
persist_session_experts(&self.cwd, |state| {
|
persist_session_experts(&self.cwd, |state| {
|
||||||
@ -583,8 +575,7 @@ impl ExpertRuntime {
|
|||||||
{
|
{
|
||||||
let mut sessions = self
|
let mut sessions = self
|
||||||
.session_experts
|
.session_experts
|
||||||
.write()
|
.write();
|
||||||
.expect("experts session_experts rwlock poisoned");
|
|
||||||
sessions.remove(session_id);
|
sessions.remove(session_id);
|
||||||
}
|
}
|
||||||
persist_session_experts(&self.cwd, |state| {
|
persist_session_experts(&self.cwd, |state| {
|
||||||
@ -597,16 +588,14 @@ impl ExpertRuntime {
|
|||||||
let name = {
|
let name = {
|
||||||
let sessions = self
|
let sessions = self
|
||||||
.session_experts
|
.session_experts
|
||||||
.read()
|
.read();
|
||||||
.expect("experts session_experts rwlock poisoned");
|
|
||||||
sessions.get(session_id).cloned()
|
sessions.get(session_id).cloned()
|
||||||
}?;
|
}?;
|
||||||
|
|
||||||
// Filter out disabled experts.
|
// Filter out disabled experts.
|
||||||
let state = self
|
let state = self
|
||||||
.disable_state
|
.disable_state
|
||||||
.read()
|
.read();
|
||||||
.expect("experts disable_state rwlock poisoned");
|
|
||||||
if state.is_disabled(&name) {
|
if state.is_disabled(&name) {
|
||||||
return None;
|
return None;
|
||||||
}
|
}
|
||||||
|
|||||||
@ -1,5 +1,5 @@
|
|||||||
use std::collections::HashMap;
|
use std::collections::HashMap;
|
||||||
use std::sync::RwLock;
|
use parking_lot::RwLock;
|
||||||
|
|
||||||
/// per-session 的用户模型覆盖选择存储。
|
/// per-session 的用户模型覆盖选择存储。
|
||||||
///
|
///
|
||||||
@ -19,8 +19,7 @@ impl ModelSelectionStore {
|
|||||||
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
|
let mut selections = self
|
||||||
.selections
|
.selections
|
||||||
.write()
|
.write();
|
||||||
.expect("model selections rwlock poisoned");
|
|
||||||
if provider.is_none() && model.is_none() {
|
if provider.is_none() && model.is_none() {
|
||||||
selections.remove(session_id);
|
selections.remove(session_id);
|
||||||
} else {
|
} else {
|
||||||
@ -32,7 +31,6 @@ impl ModelSelectionStore {
|
|||||||
pub fn get(&self, session_id: &str) -> Option<(Option<String>, Option<String>)> {
|
pub fn get(&self, session_id: &str) -> Option<(Option<String>, Option<String>)> {
|
||||||
self.selections
|
self.selections
|
||||||
.read()
|
.read()
|
||||||
.expect("model selections rwlock poisoned")
|
|
||||||
.get(session_id)
|
.get(session_id)
|
||||||
.cloned()
|
.cloned()
|
||||||
}
|
}
|
||||||
|
|||||||
@ -1,5 +1,6 @@
|
|||||||
use std::collections::HashSet;
|
use std::collections::HashSet;
|
||||||
use std::sync::{Arc, Mutex};
|
use std::sync::Arc;
|
||||||
|
use parking_lot::Mutex;
|
||||||
|
|
||||||
use tokio::sync::Semaphore;
|
use tokio::sync::Semaphore;
|
||||||
|
|
||||||
@ -344,7 +345,7 @@ impl InboundProcessor {
|
|||||||
// 检查并设置"生成中"守卫,防止竞态条件导致重复生成
|
// 检查并设置"生成中"守卫,防止竞态条件导致重复生成
|
||||||
let should_generate = {
|
let should_generate = {
|
||||||
let mut in_flight =
|
let mut in_flight =
|
||||||
self.description_generation_in_flight.lock().unwrap();
|
self.description_generation_in_flight.lock();
|
||||||
if in_flight.contains(topic_id) {
|
if in_flight.contains(topic_id) {
|
||||||
false
|
false
|
||||||
} else {
|
} else {
|
||||||
@ -373,7 +374,7 @@ impl InboundProcessor {
|
|||||||
Some(content) => content,
|
Some(content) => content,
|
||||||
None => {
|
None => {
|
||||||
tracing::warn!(topic_id = %topic_id_clone, "No user message found for topic, skipping description generation");
|
tracing::warn!(topic_id = %topic_id_clone, "No user message found for topic, skipping description generation");
|
||||||
in_flight.lock().unwrap().remove(&topic_id_clone);
|
in_flight.lock().remove(&topic_id_clone);
|
||||||
return;
|
return;
|
||||||
}
|
}
|
||||||
};
|
};
|
||||||
@ -405,7 +406,7 @@ impl InboundProcessor {
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
// 无论成功失败,释放生成守卫
|
// 无论成功失败,释放生成守卫
|
||||||
in_flight.lock().unwrap().remove(&topic_id_clone);
|
in_flight.lock().remove(&topic_id_clone);
|
||||||
});
|
});
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|||||||
@ -60,7 +60,7 @@ pub struct BusToolCallEmitter {
|
|||||||
chat_id: String,
|
chat_id: String,
|
||||||
metadata: HashMap<String, String>,
|
metadata: HashMap<String, String>,
|
||||||
store: Arc<SessionStore>,
|
store: Arc<SessionStore>,
|
||||||
stream_message_id: std::sync::Mutex<Option<String>>,
|
stream_message_id: parking_lot::Mutex<Option<String>>,
|
||||||
}
|
}
|
||||||
|
|
||||||
impl BusToolCallEmitter {
|
impl BusToolCallEmitter {
|
||||||
@ -77,7 +77,7 @@ impl BusToolCallEmitter {
|
|||||||
chat_id: chat_id.into(),
|
chat_id: chat_id.into(),
|
||||||
metadata,
|
metadata,
|
||||||
store,
|
store,
|
||||||
stream_message_id: std::sync::Mutex::new(None),
|
stream_message_id: parking_lot::Mutex::new(None),
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
@ -140,7 +140,7 @@ impl EmittedMessageHandler for BusToolCallEmitter {
|
|||||||
async fn handle_stream_delta(&self, delta: &StreamDelta) {
|
async fn handle_stream_delta(&self, delta: &StreamDelta) {
|
||||||
// Get or create the stream message ID
|
// Get or create the stream message ID
|
||||||
let message_id = {
|
let message_id = {
|
||||||
let mut guard = self.stream_message_id.lock().unwrap();
|
let mut guard = self.stream_message_id.lock();
|
||||||
guard
|
guard
|
||||||
.get_or_insert_with(|| Uuid::new_v4().to_string())
|
.get_or_insert_with(|| Uuid::new_v4().to_string())
|
||||||
.clone()
|
.clone()
|
||||||
@ -180,7 +180,7 @@ impl EmittedMessageHandler for BusToolCallEmitter {
|
|||||||
}
|
}
|
||||||
|
|
||||||
async fn set_stream_message_id(&self, id: &str) {
|
async fn set_stream_message_id(&self, id: &str) {
|
||||||
*self.stream_message_id.lock().unwrap() = Some(id.to_string());
|
*self.stream_message_id.lock() = Some(id.to_string());
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|||||||
@ -7,6 +7,11 @@ use crate::storage::{
|
|||||||
ConversationRepository, SessionRecord, SkillEventRepository, persistent_session_id,
|
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 {
|
fn preview_text(content: &str, max_chars: usize) -> String {
|
||||||
let mut preview = content.chars().take(max_chars).collect::<String>();
|
let mut preview = content.chars().take(max_chars).collect::<String>();
|
||||||
if content.chars().count() > max_chars {
|
if content.chars().count() > max_chars {
|
||||||
@ -19,6 +24,7 @@ pub(crate) struct SessionHistory {
|
|||||||
channel_name: String,
|
channel_name: String,
|
||||||
/// 按 topic_id 键化的内存历史缓存。
|
/// 按 topic_id 键化的内存历史缓存。
|
||||||
/// 不同 topic 的历史独立存储,互不干扰,支持多话题并发执行。
|
/// 不同 topic 的历史独立存储,互不干扰,支持多话题并发执行。
|
||||||
|
/// 超过 `MAX_CACHED_TOPICS` 时自动驱逐非活跃 topic。
|
||||||
topic_histories: HashMap<String, Vec<ChatMessage>>,
|
topic_histories: HashMap<String, Vec<ChatMessage>>,
|
||||||
/// UI 状态:每个 chat 当前活跃的 topic(按 chat_id 键)。
|
/// UI 状态:每个 chat 当前活跃的 topic(按 chat_id 键)。
|
||||||
chat_topic_ids: HashMap<String, String>,
|
chat_topic_ids: HashMap<String, String>,
|
||||||
@ -34,6 +40,56 @@ pub(crate) struct SessionHistory {
|
|||||||
}
|
}
|
||||||
|
|
||||||
impl 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(
|
pub(crate) fn new(
|
||||||
channel_name: impl Into<String>,
|
channel_name: impl Into<String>,
|
||||||
conversations: Arc<dyn ConversationRepository>,
|
conversations: Arc<dyn ConversationRepository>,
|
||||||
@ -103,6 +159,7 @@ impl SessionHistory {
|
|||||||
}
|
}
|
||||||
|
|
||||||
self.topic_histories.insert(tid.to_string(), history);
|
self.topic_histories.insert(tid.to_string(), history);
|
||||||
|
self.evict_inactive_if_needed();
|
||||||
Ok(())
|
Ok(())
|
||||||
}
|
}
|
||||||
|
|
||||||
@ -126,6 +183,7 @@ impl SessionHistory {
|
|||||||
|
|
||||||
pub(crate) fn set_history(&mut self, topic_id: &str, history: Vec<ChatMessage>) {
|
pub(crate) fn set_history(&mut self, topic_id: &str, history: Vec<ChatMessage>) {
|
||||||
self.topic_histories.insert(topic_id.to_string(), history);
|
self.topic_histories.insert(topic_id.to_string(), history);
|
||||||
|
self.evict_inactive_if_needed();
|
||||||
}
|
}
|
||||||
|
|
||||||
/// 设置指定 chat 的当前 topic(UI 状态)
|
/// 设置指定 chat 的当前 topic(UI 状态)
|
||||||
@ -150,6 +208,10 @@ impl SessionHistory {
|
|||||||
pub(crate) fn remove_history(&mut self, topic_id: &str) {
|
pub(crate) fn remove_history(&mut self, topic_id: &str) {
|
||||||
self.topic_histories.remove(topic_id);
|
self.topic_histories.remove(topic_id);
|
||||||
self.compression_in_flight.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 消息。
|
/// 清空指定 chat/topic 的内存历史和 DB 消息。
|
||||||
@ -278,6 +340,7 @@ impl SessionHistory {
|
|||||||
.load_messages_for_topic(topic_id, Some(&sid))
|
.load_messages_for_topic(topic_id, Some(&sid))
|
||||||
.map_err(|err| AgentError::Other(format!("session history reload error: {}", err)))?;
|
.map_err(|err| AgentError::Other(format!("session history reload error: {}", err)))?;
|
||||||
self.topic_histories.insert(topic_id.to_string(), history);
|
self.topic_histories.insert(topic_id.to_string(), history);
|
||||||
|
self.evict_inactive_if_needed();
|
||||||
Ok(())
|
Ok(())
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|||||||
@ -7,7 +7,8 @@
|
|||||||
//! - Dynamically registers MCP tools via the Tool trait adapter
|
//! - Dynamically registers MCP tools via the Tool trait adapter
|
||||||
|
|
||||||
use std::collections::HashMap;
|
use std::collections::HashMap;
|
||||||
use std::sync::{Arc, Mutex};
|
use std::sync::Arc;
|
||||||
|
use parking_lot::Mutex;
|
||||||
use tokio::sync::RwLock;
|
use tokio::sync::RwLock;
|
||||||
|
|
||||||
use http::{HeaderName, HeaderValue};
|
use http::{HeaderName, HeaderValue};
|
||||||
@ -402,7 +403,8 @@ impl McpClientManager {
|
|||||||
);
|
);
|
||||||
}
|
}
|
||||||
// Also collect into the shared buffer (cap at 50 lines)
|
// Also collect into the shared buffer (cap at 50 lines)
|
||||||
if let Ok(mut buf) = stderr_lines_for_task.lock() {
|
{
|
||||||
|
let mut buf = stderr_lines_for_task.lock();
|
||||||
if buf.len() < 50 {
|
if buf.len() < 50 {
|
||||||
buf.push(line);
|
buf.push(line);
|
||||||
}
|
}
|
||||||
@ -414,17 +416,14 @@ impl McpClientManager {
|
|||||||
// Use default client handler (empty tuple)
|
// Use default client handler (empty tuple)
|
||||||
let client = ().serve(transport).await.map_err(|e| {
|
let client = ().serve(transport).await.map_err(|e| {
|
||||||
// Include stderr summary in error if available
|
// Include stderr summary in error if available
|
||||||
let stderr_summary = stderr_lines
|
let stderr_summary = {
|
||||||
.lock()
|
let buf = stderr_lines.lock();
|
||||||
.ok()
|
if buf.is_empty() {
|
||||||
.map(|buf| {
|
String::new()
|
||||||
if buf.is_empty() {
|
} else {
|
||||||
String::new()
|
format!("\nstderr:\n {}", buf.join("\n "))
|
||||||
} else {
|
}
|
||||||
format!("\nstderr:\n {}", buf.join("\n "))
|
};
|
||||||
}
|
|
||||||
})
|
|
||||||
.unwrap_or_default();
|
|
||||||
anyhow::anyhow!(
|
anyhow::anyhow!(
|
||||||
"Failed to establish MCP stdio connection '{}': {}{}",
|
"Failed to establish MCP stdio connection '{}': {}{}",
|
||||||
effective_command.display(),
|
effective_command.display(),
|
||||||
|
|||||||
@ -6,7 +6,8 @@ use serde_json::json;
|
|||||||
use std::collections::{HashMap, HashSet};
|
use std::collections::{HashMap, HashSet};
|
||||||
use std::fs;
|
use std::fs;
|
||||||
use std::path::{Path, PathBuf};
|
use std::path::{Path, PathBuf};
|
||||||
use std::sync::{Arc, RwLock};
|
use std::sync::Arc;
|
||||||
|
use parking_lot::RwLock;
|
||||||
|
|
||||||
#[cfg(test)]
|
#[cfg(test)]
|
||||||
static SKILL_TEST_ENV_LOCK: std::sync::Mutex<()> = std::sync::Mutex::new(());
|
static SKILL_TEST_ENV_LOCK: std::sync::Mutex<()> = std::sync::Mutex::new(());
|
||||||
@ -136,7 +137,7 @@ impl SkillRuntime {
|
|||||||
|
|
||||||
pub fn reload(&self) -> Result<SkillCatalog, String> {
|
pub fn reload(&self) -> Result<SkillCatalog, String> {
|
||||||
let catalog = SkillCatalog::discover(&self.config);
|
let catalog = SkillCatalog::discover(&self.config);
|
||||||
let mut guard = self.catalog.write().expect("skills rwlock poisoned");
|
let mut guard = self.catalog.write();
|
||||||
*guard = catalog.clone();
|
*guard = catalog.clone();
|
||||||
Ok(catalog)
|
Ok(catalog)
|
||||||
}
|
}
|
||||||
@ -144,18 +145,16 @@ impl SkillRuntime {
|
|||||||
pub fn is_empty(&self) -> bool {
|
pub fn is_empty(&self) -> bool {
|
||||||
self.catalog
|
self.catalog
|
||||||
.read()
|
.read()
|
||||||
.expect("skills rwlock poisoned")
|
|
||||||
.is_empty()
|
.is_empty()
|
||||||
}
|
}
|
||||||
|
|
||||||
pub fn len(&self) -> usize {
|
pub fn len(&self) -> usize {
|
||||||
self.catalog.read().expect("skills rwlock poisoned").len()
|
self.catalog.read().len()
|
||||||
}
|
}
|
||||||
|
|
||||||
pub fn system_index_prompt(&self) -> Option<String> {
|
pub fn system_index_prompt(&self) -> Option<String> {
|
||||||
self.catalog
|
self.catalog
|
||||||
.read()
|
.read()
|
||||||
.expect("skills rwlock poisoned")
|
|
||||||
.system_index_prompt()
|
.system_index_prompt()
|
||||||
}
|
}
|
||||||
|
|
||||||
@ -167,42 +166,36 @@ impl SkillRuntime {
|
|||||||
) -> Option<String> {
|
) -> Option<String> {
|
||||||
self.catalog
|
self.catalog
|
||||||
.read()
|
.read()
|
||||||
.expect("skills rwlock poisoned")
|
|
||||||
.system_index_prompt_filtered(allowed, denied)
|
.system_index_prompt_filtered(allowed, denied)
|
||||||
}
|
}
|
||||||
|
|
||||||
pub fn discovery_event_payload(&self) -> serde_json::Value {
|
pub fn discovery_event_payload(&self) -> serde_json::Value {
|
||||||
self.catalog
|
self.catalog
|
||||||
.read()
|
.read()
|
||||||
.expect("skills rwlock poisoned")
|
|
||||||
.discovery_event_payload()
|
.discovery_event_payload()
|
||||||
}
|
}
|
||||||
|
|
||||||
pub fn offered_event_payload(&self) -> serde_json::Value {
|
pub fn offered_event_payload(&self) -> serde_json::Value {
|
||||||
self.catalog
|
self.catalog
|
||||||
.read()
|
.read()
|
||||||
.expect("skills rwlock poisoned")
|
|
||||||
.offered_event_payload()
|
.offered_event_payload()
|
||||||
}
|
}
|
||||||
|
|
||||||
pub fn activation_payload(&self, name: &str) -> Result<String, String> {
|
pub fn activation_payload(&self, name: &str) -> Result<String, String> {
|
||||||
self.catalog
|
self.catalog
|
||||||
.read()
|
.read()
|
||||||
.expect("skills rwlock poisoned")
|
|
||||||
.activation_payload(name)
|
.activation_payload(name)
|
||||||
}
|
}
|
||||||
|
|
||||||
pub fn activation_event_payload(&self, name: &str) -> Result<serde_json::Value, String> {
|
pub fn activation_event_payload(&self, name: &str) -> Result<serde_json::Value, String> {
|
||||||
self.catalog
|
self.catalog
|
||||||
.read()
|
.read()
|
||||||
.expect("skills rwlock poisoned")
|
|
||||||
.activation_event_payload(name)
|
.activation_event_payload(name)
|
||||||
}
|
}
|
||||||
|
|
||||||
pub fn list_skills(&self) -> Vec<Skill> {
|
pub fn list_skills(&self) -> Vec<Skill> {
|
||||||
self.catalog
|
self.catalog
|
||||||
.read()
|
.read()
|
||||||
.expect("skills rwlock poisoned")
|
|
||||||
.skills
|
.skills
|
||||||
.clone()
|
.clone()
|
||||||
}
|
}
|
||||||
@ -235,7 +228,6 @@ impl SkillRuntime {
|
|||||||
pub fn get_skill(&self, name: &str) -> Option<Skill> {
|
pub fn get_skill(&self, name: &str) -> Option<Skill> {
|
||||||
self.catalog
|
self.catalog
|
||||||
.read()
|
.read()
|
||||||
.expect("skills rwlock poisoned")
|
|
||||||
.find_skill(name)
|
.find_skill(name)
|
||||||
.cloned()
|
.cloned()
|
||||||
}
|
}
|
||||||
|
|||||||
@ -149,113 +149,121 @@ impl Tool for FileReadTool {
|
|||||||
});
|
});
|
||||||
}
|
}
|
||||||
|
|
||||||
// Try to read as text
|
// 文件读取与后续行处理/截断/base64 编码均为阻塞操作(大文件可能数 MB),
|
||||||
match std::fs::read_to_string(&resolved) {
|
// 统一放到 blocking 线程池执行,避免阻塞 tokio worker。
|
||||||
Ok(content) => {
|
let result = tokio::task::spawn_blocking(move || {
|
||||||
let all_lines: Vec<&str> = content.lines().collect();
|
// Try to read as text
|
||||||
let total = all_lines.len();
|
match std::fs::read_to_string(&resolved) {
|
||||||
|
Ok(content) => {
|
||||||
|
let all_lines: Vec<&str> = content.lines().collect();
|
||||||
|
let total = all_lines.len();
|
||||||
|
|
||||||
if offset < 1 {
|
if offset < 1 {
|
||||||
return Ok(ToolResult {
|
return ToolResult {
|
||||||
success: false,
|
success: false,
|
||||||
output: String::new(),
|
output: String::new(),
|
||||||
error: Some(format!("offset must be at least 1, got {}", offset)),
|
error: Some(format!("offset must be at least 1, got {}", offset)),
|
||||||
});
|
};
|
||||||
}
|
}
|
||||||
|
|
||||||
if offset > total {
|
if offset > total {
|
||||||
return Ok(ToolResult {
|
return ToolResult {
|
||||||
success: false,
|
success: false,
|
||||||
output: String::new(),
|
output: String::new(),
|
||||||
error: Some(format!(
|
error: Some(format!(
|
||||||
"offset {} is beyond end of file ({} lines)",
|
"offset {} is beyond end of file ({} lines)",
|
||||||
offset, total
|
offset, total
|
||||||
)),
|
)),
|
||||||
});
|
};
|
||||||
}
|
}
|
||||||
|
|
||||||
let start = offset - 1;
|
let start = offset - 1;
|
||||||
let end = std::cmp::min(start + limit, total);
|
let end = std::cmp::min(start + limit, total);
|
||||||
let lines: Vec<String> = all_lines[start..end]
|
let lines: Vec<String> = all_lines[start..end]
|
||||||
.iter()
|
.iter()
|
||||||
.enumerate()
|
.enumerate()
|
||||||
.map(|(i, line)| format!("{}| {}", start + i + 1, line))
|
.map(|(i, line)| format!("{}| {}", start + i + 1, line))
|
||||||
.collect();
|
.collect();
|
||||||
|
|
||||||
let mut result = lines.join("\n");
|
let mut result = lines.join("\n");
|
||||||
|
|
||||||
// Truncate if too long
|
// Truncate if too long
|
||||||
if result.len() > MAX_CHARS {
|
if result.len() > MAX_CHARS {
|
||||||
let original_len = result.len();
|
let original_len = result.len();
|
||||||
let mut truncated_chars = 0;
|
let mut truncated_chars = 0;
|
||||||
let mut end_idx = 0;
|
let mut end_idx = 0;
|
||||||
for (i, line) in lines.iter().enumerate() {
|
for (i, line) in lines.iter().enumerate() {
|
||||||
truncated_chars += line.len() + 1;
|
truncated_chars += line.len() + 1;
|
||||||
if truncated_chars > MAX_CHARS {
|
if truncated_chars > MAX_CHARS {
|
||||||
end_idx = i;
|
end_idx = i;
|
||||||
break;
|
break;
|
||||||
|
}
|
||||||
|
end_idx = i + 1;
|
||||||
}
|
}
|
||||||
end_idx = i + 1;
|
if end_idx == 0 && !lines.is_empty() {
|
||||||
|
// First line alone exceeds MAX_CHARS — take its prefix
|
||||||
|
result = take_prefix_chars(&lines[0], MAX_CHARS.saturating_sub(100));
|
||||||
|
} else {
|
||||||
|
result = lines[..end_idx].join("\n");
|
||||||
|
}
|
||||||
|
let truncated_amount = original_len.saturating_sub(result.len());
|
||||||
|
result.push_str(&format!(
|
||||||
|
"\n\n... ({} chars truncated) ...",
|
||||||
|
truncated_amount
|
||||||
|
));
|
||||||
}
|
}
|
||||||
if end_idx == 0 && !lines.is_empty() {
|
|
||||||
// First line alone exceeds MAX_CHARS — take its prefix
|
if end < total {
|
||||||
result = take_prefix_chars(&lines[0], MAX_CHARS.saturating_sub(100));
|
result.push_str(&format!(
|
||||||
|
"\n\n(Showing lines {}-{} of {}. Use offset={} to continue.)",
|
||||||
|
offset,
|
||||||
|
end,
|
||||||
|
total,
|
||||||
|
end + 1
|
||||||
|
));
|
||||||
} else {
|
} else {
|
||||||
result = lines[..end_idx].join("\n");
|
result.push_str(&format!("\n\n(End of file — {} lines total)", total));
|
||||||
}
|
}
|
||||||
let truncated_amount = original_len.saturating_sub(result.len());
|
|
||||||
result.push_str(&format!(
|
|
||||||
"\n\n... ({} chars truncated) ...",
|
|
||||||
truncated_amount
|
|
||||||
));
|
|
||||||
}
|
|
||||||
|
|
||||||
if end < total {
|
ToolResult {
|
||||||
result.push_str(&format!(
|
success: true,
|
||||||
"\n\n(Showing lines {}-{} of {}. Use offset={} to continue.)",
|
output: result,
|
||||||
offset,
|
error: None,
|
||||||
end,
|
}
|
||||||
total,
|
}
|
||||||
end + 1
|
Err(e) => {
|
||||||
));
|
// Try to read as binary and encode as base64
|
||||||
} else {
|
match std::fs::read(&resolved) {
|
||||||
result.push_str(&format!("\n\n(End of file — {} lines total)", total));
|
Ok(bytes) => {
|
||||||
}
|
use base64::{Engine, engine::general_purpose::STANDARD};
|
||||||
|
let encoded = STANDARD.encode(&bytes);
|
||||||
Ok(ToolResult {
|
let mime = mime_guess::from_path(&resolved)
|
||||||
success: true,
|
.first_or_octet_stream()
|
||||||
output: result,
|
.to_string();
|
||||||
error: None,
|
ToolResult {
|
||||||
})
|
success: true,
|
||||||
}
|
output: format!(
|
||||||
Err(e) => {
|
"(Binary file: {}, {} bytes, base64 encoded)\n{}",
|
||||||
// Try to read as binary and encode as base64
|
mime,
|
||||||
match std::fs::read(&resolved) {
|
bytes.len(),
|
||||||
Ok(bytes) => {
|
encoded
|
||||||
use base64::{Engine, engine::general_purpose::STANDARD};
|
),
|
||||||
let encoded = STANDARD.encode(&bytes);
|
error: None,
|
||||||
let mime = mime_guess::from_path(&resolved)
|
}
|
||||||
.first_or_octet_stream()
|
}
|
||||||
.to_string();
|
Err(_) => ToolResult {
|
||||||
Ok(ToolResult {
|
success: false,
|
||||||
success: true,
|
output: String::new(),
|
||||||
output: format!(
|
error: Some(format!("Failed to read file: {}", e)),
|
||||||
"(Binary file: {}, {} bytes, base64 encoded)\n{}",
|
},
|
||||||
mime,
|
|
||||||
bytes.len(),
|
|
||||||
encoded
|
|
||||||
),
|
|
||||||
error: None,
|
|
||||||
})
|
|
||||||
}
|
}
|
||||||
Err(_) => Ok(ToolResult {
|
|
||||||
success: false,
|
|
||||||
output: String::new(),
|
|
||||||
error: Some(format!("Failed to read file: {}", e)),
|
|
||||||
}),
|
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
}
|
})
|
||||||
|
.await
|
||||||
|
.map_err(|e| anyhow::anyhow!("file_read blocking task failed: {}", e))?;
|
||||||
|
|
||||||
|
Ok(result)
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|||||||
@ -1,5 +1,6 @@
|
|||||||
use std::collections::HashMap;
|
use std::collections::HashMap;
|
||||||
use std::sync::{Arc, RwLock};
|
use std::sync::Arc;
|
||||||
|
use parking_lot::RwLock;
|
||||||
|
|
||||||
use crate::domain::tools::{Tool, ToolFunction};
|
use crate::domain::tools::{Tool, ToolFunction};
|
||||||
|
|
||||||
@ -19,14 +20,12 @@ impl ToolRegistry {
|
|||||||
pub fn register<T: ToolTrait + 'static>(&self, tool: T) {
|
pub fn register<T: ToolTrait + 'static>(&self, tool: T) {
|
||||||
self.tools
|
self.tools
|
||||||
.write()
|
.write()
|
||||||
.expect("ToolRegistry lock poisoned")
|
|
||||||
.insert(tool.name().to_string(), Arc::new(tool));
|
.insert(tool.name().to_string(), Arc::new(tool));
|
||||||
}
|
}
|
||||||
|
|
||||||
pub fn get(&self, name: &str) -> Option<Arc<dyn ToolTrait>> {
|
pub fn get(&self, name: &str) -> Option<Arc<dyn ToolTrait>> {
|
||||||
self.tools
|
self.tools
|
||||||
.read()
|
.read()
|
||||||
.expect("ToolRegistry lock poisoned")
|
|
||||||
.get(name)
|
.get(name)
|
||||||
.cloned()
|
.cloned()
|
||||||
}
|
}
|
||||||
@ -36,7 +35,6 @@ impl ToolRegistry {
|
|||||||
pub fn get_all(&self) -> Vec<Arc<dyn ToolTrait>> {
|
pub fn get_all(&self) -> Vec<Arc<dyn ToolTrait>> {
|
||||||
self.tools
|
self.tools
|
||||||
.read()
|
.read()
|
||||||
.expect("ToolRegistry lock poisoned")
|
|
||||||
.values()
|
.values()
|
||||||
.cloned()
|
.cloned()
|
||||||
.collect()
|
.collect()
|
||||||
@ -45,7 +43,6 @@ impl ToolRegistry {
|
|||||||
pub fn get_definitions(&self) -> Vec<Tool> {
|
pub fn get_definitions(&self) -> Vec<Tool> {
|
||||||
self.tools
|
self.tools
|
||||||
.read()
|
.read()
|
||||||
.expect("ToolRegistry lock poisoned")
|
|
||||||
.values()
|
.values()
|
||||||
.map(|tool| Tool {
|
.map(|tool| Tool {
|
||||||
tool_type: "function".to_string(),
|
tool_type: "function".to_string(),
|
||||||
@ -62,14 +59,12 @@ impl ToolRegistry {
|
|||||||
!self
|
!self
|
||||||
.tools
|
.tools
|
||||||
.read()
|
.read()
|
||||||
.expect("ToolRegistry lock poisoned")
|
|
||||||
.is_empty()
|
.is_empty()
|
||||||
}
|
}
|
||||||
|
|
||||||
pub fn tool_names(&self) -> Vec<String> {
|
pub fn tool_names(&self) -> Vec<String> {
|
||||||
self.tools
|
self.tools
|
||||||
.read()
|
.read()
|
||||||
.expect("ToolRegistry lock poisoned")
|
|
||||||
.keys()
|
.keys()
|
||||||
.cloned()
|
.cloned()
|
||||||
.collect()
|
.collect()
|
||||||
@ -78,7 +73,7 @@ impl ToolRegistry {
|
|||||||
/// 创建一个排除指定工具的新 registry 副本
|
/// 创建一个排除指定工具的新 registry 副本
|
||||||
pub fn without(&self, exclude: &[&str]) -> Self {
|
pub fn without(&self, exclude: &[&str]) -> Self {
|
||||||
let exclude_set: std::collections::HashSet<&str> = exclude.iter().copied().collect();
|
let exclude_set: std::collections::HashSet<&str> = exclude.iter().copied().collect();
|
||||||
let tools = self.tools.read().expect("ToolRegistry lock poisoned");
|
let tools = self.tools.read();
|
||||||
let filtered: HashMap<String, Arc<dyn ToolTrait>> = tools
|
let filtered: HashMap<String, Arc<dyn ToolTrait>> = tools
|
||||||
.iter()
|
.iter()
|
||||||
.filter(|(name, _)| !exclude_set.contains(name.as_str()))
|
.filter(|(name, _)| !exclude_set.contains(name.as_str()))
|
||||||
@ -87,8 +82,7 @@ impl ToolRegistry {
|
|||||||
let new_registry = ToolRegistry::new();
|
let new_registry = ToolRegistry::new();
|
||||||
*new_registry
|
*new_registry
|
||||||
.tools
|
.tools
|
||||||
.write()
|
.write() = filtered;
|
||||||
.expect("ToolRegistry lock poisoned") = filtered;
|
|
||||||
new_registry
|
new_registry
|
||||||
}
|
}
|
||||||
|
|
||||||
@ -96,7 +90,7 @@ impl ToolRegistry {
|
|||||||
/// include 中不存在于当前 registry 的名称会被静默跳过(取交集语义)。
|
/// include 中不存在于当前 registry 的名称会被静默跳过(取交集语义)。
|
||||||
pub fn only(&self, include: &[&str]) -> Self {
|
pub fn only(&self, include: &[&str]) -> Self {
|
||||||
let include_set: std::collections::HashSet<&str> = include.iter().copied().collect();
|
let include_set: std::collections::HashSet<&str> = include.iter().copied().collect();
|
||||||
let tools = self.tools.read().expect("ToolRegistry lock poisoned");
|
let tools = self.tools.read();
|
||||||
let filtered: HashMap<String, Arc<dyn ToolTrait>> = tools
|
let filtered: HashMap<String, Arc<dyn ToolTrait>> = tools
|
||||||
.iter()
|
.iter()
|
||||||
.filter(|(name, _)| include_set.contains(name.as_str()))
|
.filter(|(name, _)| include_set.contains(name.as_str()))
|
||||||
@ -105,8 +99,7 @@ impl ToolRegistry {
|
|||||||
let new_registry = ToolRegistry::new();
|
let new_registry = ToolRegistry::new();
|
||||||
*new_registry
|
*new_registry
|
||||||
.tools
|
.tools
|
||||||
.write()
|
.write() = filtered;
|
||||||
.expect("ToolRegistry lock poisoned") = filtered;
|
|
||||||
new_registry
|
new_registry
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|||||||
@ -1,5 +1,5 @@
|
|||||||
use std::collections::HashMap;
|
use std::collections::HashMap;
|
||||||
use std::sync::RwLock;
|
use parking_lot::RwLock;
|
||||||
|
|
||||||
use async_trait::async_trait;
|
use async_trait::async_trait;
|
||||||
|
|
||||||
@ -65,18 +65,17 @@ impl TaskRepository for InMemoryTaskRepository {
|
|||||||
);
|
);
|
||||||
self.sessions
|
self.sessions
|
||||||
.write()
|
.write()
|
||||||
.unwrap()
|
|
||||||
.insert(session.id.clone(), session.clone());
|
.insert(session.id.clone(), session.clone());
|
||||||
tracing::debug!(
|
tracing::debug!(
|
||||||
task_id = %session.id,
|
task_id = %session.id,
|
||||||
total_tasks = self.sessions.read().unwrap().len(),
|
total_tasks = self.sessions.read().len(),
|
||||||
"Task session saved, current repository size"
|
"Task session saved, current repository size"
|
||||||
);
|
);
|
||||||
Ok(())
|
Ok(())
|
||||||
}
|
}
|
||||||
|
|
||||||
async fn load_task_session(&self, task_id: &str) -> Result<Option<TaskSession>, StorageError> {
|
async fn load_task_session(&self, task_id: &str) -> Result<Option<TaskSession>, StorageError> {
|
||||||
let sessions = self.sessions.read().unwrap();
|
let sessions = self.sessions.read();
|
||||||
let total = sessions.len();
|
let total = sessions.len();
|
||||||
let keys: Vec<&str> = sessions.keys().map(|k| k.as_str()).collect();
|
let keys: Vec<&str> = sessions.keys().map(|k| k.as_str()).collect();
|
||||||
tracing::debug!(
|
tracing::debug!(
|
||||||
@ -89,7 +88,7 @@ impl TaskRepository for InMemoryTaskRepository {
|
|||||||
}
|
}
|
||||||
|
|
||||||
async fn delete_task_session(&self, task_id: &str) -> Result<bool, StorageError> {
|
async fn delete_task_session(&self, task_id: &str) -> Result<bool, StorageError> {
|
||||||
Ok(self.sessions.write().unwrap().remove(task_id).is_some())
|
Ok(self.sessions.write().remove(task_id).is_some())
|
||||||
}
|
}
|
||||||
|
|
||||||
async fn list_tasks_for_session(
|
async fn list_tasks_for_session(
|
||||||
@ -99,7 +98,6 @@ impl TaskRepository for InMemoryTaskRepository {
|
|||||||
Ok(self
|
Ok(self
|
||||||
.sessions
|
.sessions
|
||||||
.read()
|
.read()
|
||||||
.unwrap()
|
|
||||||
.values()
|
.values()
|
||||||
.filter(|s| s.parent_session_id == parent_session_id)
|
.filter(|s| s.parent_session_id == parent_session_id)
|
||||||
.cloned()
|
.cloned()
|
||||||
@ -113,7 +111,6 @@ impl TaskRepository for InMemoryTaskRepository {
|
|||||||
Ok(self
|
Ok(self
|
||||||
.sessions
|
.sessions
|
||||||
.read()
|
.read()
|
||||||
.unwrap()
|
|
||||||
.values()
|
.values()
|
||||||
.filter(|s| s.parent_topic_id.as_ref() == Some(&parent_topic_id.to_string()))
|
.filter(|s| s.parent_topic_id.as_ref() == Some(&parent_topic_id.to_string()))
|
||||||
.cloned()
|
.cloned()
|
||||||
@ -123,7 +120,7 @@ impl TaskRepository for InMemoryTaskRepository {
|
|||||||
async fn cleanup_expired_tasks(&self, ttl_hours: u64) -> Result<usize, StorageError> {
|
async fn cleanup_expired_tasks(&self, ttl_hours: u64) -> Result<usize, StorageError> {
|
||||||
let now = current_timestamp();
|
let now = current_timestamp();
|
||||||
let ttl_millis = ttl_hours * 3600 * 1000;
|
let ttl_millis = ttl_hours * 3600 * 1000;
|
||||||
let mut sessions = self.sessions.write().unwrap();
|
let mut sessions = self.sessions.write();
|
||||||
let before = sessions.len();
|
let before = sessions.len();
|
||||||
sessions.retain(|_, s| now - s.updated_at < ttl_millis as i64);
|
sessions.retain(|_, s| now - s.updated_at < ttl_millis as i64);
|
||||||
Ok(before - sessions.len())
|
Ok(before - sessions.len())
|
||||||
|
|||||||
@ -1,7 +1,8 @@
|
|||||||
use std::collections::{HashMap, HashSet};
|
use std::collections::{HashMap, HashSet};
|
||||||
use std::fs;
|
use std::fs;
|
||||||
use std::path::{Path, PathBuf};
|
use std::path::{Path, PathBuf};
|
||||||
use std::sync::{Arc, RwLock};
|
use std::sync::Arc;
|
||||||
|
use parking_lot::RwLock;
|
||||||
use std::time::Duration;
|
use std::time::Duration;
|
||||||
|
|
||||||
use async_trait::async_trait;
|
use async_trait::async_trait;
|
||||||
@ -113,7 +114,7 @@ struct SubAgentEmitter {
|
|||||||
store: Arc<SessionStore>,
|
store: Arc<SessionStore>,
|
||||||
/// 子/孙智能体自身的 task_id,用于持久化时作为 scope_key
|
/// 子/孙智能体自身的 task_id,用于持久化时作为 scope_key
|
||||||
task_id: String,
|
task_id: String,
|
||||||
stream_message_id: std::sync::Mutex<Option<String>>,
|
stream_message_id: parking_lot::Mutex<Option<String>>,
|
||||||
}
|
}
|
||||||
|
|
||||||
#[async_trait]
|
#[async_trait]
|
||||||
@ -173,7 +174,7 @@ impl EmittedMessageHandler for SubAgentEmitter {
|
|||||||
|
|
||||||
async fn handle_stream_delta(&self, delta: &StreamDelta) {
|
async fn handle_stream_delta(&self, delta: &StreamDelta) {
|
||||||
let message_id = {
|
let message_id = {
|
||||||
let mut guard = self.stream_message_id.lock().unwrap();
|
let mut guard = self.stream_message_id.lock();
|
||||||
guard
|
guard
|
||||||
.get_or_insert_with(|| uuid::Uuid::new_v4().to_string())
|
.get_or_insert_with(|| uuid::Uuid::new_v4().to_string())
|
||||||
.clone()
|
.clone()
|
||||||
@ -212,7 +213,7 @@ impl EmittedMessageHandler for SubAgentEmitter {
|
|||||||
}
|
}
|
||||||
|
|
||||||
async fn set_stream_message_id(&self, id: &str) {
|
async fn set_stream_message_id(&self, id: &str) {
|
||||||
*self.stream_message_id.lock().unwrap() = Some(id.to_string());
|
*self.stream_message_id.lock() = Some(id.to_string());
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
@ -548,7 +549,7 @@ impl DefaultSubAgentRuntime {
|
|||||||
metadata,
|
metadata,
|
||||||
store: self.store.clone(),
|
store: self.store.clone(),
|
||||||
task_id: session.id.clone(),
|
task_id: session.id.clone(),
|
||||||
stream_message_id: std::sync::Mutex::new(None),
|
stream_message_id: parking_lot::Mutex::new(None),
|
||||||
},
|
},
|
||||||
self.conversation_repository.clone(),
|
self.conversation_repository.clone(),
|
||||||
session.session_id.clone(),
|
session.session_id.clone(),
|
||||||
@ -1298,7 +1299,7 @@ impl SubagentRuntime {
|
|||||||
let mut guard = self
|
let mut guard = self
|
||||||
.catalog
|
.catalog
|
||||||
.write()
|
.write()
|
||||||
.expect("subagent catalog rwlock poisoned");
|
;
|
||||||
*guard = new_catalog;
|
*guard = new_catalog;
|
||||||
Ok(())
|
Ok(())
|
||||||
}
|
}
|
||||||
@ -1308,11 +1309,11 @@ impl SubagentRuntime {
|
|||||||
let state = self
|
let state = self
|
||||||
.disable_state
|
.disable_state
|
||||||
.read()
|
.read()
|
||||||
.expect("subagent state rwlock poisoned");
|
;
|
||||||
let catalog = self
|
let catalog = self
|
||||||
.catalog
|
.catalog
|
||||||
.read()
|
.read()
|
||||||
.expect("subagent catalog rwlock poisoned");
|
;
|
||||||
let mut items: Vec<SubagentWithStatus> = catalog
|
let mut items: Vec<SubagentWithStatus> = catalog
|
||||||
.all()
|
.all()
|
||||||
.iter()
|
.iter()
|
||||||
@ -1339,11 +1340,11 @@ impl SubagentRuntime {
|
|||||||
let state = self
|
let state = self
|
||||||
.disable_state
|
.disable_state
|
||||||
.read()
|
.read()
|
||||||
.expect("subagent state rwlock poisoned");
|
;
|
||||||
let catalog = self
|
let catalog = self
|
||||||
.catalog
|
.catalog
|
||||||
.read()
|
.read()
|
||||||
.expect("subagent catalog rwlock poisoned");
|
;
|
||||||
catalog
|
catalog
|
||||||
.names()
|
.names()
|
||||||
.into_iter()
|
.into_iter()
|
||||||
@ -1356,13 +1357,13 @@ impl SubagentRuntime {
|
|||||||
let state = self
|
let state = self
|
||||||
.disable_state
|
.disable_state
|
||||||
.read()
|
.read()
|
||||||
.expect("subagent state rwlock poisoned");
|
;
|
||||||
if state.is_disabled(name) {
|
if state.is_disabled(name) {
|
||||||
return None;
|
return None;
|
||||||
}
|
}
|
||||||
self.catalog
|
self.catalog
|
||||||
.read()
|
.read()
|
||||||
.expect("subagent catalog rwlock poisoned")
|
|
||||||
.find(name)
|
.find(name)
|
||||||
.cloned()
|
.cloned()
|
||||||
}
|
}
|
||||||
@ -1372,11 +1373,11 @@ impl SubagentRuntime {
|
|||||||
let state = self
|
let state = self
|
||||||
.disable_state
|
.disable_state
|
||||||
.read()
|
.read()
|
||||||
.expect("subagent state rwlock poisoned");
|
;
|
||||||
let catalog = self
|
let catalog = self
|
||||||
.catalog
|
.catalog
|
||||||
.read()
|
.read()
|
||||||
.expect("subagent catalog rwlock poisoned");
|
;
|
||||||
let available_defs: Vec<&SubagentDef> = catalog
|
let available_defs: Vec<&SubagentDef> = catalog
|
||||||
.all()
|
.all()
|
||||||
.into_iter()
|
.into_iter()
|
||||||
@ -1417,11 +1418,11 @@ impl SubagentRuntime {
|
|||||||
let state = self
|
let state = self
|
||||||
.disable_state
|
.disable_state
|
||||||
.read()
|
.read()
|
||||||
.expect("subagent state rwlock poisoned");
|
;
|
||||||
let catalog = self
|
let catalog = self
|
||||||
.catalog
|
.catalog
|
||||||
.read()
|
.read()
|
||||||
.expect("subagent catalog rwlock poisoned");
|
;
|
||||||
let available_defs: Vec<&SubagentDef> = catalog
|
let available_defs: Vec<&SubagentDef> = catalog
|
||||||
.all()
|
.all()
|
||||||
.into_iter()
|
.into_iter()
|
||||||
@ -1487,7 +1488,7 @@ impl SubagentRuntime {
|
|||||||
if self
|
if self
|
||||||
.catalog
|
.catalog
|
||||||
.read()
|
.read()
|
||||||
.expect("subagent catalog rwlock poisoned")
|
|
||||||
.find(name)
|
.find(name)
|
||||||
.is_none()
|
.is_none()
|
||||||
{
|
{
|
||||||
@ -1514,7 +1515,7 @@ impl SubagentRuntime {
|
|||||||
let mut state = self
|
let mut state = self
|
||||||
.disable_state
|
.disable_state
|
||||||
.write()
|
.write()
|
||||||
.expect("subagent state rwlock poisoned");
|
;
|
||||||
match scope {
|
match scope {
|
||||||
SubagentScope::User => {
|
SubagentScope::User => {
|
||||||
if enabled {
|
if enabled {
|
||||||
@ -1537,7 +1538,7 @@ impl SubagentRuntime {
|
|||||||
let state = self
|
let state = self
|
||||||
.disable_state
|
.disable_state
|
||||||
.read()
|
.read()
|
||||||
.expect("subagent state rwlock poisoned");
|
;
|
||||||
let disabled_in_scopes = state.disabled_scopes_for(name);
|
let disabled_in_scopes = state.disabled_scopes_for(name);
|
||||||
|
|
||||||
Ok(SubagentAvailabilityChange {
|
Ok(SubagentAvailabilityChange {
|
||||||
@ -1568,7 +1569,7 @@ impl SubagentRuntime {
|
|||||||
let catalog = self
|
let catalog = self
|
||||||
.catalog
|
.catalog
|
||||||
.read()
|
.read()
|
||||||
.expect("subagent catalog rwlock poisoned");
|
;
|
||||||
catalog
|
catalog
|
||||||
.find(name)
|
.find(name)
|
||||||
.ok_or_else(|| format!("subagent '{}' not found", name))?
|
.ok_or_else(|| format!("subagent '{}' not found", name))?
|
||||||
@ -1634,7 +1635,7 @@ impl SubagentRuntime {
|
|||||||
let catalog = self
|
let catalog = self
|
||||||
.catalog
|
.catalog
|
||||||
.read()
|
.read()
|
||||||
.expect("subagent catalog rwlock poisoned");
|
;
|
||||||
if catalog.find(name).is_some() {
|
if catalog.find(name).is_some() {
|
||||||
return Err(format!("subagent '{}' already exists", name));
|
return Err(format!("subagent '{}' already exists", name));
|
||||||
}
|
}
|
||||||
@ -1688,7 +1689,7 @@ impl SubagentRuntime {
|
|||||||
let catalog = self
|
let catalog = self
|
||||||
.catalog
|
.catalog
|
||||||
.read()
|
.read()
|
||||||
.expect("subagent catalog rwlock poisoned");
|
;
|
||||||
let def = catalog
|
let def = catalog
|
||||||
.find(name)
|
.find(name)
|
||||||
.ok_or_else(|| format!("subagent '{}' not found", name))?;
|
.ok_or_else(|| format!("subagent '{}' not found", name))?;
|
||||||
|
|||||||
Loading…
x
Reference in New Issue
Block a user