Compare commits

...

2 Commits

View File

@ -644,8 +644,10 @@ pub struct AgentLoop {
observer: Option<Arc<dyn Observer>>, observer: Option<Arc<dyn Observer>>,
emitted_message_handler: Option<Arc<dyn EmittedMessageHandler>>, emitted_message_handler: Option<Arc<dyn EmittedMessageHandler>>,
max_iterations: usize, max_iterations: usize,
/// 取消信号接收端Agent 在每次迭代开始时检查是否被取消 /// 取消信号接收端Agent 在每次迭代开始时检查是否被取消。
cancel_token: Option<tokio::sync::watch::Receiver<()>>, /// 包装在 Mutex 中以支持 interior mutability ——
/// watch::Receiver::changed() 需要 &mut self但 process() 持有 &self。
cancel_token: Option<tokio::sync::Mutex<tokio::sync::watch::Receiver<()>>>,
} }
#[derive(Debug, Clone)] #[derive(Debug, Clone)]
@ -862,8 +864,9 @@ impl AgentLoop {
/// ///
/// Agent 在每次迭代开始时检查 `cancel_token.has_changed()` /// Agent 在每次迭代开始时检查 `cancel_token.has_changed()`
/// 如果已收到取消信号则提前返回。 /// 如果已收到取消信号则提前返回。
/// 同时LLM 调用和工具执行期间通过 tokio::select! 与 cancel_signal() 竞速。
pub fn with_cancel_token(mut self, token: tokio::sync::watch::Receiver<()>) -> Self { pub fn with_cancel_token(mut self, token: tokio::sync::watch::Receiver<()>) -> Self {
self.cancel_token = Some(token); self.cancel_token = Some(tokio::sync::Mutex::new(token));
self self
} }
@ -907,21 +910,14 @@ impl AgentLoop {
tracing::debug!(iteration, "Agent iteration started"); tracing::debug!(iteration, "Agent iteration started");
// 检查取消信号 // 检查取消信号
if let Some(ref token) = self.cancel_token { // 使用 unwrap_or(true):即使 watch channel 因异常关闭sender drop 但未 send
if token.has_changed().unwrap_or(false) { // 也视为取消信号。defense-in-depth —— fail safe 而非 fail silent。
if let Some(ref mutex) = self.cancel_token {
if mutex.lock().await.has_changed().unwrap_or(true) {
tracing::info!(iteration, "Agent execution cancelled by user"); tracing::info!(iteration, "Agent execution cancelled by user");
let cancel_message = format!( let cancel = Self::build_cancel_result(iteration, emitted_messages);
"\n\n[用户已取消执行。已迭代 {} 次,取消前共生成了 {} 条消息。]", self.emit_live_tool_call_message(cancel.final_response.clone()).await;
iteration, return Ok(cancel);
emitted_messages.len()
);
let assistant_message = ChatMessage::assistant(cancel_message);
emitted_messages.push(assistant_message.clone());
self.emit_live_tool_call_message(assistant_message.clone()).await;
return Ok(AgentProcessResult {
final_response: assistant_message,
emitted_messages,
});
} }
} }
@ -1016,11 +1012,45 @@ impl AgentLoop {
let _ = delta_tx.try_send(delta); let _ = delta_tx.try_send(delta);
}); });
let response = match (*self.provider).chat_with_streaming(request, stream_callback).await { // LLM 调用与取消信号竞速若取消信号到达drop LLM future 以 abort HTTP 请求。
// stream_callback 是 Arc<...>LLM future 持有其 clone。
// 取消时需显式 drop 外部 stream_callback 以关闭 mpsc channel
// 让 consumer_task 自然退出。
let llm_result: Result<
crate::providers::ChatCompletionResponse,
Box<dyn std::error::Error + Send + Sync>,
>;
if self.cancel_token.is_some() {
tokio::select! {
_ = self.cancel_signal() => {
// LLM future 已被 select! drop → stream_callback clone 已释放。
// 显式 drop 外部 stream_callback → delta_tx 释放 → channel 关闭。
drop(stream_callback);
let _ = consumer_task.await;
let cancel = Self::build_cancel_result(iteration, emitted_messages);
self.emit_live_tool_call_message(cancel.final_response.clone()).await;
return Ok(cancel);
}
result = self.provider.chat_with_streaming(request, stream_callback.clone()) => {
llm_result = result;
}
}
// LLM 调用正常完成clone 的 Arc 已被 select! drop
// 但外部 stream_callback 仍存活。显式 drop 以关闭 mpsc channel
// 让 consumer_task 能自然退出。
drop(stream_callback);
} else {
// 无取消令牌stream_callback 被 move 进 chat_with_streaming调用完成即释放。
llm_result = self.provider.chat_with_streaming(request, stream_callback).await;
}
// Close delta channel and wait for consumer to finish processing
// (delta_tx is dropped when the callback closure is dropped)
let _ = consumer_task.await;
let response = match llm_result {
Ok(response) => response, Ok(response) => response,
Err(e) => { Err(e) => {
// delta_tx is dropped with the callback; await consumer to finish
let _ = consumer_task.await;
tracing::error!( tracing::error!(
provider = %self.provider.name(), provider = %self.provider.name(),
model = %self.provider.model_id(), model = %self.provider.model_id(),
@ -1039,10 +1069,6 @@ impl AgentLoop {
} }
}; };
// Close delta channel and wait for consumer to finish processing
// (delta_tx is dropped when the callback closure is dropped)
let _ = consumer_task.await;
// 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();
if had_streaming { if had_streaming {
@ -1121,7 +1147,22 @@ impl AgentLoop {
.await; .await;
// Execute tools and add results to messages // Execute tools and add results to messages
let tool_results = self.execute_tools(&response.tool_calls).await; // 工具执行与取消信号竞速:取消时 drop join_all 或 sequential future
// 未完成的工具调用被丢弃。
let tool_results = if self.cancel_token.is_some() {
tokio::select! {
_ = self.cancel_signal() => {
let cancel = Self::build_cancel_result(iteration, emitted_messages);
self.emit_live_tool_call_message(cancel.final_response.clone()).await;
return Ok(cancel);
}
results = self.execute_tools(&response.tool_calls) => {
results
}
}
} else {
self.execute_tools(&response.tool_calls).await
};
for (tool_call, result) in response.tool_calls.iter().zip(tool_results.iter()) { for (tool_call, result) in response.tool_calls.iter().zip(tool_results.iter()) {
// Truncate tool result if too large // Truncate tool result if too large
@ -1238,7 +1279,27 @@ impl AgentLoop {
tools: None, // No tools in final summary call tools: None, // No tools in final summary call
}; };
match (*self.provider).chat(request).await { // 最终 summary 调用也与取消信号竞速
let final_result: Result<
crate::providers::ChatCompletionResponse,
Box<dyn std::error::Error + Send + Sync>,
>;
if self.cancel_token.is_some() {
tokio::select! {
_ = self.cancel_signal() => {
let cancel = Self::build_cancel_result(self.max_iterations, emitted_messages);
self.emit_live_tool_call_message(cancel.final_response.clone()).await;
return Ok(cancel);
}
result = self.provider.chat(request) => {
final_result = result;
}
}
} else {
final_result = self.provider.chat(request).await;
}
match final_result {
Ok(response) => { Ok(response) => {
let assistant_message = if let Some(reasoning_content) = response.reasoning_content let assistant_message = if let Some(reasoning_content) = response.reasoning_content
{ {
@ -1272,6 +1333,39 @@ impl AgentLoop {
} }
} }
/// 等待取消信号。若未配置 cancel_token永远不返回。
///
/// 封装了与 watch channel 的交互changed() 返回 Ok 表示收到信号,
/// 返回 Err(Closed) 表示 sender 已 drop两种情况都视为取消。
async fn cancel_signal(&self) {
if let Some(ref mutex) = self.cancel_token {
let mut token = mutex.lock().await;
let _ = token.changed().await;
} else {
std::future::pending::<()>().await;
}
}
/// 构建取消响应,包含已完成的迭代次数和已生成的消息数量。
/// 构建取消响应,将取消通知追加到 `emitted_messages` 末尾后一并返回。
/// 这样 finalize_result 会把中间消息加入内存历史,确保下一个 LLM 调用有完整上下文。
fn build_cancel_result(
iteration: usize,
mut emitted_messages: Vec<ChatMessage>,
) -> AgentProcessResult {
let emitted_count = emitted_messages.len();
let cancel_message = format!(
"\n\n[用户已取消执行。已迭代 {} 次,取消前共生成了 {} 条消息。]",
iteration, emitted_count
);
let assistant_message = ChatMessage::assistant(cancel_message);
emitted_messages.push(assistant_message.clone());
AgentProcessResult {
final_response: assistant_message,
emitted_messages,
}
}
async fn emit_live_tool_call_message(&self, message: ChatMessage) { async fn emit_live_tool_call_message(&self, message: ChatMessage) {
if let Some(handler) = &self.emitted_message_handler { if let Some(handler) = &self.emitted_message_handler {
handler.handle(message).await; handler.handle(message).await;