diff --git a/src/agent/agent_loop.rs b/src/agent/agent_loop.rs index a6668e4..d624386 100644 --- a/src/agent/agent_loop.rs +++ b/src/agent/agent_loop.rs @@ -155,10 +155,38 @@ fn estimate_tokens_from_serialized_json(value: &T) -> usize as usize } +/// 轻量 token 估算：直接从 ChatMessage 累加文本字段长度， +/// 避免对完整 Vec 做 serde_json::to_string（O(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( runtime_config: &AgentRuntimeConfig, - text_only_messages: &[Message], - tools: Option<&Vec>, + text_tokens: usize, + tools_tokens: usize, ) -> usize { let completion_reserve = runtime_config .provider @@ -170,12 +198,7 @@ fn image_token_budget_for_request( .saturating_sub(completion_reserve); 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) - + tools - .map(estimate_tokens_from_serialized_json) - .unwrap_or_default(); - - safe_input_window.saturating_sub(text_tokens) + safe_input_window.saturating_sub(text_tokens + tools_tokens) } fn count_supported_image_media_refs(messages: &[ChatMessage]) -> usize { @@ -700,23 +723,46 @@ impl PersistingEmittedMessageHandler { #[async_trait] impl EmittedMessageHandler for PersistingEmittedMessageHandler { async fn handle(&self, message: ChatMessage) { - if let Err(e) = self.conversation_repository - .append_message_with_topic(&self.session_id, self.topic_id.as_deref(), &message) - { - tracing::error!(error = %e, session_id = %self.session_id, - "Failed to persist emitted message"); - } + // SQLite 是同步 I/O，直接调用会阻塞 tokio worker 线程。 + // 用 spawn_blocking 把写操作挪到 blocking 线程池，不改变语义： + // 失败仅记日志不阻断主流程（与原 if let Err 一致）。 + let repo = self.conversation_repository.clone(); + 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; } async fn handle_tool_result(&self, message: ChatMessage, duration_ms: Option) { // Persist the ChatMessage first (no duration field, same as before) - if let Err(e) = self.conversation_repository - .append_message_with_topic(&self.session_id, self.topic_id.as_deref(), &message) - { - tracing::error!(error = %e, session_id = %self.session_id, - "Failed to persist emitted message"); - } + let repo = self.conversation_repository.clone(); + 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(); 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 emitted_messages = Vec::new(); + // 工具集在 create_agent 时按当前专家 capability 过滤后注入（快照语义）。 + // 单次 process() 期间专家切换不影响当前 Arc 引用， + // 因此循环外构建一次 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 { #[cfg(debug_assertions)] tracing::debug!(iteration, "Agent iteration started"); @@ -934,14 +997,6 @@ impl AgentLoop { 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 // before EVERY LLM request, not just once at process() entry. // 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 // Pre-generate the message ID so stream deltas and the final assistant @@ -1182,11 +1237,16 @@ impl AgentLoop { /// 构建 LLM 请求：过滤图片、构建系统提示、转换消息格式。 /// 统一了循环内调用和最终 summary 调用的请求构建逻辑。 + /// + /// 优化：token 估算直接基于 ChatMessage 累加字段长度（轻量）， + /// 不再构造中间 text_only_messages: Vec（避免 N 条消息 × 5 字段双克隆）。 + /// 工具 token 估算在循环外预算后传入（tools_tokens）。 fn build_llm_request( &self, messages: &[ChatMessage], system_prompt_context: Option<&SystemPromptContext>, tools: Option>, + tools_tokens: usize, ) -> ChatCompletionRequest { let filtered_messages = filter_images_by_age_and_count( messages, @@ -1201,18 +1261,18 @@ impl AgentLoop { .and_then(|provider| provider.build(ctx)) }); - let mut text_only_messages: Vec = Vec::with_capacity(filtered_messages.len() + 2); - if let Some(ref prompt) = system_prompt { - text_only_messages.push(Message::system(prompt.content.clone())); - } - text_only_messages.extend(filtered_messages.iter().map(chat_message_to_text_only_llm_message)); - + // 轻量 token 估算：直接从 ChatMessage 累加，无需构造中间 Vec + let text_tokens = estimate_text_tokens_from_chat_messages( + &filtered_messages, + system_prompt.as_ref().map(|p| p.content.as_str()), + ); let image_tokens = image_token_budget_for_request( &self.runtime_config, - &text_only_messages, - tools.as_ref(), + text_tokens, + tools_tokens, ); let mut image_budget = ImageInlineBudget::new(image_tokens, image_count); + let mut messages_for_llm: Vec = Vec::with_capacity(filtered_messages.len() + 2); if let Some(ref prompt) = system_prompt { messages_for_llm.push(Message::system(prompt.content.clone())); @@ -1321,7 +1381,7 @@ impl AgentLoop { ); 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 调用也与取消信号竞速 let final_result: Result<