1019dbe8cc
- 移除无用克隆与冗余引用，减少不必要内存分配 - 规范 unwrap/expect 使用，修复可提前失败路径 - 修复 anthropic provider llm_timeout_secs 死代码并补全超时日志 - cargo fmt 统一格式
288 lines
12 KiB
Rust
288 lines
12 KiB
Rust
use std::sync::Arc;
|
||
use std::time::Duration;
|
||
|
||
use async_trait::async_trait;
|
||
use tokio::sync::{Mutex, OwnedMutexGuard, watch};
|
||
use tokio::time::sleep;
|
||
|
||
use crate::gateway::session::Session;
|
||
use crate::storage::SessionStore;
|
||
use crate::tools::task::SubagentResult;
|
||
use crate::tools::{WaitCoordinator, WaitEvent};
|
||
|
||
/// 基于 Session 的 wait 协调器实现。
|
||
///
|
||
/// 持有 serial_lock 的 guard slot（可通过 take/drop 释放锁，通过 put 回填新 guard），
|
||
/// 以及 Session 引用（用于管理 sub_done_receiver、wait_wakeup、waiting_flag）。
|
||
///
|
||
/// wait() 流程：
|
||
/// 1. 设置 waiting=true（让 process_one 走注入路径）
|
||
/// 2. 取出 sub_done_receiver（从 Session 中 take，select! 期间不持有 Session 锁）
|
||
/// 3. 获取 wait_wakeup（Arc<Notify>，clone 后不持有 Session 锁）
|
||
/// 4. 释放 serial_lock（从 guard_slot take 并 drop guard）
|
||
/// 5. select! { sub_done_q.recv(), wakeup.notified(), timeout }
|
||
/// 6. 重新获取 serial_lock（serial_lock.lock_owned().await）
|
||
/// 7. 回填 guard 到 guard_slot
|
||
/// 8. 设置 waiting=false（先获取锁后清除，避免 TOCTOU）
|
||
/// 9. 归还 sub_done_receiver
|
||
pub struct SessionWaitCoordinator {
|
||
/// Session 引用（Arc<Mutex<Session>>），用于访问 SessionHistory 的队列和状态
|
||
session: Arc<Mutex<Session>>,
|
||
/// serial_lock guard 的存储槽。
|
||
/// 执行路径（execution.rs）获取锁后将 guard 存入此槽；
|
||
/// wait() 取出并 drop 以释放锁，重获取后回填新 guard。
|
||
guard_slot: Arc<Mutex<Option<OwnedMutexGuard<()>>>>,
|
||
/// serial_lock 本体（Arc<tokio::sync::Mutex<()>>），用于重获取锁
|
||
serial_lock: Arc<tokio::sync::Mutex<()>>,
|
||
/// SessionStore 引用，用于查询 pending_subagents
|
||
store: Arc<SessionStore>,
|
||
/// 当前 topic_id
|
||
topic_id: String,
|
||
}
|
||
|
||
impl SessionWaitCoordinator {
|
||
pub fn new(
|
||
session: Arc<Mutex<Session>>,
|
||
guard_slot: Arc<Mutex<Option<OwnedMutexGuard<()>>>>,
|
||
serial_lock: Arc<tokio::sync::Mutex<()>>,
|
||
store: Arc<SessionStore>,
|
||
topic_id: String,
|
||
) -> Self {
|
||
Self {
|
||
session,
|
||
guard_slot,
|
||
serial_lock,
|
||
store,
|
||
topic_id,
|
||
}
|
||
}
|
||
}
|
||
|
||
#[async_trait]
|
||
impl WaitCoordinator for SessionWaitCoordinator {
|
||
fn query_pending_task_ids(&self) -> Vec<String> {
|
||
self.store
|
||
.list_pending_subagents(&self.topic_id, Some("running"))
|
||
.map(|records| records.into_iter().map(|r| r.task_id).collect())
|
||
.unwrap_or_default()
|
||
}
|
||
|
||
async fn try_drain_queued_results(&self) -> Vec<SubagentResult> {
|
||
// 取出 receiver → try_recv 排空 → 归还 receiver
|
||
// 安全性：此方法在执行路径中被调用（serial_lock 已持有），
|
||
// 无其他代码并发访问 receiver。
|
||
let rx = {
|
||
let mut session = self.session.lock().await;
|
||
session.take_sub_done_receiver(&self.topic_id)
|
||
};
|
||
|
||
let mut results = Vec::new();
|
||
if let Some(mut rx) = rx {
|
||
while let Ok(result) = rx.try_recv() {
|
||
results.push(result);
|
||
}
|
||
// 归还 receiver（即使已排空，仍需放回供后续 wait() 使用）
|
||
let mut session = self.session.lock().await;
|
||
session.restore_sub_done_receiver(&self.topic_id, rx);
|
||
}
|
||
|
||
if !results.is_empty() {
|
||
tracing::debug!(
|
||
topic_id = %self.topic_id,
|
||
drained_count = results.len(),
|
||
"Drained buffered subagent results from sub_done_q"
|
||
);
|
||
}
|
||
|
||
results
|
||
}
|
||
|
||
async fn wait(&self, timeout: Duration, cancel_rx: Option<watch::Receiver<()>>) -> WaitEvent {
|
||
// 1. 设置 waiting=true
|
||
{
|
||
let mut session = self.session.lock().await;
|
||
session.set_waiting(&self.topic_id, true);
|
||
}
|
||
|
||
// 2. 取出 sub_done_receiver（select! 期间不持有 Session 锁）
|
||
let receiver = {
|
||
let mut session = self.session.lock().await;
|
||
session.take_sub_done_receiver(&self.topic_id)
|
||
};
|
||
|
||
// 3. 获取 wait_wakeup（Arc<Notify>，clone 后不持有 Session 锁）
|
||
let wakeup = {
|
||
let mut session = self.session.lock().await;
|
||
session.wait_wakeup(&self.topic_id)
|
||
};
|
||
|
||
// 3.5. 记录等待前的用户消息数量（用于 wakeup 后提取新注入的消息）
|
||
// 直接从 SQLite 读取，不持有任何锁
|
||
let user_msg_count_before = self
|
||
.store
|
||
.load_messages_for_topic(&self.topic_id, None)
|
||
.map(|msgs| msgs.iter().filter(|m| m.role == "user").count())
|
||
.unwrap_or(0);
|
||
|
||
// 4. 释放 serial_lock（取出 guard 并 drop）
|
||
{
|
||
let mut slot = self.guard_slot.lock().await;
|
||
let _ = slot.take(); // drop guard → 释放 serial_lock
|
||
}
|
||
|
||
tracing::debug!(
|
||
topic_id = %self.topic_id,
|
||
timeout_secs = timeout.as_secs(),
|
||
user_msg_count_before,
|
||
has_cancel_rx = cancel_rx.is_some(),
|
||
"SessionWaitCoordinator: lock released, entering select!"
|
||
);
|
||
|
||
// 5. select! 等待（不持有任何锁）
|
||
//
|
||
// cancel 分支放在最后（biased 排序中最后被 poll），
|
||
// 确保子代理结果和用户消息优先于取消信号被处理。
|
||
// 场景：/stop 后用户立即发消息 → process_one 注入消息 + wakeup，
|
||
// select! 优先消费 wakeup（UserMessage），而非 cancel（Cancelled），
|
||
// 使已注入的用户消息能被 Agent 处理而非丢失。
|
||
//
|
||
// 但 cancel_rx.changed() 不会无限阻塞——若无子代理结果、无用户消息，
|
||
// cancel 仍是唯一就绪分支，等待被优雅终止。
|
||
let event = if let Some(mut rx) = receiver {
|
||
let mut cancel_rx = cancel_rx;
|
||
tokio::select! {
|
||
biased;
|
||
result = rx.recv() => {
|
||
match result {
|
||
Some(subagent_result) => {
|
||
let event = WaitEvent::SubagentResult(subagent_result);
|
||
let mut session = self.session.lock().await;
|
||
session.restore_sub_done_receiver(&self.topic_id, rx);
|
||
event
|
||
}
|
||
None => {
|
||
// sender 全部 drop（所有 sub_done_sender 被释放）
|
||
WaitEvent::Timeout
|
||
}
|
||
}
|
||
}
|
||
_ = wakeup.notified() => {
|
||
// 用户消息到达（process_one 已注入 history 并 wakeup）
|
||
let mut session = self.session.lock().await;
|
||
session.restore_sub_done_receiver(&self.topic_id, rx);
|
||
// 提取等待期间新注入的用户消息内容
|
||
let new_messages = self.fetch_new_user_messages(user_msg_count_before);
|
||
tracing::info!(
|
||
topic_id = %self.topic_id,
|
||
new_msg_count = new_messages.len(),
|
||
"SessionWaitCoordinator: woke up by user message"
|
||
);
|
||
WaitEvent::UserMessage(new_messages)
|
||
}
|
||
_ = async {
|
||
if let Some(ref mut crx) = cancel_rx {
|
||
let _ = crx.changed().await;
|
||
} else {
|
||
std::future::pending::<()>().await;
|
||
}
|
||
} => {
|
||
// 取消信号到达（/stop）→ 归还 receiver，返回 Cancelled
|
||
tracing::info!(
|
||
topic_id = %self.topic_id,
|
||
"SessionWaitCoordinator: cancelled by /stop during wait"
|
||
);
|
||
let mut session = self.session.lock().await;
|
||
session.restore_sub_done_receiver(&self.topic_id, rx);
|
||
WaitEvent::Cancelled
|
||
}
|
||
_ = sleep(timeout) => {
|
||
let mut session = self.session.lock().await;
|
||
session.restore_sub_done_receiver(&self.topic_id, rx);
|
||
WaitEvent::Timeout
|
||
}
|
||
}
|
||
} else {
|
||
// 无 receiver（topic 无 sub_done 队列），直接等待 timeout 或 wakeup
|
||
let mut cancel_rx = cancel_rx;
|
||
tokio::select! {
|
||
biased;
|
||
_ = wakeup.notified() => {
|
||
let new_messages = self.fetch_new_user_messages(user_msg_count_before);
|
||
tracing::info!(
|
||
topic_id = %self.topic_id,
|
||
new_msg_count = new_messages.len(),
|
||
"SessionWaitCoordinator: woke up by user message (no receiver)"
|
||
);
|
||
WaitEvent::UserMessage(new_messages)
|
||
}
|
||
_ = async {
|
||
if let Some(ref mut crx) = cancel_rx {
|
||
let _ = crx.changed().await;
|
||
} else {
|
||
std::future::pending::<()>().await;
|
||
}
|
||
} => {
|
||
tracing::info!(
|
||
topic_id = %self.topic_id,
|
||
"SessionWaitCoordinator: cancelled by /stop during wait (no receiver)"
|
||
);
|
||
WaitEvent::Cancelled
|
||
}
|
||
_ = sleep(timeout) => WaitEvent::Timeout,
|
||
}
|
||
};
|
||
|
||
// 6. 重新获取 serial_lock
|
||
// lock_owned 消费 Arc<Self>，需 clone 保留 self.serial_lock 供后续可能的重入
|
||
//
|
||
// 取消场景下此处可能阻塞——如果 process_one 正持有锁注入用户消息，
|
||
// 需等其释放后才能重获取。这是正确行为：确保 is_waiting 清除与
|
||
// process_one 的注入互斥，避免 TOCTOU。
|
||
let new_guard = self.serial_lock.clone().lock_owned().await;
|
||
|
||
// 7. 回填 guard 到 guard_slot
|
||
{
|
||
let mut slot = self.guard_slot.lock().await;
|
||
*slot = Some(new_guard);
|
||
}
|
||
|
||
// 8. 设置 waiting=false（先获取锁后清除，避免 TOCTOU）
|
||
// 注意：serial_lock 已在步骤 6 获取，此时 process_one 无法获取锁，
|
||
// 所以清除 waiting 是安全的。
|
||
{
|
||
let mut session = self.session.lock().await;
|
||
session.set_waiting(&self.topic_id, false);
|
||
}
|
||
|
||
tracing::debug!(
|
||
topic_id = %self.topic_id,
|
||
"SessionWaitCoordinator: lock reacquired, waiting cleared"
|
||
);
|
||
|
||
event
|
||
}
|
||
}
|
||
|
||
impl SessionWaitCoordinator {
|
||
/// 提取等待期间新注入的用户消息内容。
|
||
/// 通过对比等待前的用户消息数量，从 SQLite 中取出新增的用户消息。
|
||
fn fetch_new_user_messages(&self, count_before: usize) -> Vec<String> {
|
||
match self.store.load_messages_for_topic(&self.topic_id, None) {
|
||
Ok(msgs) => msgs
|
||
.iter()
|
||
.filter(|m| m.role == "user")
|
||
.skip(count_before)
|
||
.map(|m| m.content.clone())
|
||
.collect(),
|
||
Err(e) => {
|
||
tracing::warn!(
|
||
error = %e,
|
||
topic_id = %self.topic_id,
|
||
"Failed to load messages for fetching new user messages"
|
||
);
|
||
Vec::new()
|
||
}
|
||
}
|
||
}
|
||
}
|