perf(agent-loop): 缓存工具定义与轻量 token 估算,DB 写入 offload 至 blocking 池

- 工具定义在 process() 循环外构建一次(快照语义:单次 process 内 Arc<ToolRegistry> 不变)
- 工具 token 估算循环外预算一次,避免每轮 serde_json::to_string 全量序列化
- 新增 estimate_text_tokens_from_chat_messages 直接基于 ChatMessage 累加字段长度,
  消除中间 text_only_messages: Vec<Message> 构造(N 条消息 × 5 字段双克隆 → 0)
- PersistingEmittedMessageHandler 用 spawn_blocking 包装 SQLite 同步写,
  解除 tokio worker 线程阻塞,不改变失败语义(仅记日志不阻断主流程)
This commit is contained in:
oudecheng 2026-07-31 08:22:58 +08:00
parent 649ea2998d
commit 92db80dc3f

View File

@ -155,10 +155,38 @@ fn estimate_tokens_from_serialized_json<T: serde::Serialize>(value: &T) -> usize
as usize as usize
} }
/// 轻量 token 估算:直接从 ChatMessage 累加文本字段长度,
/// 避免对完整 Vec<Message> 做 serde_json::to_stringO(N) 序列化)。
/// 精度略低于完整 JSON 序列化(忽略 role/tool_call_id 等 overhead
/// 但作为图片预算估算已足够——上下文窗口有 CONTEXT_INPUT_SAFETY_RATIO 兜底。
/// 与原 text_only_messages 路径等价覆盖content + reasoning_content + tool_calls。
fn estimate_text_tokens_from_chat_messages(
messages: &[ChatMessage],
system_prompt: Option<&str>,
) -> usize {
let mut total_chars: usize = system_prompt.map(|s| s.len()).unwrap_or(0);
for m in messages {
total_chars += m.content.len();
if let Some(rc) = m.reasoning_content.as_ref() {
total_chars += rc.len();
}
if let Some(tcs) = m.tool_calls.as_ref() {
for tc in tcs {
total_chars += tc.name.len();
total_chars += serde_json::to_string(&tc.arguments)
.map(|s| s.len())
.unwrap_or(0);
}
}
}
((total_chars.div_ceil(TOKEN_ESTIMATE_CHARS_PER_TOKEN) as f64)
* TOKEN_ESTIMATE_SAFETY_MULTIPLIER) as usize
}
fn image_token_budget_for_request( fn image_token_budget_for_request(
runtime_config: &AgentRuntimeConfig, runtime_config: &AgentRuntimeConfig,
text_only_messages: &[Message], text_tokens: usize,
tools: Option<&Vec<crate::domain::tools::Tool>>, tools_tokens: usize,
) -> usize { ) -> usize {
let completion_reserve = runtime_config let completion_reserve = runtime_config
.provider .provider
@ -170,12 +198,7 @@ fn image_token_budget_for_request(
.saturating_sub(completion_reserve); .saturating_sub(completion_reserve);
let safe_input_window = (input_window as f64 * CONTEXT_INPUT_SAFETY_RATIO) as usize; let safe_input_window = (input_window as f64 * CONTEXT_INPUT_SAFETY_RATIO) as usize;
let text_tokens = estimate_tokens_from_serialized_json(&text_only_messages) safe_input_window.saturating_sub(text_tokens + tools_tokens)
+ tools
.map(estimate_tokens_from_serialized_json)
.unwrap_or_default();
safe_input_window.saturating_sub(text_tokens)
} }
fn count_supported_image_media_refs(messages: &[ChatMessage]) -> usize { fn count_supported_image_media_refs(messages: &[ChatMessage]) -> usize {
@ -700,23 +723,46 @@ impl<H: EmittedMessageHandler> PersistingEmittedMessageHandler<H> {
#[async_trait] #[async_trait]
impl<H: EmittedMessageHandler> EmittedMessageHandler for PersistingEmittedMessageHandler<H> { impl<H: EmittedMessageHandler> EmittedMessageHandler for PersistingEmittedMessageHandler<H> {
async fn handle(&self, message: ChatMessage) { async fn handle(&self, message: ChatMessage) {
if let Err(e) = self.conversation_repository // SQLite 是同步 I/O直接调用会阻塞 tokio worker 线程。
.append_message_with_topic(&self.session_id, self.topic_id.as_deref(), &message) // 用 spawn_blocking 把写操作挪到 blocking 线程池,不改变语义:
{ // 失败仅记日志不阻断主流程(与原 if let Err 一致)。
tracing::error!(error = %e, session_id = %self.session_id, let repo = self.conversation_repository.clone();
"Failed to persist emitted message"); let session_id = self.session_id.clone();
} let topic_id = self.topic_id.clone();
let msg_for_persist = message.clone();
tokio::task::spawn_blocking(move || {
if let Err(e) = repo.append_message_with_topic(
&session_id,
topic_id.as_deref(),
&msg_for_persist,
) {
tracing::error!(error = %e, session_id = %session_id,
"Failed to persist emitted message");
}
})
.await
.ok(); // JoinError 不影响主流程
self.inner.handle(message).await; self.inner.handle(message).await;
} }
async fn handle_tool_result(&self, message: ChatMessage, duration_ms: Option<u64>) { async fn handle_tool_result(&self, message: ChatMessage, duration_ms: Option<u64>) {
// Persist the ChatMessage first (no duration field, same as before) // Persist the ChatMessage first (no duration field, same as before)
if let Err(e) = self.conversation_repository let repo = self.conversation_repository.clone();
.append_message_with_topic(&self.session_id, self.topic_id.as_deref(), &message) let session_id = self.session_id.clone();
{ let topic_id = self.topic_id.clone();
tracing::error!(error = %e, session_id = %self.session_id, let msg_for_persist = message.clone();
"Failed to persist emitted message"); tokio::task::spawn_blocking(move || {
} if let Err(e) = repo.append_message_with_topic(
&session_id,
topic_id.as_deref(),
&msg_for_persist,
) {
tracing::error!(error = %e, session_id = %session_id,
"Failed to persist emitted message");
}
})
.await
.ok();
self.inner.handle_tool_result(message, duration_ms).await; self.inner.handle_tool_result(message, duration_ms).await;
} }
@ -920,6 +966,23 @@ impl AgentLoop {
let mut loop_detector = LoopDetector::new(LoopDetectorConfig::default()); let mut loop_detector = LoopDetector::new(LoopDetectorConfig::default());
let mut emitted_messages = Vec::new(); let mut emitted_messages = Vec::new();
// 工具集在 create_agent 时按当前专家 capability 过滤后注入(快照语义)。
// 单次 process() 期间专家切换不影响当前 Arc<ToolRegistry> 引用,
// 因此循环外构建一次 tool_defs避免每轮重复读锁 + 全量 clone。
// 注意:若未来支持 process() 中途热切换专家(当前不支持),需重新评估。
let tool_defs = self.tools.get_definitions();
let tools = if tool_defs.is_empty() {
None
} else {
Some(tool_defs)
};
// 工具 token 估算在循环外算一次tool_defs 在 process() 期间不变),
// 避免每轮 serde_json::to_string 全量序列化工具定义。
let tools_tokens = tools
.as_ref()
.map(|t| estimate_tokens_from_serialized_json(t))
.unwrap_or_default();
for iteration in 0..self.max_iterations { for iteration in 0..self.max_iterations {
#[cfg(debug_assertions)] #[cfg(debug_assertions)]
tracing::debug!(iteration, "Agent iteration started"); tracing::debug!(iteration, "Agent iteration started");
@ -934,14 +997,6 @@ impl AgentLoop {
return Ok(cancel); return Ok(cancel);
} }
// Build request
let tool_defs = self.tools.get_definitions();
let tools = if tool_defs.is_empty() {
None
} else {
Some(tool_defs)
};
// Defense-in-depth: sanitize incomplete tool call sequences // Defense-in-depth: sanitize incomplete tool call sequences
// before EVERY LLM request, not just once at process() entry. // before EVERY LLM request, not just once at process() entry.
// This catches edge cases where compression, persistence races, // This catches edge cases where compression, persistence races,
@ -956,7 +1011,7 @@ impl AgentLoop {
); );
} }
let request = self.build_llm_request(&messages, system_prompt_context, tools); let request = self.build_llm_request(&messages, system_prompt_context, tools.clone(), tools_tokens);
// 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
@ -1182,11 +1237,16 @@ impl AgentLoop {
/// 构建 LLM 请求:过滤图片、构建系统提示、转换消息格式。 /// 构建 LLM 请求:过滤图片、构建系统提示、转换消息格式。
/// 统一了循环内调用和最终 summary 调用的请求构建逻辑。 /// 统一了循环内调用和最终 summary 调用的请求构建逻辑。
///
/// 优化token 估算直接基于 ChatMessage 累加字段长度(轻量),
/// 不再构造中间 text_only_messages: Vec<Message>(避免 N 条消息 × 5 字段双克隆)。
/// 工具 token 估算在循环外预算后传入tools_tokens
fn build_llm_request( fn build_llm_request(
&self, &self,
messages: &[ChatMessage], messages: &[ChatMessage],
system_prompt_context: Option<&SystemPromptContext>, system_prompt_context: Option<&SystemPromptContext>,
tools: Option<Vec<crate::domain::tools::Tool>>, tools: Option<Vec<crate::domain::tools::Tool>>,
tools_tokens: usize,
) -> ChatCompletionRequest { ) -> ChatCompletionRequest {
let filtered_messages = filter_images_by_age_and_count( let filtered_messages = filter_images_by_age_and_count(
messages, messages,
@ -1201,18 +1261,18 @@ impl AgentLoop {
.and_then(|provider| provider.build(ctx)) .and_then(|provider| provider.build(ctx))
}); });
let mut text_only_messages: Vec<Message> = Vec::with_capacity(filtered_messages.len() + 2); // 轻量 token 估算:直接从 ChatMessage 累加,无需构造中间 Vec<Message>
if let Some(ref prompt) = system_prompt { let text_tokens = estimate_text_tokens_from_chat_messages(
text_only_messages.push(Message::system(prompt.content.clone())); &filtered_messages,
} system_prompt.as_ref().map(|p| p.content.as_str()),
text_only_messages.extend(filtered_messages.iter().map(chat_message_to_text_only_llm_message)); );
let image_tokens = image_token_budget_for_request( let image_tokens = image_token_budget_for_request(
&self.runtime_config, &self.runtime_config,
&text_only_messages, text_tokens,
tools.as_ref(), tools_tokens,
); );
let mut image_budget = ImageInlineBudget::new(image_tokens, image_count); let mut image_budget = ImageInlineBudget::new(image_tokens, image_count);
let mut messages_for_llm: Vec<Message> = Vec::with_capacity(filtered_messages.len() + 2); let mut messages_for_llm: Vec<Message> = Vec::with_capacity(filtered_messages.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()));
@ -1321,7 +1381,7 @@ impl AgentLoop {
); );
messages.push(summary_request); messages.push(summary_request);
let request = self.build_llm_request(messages, system_prompt_context, None); let request = self.build_llm_request(messages, system_prompt_context, None, 0);
// 最终 summary 调用也与取消信号竞速 // 最终 summary 调用也与取消信号竞速
let final_result: Result< let final_result: Result<