feat: 异步子代理 v8 实现与 /stop 取消流程修复
后端核心实现： - 新增 wait_coordinator：释放/重获取 serial_lock 的 in-tool waiting 模式，替代旧的 break-exit 方案 - 新增 wait_for_subagents 工具：select! 等待子代理完成/用户消息/超时，支持 try_drain 批量消费 - task 工具异步 spawn 路径：CancellationToken 注册表 + RAII guard + Semaphore 并发限流 - process exit 安全检查：pending 子代理存在时抑制 ExecutionCompleted - 崩溃恢复：启动时标记 running->interrupted，history 加载时对账占位 /stop 取消流程修复： - wait_coordinator select! 添加 cancel 分支，完整清理状态 - ToolContext 注入 cancel_rx（watch::Receiver clone） - agent_loop 工具执行 select! 对 wait 工具跳过竞速，防止 drop coordinator 清理逻辑 - 子代理完成状态通过 execution_completed metadata 传播 存储层： - pending_subagents 表 + 条件 UPDATE - mark_all_running_as_interrupted 崩溃恢复
This commit is contained in:
oudecheng
parent
e0313ab8f3
commit
fc3a95b152
@ -1447,10 +1447,29 @@ impl AgentLoop {
|
||||
.await;
|
||||
|
||||
// Execute tools and add results to messages
|
||||
// 工具执行与取消信号竞速：取消时 drop join_all 或 sequential future，
|
||||
// 未完成的工具调用被丢弃。
|
||||
let tool_results = if self.cancel_token.is_some() {
|
||||
//
|
||||
// 取消竞速策略：
|
||||
// - 包含 wait_for_subagents 时：不使用 select!，直接 await execute_tools。
|
||||
// 原因：coordinator.wait() 在 select! 返回后有不可中断的清理逻辑
|
||||
// （步骤 6-8：重获取 serial_lock、回填 guard_slot、清除 is_waiting）。
|
||||
// 若 agent_loop 的 select! 在清理期间 drop execute_tools，
|
||||
// 会导致 is_waiting=true 永久残留、guard_slot 为空、serial_lock 未持有，
|
||||
// 后续所有用户消息走注入路径但 wakeup 无接收者 → 系统永久卡死。
|
||||
// cancel 由 coordinator 内部 select! 的 cancel 分支处理，清理不会被打断。
|
||||
//
|
||||
// - 不含 wait_for_subagents 时：保留 select! 竞速，允许 /stop 中断
|
||||
// 长时间运行的工具（如 MCP HTTP 请求）。
|
||||
let has_wait_tool = response
|
||||
.tool_calls
|
||||
.iter()
|
||||
.any(|tc| tc.name == "wait_for_subagents");
|
||||
|
||||
let tool_results = if self.cancel_token.is_some() && !has_wait_tool {
|
||||
tokio::select! {
|
||||
biased;
|
||||
results = self.execute_tools(&response.tool_calls) => {
|
||||
results
|
||||
}
|
||||
_ = self.cancel_signal() => {
|
||||
// 为所有 tool_calls 补充取消结果，避免孤立 assistant(tool_calls)
|
||||
for tool_call in &response.tool_calls {
|
||||
@ -1468,9 +1487,6 @@ impl AgentLoop {
|
||||
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
|
||||
|
||||
@ -1,4 +1,5 @@
|
||||
use async_trait::async_trait;
|
||||
use std::sync::Arc;
|
||||
|
||||
use crate::command::Command;
|
||||
use crate::command::context::CommandContext;
|
||||
@ -6,18 +7,27 @@ use crate::command::handler::{CommandHandler, CommandMetadata};
|
||||
use crate::command::response::{CommandError, CommandResponse, MessageKind};
|
||||
use crate::gateway::cancel_manager::CancelManager;
|
||||
use crate::gateway::session::SessionManager;
|
||||
use crate::tools::SubAgentRuntime;
|
||||
|
||||
/// 处理 StopExecution 命令：按话题取消当前正在执行的 Agent。
|
||||
///
|
||||
/// 取消传播：同时取消该 topic 下所有正在运行的异步子代理（通过 CancellationToken）。
|
||||
pub struct StopExecutionCommandHandler {
|
||||
cancel_manager: CancelManager,
|
||||
session_manager: SessionManager,
|
||||
subagent_executor: Option<Arc<dyn SubAgentRuntime>>,
|
||||
}
|
||||
|
||||
impl StopExecutionCommandHandler {
|
||||
pub fn new(cancel_manager: CancelManager, session_manager: SessionManager) -> Self {
|
||||
pub fn new(
|
||||
cancel_manager: CancelManager,
|
||||
session_manager: SessionManager,
|
||||
subagent_executor: Option<Arc<dyn SubAgentRuntime>>,
|
||||
) -> Self {
|
||||
Self {
|
||||
cancel_manager,
|
||||
session_manager,
|
||||
subagent_executor,
|
||||
}
|
||||
}
|
||||
}
|
||||
@ -89,9 +99,26 @@ impl CommandHandler for StopExecutionCommandHandler {
|
||||
|
||||
let cancelled = self.cancel_manager.cancel_by_topic(&topic_id).await;
|
||||
|
||||
if cancelled {
|
||||
// 取消传播：同时取消该 topic 下所有正在运行的异步子代理
|
||||
let cancelled_subagents = if let Some(ref executor) = self.subagent_executor {
|
||||
executor.cancel_pending_for_topic(&topic_id).await
|
||||
} else {
|
||||
0
|
||||
};
|
||||
|
||||
if cancelled || cancelled_subagents > 0 {
|
||||
let msg = if cancelled && cancelled_subagents > 0 {
|
||||
format!(
|
||||
"正在停止当前任务及 {} 个后台子代理...",
|
||||
cancelled_subagents
|
||||
)
|
||||
} else if cancelled {
|
||||
"正在停止当前任务...".to_string()
|
||||
} else {
|
||||
format!("正在停止 {} 个后台子代理...", cancelled_subagents)
|
||||
};
|
||||
Ok(CommandResponse::success(ctx.request_id)
|
||||
.with_message(MessageKind::Notification, "正在停止当前任务..."))
|
||||
.with_message(MessageKind::Notification, msg))
|
||||
} else {
|
||||
Ok(CommandResponse::success(ctx.request_id)
|
||||
.with_message(MessageKind::Notification, "当前没有正在执行的任务"))
|
||||
|
||||
@ -337,6 +337,12 @@ pub struct TaskConfig {
|
||||
pub allowed_tools: Vec<String>,
|
||||
#[serde(default = "default_task_max_nesting_depth")]
|
||||
pub max_nesting_depth: u32,
|
||||
/// 异步子代理最大并发数（Semaphore 限流，仅主 agent 顶层 spawn 生效）
|
||||
#[serde(default = "default_task_max_concurrent")]
|
||||
pub max_concurrent: usize,
|
||||
/// wait_for_subagents 工具默认超时（秒），LLM 可通过参数覆盖
|
||||
#[serde(default = "default_task_wait_default_timeout_secs")]
|
||||
pub wait_default_timeout_secs: u64,
|
||||
}
|
||||
|
||||
fn default_task_enabled() -> bool {
|
||||
@ -355,6 +361,14 @@ fn default_task_max_nesting_depth() -> u32 {
|
||||
2
|
||||
}
|
||||
|
||||
fn default_task_max_concurrent() -> usize {
|
||||
8
|
||||
}
|
||||
|
||||
fn default_task_wait_default_timeout_secs() -> u64 {
|
||||
60
|
||||
}
|
||||
|
||||
fn default_task_allowed_tools() -> Vec<String> {
|
||||
vec![
|
||||
"read".to_string(),
|
||||
@ -380,6 +394,8 @@ impl Default for TaskConfig {
|
||||
ttl_hours: default_task_ttl_hours(),
|
||||
allowed_tools: default_task_allowed_tools(),
|
||||
max_nesting_depth: default_task_max_nesting_depth(),
|
||||
max_concurrent: default_task_max_concurrent(),
|
||||
wait_default_timeout_secs: default_task_wait_default_timeout_secs(),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@ -1,5 +1,7 @@
|
||||
use std::sync::Arc;
|
||||
|
||||
use tokio::sync::mpsc;
|
||||
|
||||
use crate::agent::context_compressor::ContextCompressor;
|
||||
use crate::agent::{AgentError, AgentLoop, AgentRuntimeConfig, CompositeSystemPromptProvider, SystemPromptProvider};
|
||||
use crate::config::{CompactionConfig, LLMProviderConfig, ModelResolver};
|
||||
@ -14,7 +16,8 @@ use crate::skills::{SkillPromptProvider, SkillRuntime};
|
||||
use crate::storage::PromptInjectionRepository;
|
||||
use crate::storage::persistent_session_id;
|
||||
use crate::tools::task::runtime::{SubagentPromptProvider, SubagentRuntime};
|
||||
use crate::tools::{ToolContext, ToolRegistry};
|
||||
use crate::tools::task::SubagentResult;
|
||||
use crate::tools::{ToolContext, ToolRegistry, WaitCoordinator};
|
||||
|
||||
/// 构建与 Agent 实际使用的完全一致的组合系统提示词 Provider。
|
||||
///
|
||||
@ -76,6 +79,11 @@ pub(crate) struct AgentBuildRequest<'a> {
|
||||
pub(crate) cancel_token: Option<tokio::sync::watch::Receiver<()>>,
|
||||
/// 端到端追踪 ID（从 InboundMessage 继承，注入 ToolContext 供 tool 执行路径日志关联）
|
||||
pub(crate) trace_id: Option<String>,
|
||||
/// 异步子代理完成队列的 sender（按 topic 隔离）。
|
||||
/// 仅主 agent 有值：TaskTool 据此在子代理完成时发送 SubagentResult。
|
||||
pub(crate) sub_done_sender: Option<mpsc::Sender<SubagentResult>>,
|
||||
/// wait_for_subagents 工具的协调器（仅主 agent 有值）。
|
||||
pub(crate) wait_coordinator: Option<Arc<dyn WaitCoordinator>>,
|
||||
}
|
||||
|
||||
impl AgentFactory {
|
||||
@ -234,6 +242,12 @@ impl AgentFactory {
|
||||
.notification_chat_id
|
||||
.unwrap_or(request.session_chat_id);
|
||||
// 构建上下文压缩器（参数内聚到 ContextCompressor，CompactionConfig 注入）
|
||||
// 注入取消信号 receiver 的 clone 到 ToolContext，
|
||||
// 供 wait_for_subagents 工具传递给 coordinator.wait() 的 select!。
|
||||
// watch::Receiver::clone() 创建共享同一 sender 的新 receiver，
|
||||
// 各 receiver 的 has_changed()/changed() 状态独立，互不影响。
|
||||
let cancel_rx_for_context = request.cancel_token.as_ref().map(|rx| rx.clone());
|
||||
|
||||
let runtime_config = AgentRuntimeConfig::from(effective_provider_config.clone());
|
||||
let compressor = Arc::new(self.build_compressor(&runtime_config));
|
||||
let mut agent = agent
|
||||
@ -253,6 +267,12 @@ impl AgentFactory {
|
||||
// 注入专家 capability，TaskTool 据此强制校验子代理白/黑名单
|
||||
parent_capability: expert_capability.clone(),
|
||||
trace_id: request.trace_id.clone(),
|
||||
// 注入异步子代理完成队列 sender（按 topic 隔离）
|
||||
sub_done_sender: request.sub_done_sender.clone(),
|
||||
// 注入 wait 协调器（封装释放/重获取 serial_lock 逻辑）
|
||||
wait_coordinator: request.wait_coordinator.clone(),
|
||||
// 注入取消信号 receiver clone（供 wait 工具的 cancel 检查）
|
||||
cancel_rx: cancel_rx_for_context,
|
||||
})
|
||||
.with_compressor(Some(compressor));
|
||||
// 注入观测器（依赖注入，agent_loop 只认 Observer trait）
|
||||
|
||||
@ -15,6 +15,8 @@ use tokio::sync::Mutex;
|
||||
use super::compaction::schedule_background_history_compaction;
|
||||
use super::message_prepare::enrich_user_content_with_media_refs;
|
||||
use super::session::Session;
|
||||
use super::wait_coordinator::SessionWaitCoordinator;
|
||||
use crate::tools::WaitCoordinator;
|
||||
|
||||
/// 空的 EmittedMessageHandler，不转发消息，仅配合 PersistingEmittedMessageHandler 做持久化。
|
||||
struct NoOpEmittedMessageHandler;
|
||||
@ -284,15 +286,32 @@ impl AgentExecutionService {
|
||||
// 获取该 topic 的串行锁（通过短暂获取 session 锁）
|
||||
// 同一 topic 的消息处理必须串行执行，防止并发 loop 操作同一历史的不同快照
|
||||
// 不同 topic 之间互不阻塞，支持多话题并发执行
|
||||
let serial_lock = {
|
||||
let (serial_lock, store, lock_key) = {
|
||||
let mut session_guard = request.session.lock().await;
|
||||
let lock_key = request.topic_id.as_deref().unwrap_or(request.chat_id);
|
||||
session_guard.topic_serial_lock(lock_key)
|
||||
let lock_key = request
|
||||
.topic_id
|
||||
.as_deref()
|
||||
.unwrap_or(request.chat_id)
|
||||
.to_string();
|
||||
session_guard.ensure_sub_done_channel(&lock_key);
|
||||
(
|
||||
session_guard.topic_serial_lock(&lock_key),
|
||||
session_guard.session_store(),
|
||||
lock_key,
|
||||
)
|
||||
};
|
||||
|
||||
// 等待该 topic 的前一条消息处理完成（含压缩）
|
||||
// await 串行锁时不持有 session 锁，其他 topic 的消息可以正常处理
|
||||
let _serial_guard = serial_lock.lock().await;
|
||||
// 使用 lock_owned 获取 OwnedMutexGuard，存入 guard_slot 供 wait_coordinator 释放/重获取
|
||||
// 注意：lock_owned 消费 Arc<Self>，需 clone 保留 serial_lock 供 coordinator 使用
|
||||
let serial_guard = serial_lock.clone().lock_owned().await;
|
||||
|
||||
// guard_slot：wait_coordinator 通过此 slot 释放/重获取 serial_lock。
|
||||
// 正常执行时 guard 留在 slot 中（锁持有）；wait 工具调用时 take guard 释放锁，
|
||||
// select! 等待结束后重获取锁并回填新 guard。
|
||||
// guard_slot 作为 Arc 共享于执行路径与 coordinator，二者全部 drop 时 guard 才释放锁。
|
||||
let guard_slot = Arc::new(Mutex::new(Some(serial_guard)));
|
||||
|
||||
let (history, agent, user_message, user_message_count, original_topic_id) = {
|
||||
let mut session_guard = request.session.lock().await;
|
||||
@ -340,12 +359,27 @@ impl AgentExecutionService {
|
||||
let history = session_guard.get_or_create_history(history_key).clone();
|
||||
session_guard.record_skill_offer(request.chat_id)?;
|
||||
|
||||
// 创建 wait 协调器（封装释放/重获取 serial_lock + select! 等待逻辑）。
|
||||
// 仅主 agent 注入；coordinator 通过 guard_slot 释放/重获取 serial_lock，
|
||||
// 使 wait_for_subagents 工具能在等待期间让 process_one 注入用户消息。
|
||||
let wait_coordinator: Option<Arc<dyn WaitCoordinator>> = {
|
||||
let coordinator = SessionWaitCoordinator::new(
|
||||
request.session.clone(),
|
||||
guard_slot.clone(),
|
||||
serial_lock.clone(),
|
||||
store.clone(),
|
||||
lock_key.clone(),
|
||||
);
|
||||
Some(Arc::new(coordinator))
|
||||
};
|
||||
|
||||
let mut agent = session_guard.create_agent(
|
||||
request.chat_id,
|
||||
Some(request.sender_id),
|
||||
Some(&user_message.id),
|
||||
original_topic_id.as_deref(),
|
||||
request.trace_id,
|
||||
wait_coordinator,
|
||||
)?;
|
||||
if let Some(handler) = request.live_emitter.clone() {
|
||||
agent = agent.with_emitted_message_handler(handler);
|
||||
@ -408,17 +442,26 @@ impl AgentExecutionService {
|
||||
// 获取该 topic 的串行锁（与普通消息路径共享，保证串行执行）
|
||||
// 定时任务由调度器触发，无用户消息竞态；在锁前一次性捕获 topic_id，
|
||||
// 锁后复用同一值作为 original_topic_id，保证锁键与写入目标一致。
|
||||
let (serial_lock, lock_time_topic_id) = {
|
||||
let (serial_lock, session_store, lock_key, lock_time_topic_id) = {
|
||||
let mut session_guard = request.session.lock().await;
|
||||
let tid = session_guard
|
||||
.current_topic(request.chat_id)
|
||||
.map(|s| s.to_string());
|
||||
let lock_key = tid.as_deref().unwrap_or(request.chat_id);
|
||||
(session_guard.topic_serial_lock(lock_key), tid)
|
||||
let lock_key = tid.as_deref().unwrap_or(request.chat_id).to_string();
|
||||
session_guard.ensure_sub_done_channel(&lock_key);
|
||||
(
|
||||
session_guard.topic_serial_lock(&lock_key),
|
||||
session_guard.session_store(),
|
||||
lock_key,
|
||||
tid,
|
||||
)
|
||||
};
|
||||
|
||||
// 等待该 topic 的前一条消息处理完成（含压缩）
|
||||
let _serial_guard = serial_lock.lock().await;
|
||||
// 使用 lock_owned 获取 OwnedMutexGuard，存入 guard_slot 供 wait_coordinator 释放/重获取
|
||||
// 注意：lock_owned 消费 Arc<Self>，需 clone 保留 serial_lock 供 coordinator 使用
|
||||
let serial_guard = serial_lock.clone().lock_owned().await;
|
||||
let guard_slot = Arc::new(Mutex::new(Some(serial_guard)));
|
||||
|
||||
let (
|
||||
history,
|
||||
@ -475,6 +518,18 @@ impl AgentExecutionService {
|
||||
let history = session_guard.get_or_create_history(history_key).clone();
|
||||
session_guard.record_skill_offer(request.chat_id)?;
|
||||
|
||||
// 创建 wait 协调器（与普通消息路径一致，支持定时任务中 spawn 异步子代理）
|
||||
let wait_coordinator: Option<Arc<dyn WaitCoordinator>> = {
|
||||
let coordinator = SessionWaitCoordinator::new(
|
||||
request.session.clone(),
|
||||
guard_slot.clone(),
|
||||
serial_lock.clone(),
|
||||
session_store.clone(),
|
||||
lock_key.clone(),
|
||||
);
|
||||
Some(Arc::new(coordinator))
|
||||
};
|
||||
|
||||
let agent = session_guard.create_agent_with_provider_config(
|
||||
request.chat_id,
|
||||
request.notification_chat_id, // 传入真实 chat_id
|
||||
@ -483,6 +538,7 @@ impl AgentExecutionService {
|
||||
request.provider_config.clone(),
|
||||
original_topic_id.as_deref(),
|
||||
&request.trace_id,
|
||||
wait_coordinator,
|
||||
)?;
|
||||
|
||||
// 获取 store 和 session_id，用于构造消息持久化 handler
|
||||
|
||||
@ -28,6 +28,7 @@ pub mod session_pool;
|
||||
pub mod static_files;
|
||||
pub mod tool_prompt_provider;
|
||||
pub mod tool_registry_factory;
|
||||
pub mod wait_coordinator;
|
||||
pub mod ws;
|
||||
|
||||
use axum::{Router, middleware, routing};
|
||||
@ -70,6 +71,8 @@ pub struct GatewayState {
|
||||
pub skills: Arc<SkillRuntime>,
|
||||
pub experts: Arc<crate::experts::ExpertRuntime>,
|
||||
pub subagent_runtime: Arc<SubagentRuntime>,
|
||||
/// 异步子代理执行器（DefaultSubAgentRuntime），用于取消传播等操作
|
||||
pub subagent_executor: Option<Arc<dyn crate::tools::SubAgentRuntime>>,
|
||||
/// per-session 的用户模型选择（覆盖专家配置）
|
||||
pub model_selections: Arc<model_selection::ModelSelectionStore>,
|
||||
/// Prometheus metrics handle（/metrics 端点渲染用）。
|
||||
@ -105,7 +108,7 @@ impl GatewayState {
|
||||
mcp_servers: config.mcp_servers.clone(),
|
||||
};
|
||||
|
||||
let (session_manager, task_repository, mcp_manager, subagent_runtime, model_selections) =
|
||||
let (session_manager, task_repository, mcp_manager, subagent_runtime, model_selections, subagent_executor) =
|
||||
build_session_manager_with_sender(
|
||||
agent_prompt_reinject_every,
|
||||
show_tool_results,
|
||||
@ -150,6 +153,7 @@ impl GatewayState {
|
||||
skills,
|
||||
experts,
|
||||
subagent_runtime,
|
||||
subagent_executor,
|
||||
model_selections,
|
||||
prometheus_handle,
|
||||
})
|
||||
@ -178,6 +182,7 @@ impl GatewayState {
|
||||
semaphore,
|
||||
provider_config,
|
||||
self.cancel_manager.clone(),
|
||||
self.subagent_executor.clone(),
|
||||
);
|
||||
tokio::spawn(inbound_processor.run());
|
||||
|
||||
@ -213,6 +218,30 @@ pub async fn run(
|
||||
|
||||
let state = Arc::new(GatewayState::from_config(config, restart_tx)?);
|
||||
|
||||
// ── 崩溃恢复：标记中断的异步子代理 ──
|
||||
// 服务器重启后，之前 spawn 的异步子代理进程已不存在，
|
||||
// 将 pending_subagents 表中所有 status='running' 的记录标记为 'interrupted'。
|
||||
// 下次 wait_for_subagents 调用时，这些 task_id 不会出现在 pending 列表中，
|
||||
// agent 可据此判断子代理未正常完成。
|
||||
match state.session_manager.store().mark_all_running_as_interrupted() {
|
||||
Ok(0) => {
|
||||
tracing::info!("Crash recovery: no interrupted subagents to recover");
|
||||
}
|
||||
Ok(n) => {
|
||||
tracing::info!(
|
||||
recovered_count = n,
|
||||
"Crash recovery: marked {} running subagents as interrupted (server restarted)",
|
||||
n
|
||||
);
|
||||
}
|
||||
Err(e) => {
|
||||
tracing::error!(
|
||||
error = %e,
|
||||
"Crash recovery: failed to mark interrupted subagents"
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
// Get provider config for channels
|
||||
let cfg = state.config.read().await;
|
||||
let provider_config = cfg.get_provider_config("default")?;
|
||||
|
||||
@ -29,6 +29,7 @@ use crate::storage::persistent_session_id;
|
||||
use crate::topic_description::generate_topic_description;
|
||||
|
||||
use super::session::{BusToolCallEmitter, SessionManager};
|
||||
use super::message_prepare::enrich_user_content_with_media_refs;
|
||||
|
||||
#[derive(Clone)]
|
||||
pub struct InboundProcessor {
|
||||
@ -48,6 +49,7 @@ impl InboundProcessor {
|
||||
semaphore: Arc<Semaphore>,
|
||||
provider_config: LLMProviderConfig,
|
||||
cancel_manager: CancelManager,
|
||||
subagent_executor: Option<Arc<dyn crate::tools::SubAgentRuntime>>,
|
||||
) -> Self {
|
||||
// 创建命令路由器并注册处理器
|
||||
let mut command_router = CommandRouter::new();
|
||||
@ -118,6 +120,7 @@ impl InboundProcessor {
|
||||
command_router.register(Box::new(StopExecutionCommandHandler::new(
|
||||
cancel_manager.clone(),
|
||||
session_manager.clone(),
|
||||
subagent_executor,
|
||||
)));
|
||||
|
||||
Self {
|
||||
@ -312,6 +315,86 @@ impl InboundProcessor {
|
||||
if let Some(ref topic_id) = current_topic {
|
||||
emitter_metadata.insert("topic_id".to_string(), topic_id.clone());
|
||||
}
|
||||
|
||||
// ── 异步子代理等待注入路径 ──
|
||||
// 当主 agent 正在 wait_for_subagents 中等待（已释放 serial_lock、is_waiting=true）时，
|
||||
// 新用户消息不应启动新的 agent loop，而应注入 history 并唤醒等待中的 agent。
|
||||
//
|
||||
// 流程：
|
||||
// 1. 获取 serial_lock（若 agent 正常运行则阻塞；若 agent 在 wait 中则立即获取）
|
||||
// 2. 检查 is_waiting：true → 注入 + wakeup + return；false → 释放锁走正常路径
|
||||
//
|
||||
// 安全性：is_waiting 在持锁状态下检查，wait_coordinator 清除 is_waiting 需先重获取锁，
|
||||
// 两者互斥，无 TOCTOU。
|
||||
if let Some(ref topic_id) = current_topic {
|
||||
if let Some(session) = self.session_manager.get(&inbound.channel).await {
|
||||
let lock_key = topic_id.clone();
|
||||
|
||||
// 获取 serial_lock Arc（短暂持有 session 锁）
|
||||
let serial_lock = {
|
||||
let mut g = session.lock().await;
|
||||
g.ensure_sub_done_channel(&lock_key);
|
||||
g.topic_serial_lock(&lock_key)
|
||||
};
|
||||
|
||||
// 阻塞获取 serial_lock
|
||||
// - agent 正常运行：阻塞至其完成（天然串行化）
|
||||
// - agent 在 wait 中：wait 已释放锁，可立即获取
|
||||
let _inject_guard = serial_lock.clone().lock_owned().await;
|
||||
|
||||
// 检查 is_waiting（持锁状态下安全）
|
||||
let is_waiting = {
|
||||
let g = session.lock().await;
|
||||
g.is_waiting(&lock_key)
|
||||
};
|
||||
|
||||
if is_waiting {
|
||||
// Agent 正在 wait_for_subagents 中等待 → 注入用户消息 + 唤醒
|
||||
tracing::info!(
|
||||
topic_id = %lock_key,
|
||||
"Topic is in waiting state, injecting user message and waking up agent"
|
||||
);
|
||||
|
||||
let wakeup = {
|
||||
let mut g = session.lock().await;
|
||||
// 确保 session 和 chat 已加载
|
||||
g.ensure_persistent_session(&inbound.chat_id)?;
|
||||
g.ensure_chat_loaded(&inbound.chat_id, Some(&lock_key))?;
|
||||
|
||||
// 构造用户消息（与 prepare_and_execute_message 一致的处理流程）
|
||||
let media_refs: Vec<String> = inbound
|
||||
.media
|
||||
.iter()
|
||||
.map(|m| m.path.clone())
|
||||
.collect();
|
||||
let enriched_content =
|
||||
enrich_user_content_with_media_refs(&inbound.content, &media_refs)?;
|
||||
let user_message =
|
||||
g.create_user_message(&enriched_content, media_refs);
|
||||
g.append_persisted_message(
|
||||
&inbound.chat_id,
|
||||
Some(&lock_key),
|
||||
user_message,
|
||||
)?;
|
||||
|
||||
// 获取 wakeup 信号
|
||||
g.wait_wakeup(&lock_key)
|
||||
};
|
||||
|
||||
// 唤醒等待中的 agent（wait_coordinator 的 select! 会捕获此通知）
|
||||
wakeup.notify_one();
|
||||
|
||||
// _inject_guard 在此处 drop → 释放 serial_lock
|
||||
// wait_coordinator 重获取锁后继续处理（history 已包含新用户消息）
|
||||
//
|
||||
// 跳过 handle_message / cancel 注册 / execution_completed，
|
||||
// 因为等待中的 agent 会处理这条消息。
|
||||
return Ok(());
|
||||
}
|
||||
// is_waiting=false：_inject_guard drop 释放锁，走正常 handle_message 路径
|
||||
}
|
||||
}
|
||||
|
||||
let live_emitter = Arc::new(PersistingEmittedMessageHandler::new(
|
||||
BusToolCallEmitter::new(
|
||||
self.bus.clone(),
|
||||
@ -497,31 +580,58 @@ impl InboundProcessor {
|
||||
self.cancel_manager.remove_by_topic(topic_id).await;
|
||||
}
|
||||
|
||||
// 发送执行完成信号，通知前端可以停止 loading 状态
|
||||
// 无论成功还是失败都发送，确保前端状态正确
|
||||
let mut completion_metadata = inbound.forwarded_metadata.clone();
|
||||
if let Some(ref topic_id) = current_topic {
|
||||
completion_metadata.insert("topic_id".to_string(), topic_id.clone());
|
||||
}
|
||||
if let Err(error) = self
|
||||
.bus
|
||||
.publish_outbound(
|
||||
OutboundMessage::execution_completed(
|
||||
channel,
|
||||
chat_id,
|
||||
Some(session_id),
|
||||
completion_metadata,
|
||||
// 发送执行完成信号，通知前端可以停止 loading 状态。
|
||||
//
|
||||
// 退出兜底（safety net）：如果当前 topic 仍有 running 状态的子代理，
|
||||
// 不发送 ExecutionCompleted。这防止 LLM 未调用 wait_for_subagents 就退出时，
|
||||
// 前端过早停止 loading 导致子代理结果"丢失"的观感。
|
||||
// 恢复路径：下一条用户消息触发新的 process_one → 加载 history →
|
||||
// LLM 看到 "running" 占位 → 调用 wait_for_subagents → 消费 sub_done_q 结果。
|
||||
let has_pending_subagents = if let Some(ref topic_id) = current_topic {
|
||||
let pending = self
|
||||
.session_manager
|
||||
.store()
|
||||
.list_pending_subagents(topic_id, Some("running"))
|
||||
.unwrap_or_default();
|
||||
if !pending.is_empty() {
|
||||
tracing::info!(
|
||||
topic_id = %topic_id,
|
||||
pending_count = pending.len(),
|
||||
"Skipping ExecutionCompleted: pending subagents still running"
|
||||
);
|
||||
true
|
||||
} else {
|
||||
false
|
||||
}
|
||||
} else {
|
||||
false
|
||||
};
|
||||
|
||||
if !has_pending_subagents {
|
||||
let mut completion_metadata = inbound.forwarded_metadata.clone();
|
||||
if let Some(ref topic_id) = current_topic {
|
||||
completion_metadata.insert("topic_id".to_string(), topic_id.clone());
|
||||
}
|
||||
if let Err(error) = self
|
||||
.bus
|
||||
.publish_outbound(
|
||||
OutboundMessage::execution_completed(
|
||||
channel,
|
||||
chat_id,
|
||||
Some(session_id),
|
||||
completion_metadata,
|
||||
)
|
||||
.with_trace_id(&inbound.trace_id),
|
||||
)
|
||||
.with_trace_id(&inbound.trace_id),
|
||||
)
|
||||
.await
|
||||
{
|
||||
match error {
|
||||
crate::bus::BusError::Dropped => {
|
||||
tracing::warn!(error = %error, "Outbound dropped (bus full)");
|
||||
}
|
||||
crate::bus::BusError::Closed => {
|
||||
tracing::error!(error = %error, "Failed to publish execution_completed");
|
||||
.await
|
||||
{
|
||||
match error {
|
||||
crate::bus::BusError::Dropped => {
|
||||
tracing::warn!(error = %error, "Outbound dropped (bus full)");
|
||||
}
|
||||
crate::bus::BusError::Closed => {
|
||||
tracing::error!(error = %error, "Failed to publish execution_completed");
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@ -26,7 +26,7 @@ use crate::tools::task::runtime::SubagentRuntime;
|
||||
use crate::tools::todo_write::TodoItem;
|
||||
use crate::tools::{
|
||||
DefaultSubAgentRuntime, InMemoryTaskRepository, NoopSessionMessageSender, SessionMessageSender,
|
||||
SubAgentRuntimeConfig, SubagentCatalog, TaskTool, ToolRegistry,
|
||||
SubAgentRuntime, SubAgentRuntimeConfig, SubagentCatalog, TaskTool, ToolRegistry,
|
||||
};
|
||||
|
||||
use super::agent_factory::AgentFactory;
|
||||
@ -67,6 +67,7 @@ pub(crate) fn build_session_manager(
|
||||
Option<Arc<McpClientManager>>,
|
||||
Arc<SubagentRuntime>,
|
||||
Arc<ModelSelectionStore>,
|
||||
Option<Arc<dyn SubAgentRuntime>>,
|
||||
),
|
||||
AgentError,
|
||||
> {
|
||||
@ -119,6 +120,7 @@ pub(crate) fn build_session_manager_with_sender(
|
||||
Option<Arc<McpClientManager>>,
|
||||
Arc<SubagentRuntime>,
|
||||
Arc<ModelSelectionStore>,
|
||||
Option<Arc<dyn SubAgentRuntime>>,
|
||||
),
|
||||
AgentError,
|
||||
> {
|
||||
@ -209,10 +211,11 @@ pub(crate) fn build_session_manager_with_sender(
|
||||
}
|
||||
|
||||
// Create SubAgentRuntime (if task tool is enabled)
|
||||
let (factory, task_repository, subagent_runtime): (
|
||||
let (factory, task_repository, subagent_runtime, subagent_executor): (
|
||||
_,
|
||||
Arc<dyn TaskRepository>,
|
||||
Arc<SubagentRuntime>,
|
||||
Option<Arc<dyn SubAgentRuntime>>,
|
||||
) = if task_config.enabled {
|
||||
let task_repository = Arc::new(InMemoryTaskRepository::new());
|
||||
// Build subagent tools with MCP tools (task tool registered separately below)
|
||||
@ -237,6 +240,7 @@ pub(crate) fn build_session_manager_with_sender(
|
||||
default_max_execution_secs: task_config.max_execution_secs,
|
||||
ttl_hours: task_config.ttl_hours,
|
||||
max_nesting_depth: task_config.max_nesting_depth,
|
||||
max_concurrent: task_config.max_concurrent,
|
||||
};
|
||||
|
||||
let default_subagent_runtime = Arc::new(DefaultSubAgentRuntime::new(
|
||||
@ -260,10 +264,14 @@ pub(crate) fn build_session_manager_with_sender(
|
||||
));
|
||||
}
|
||||
|
||||
let subagent_executor: Option<Arc<dyn SubAgentRuntime>> =
|
||||
Some(default_subagent_runtime.clone());
|
||||
|
||||
(
|
||||
factory.with_subagent_runtime(default_subagent_runtime),
|
||||
task_repository,
|
||||
subagent_runtime,
|
||||
subagent_executor,
|
||||
)
|
||||
} else {
|
||||
// task_config 未启用时仍创建 subagent_runtime（供 API 使用）
|
||||
@ -272,6 +280,7 @@ pub(crate) fn build_session_manager_with_sender(
|
||||
factory,
|
||||
Arc::new(InMemoryTaskRepository::new()),
|
||||
subagent_runtime,
|
||||
None,
|
||||
)
|
||||
};
|
||||
|
||||
@ -367,5 +376,6 @@ pub(crate) fn build_session_manager_with_sender(
|
||||
mcp_manager,
|
||||
subagent_runtime,
|
||||
model_selections,
|
||||
subagent_executor,
|
||||
))
|
||||
}
|
||||
|
||||
@ -14,10 +14,12 @@ use crate::storage::{
|
||||
use crate::tools::ToolRegistry;
|
||||
use crate::tools::task::repository::TaskRepository;
|
||||
use crate::tools::task::runtime::SubagentRuntime;
|
||||
use crate::tools::task::SubagentResult;
|
||||
use crate::tools::WaitCoordinator;
|
||||
use async_trait::async_trait;
|
||||
use std::collections::HashMap;
|
||||
use std::sync::Arc;
|
||||
use tokio::sync::{Mutex, mpsc};
|
||||
use tokio::sync::{Mutex, Notify, mpsc};
|
||||
use uuid::Uuid;
|
||||
|
||||
use super::agent_factory::{AgentBuildRequest, AgentFactory};
|
||||
@ -413,12 +415,17 @@ impl Session {
|
||||
|
||||
/// 确保指定 topic 的历史已加载到内存。
|
||||
/// 按 topic_id 键化查找，已存在则直接返回，否则从 DB 加载。
|
||||
/// 加载后对账 pending_subagents DB，替换过时的 "running" 占位。
|
||||
pub fn ensure_chat_loaded(
|
||||
&mut self,
|
||||
chat_id: &str,
|
||||
topic_id: Option<&str>,
|
||||
) -> Result<(), AgentError> {
|
||||
self.history.ensure_chat_loaded(chat_id, topic_id)
|
||||
self.history.ensure_chat_loaded(chat_id, topic_id)?;
|
||||
if let Some(tid) = topic_id {
|
||||
self.reconcile_running_placeholders(tid);
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub fn ensure_agent_prompt_before_user_message(
|
||||
@ -587,19 +594,147 @@ impl Session {
|
||||
self.history.topic_serial_lock(topic_id)
|
||||
}
|
||||
|
||||
/// 确保该 topic 的 sub_done 队列已创建（与 topic_serial_lock 同步初始化）。
|
||||
pub(crate) fn ensure_sub_done_channel(&mut self, topic_id: &str) {
|
||||
self.history.ensure_sub_done_channel(topic_id);
|
||||
}
|
||||
|
||||
/// 获取该 topic 的 sub_done 队列 sender（用于后台子代理发送结果）。
|
||||
#[allow(dead_code)]
|
||||
pub(crate) fn sub_done_sender(
|
||||
&mut self,
|
||||
topic_id: &str,
|
||||
) -> Option<mpsc::Sender<SubagentResult>> {
|
||||
self.history.sub_done_sender(topic_id)
|
||||
}
|
||||
|
||||
/// 取出该 topic 的 sub_done 队列 receiver（wait 工具进入 select! 前调用）。
|
||||
pub(crate) fn take_sub_done_receiver(
|
||||
&mut self,
|
||||
topic_id: &str,
|
||||
) -> Option<mpsc::Receiver<SubagentResult>> {
|
||||
self.history.take_sub_done_receiver(topic_id)
|
||||
}
|
||||
|
||||
/// 归还该 topic 的 sub_done 队列 receiver（wait 工具 select! 结束后调用）。
|
||||
pub(crate) fn restore_sub_done_receiver(
|
||||
&mut self,
|
||||
topic_id: &str,
|
||||
rx: mpsc::Receiver<SubagentResult>,
|
||||
) {
|
||||
self.history.restore_sub_done_receiver(topic_id, rx);
|
||||
}
|
||||
|
||||
/// 获取或创建该 topic 的 wait 唤醒信号。
|
||||
pub(crate) fn wait_wakeup(&mut self, topic_id: &str) -> Arc<Notify> {
|
||||
self.history.wait_wakeup(topic_id)
|
||||
}
|
||||
|
||||
/// 设置该 topic 的等待状态。
|
||||
pub(crate) fn set_waiting(&mut self, topic_id: &str, waiting: bool) {
|
||||
self.history.set_waiting(topic_id, waiting);
|
||||
}
|
||||
|
||||
/// 检查该 topic 是否处于等待状态。
|
||||
pub(crate) fn is_waiting(&self, topic_id: &str) -> bool {
|
||||
self.history.is_waiting(topic_id)
|
||||
}
|
||||
|
||||
/// 按 topic_id 从 DB 重新加载历史到内存
|
||||
pub(crate) fn reload_topic_history(
|
||||
&mut self,
|
||||
chat_id: &str,
|
||||
topic_id: &str,
|
||||
) -> Result<(), AgentError> {
|
||||
self.history.reload_topic_history(chat_id, topic_id)
|
||||
self.history.reload_topic_history(chat_id, topic_id)?;
|
||||
self.reconcile_running_placeholders(topic_id);
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// 对账 pending_subagents DB，将内存 history 中过时的 "running" 占位
|
||||
/// 替换为 DB 中的实际状态。
|
||||
///
|
||||
/// 场景：服务器崩溃重启后，history 中仍保留 "running" 占位，
|
||||
/// 但 DB 中该子代理状态可能已被启动扫描标记为 "interrupted"。
|
||||
/// 若不对账，LLM 会看到 "running" → 调用 wait_for_subagents →
|
||||
/// query_pending_task_ids 返回空（DB 已非 running）→ 返回 "no pending" →
|
||||
/// LLM 困惑：history 说 running 但 wait 说无 pending。
|
||||
///
|
||||
/// 支持两种 content 格式：
|
||||
/// - JSON: `{"status":"running","task_id":"xxx",...}`（当前 task 工具返回格式）
|
||||
/// - 纯文本: `running, task_id=xxx. ...`（旧格式，向后兼容）
|
||||
///
|
||||
/// 仅修改内存缓存，不持久化到 DB（每次从 DB 加载时重新对账，幂等）。
|
||||
fn reconcile_running_placeholders(&mut self, topic_id: &str) {
|
||||
let pending = match self.store.list_pending_subagents(topic_id, None) {
|
||||
Ok(records) => records,
|
||||
Err(e) => {
|
||||
tracing::warn!(
|
||||
error = %e,
|
||||
topic_id = %topic_id,
|
||||
"Failed to query pending_subagents for reconciliation"
|
||||
);
|
||||
return;
|
||||
}
|
||||
};
|
||||
|
||||
if pending.is_empty() {
|
||||
return;
|
||||
}
|
||||
|
||||
let status_map: std::collections::HashMap<&str, &str> = pending
|
||||
.iter()
|
||||
.map(|r| (r.task_id.as_str(), r.status.as_str()))
|
||||
.collect();
|
||||
|
||||
let history = self.history.get_or_create_history(topic_id);
|
||||
let mut reconciled = 0;
|
||||
|
||||
for msg in history.iter_mut() {
|
||||
if msg.role != "tool" {
|
||||
continue;
|
||||
}
|
||||
|
||||
// 尝试提取 task_id（同时支持 JSON 和纯文本格式）
|
||||
let (task_id, is_json) = match extract_task_id_from_content(&msg.content) {
|
||||
Some(id) => id,
|
||||
None => continue,
|
||||
};
|
||||
|
||||
let actual_status = match status_map.get(task_id.as_str()) {
|
||||
Some(s) => *s,
|
||||
None => continue, // 记录不存在（已清理），保留原占位
|
||||
};
|
||||
|
||||
if actual_status == "running" {
|
||||
continue; // 仍在运行，保留占位
|
||||
}
|
||||
|
||||
// 替换为实际状态
|
||||
msg.content = format_reconciled_content(&task_id, actual_status, is_json);
|
||||
reconciled += 1;
|
||||
}
|
||||
|
||||
if reconciled > 0 {
|
||||
tracing::info!(
|
||||
topic_id = %topic_id,
|
||||
reconciled_count = reconciled,
|
||||
"Reconciled stale 'running' placeholders with DB status"
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) fn store(&self) -> Arc<dyn ConversationRepository> {
|
||||
self.history.conversations()
|
||||
}
|
||||
|
||||
/// 获取底层 SessionStore（用于 pending_subagents 查询等）。
|
||||
/// 与 `store()` 不同：后者返回 ConversationRepository trait object，
|
||||
/// 此方法返回具体的 SessionStore 类型，暴露 pending_subagents 等 CRUD。
|
||||
pub(crate) fn session_store(&self) -> Arc<SessionStore> {
|
||||
self.store.clone()
|
||||
}
|
||||
|
||||
pub fn record_skill_offer(&self, chat_id: &str) -> Result<(), AgentError> {
|
||||
if self.skills.is_empty() {
|
||||
return Ok(());
|
||||
@ -621,6 +756,7 @@ impl Session {
|
||||
message_id: Option<&str>,
|
||||
explicit_topic_id: Option<&str>,
|
||||
trace_id: &str,
|
||||
wait_coordinator: Option<Arc<dyn WaitCoordinator>>,
|
||||
) -> Result<AgentLoop, AgentError> {
|
||||
self.create_agent_with_provider_config(
|
||||
chat_id,
|
||||
@ -630,6 +766,7 @@ impl Session {
|
||||
self.provider_config.clone(),
|
||||
explicit_topic_id,
|
||||
trace_id,
|
||||
wait_coordinator,
|
||||
)
|
||||
}
|
||||
|
||||
@ -642,6 +779,7 @@ impl Session {
|
||||
provider_config: LLMProviderConfig,
|
||||
explicit_topic_id: Option<&str>,
|
||||
trace_id: &str,
|
||||
wait_coordinator: Option<Arc<dyn WaitCoordinator>>,
|
||||
) -> Result<AgentLoop, AgentError> {
|
||||
// 优先使用显式传入的 topic_id；回退到当前 chat 的活跃 topic
|
||||
let topic_id = explicit_topic_id
|
||||
@ -658,6 +796,12 @@ impl Session {
|
||||
None => self.pending_cancel_tokens.remove(session_chat_id),
|
||||
};
|
||||
|
||||
// 获取该 topic 的 sub_done 队列 sender（用于异步子代理完成回调）
|
||||
// 仅主 agent 注入；无 topic 时为 None（走同步路径）
|
||||
let sub_done_sender = topic_id
|
||||
.as_deref()
|
||||
.and_then(|tid| self.history.sub_done_sender(tid));
|
||||
|
||||
self.agent_factory.create(AgentBuildRequest {
|
||||
channel_name: &self.channel_name,
|
||||
session_chat_id,
|
||||
@ -668,10 +812,97 @@ impl Session {
|
||||
provider_config,
|
||||
cancel_token,
|
||||
trace_id: Some(trace_id.to_string()),
|
||||
sub_done_sender,
|
||||
wait_coordinator,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
/// 从 tool result content 中提取 task_id，同时支持 JSON 和纯文本格式。
|
||||
///
|
||||
/// JSON 格式: `{"status":"running","task_id":"xxx",...}`
|
||||
/// 纯文本格式: `running, task_id=xxx. ...`
|
||||
///
|
||||
/// 返回 (task_id, is_json)。仅当 status=="running" 时才提取（只对账 running 占位）。
|
||||
fn extract_task_id_from_content(content: &str) -> Option<(String, bool)> {
|
||||
// 尝试 JSON 格式
|
||||
let json_start = content.find('{')?;
|
||||
if let Ok(parsed) = serde_json::from_str::<serde_json::Value>(&content[json_start..]) {
|
||||
let status = parsed.get("status").and_then(|v| v.as_str())?;
|
||||
if status != "running" {
|
||||
return None;
|
||||
}
|
||||
let task_id = parsed.get("task_id").and_then(|v| v.as_str())?;
|
||||
if task_id.is_empty() {
|
||||
return None;
|
||||
}
|
||||
return Some((task_id.to_string(), true));
|
||||
}
|
||||
|
||||
// 回退到纯文本格式（向后兼容）
|
||||
let prefix = "running, task_id=";
|
||||
let rest = content.strip_prefix(prefix)?;
|
||||
let end = rest.find('.').unwrap_or(rest.len());
|
||||
let task_id = &rest[..end];
|
||||
if task_id.is_empty() {
|
||||
None
|
||||
} else {
|
||||
Some((task_id.to_string(), false))
|
||||
}
|
||||
}
|
||||
|
||||
/// 根据子代理的实际状态生成替换内容。
|
||||
///
|
||||
/// JSON 格式：更新 JSON 中的 status 字段（保持前端 parseTaskResult 兼容）。
|
||||
/// 纯文本格式：替换为描述性文本（向后兼容）。
|
||||
fn format_reconciled_content(task_id: &str, status: &str, is_json: bool) -> String {
|
||||
if is_json {
|
||||
// 更新 JSON 中的 status 字段，保持前端 parseTaskResult 能正确解析
|
||||
let placeholder = match status {
|
||||
"interrupted" => format!(
|
||||
"Subagent {} was interrupted (server restart). Result is unavailable.",
|
||||
task_id
|
||||
),
|
||||
"completed" => format!(
|
||||
"Subagent {} has completed. Call wait_for_subagents to retrieve the result.",
|
||||
task_id
|
||||
),
|
||||
"failed" => format!(
|
||||
"Subagent {} has failed. Call wait_for_subagents to retrieve error details.",
|
||||
task_id
|
||||
),
|
||||
"timeout" => format!("Subagent {} timed out.", task_id),
|
||||
"cancelled" => format!("Subagent {} was cancelled.", task_id),
|
||||
other => format!("Subagent {} status: {}.", task_id, other),
|
||||
};
|
||||
serde_json::json!({
|
||||
"status": status,
|
||||
"summary": placeholder,
|
||||
"output": placeholder,
|
||||
"task_id": task_id,
|
||||
})
|
||||
.to_string()
|
||||
} else {
|
||||
match status {
|
||||
"interrupted" => format!(
|
||||
"Subagent {} was interrupted (server restart). Result is unavailable.",
|
||||
task_id
|
||||
),
|
||||
"completed" => format!(
|
||||
"Subagent {} has completed. Call wait_for_subagents to retrieve the result.",
|
||||
task_id
|
||||
),
|
||||
"failed" => format!(
|
||||
"Subagent {} has failed. Call wait_for_subagents to retrieve error details.",
|
||||
task_id
|
||||
),
|
||||
"timeout" => format!("Subagent {} timed out.", task_id),
|
||||
"cancelled" => format!("Subagent {} was cancelled.", task_id),
|
||||
other => format!("Subagent {} status: {}.", task_id, other),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// SessionManager 管理所有 Session，按 channel_name 路由
|
||||
#[derive(Clone)]
|
||||
pub struct SessionManager {
|
||||
@ -763,7 +994,7 @@ impl SessionManager {
|
||||
model_resolver,
|
||||
crate::config::CompactionConfig::default(),
|
||||
)
|
||||
.map(|(session_manager, _, _, _, _)| session_manager)
|
||||
.map(|(session_manager, _, _, _, _, _)| session_manager)
|
||||
}
|
||||
|
||||
pub fn tools(&self) -> Arc<ToolRegistry> {
|
||||
|
||||
@ -1,11 +1,14 @@
|
||||
use std::collections::{HashMap, HashSet};
|
||||
use std::sync::Arc;
|
||||
|
||||
use tokio::sync::{Notify, mpsc};
|
||||
|
||||
use crate::agent::AgentError;
|
||||
use crate::bus::ChatMessage;
|
||||
use crate::storage::{
|
||||
ConversationRepository, SessionRecord, SkillEventRepository, persistent_session_id,
|
||||
};
|
||||
use crate::tools::task::SubagentResult;
|
||||
|
||||
/// 内存中缓存的 topic 历史上限。
|
||||
/// 超过此值时，驱逐非活跃 topic（不在 chat_topic_ids 当前引用中的 topic）。
|
||||
@ -35,6 +38,23 @@ pub(crate) struct SessionHistory {
|
||||
/// 防止并发 loop 操作同一历史的不同快照产生交错序列。
|
||||
/// 不同 topic 之间互不阻塞，支持多话题并发执行。
|
||||
topic_serial_locks: HashMap<String, Arc<tokio::sync::Mutex<()>>>,
|
||||
/// per-topic 子代理完成队列的 sender 端。
|
||||
/// 后台子代理完成后 clone 此 sender 发送 SubagentResult，
|
||||
/// 由 wait_for_subagents 工具的 receiver 端消费。
|
||||
/// 生命周期与 topic_serial_locks 完全一致。
|
||||
sub_done_senders: HashMap<String, mpsc::Sender<SubagentResult>>,
|
||||
/// per-topic 子代理完成队列的 receiver 端。
|
||||
/// 使用 Option 包装以便 wait_for_subagents 工具在 select! 期间
|
||||
/// 将其取出（不持有 Session 锁），结束后归还。
|
||||
sub_done_receivers: HashMap<String, Option<mpsc::Receiver<SubagentResult>>>,
|
||||
/// per-topic wait 唤醒信号。
|
||||
/// wait 释放锁进入 select! 后，process_one 在注入用户消息到 history 后
|
||||
/// 调用 notify_one() 唤醒 wait。
|
||||
wait_wakeups: HashMap<String, Arc<Notify>>,
|
||||
/// per-topic 等待状态标志。
|
||||
/// true 表示该 topic 的 agent 正在 wait_for_subagents 中等待（已释放 serial_lock）。
|
||||
/// process_one 持锁后检查此标志：true 则注入用户消息 + wakeup，false 则正常处理。
|
||||
waiting_flags: HashMap<String, bool>,
|
||||
conversations: Arc<dyn ConversationRepository>,
|
||||
skill_events: Arc<dyn SkillEventRepository>,
|
||||
}
|
||||
@ -67,6 +87,20 @@ impl SessionHistory {
|
||||
if active.contains(tid.as_str()) || self.compression_in_flight.contains(*tid) {
|
||||
return false;
|
||||
}
|
||||
// 不变量 2：lock 实例单射性 — 驱逐会移除 topic_serial_locks/wakeup/sender 等 entry，
|
||||
// 下次访问会创建新实例，破坏原 Arc 持有者（wait_coordinator）与新建者（process_one）
|
||||
// 的串行化。因此必须充分覆盖所有"活跃"语义：
|
||||
// a) waiting_flag=true → wait_coordinator 正在 select! 等待
|
||||
// b) sub_done_receivers[tid]=None → receiver 被 take 走（wait_coordinator 持有）
|
||||
// c) topic_serial_lock 被持有 → 有 agent 任务正在处理
|
||||
// 任一为真都不能驱逐。
|
||||
if self.waiting_flags.get(*tid).copied().unwrap_or(false) {
|
||||
return false;
|
||||
}
|
||||
if matches!(self.sub_done_receivers.get(*tid), Some(None)) {
|
||||
// receiver 被 take 走 = wait_coordinator 正在 select! 中
|
||||
return false;
|
||||
}
|
||||
// 检查是否有活跃 agent 任务（serial lock 被持有）
|
||||
// try_lock 成功 = 锁空闲 = 无活跃任务 = 可驱逐
|
||||
// try_lock 失败 = 锁被持有 = 有活跃任务 = 不驱逐
|
||||
@ -81,6 +115,11 @@ impl SessionHistory {
|
||||
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);
|
||||
self.topic_serial_locks.remove(&tid);
|
||||
self.sub_done_senders.remove(&tid);
|
||||
self.sub_done_receivers.remove(&tid);
|
||||
self.wait_wakeups.remove(&tid);
|
||||
self.waiting_flags.remove(&tid);
|
||||
tracing::info!(
|
||||
topic_id = %tid,
|
||||
evicted_messages = msg_count,
|
||||
@ -101,6 +140,10 @@ impl SessionHistory {
|
||||
chat_topic_ids: HashMap::new(),
|
||||
compression_in_flight: HashSet::new(),
|
||||
topic_serial_locks: HashMap::new(),
|
||||
sub_done_senders: HashMap::new(),
|
||||
sub_done_receivers: HashMap::new(),
|
||||
wait_wakeups: HashMap::new(),
|
||||
waiting_flags: HashMap::new(),
|
||||
conversations,
|
||||
skill_events,
|
||||
}
|
||||
@ -109,6 +152,9 @@ impl SessionHistory {
|
||||
/// 获取或创建该 topic 的串行化锁。
|
||||
/// 同一 topic 的所有消息处理共享同一个锁，保证串行执行；
|
||||
/// 不同 topic 之间互不阻塞，支持多话题并发执行。
|
||||
///
|
||||
/// 同时同步创建该 topic 的 sub_done 队列、wait 唤醒信号和等待标志，
|
||||
/// 生命周期与 serial_lock 完全一致。
|
||||
pub(crate) fn topic_serial_lock(&mut self, topic_id: &str) -> Arc<tokio::sync::Mutex<()>> {
|
||||
self.topic_serial_locks
|
||||
.entry(topic_id.to_string())
|
||||
@ -116,6 +162,68 @@ impl SessionHistory {
|
||||
.clone()
|
||||
}
|
||||
|
||||
/// 确保该 topic 的 sub_done 队列已创建。
|
||||
/// 应在 topic 初始化时（与 topic_serial_lock 同步）调用。
|
||||
/// 队列容量为 32，足够缓存多个子代理同时完成的结果。
|
||||
pub(crate) fn ensure_sub_done_channel(&mut self, topic_id: &str) {
|
||||
if !self.sub_done_senders.contains_key(topic_id) {
|
||||
let (tx, rx) = mpsc::channel::<SubagentResult>(32);
|
||||
self.sub_done_senders.insert(topic_id.to_string(), tx);
|
||||
self.sub_done_receivers
|
||||
.insert(topic_id.to_string(), Some(rx));
|
||||
}
|
||||
}
|
||||
|
||||
/// 获取该 topic 的 sub_done 队列 sender（用于后台子代理发送结果）。
|
||||
/// 调用前应已通过 `ensure_sub_done_channel` 创建队列。
|
||||
pub(crate) fn sub_done_sender(&mut self, topic_id: &str) -> Option<mpsc::Sender<SubagentResult>> {
|
||||
self.ensure_sub_done_channel(topic_id);
|
||||
self.sub_done_senders.get(topic_id).cloned()
|
||||
}
|
||||
|
||||
/// 取出该 topic 的 sub_done 队列 receiver。
|
||||
/// 由 wait_for_subagents 工具在进入 select! 前调用（需释放 Session 锁）。
|
||||
/// 调用后 receiver 不在 map 中，需通过 `restore_sub_done_receiver` 归还。
|
||||
pub(crate) fn take_sub_done_receiver(
|
||||
&mut self,
|
||||
topic_id: &str,
|
||||
) -> Option<mpsc::Receiver<SubagentResult>> {
|
||||
self.sub_done_receivers
|
||||
.get_mut(topic_id)
|
||||
.and_then(|opt| opt.take())
|
||||
}
|
||||
|
||||
/// 归还该 topic 的 sub_done 队列 receiver。
|
||||
/// 由 wait_for_subagents 工具在 select! 结束后调用。
|
||||
pub(crate) fn restore_sub_done_receiver(
|
||||
&mut self,
|
||||
topic_id: &str,
|
||||
rx: mpsc::Receiver<SubagentResult>,
|
||||
) {
|
||||
self.sub_done_receivers
|
||||
.insert(topic_id.to_string(), Some(rx));
|
||||
}
|
||||
|
||||
/// 获取或创建该 topic 的 wait 唤醒信号。
|
||||
pub(crate) fn wait_wakeup(&mut self, topic_id: &str) -> Arc<Notify> {
|
||||
self.wait_wakeups
|
||||
.entry(topic_id.to_string())
|
||||
.or_insert_with(|| Arc::new(Notify::new()))
|
||||
.clone()
|
||||
}
|
||||
|
||||
/// 设置该 topic 的等待状态。
|
||||
/// true = agent 正在 wait_for_subagents 中等待（已释放 serial_lock）。
|
||||
pub(crate) fn set_waiting(&mut self, topic_id: &str, waiting: bool) {
|
||||
self.waiting_flags.insert(topic_id.to_string(), waiting);
|
||||
}
|
||||
|
||||
/// 检查该 topic 是否处于等待状态。
|
||||
/// process_one 持锁后调用：true 则走注入路径，false 则正常处理。
|
||||
pub(crate) fn is_waiting(&self, topic_id: &str) -> bool {
|
||||
self.waiting_flags.get(topic_id).copied().unwrap_or(false)
|
||||
}
|
||||
|
||||
pub(crate) fn persistent_session_id(&self, chat_id: &str) -> String {
|
||||
persistent_session_id(&self.channel_name, chat_id)
|
||||
}
|
||||
@ -212,6 +320,11 @@ impl SessionHistory {
|
||||
// （仅在无活跃任务时安全移除；有活跃任务时 lock 被 Arc clone 持有，
|
||||
// 移除 HashMap entry 不影响正在使用 lock 的任务）
|
||||
self.topic_serial_locks.remove(topic_id);
|
||||
// 同步清理 sub_done 队列、wait 唤醒信号和等待标志
|
||||
self.sub_done_senders.remove(topic_id);
|
||||
self.sub_done_receivers.remove(topic_id);
|
||||
self.wait_wakeups.remove(topic_id);
|
||||
self.waiting_flags.remove(topic_id);
|
||||
}
|
||||
|
||||
/// 清空指定 chat/topic 的内存历史和 DB 消息。
|
||||
|
||||
@ -14,7 +14,7 @@ use crate::tools::{
|
||||
BashTool, CalculatorTool, FileEditTool, FileReadTool, FileWriteTool, HttpRequestTool,
|
||||
MemoryManageTool, MemorySearchTool, SchedulerManageTool, SessionMessageSender, SessionSendTool,
|
||||
ShellSessionManager, SkillActivateTool, SkillManageTool, SubAgentRuntime, TaskTool, TimeTool,
|
||||
TodoReadTool, TodoWriteTool, ToolRegistry, WebFetchTool,
|
||||
TodoReadTool, TodoWriteTool, ToolRegistry, WaitForSubagentsTool, WebFetchTool,
|
||||
};
|
||||
|
||||
pub(crate) struct ToolRegistryFactory {
|
||||
@ -160,6 +160,11 @@ impl ToolRegistryFactory {
|
||||
if self.is_enabled("task") && self.task_config.enabled {
|
||||
if let Some(runtime) = &self.subagent_runtime {
|
||||
registry.register(TaskTool::new(runtime.clone(), None));
|
||||
// 注册 wait_for_subagents 工具（仅主 agent，用于等待异步子代理完成）
|
||||
// 默认超时从配置读取，LLM 可通过 timeout_secs 参数覆盖
|
||||
registry.register(WaitForSubagentsTool::new(
|
||||
self.task_config.wait_default_timeout_secs,
|
||||
));
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
291
src/gateway/wait_coordinator.rs
Normal file
291
src/gateway/wait_coordinator.rs
Normal file
@ -0,0 +1,291 @@
|
||||
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()
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
@ -535,6 +535,7 @@ async fn handle_inbound(
|
||||
router.register(Box::new(StopExecutionCommandHandler::new(
|
||||
state.cancel_manager.clone(),
|
||||
state.session_manager.clone(),
|
||||
state.subagent_executor.clone(),
|
||||
)));
|
||||
|
||||
// 构建命令上下文
|
||||
|
||||
@ -330,6 +330,12 @@ pub enum WsOutbound {
|
||||
timestamp: Option<i64>,
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
subagent_task_id: Option<String>,
|
||||
/// 子代理最终状态（completed/failed/timeout/cancelled/interrupted），
|
||||
/// 供前端更新主视图中 task tool result 占位消息的显示状态。
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
subagent_status: Option<String>,
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
subagent_summary: Option<String>,
|
||||
},
|
||||
#[serde(rename = "todo_list")]
|
||||
TodoList {
|
||||
|
||||
@ -221,6 +221,8 @@ pub(crate) fn ws_outbound_from_outbound_message(message: &OutboundMessage) -> Ve
|
||||
topic_id: message.metadata.get("topic_id").cloned(),
|
||||
timestamp: Some(crate::protocol::now_timestamp()),
|
||||
subagent_task_id: message.metadata.get("subagent_task_id").cloned(),
|
||||
subagent_status: message.metadata.get("subagent_status").cloned(),
|
||||
subagent_summary: message.metadata.get("subagent_summary").cloned(),
|
||||
}],
|
||||
}
|
||||
}
|
||||
|
||||
@ -323,3 +323,28 @@ pub(super) fn add_column_if_missing(conn: &Connection, sql: &str) -> Result<(),
|
||||
Err(error) => Err(StorageError::Database(error)),
|
||||
}
|
||||
}
|
||||
|
||||
/// pending_subagents 表：跟踪异步子代理执行状态，用于崩溃恢复。
|
||||
pub(super) fn ensure_pending_subagents_schema(conn: &Connection) -> Result<(), StorageError> {
|
||||
conn.execute_batch(
|
||||
"
|
||||
CREATE TABLE IF NOT EXISTS pending_subagents (
|
||||
task_id TEXT PRIMARY KEY,
|
||||
parent_session_id TEXT NOT NULL,
|
||||
parent_topic_id TEXT NOT NULL,
|
||||
parent_chat_id TEXT NOT NULL,
|
||||
parent_channel TEXT NOT NULL,
|
||||
def_name TEXT,
|
||||
spawned_at INTEGER NOT NULL,
|
||||
status TEXT NOT NULL DEFAULT 'running'
|
||||
);
|
||||
|
||||
CREATE INDEX IF NOT EXISTS idx_pending_subagents_topic
|
||||
ON pending_subagents(parent_topic_id, status);
|
||||
|
||||
CREATE INDEX IF NOT EXISTS idx_pending_subagents_session
|
||||
ON pending_subagents(parent_session_id);
|
||||
",
|
||||
)?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
@ -28,10 +28,10 @@ pub use ports::{
|
||||
SkillEventRepository, TodoRepository,
|
||||
};
|
||||
pub use records::{
|
||||
ALLOWED_MEMORY_NAMESPACES, GLOBAL_SCOPE_KEY, MemoryRecord, MemoryUpsert, SchedulerJobRecord,
|
||||
SchedulerJobState, SchedulerJobStatus, SchedulerJobUpsert, SessionRecord, SessionTokenStats,
|
||||
SkillEventRecord, TodoRecord, TopicRecord, allowed_namespace_names, get_namespace_description,
|
||||
is_valid_namespace,
|
||||
ALLOWED_MEMORY_NAMESPACES, GLOBAL_SCOPE_KEY, MemoryRecord, MemoryUpsert, PendingSubagentRecord,
|
||||
SchedulerJobRecord, SchedulerJobState, SchedulerJobStatus, SchedulerJobUpsert, SessionRecord,
|
||||
SessionTokenStats, SkillEventRecord, TodoRecord, TopicRecord, allowed_namespace_names,
|
||||
get_namespace_description, is_valid_namespace,
|
||||
};
|
||||
|
||||
#[derive(Clone)]
|
||||
@ -234,6 +234,7 @@ impl SessionStore {
|
||||
ensure_scheduler_schema(&conn)?;
|
||||
ensure_memory_scope_key_migration(&conn)?;
|
||||
ensure_todos_schema(&conn)?;
|
||||
ensure_pending_subagents_schema(&conn)?;
|
||||
|
||||
drop(conn);
|
||||
|
||||
@ -2052,6 +2053,140 @@ impl SessionStore {
|
||||
}
|
||||
Ok(todos)
|
||||
}
|
||||
|
||||
// ==================== pending_subagents ====================
|
||||
|
||||
/// 插入一条 pending_subagent 记录（task 工具 spawn 时调用）。
|
||||
pub fn insert_pending_subagent(
|
||||
&self,
|
||||
record: &PendingSubagentRecord,
|
||||
) -> Result<(), StorageError> {
|
||||
let conn = self.pool.get()?;
|
||||
conn.execute(
|
||||
"INSERT OR REPLACE INTO pending_subagents
|
||||
(task_id, parent_session_id, parent_topic_id, parent_chat_id, parent_channel, def_name, spawned_at, status)
|
||||
VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8)",
|
||||
params![
|
||||
record.task_id,
|
||||
record.parent_session_id,
|
||||
record.parent_topic_id,
|
||||
record.parent_chat_id,
|
||||
record.parent_channel,
|
||||
record.def_name,
|
||||
record.spawned_at,
|
||||
record.status,
|
||||
],
|
||||
)?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// 查询指定 topic 下匹配状态的 pending_subagent 记录。
|
||||
/// `status` 为 None 时查询所有状态。
|
||||
pub fn list_pending_subagents(
|
||||
&self,
|
||||
topic_id: &str,
|
||||
status: Option<&str>,
|
||||
) -> Result<Vec<PendingSubagentRecord>, StorageError> {
|
||||
let conn = self.pool.get()?;
|
||||
let sql = if status.is_some() {
|
||||
"SELECT task_id, parent_session_id, parent_topic_id, parent_chat_id, parent_channel, def_name, spawned_at, status
|
||||
FROM pending_subagents
|
||||
WHERE parent_topic_id = ?1 AND status = ?2
|
||||
ORDER BY spawned_at ASC"
|
||||
} else {
|
||||
"SELECT task_id, parent_session_id, parent_topic_id, parent_chat_id, parent_channel, def_name, spawned_at, status
|
||||
FROM pending_subagents
|
||||
WHERE parent_topic_id = ?1
|
||||
ORDER BY spawned_at ASC"
|
||||
};
|
||||
let mut stmt = conn.prepare(sql)?;
|
||||
let rows = if let Some(s) = status {
|
||||
stmt.query_map(params![topic_id, s], map_pending_subagent_record)?
|
||||
} else {
|
||||
stmt.query_map(params![topic_id], map_pending_subagent_record)?
|
||||
};
|
||||
let mut result = Vec::new();
|
||||
for row in rows {
|
||||
result.push(row?);
|
||||
}
|
||||
Ok(result)
|
||||
}
|
||||
|
||||
/// 获取指定 task_id 的 pending_subagent 记录。
|
||||
pub fn get_pending_subagent(
|
||||
&self,
|
||||
task_id: &str,
|
||||
) -> Result<Option<PendingSubagentRecord>, StorageError> {
|
||||
let conn = self.pool.get()?;
|
||||
let mut stmt = conn.prepare(
|
||||
"SELECT task_id, parent_session_id, parent_topic_id, parent_chat_id, parent_channel, def_name, spawned_at, status
|
||||
FROM pending_subagents
|
||||
WHERE task_id = ?1",
|
||||
)?;
|
||||
let mut rows = stmt.query_map(params![task_id], map_pending_subagent_record)?;
|
||||
match rows.next() {
|
||||
Some(row) => Ok(Some(row?)),
|
||||
None => Ok(None),
|
||||
}
|
||||
}
|
||||
|
||||
/// 更新指定 task_id 的状态（子代理完成或取消时调用）。
|
||||
pub fn update_pending_subagent_status(
|
||||
&self,
|
||||
task_id: &str,
|
||||
new_status: &str,
|
||||
) -> Result<(), StorageError> {
|
||||
let conn = self.pool.get()?;
|
||||
conn.execute(
|
||||
"UPDATE pending_subagents SET status = ?1 WHERE task_id = ?2",
|
||||
params![new_status, task_id],
|
||||
)?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// 条件更新状态：仅在当前状态为 `expected_current` 时才更新为 `new_status`。
|
||||
///
|
||||
/// 实现状态机不可逆性不变量：避免 cancel 路径覆盖 spawn 已写入的终态
|
||||
/// （completed → cancelled 是非法转换）。
|
||||
///
|
||||
/// 返回是否实际更新（affected rows > 0）。false 表示状态已被其他路径更新，
|
||||
/// 调用方应跳过后续基于该假设的操作。
|
||||
pub fn try_update_pending_subagent_status(
|
||||
&self,
|
||||
task_id: &str,
|
||||
expected_current: &str,
|
||||
new_status: &str,
|
||||
) -> Result<bool, StorageError> {
|
||||
let conn = self.pool.get()?;
|
||||
let affected = conn.execute(
|
||||
"UPDATE pending_subagents SET status = ?1 WHERE task_id = ?2 AND status = ?3",
|
||||
params![new_status, task_id, expected_current],
|
||||
)?;
|
||||
Ok(affected > 0)
|
||||
}
|
||||
|
||||
/// 将所有 running 状态的 pending_subagent 标记为 interrupted（启动时崩溃恢复调用）。
|
||||
pub fn mark_all_running_as_interrupted(&self) -> Result<usize, StorageError> {
|
||||
let conn = self.pool.get()?;
|
||||
let affected = conn.execute(
|
||||
"UPDATE pending_subagents SET status = 'interrupted' WHERE status = 'running'",
|
||||
[],
|
||||
)?;
|
||||
Ok(affected)
|
||||
}
|
||||
}
|
||||
|
||||
fn map_pending_subagent_record(row: &rusqlite::Row<'_>) -> rusqlite::Result<PendingSubagentRecord> {
|
||||
Ok(PendingSubagentRecord {
|
||||
task_id: row.get(0)?,
|
||||
parent_session_id: row.get(1)?,
|
||||
parent_topic_id: row.get(2)?,
|
||||
parent_chat_id: row.get(3)?,
|
||||
parent_channel: row.get(4)?,
|
||||
def_name: row.get(5)?,
|
||||
spawned_at: row.get(6)?,
|
||||
status: row.get(7)?,
|
||||
})
|
||||
}
|
||||
|
||||
pub fn persistent_session_id(channel_name: &str, chat_id: &str) -> String {
|
||||
|
||||
@ -111,6 +111,31 @@ pub struct TopicRecord {
|
||||
pub message_count: i64,
|
||||
}
|
||||
|
||||
/// pending_subagents 表的记录，跟踪异步子代理执行状态。
|
||||
///
|
||||
/// 生命周期：task 工具 spawn 时插入（status=running），
|
||||
/// 子代理完成时更新为 completed/failed/timeout，
|
||||
/// 进程重启时 running 状态被标记为 interrupted。
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct PendingSubagentRecord {
|
||||
/// 子代理 task_id（主键，与 TaskSession.id 一致）
|
||||
pub task_id: String,
|
||||
/// 父会话 session_id
|
||||
pub parent_session_id: String,
|
||||
/// 父会话 topic_id（用于按 topic 查询未完成子代理）
|
||||
pub parent_topic_id: String,
|
||||
/// 父会话 chat_id
|
||||
pub parent_chat_id: String,
|
||||
/// 父会话 channel_name
|
||||
pub parent_channel: String,
|
||||
/// 子代理定义名称（可选，用于诊断）
|
||||
pub def_name: Option<String>,
|
||||
/// 子代理启动时间戳
|
||||
pub spawned_at: i64,
|
||||
/// 执行状态：running / completed / failed / interrupted / cancelled / timeout
|
||||
pub status: String,
|
||||
}
|
||||
|
||||
/// 单个 session 的 token 用量统计（聚合结果）。
|
||||
#[derive(Debug, Clone, Default, Serialize, Deserialize)]
|
||||
pub struct SessionTokenStats {
|
||||
|
||||
@ -18,6 +18,7 @@ pub mod time;
|
||||
pub mod todo_read;
|
||||
pub mod todo_write;
|
||||
pub mod traits;
|
||||
pub mod wait_tool;
|
||||
pub mod web_fetch;
|
||||
|
||||
pub use bash::BashTool;
|
||||
@ -45,7 +46,8 @@ pub use task::{
|
||||
pub use time::TimeTool;
|
||||
pub use todo_read::TodoReadTool;
|
||||
pub use todo_write::TodoWriteTool;
|
||||
pub use traits::{Tool, ToolContext, ToolResult};
|
||||
pub use traits::{Tool, ToolContext, ToolResult, WaitCoordinator, WaitEvent};
|
||||
pub use wait_tool::WaitForSubagentsTool;
|
||||
pub use web_fetch::WebFetchTool;
|
||||
|
||||
/// Extract a string parameter from JSON args.
|
||||
|
||||
@ -18,6 +18,9 @@ pub enum TaskError {
|
||||
#[error("Task execution timed out")]
|
||||
Timeout,
|
||||
|
||||
#[error("Task cancelled by user")]
|
||||
Cancelled,
|
||||
|
||||
#[error("Repository error: {0}")]
|
||||
RepositoryError(#[from] StorageError),
|
||||
|
||||
@ -35,6 +38,7 @@ impl TaskError {
|
||||
pub fn as_status(&self) -> &'static str {
|
||||
match self {
|
||||
Self::Timeout => "timeout",
|
||||
Self::Cancelled => "cancelled",
|
||||
Self::SessionNotFound(_) => "failed",
|
||||
Self::InvalidParentSession => "failed",
|
||||
Self::AgentCreationFailed(_) => "failed",
|
||||
|
||||
@ -14,6 +14,6 @@ pub use runtime::{
|
||||
};
|
||||
pub use tool::TaskTool;
|
||||
pub use types::{
|
||||
SubagentDef, SubagentSource, SubagentType, TaskDefinition, TaskHandle, TaskSession,
|
||||
TaskSessionState, TaskToolArgs, TaskToolResult,
|
||||
SubagentDef, SubagentResult, SubagentSource, SubagentStatus, SubagentType, TaskDefinition,
|
||||
TaskHandle, TaskSession, TaskSessionState, TaskToolArgs, TaskToolResult,
|
||||
};
|
||||
|
||||
@ -8,6 +8,49 @@ use std::time::Duration;
|
||||
use async_trait::async_trait;
|
||||
use serde::Deserialize;
|
||||
|
||||
/// RAII guard：spawn 任务退出时（正常/early return/panic）确定性清理 cancel_registry 条目。
|
||||
///
|
||||
/// 实现不变量 3（资源生命周期与作用域严格绑定）：
|
||||
/// spawn 块末行的手动清理是脆弱的，panic/early return 会绕过；
|
||||
/// 用 Drop impl 把清理封进作用域语义，由编译器保证执行。
|
||||
struct CancelRegistryGuard {
|
||||
task_id: String,
|
||||
registry: Arc<parking_lot::Mutex<HashMap<String, tokio_util::sync::CancellationToken>>>,
|
||||
/// 标记是否已显式释放（例如 spawn 块成功路径末尾主动 disarm）。
|
||||
/// 默认 false，drop 时执行清理。
|
||||
disarmed: bool,
|
||||
}
|
||||
|
||||
impl CancelRegistryGuard {
|
||||
fn new(
|
||||
task_id: String,
|
||||
registry: Arc<parking_lot::Mutex<HashMap<String, tokio_util::sync::CancellationToken>>>,
|
||||
) -> Self {
|
||||
Self {
|
||||
task_id,
|
||||
registry,
|
||||
disarmed: false,
|
||||
}
|
||||
}
|
||||
|
||||
/// 显式释放：成功路径末尾调用，避免重复清理。
|
||||
/// （实际 drop 也会幂等移除，但 disarm 让语义更清晰。）
|
||||
#[allow(dead_code)]
|
||||
fn disarm(&mut self) {
|
||||
self.disarmed = true;
|
||||
}
|
||||
}
|
||||
|
||||
impl Drop for CancelRegistryGuard {
|
||||
fn drop(&mut self) {
|
||||
if self.disarmed {
|
||||
return;
|
||||
}
|
||||
// 幂等：条目可能已被 cancel_pending_for_topic 移除或先前已 drop
|
||||
self.registry.lock().remove(&self.task_id);
|
||||
}
|
||||
}
|
||||
|
||||
use crate::agent::{
|
||||
AgentLoop, AgentRuntimeConfig, EmittedMessageHandler, PersistingEmittedMessageHandler,
|
||||
SystemPrompt, SystemPromptContext, SystemPromptProvider,
|
||||
@ -20,14 +63,18 @@ use crate::domain::CapabilityPolicy;
|
||||
use crate::experts::ExpertRuntime;
|
||||
use crate::providers::StreamDelta;
|
||||
use crate::skills::SkillRuntime;
|
||||
use crate::storage::{ConversationRepository, SessionStore};
|
||||
use crate::storage::{ConversationRepository, PendingSubagentRecord, SessionStore};
|
||||
use crate::tools::{ToolContext, ToolRegistry};
|
||||
use crate::utils::current_timestamp;
|
||||
|
||||
use super::error::TaskError;
|
||||
use super::prompt::{SubagentPromptBuilder, extract_summary};
|
||||
use super::repository::TaskRepository;
|
||||
use super::tool::TaskTool;
|
||||
use super::types::{SubagentDef, SubagentSource, TaskDefinition, TaskSession, TaskToolResult};
|
||||
use super::types::{
|
||||
SubagentDef, SubagentResult, SubagentSource, SubagentStatus, TaskDefinition, TaskSession,
|
||||
TaskToolResult,
|
||||
};
|
||||
|
||||
/// 子代理运行时配置
|
||||
#[derive(Debug, Clone)]
|
||||
@ -40,6 +87,8 @@ pub struct SubAgentRuntimeConfig {
|
||||
pub ttl_hours: u64,
|
||||
/// 子代理最大嵌套深度（0 = 禁止嵌套，1 = 允许 1 层孙代理）
|
||||
pub max_nesting_depth: u32,
|
||||
/// 异步子代理最大并发数（Semaphore 限流）
|
||||
pub max_concurrent: usize,
|
||||
}
|
||||
|
||||
impl Default for SubAgentRuntimeConfig {
|
||||
@ -62,6 +111,7 @@ impl Default for SubAgentRuntimeConfig {
|
||||
default_max_execution_secs: 3600, // 60分钟
|
||||
ttl_hours: 24,
|
||||
max_nesting_depth: 1,
|
||||
max_concurrent: 8,
|
||||
}
|
||||
}
|
||||
}
|
||||
@ -92,6 +142,12 @@ pub trait SubAgentRuntime: Send + Sync + 'static {
|
||||
|
||||
/// 获取可用的子代理类型列表
|
||||
fn available_subagent_names(&self) -> Vec<String>;
|
||||
|
||||
/// 取消指定 topic 下所有正在运行的异步子代理。
|
||||
///
|
||||
/// 用于 /stop 命令传播：用户取消主 agent 时，同步取消其后台子代理。
|
||||
/// 返回被触发取消的子代理数量。
|
||||
async fn cancel_pending_for_topic(&self, topic_id: &str) -> usize;
|
||||
}
|
||||
|
||||
/// 静态系统提示词提供者（用于子代理）
|
||||
@ -312,6 +368,15 @@ fn build_subagent_event_metadata(session: &TaskSession) -> HashMap<String, Strin
|
||||
"topic_id".to_string(),
|
||||
session.parent_topic_id.clone().unwrap_or_default(),
|
||||
);
|
||||
// 子代理最终状态（completed/failed/timeout/cancelled/interrupted），
|
||||
// 供前端更新主视图中 task tool result 占位消息的显示状态。
|
||||
metadata.insert(
|
||||
"subagent_status".to_string(),
|
||||
session.state.as_str().to_string(),
|
||||
);
|
||||
if let Some(ref summary) = session.summary {
|
||||
metadata.insert("subagent_summary".to_string(), summary.clone());
|
||||
}
|
||||
metadata
|
||||
}
|
||||
|
||||
@ -392,6 +457,11 @@ pub struct DefaultSubAgentRuntime {
|
||||
store: Arc<SessionStore>,
|
||||
/// 技能运行时（实时计算技能索引，替代冻结快照）
|
||||
skills: Arc<SkillRuntime>,
|
||||
/// 异步子代理并发限流（按 config.max_concurrent 初始化）
|
||||
semaphore: Arc<tokio::sync::Semaphore>,
|
||||
/// task_id → CancellationToken 映射，用于取消传播
|
||||
/// Arc 包装以便 spawned task 完成后清理自身条目
|
||||
cancel_registry: Arc<parking_lot::Mutex<HashMap<String, tokio_util::sync::CancellationToken>>>,
|
||||
}
|
||||
|
||||
impl DefaultSubAgentRuntime {
|
||||
@ -407,6 +477,8 @@ impl DefaultSubAgentRuntime {
|
||||
store: Arc<SessionStore>,
|
||||
skills: Arc<SkillRuntime>,
|
||||
) -> Self {
|
||||
let max_concurrent = config.max_concurrent.max(1);
|
||||
let semaphore = Arc::new(tokio::sync::Semaphore::new(max_concurrent));
|
||||
Self {
|
||||
config,
|
||||
task_repository,
|
||||
@ -418,6 +490,8 @@ impl DefaultSubAgentRuntime {
|
||||
bus,
|
||||
store,
|
||||
skills,
|
||||
semaphore,
|
||||
cancel_registry: Arc::new(parking_lot::Mutex::new(HashMap::new())),
|
||||
}
|
||||
}
|
||||
|
||||
@ -553,6 +627,12 @@ impl DefaultSubAgentRuntime {
|
||||
parent_capability: def.map(|d| d.capability.clone()),
|
||||
// 从父 ToolContext 继承 trace_id，保持端到端追踪贯通子代理
|
||||
trace_id: trace_id.clone(),
|
||||
// 子代理不注入 sub_done_sender：嵌套层不支持异步，走同步路径
|
||||
sub_done_sender: None,
|
||||
// 子代理不注入 wait_coordinator：嵌套层不支持异步 wait
|
||||
wait_coordinator: None,
|
||||
// 子代理不注入 cancel_rx：嵌套层不支持异步 wait，无需 cancel 检查
|
||||
cancel_rx: None,
|
||||
});
|
||||
|
||||
// 如果有 MessageBus，附加实时广播 emitter
|
||||
@ -596,6 +676,18 @@ impl DefaultSubAgentRuntime {
|
||||
session: &TaskSession,
|
||||
def: &SubagentDef,
|
||||
prompt: String,
|
||||
) -> Result<TaskToolResult, TaskError> {
|
||||
let max_secs = self.effective_max_execution_secs(def);
|
||||
Self::execute_task_static(agent, session, def, prompt, max_secs).await
|
||||
}
|
||||
|
||||
/// 静态执行任务（供 tokio::spawn 调用，不依赖 &self）
|
||||
async fn execute_task_static(
|
||||
agent: AgentLoop,
|
||||
session: &TaskSession,
|
||||
_def: &SubagentDef,
|
||||
prompt: String,
|
||||
max_secs: u64,
|
||||
) -> Result<TaskToolResult, TaskError> {
|
||||
// 构建初始消息
|
||||
let history = vec![ChatMessage::user(prompt)];
|
||||
@ -606,7 +698,6 @@ impl DefaultSubAgentRuntime {
|
||||
};
|
||||
|
||||
// 设置超时
|
||||
let max_secs = self.effective_max_execution_secs(def);
|
||||
let timeout_duration = Duration::from_secs(max_secs);
|
||||
|
||||
let result = tokio::time::timeout(
|
||||
@ -825,62 +916,282 @@ impl SubAgentRuntime for DefaultSubAgentRuntime {
|
||||
}
|
||||
}
|
||||
|
||||
// 6-8. 构建提示词、创建子代理、执行任务
|
||||
// 统一为单个 Result 表达式：model_resolver / create_subagent / execute_task
|
||||
// 的任何失败都流入下方 match 的 Err 分支，经 handle_task_failure 返回结构化结果。
|
||||
let result: Result<TaskToolResult, TaskError> = {
|
||||
// 6. 构建子代理系统提示词
|
||||
// 实时按 def.capability 过滤技能索引（替代冻结快照，反映运行时技能增删）
|
||||
let skills_index = if def.capability.has_skill_policy() {
|
||||
self.skills.system_index_prompt_filtered(
|
||||
def.capability.allowed_skills.as_deref(),
|
||||
&def.capability.denied_skills,
|
||||
// 6. 构建子代理系统提示词
|
||||
// 实时按 def.capability 过滤技能索引（替代冻结快照，反映运行时技能增删）
|
||||
let skills_index = if def.capability.has_skill_policy() {
|
||||
self.skills.system_index_prompt_filtered(
|
||||
def.capability.allowed_skills.as_deref(),
|
||||
&def.capability.denied_skills,
|
||||
)
|
||||
} else {
|
||||
self.skills.system_index_prompt()
|
||||
};
|
||||
// 同步解析 def 中的 provider/model 覆盖，保证环境提示中的模型名与实际使用的模型一致
|
||||
let effective_provider_config = match (def.provider.is_some(), def.model.is_some()) {
|
||||
(true, _) | (_, true) => self
|
||||
.model_resolver
|
||||
.resolve(
|
||||
def.provider.as_deref(),
|
||||
def.model.as_deref(),
|
||||
&self.provider_config,
|
||||
)
|
||||
} else {
|
||||
self.skills.system_index_prompt()
|
||||
};
|
||||
// 同步解析 def 中的 provider/model 覆盖，保证环境提示中的模型名与实际使用的模型一致
|
||||
let effective_provider_config = match (def.provider.is_some(), def.model.is_some()) {
|
||||
(true, _) | (_, true) => self
|
||||
.model_resolver
|
||||
.resolve(
|
||||
def.provider.as_deref(),
|
||||
def.model.as_deref(),
|
||||
&self.provider_config,
|
||||
)
|
||||
.map_err(|e| {
|
||||
TaskError::AgentCreationFailed(format!(
|
||||
"subagent '{}' model resolution failed: {}",
|
||||
def.name, e
|
||||
))
|
||||
})?,
|
||||
_ => self.provider_config.clone(),
|
||||
};
|
||||
let system_prompt = SubagentPromptBuilder::build(
|
||||
&def,
|
||||
&task.description,
|
||||
&task.prompt,
|
||||
&effective_provider_config,
|
||||
skills_index.as_deref(),
|
||||
);
|
||||
.map_err(|e| {
|
||||
TaskError::AgentCreationFailed(format!(
|
||||
"subagent '{}' model resolution failed: {}",
|
||||
def.name, e
|
||||
))
|
||||
})?,
|
||||
_ => self.provider_config.clone(),
|
||||
};
|
||||
let system_prompt = SubagentPromptBuilder::build(
|
||||
&def,
|
||||
&task.description,
|
||||
&task.prompt,
|
||||
&effective_provider_config,
|
||||
skills_index.as_deref(),
|
||||
);
|
||||
|
||||
// 7. 创建子代理
|
||||
let agent = self.create_subagent(
|
||||
&session,
|
||||
system_prompt,
|
||||
Some(&def),
|
||||
parent_context.nesting_depth,
|
||||
parent_context.task_id.clone(),
|
||||
parent_context.trace_id.clone(),
|
||||
)?;
|
||||
|
||||
// 8. 执行任务
|
||||
self.execute_task(agent, &session, &def, task.prompt.clone())
|
||||
.await
|
||||
// 7. 创建子代理
|
||||
let agent = match self.create_subagent(
|
||||
&session,
|
||||
system_prompt,
|
||||
Some(&def),
|
||||
parent_context.nesting_depth,
|
||||
parent_context.task_id.clone(),
|
||||
parent_context.trace_id.clone(),
|
||||
) {
|
||||
Ok(agent) => agent,
|
||||
Err(e) => {
|
||||
let trace_id = parent_context.trace_id.as_deref().unwrap_or("");
|
||||
return self.handle_task_failure(session, e, trace_id).await;
|
||||
}
|
||||
};
|
||||
|
||||
// 9. 更新会话状态并保存
|
||||
let trace_id = parent_context.trace_id.as_deref().unwrap_or("");
|
||||
|
||||
// 8. 判断执行模式：异步（主 agent + 有 sub_done_sender）或同步（子代理/无 sender）
|
||||
let is_async_mode = parent_context.nesting_depth == 0
|
||||
&& parent_context.sub_done_sender.is_some()
|
||||
&& session.parent_topic_id.is_some();
|
||||
|
||||
if is_async_mode {
|
||||
// ===== 异步路径 =====
|
||||
let topic_id = session.parent_topic_id.clone().unwrap_or_default();
|
||||
let task_id = session.id.clone();
|
||||
|
||||
// 8a. INSERT pending_subagents 记录
|
||||
let pending_record = PendingSubagentRecord {
|
||||
task_id: task_id.clone(),
|
||||
parent_session_id: session.parent_session_id.clone(),
|
||||
parent_topic_id: topic_id.clone(),
|
||||
parent_chat_id: session.parent_chat_id.clone(),
|
||||
parent_channel: session.parent_channel_name.clone(),
|
||||
def_name: Some(def.name.clone()),
|
||||
spawned_at: current_timestamp(),
|
||||
status: "running".to_string(),
|
||||
};
|
||||
if let Err(e) = self.store.insert_pending_subagent(&pending_record) {
|
||||
tracing::warn!(
|
||||
error = %e,
|
||||
task_id = %task_id,
|
||||
"Failed to insert pending_subagent record"
|
||||
);
|
||||
}
|
||||
|
||||
// 8b. tokio::spawn 后台执行子代理
|
||||
let store = self.store.clone();
|
||||
let task_repository = self.task_repository.clone();
|
||||
let bus = self.bus.clone();
|
||||
let sub_done_sender = parent_context.sub_done_sender.clone().unwrap();
|
||||
let session_clone = session.clone();
|
||||
let def_clone = def.clone();
|
||||
let prompt = task.prompt.clone();
|
||||
let trace_id_owned = trace_id.to_string();
|
||||
let max_secs = self.effective_max_execution_secs(&def);
|
||||
let task_id_for_spawn = task_id.clone();
|
||||
let semaphore = self.semaphore.clone();
|
||||
let cancel_registry = self.cancel_registry.clone();
|
||||
|
||||
// 创建 CancellationToken 并注册到 registry（供 /stop 取消传播）
|
||||
let cancel_token = tokio_util::sync::CancellationToken::new();
|
||||
cancel_registry
|
||||
.lock()
|
||||
.insert(task_id_for_spawn.clone(), cancel_token.clone());
|
||||
|
||||
// RAII guard：spawn 任务退出时（正常/early return/panic）确定性清理 registry
|
||||
// 不变量 3：清理与作用域绑定，避免末行清理被 panic 绕过
|
||||
let registry_guard =
|
||||
CancelRegistryGuard::new(task_id_for_spawn.clone(), cancel_registry.clone());
|
||||
|
||||
tokio::spawn(async move {
|
||||
// guard 在闭包退出时 drop，确定性清理 cancel_registry 条目
|
||||
let _registry_guard = registry_guard;
|
||||
|
||||
// 获取并发许可（Semaphore 限流）
|
||||
let _permit = match semaphore.acquire_owned().await {
|
||||
Ok(p) => p,
|
||||
Err(e) => {
|
||||
// 不变量 1：状态机收敛性 — early return 也必须收敛终态
|
||||
// 发送 Failed 结果让 wait 收到，更新 DB 状态，否则系统出现悬空记录
|
||||
tracing::warn!(
|
||||
error = %e,
|
||||
task_id = %task_id_for_spawn,
|
||||
"Semaphore closed, subagent cannot start; converging state machine to Failed"
|
||||
);
|
||||
let result = SubagentResult {
|
||||
task_id: task_id_for_spawn.clone(),
|
||||
status: SubagentStatus::Failed,
|
||||
output: String::new(),
|
||||
pending_task_ids: store
|
||||
.list_pending_subagents(&topic_id, Some("running"))
|
||||
.map(|records| {
|
||||
records
|
||||
.into_iter()
|
||||
.map(|r| r.task_id)
|
||||
.filter(|id| id != &task_id_for_spawn)
|
||||
.collect::<Vec<_>>()
|
||||
})
|
||||
.unwrap_or_default(),
|
||||
};
|
||||
let _ = sub_done_sender.send(result).await;
|
||||
let _ =
|
||||
store.update_pending_subagent_status(&task_id_for_spawn, "failed");
|
||||
// _registry_guard drop 时清理 registry 条目
|
||||
return;
|
||||
}
|
||||
};
|
||||
|
||||
// select! 等待执行完成或取消信号
|
||||
let exec_result = tokio::select! {
|
||||
biased;
|
||||
_ = cancel_token.cancelled() => {
|
||||
tracing::info!(
|
||||
task_id = %task_id_for_spawn,
|
||||
"Subagent cancelled by user"
|
||||
);
|
||||
Err(TaskError::Cancelled)
|
||||
}
|
||||
r = Self::execute_task_static(
|
||||
agent,
|
||||
&session_clone,
|
||||
&def_clone,
|
||||
prompt,
|
||||
max_secs,
|
||||
) => r,
|
||||
};
|
||||
|
||||
// 完成回调：查询未完成 → send SubagentResult → UPDATE status
|
||||
let (status, output, _summary) = match &exec_result {
|
||||
Ok(tool_result) => (
|
||||
SubagentStatus::Completed,
|
||||
serde_json::to_string(&tool_result).unwrap_or_default(),
|
||||
tool_result.summary.clone(),
|
||||
),
|
||||
Err(TaskError::Timeout) => (
|
||||
SubagentStatus::Timeout,
|
||||
String::new(),
|
||||
"timeout".to_string(),
|
||||
),
|
||||
Err(TaskError::Cancelled) => (
|
||||
SubagentStatus::Cancelled,
|
||||
String::new(),
|
||||
"cancelled".to_string(),
|
||||
),
|
||||
Err(e) => (
|
||||
SubagentStatus::Failed,
|
||||
String::new(),
|
||||
e.to_string(),
|
||||
),
|
||||
};
|
||||
|
||||
// 查询同 topic 下仍未完成的子代理列表
|
||||
let pending_task_ids = store
|
||||
.list_pending_subagents(&topic_id, Some("running"))
|
||||
.map(|records| {
|
||||
records
|
||||
.into_iter()
|
||||
.map(|r| r.task_id)
|
||||
.filter(|id| id != &task_id_for_spawn)
|
||||
.collect::<Vec<_>>()
|
||||
})
|
||||
.unwrap_or_default();
|
||||
|
||||
// 发送 SubagentResult 到 sub_done_q
|
||||
let result = SubagentResult {
|
||||
task_id: task_id_for_spawn.clone(),
|
||||
status,
|
||||
output,
|
||||
pending_task_ids,
|
||||
};
|
||||
if let Err(e) = sub_done_sender.send(result).await {
|
||||
tracing::warn!(
|
||||
error = %e,
|
||||
task_id = %task_id_for_spawn,
|
||||
"Failed to send SubagentResult to sub_done_q (receiver dropped?)"
|
||||
);
|
||||
}
|
||||
|
||||
// UPDATE pending_subagents 状态
|
||||
let status_str = match status {
|
||||
SubagentStatus::Completed => "completed",
|
||||
SubagentStatus::Failed => "failed",
|
||||
SubagentStatus::Timeout => "timeout",
|
||||
SubagentStatus::Cancelled => "cancelled",
|
||||
};
|
||||
if let Err(e) = store.update_pending_subagent_status(&task_id_for_spawn, status_str) {
|
||||
tracing::warn!(
|
||||
error = %e,
|
||||
task_id = %task_id_for_spawn,
|
||||
"Failed to update pending_subagent status"
|
||||
);
|
||||
}
|
||||
|
||||
// 更新 TaskSession 状态并发布完成事件
|
||||
let mut session_done = session_clone;
|
||||
match exec_result {
|
||||
Ok(tool_result) => {
|
||||
session_done.mark_completed(tool_result.summary);
|
||||
if let Err(e) = task_repository.save_task_session(&session_done).await {
|
||||
tracing::warn!(error = %e, task_id = %task_id_for_spawn, "Failed to save completed session");
|
||||
}
|
||||
publish_subagent_completion(&bus, &session_done, &trace_id_owned).await;
|
||||
}
|
||||
Err(e) => {
|
||||
let err_str = e.to_string();
|
||||
if matches!(e, TaskError::Timeout) {
|
||||
session_done.mark_timeout();
|
||||
} else if matches!(e, TaskError::Cancelled) {
|
||||
session_done.mark_cancelled();
|
||||
} else {
|
||||
session_done.mark_failed(err_str);
|
||||
}
|
||||
if let Err(e) = task_repository.save_task_session(&session_done).await {
|
||||
tracing::warn!(error = %e, task_id = %task_id_for_spawn, "Failed to save failed session");
|
||||
}
|
||||
publish_subagent_error(&bus, &session_done, &e.to_string(), &trace_id_owned).await;
|
||||
}
|
||||
}
|
||||
// _registry_guard 在此 drop，确定性清理 cancel_registry 条目
|
||||
// （替代原末行手动 remove，覆盖 panic/early return 全路径）
|
||||
});
|
||||
|
||||
// 8c. 立即返回 running 占位结果
|
||||
// 注意: summary 留空，output 只含引导信息。LLM 看到 running 后应调 wait_for_subagents。
|
||||
return Ok(TaskToolResult {
|
||||
status: "running".to_string(),
|
||||
summary: format!("Task {} spawned asynchronously", task_id),
|
||||
output: format!(
|
||||
"running, task_id={}. Call wait_for_subagents(timeout_secs) to wait for subagent completion and get results.",
|
||||
task_id
|
||||
),
|
||||
task_id,
|
||||
});
|
||||
}
|
||||
|
||||
// ===== 同步路径（子代理嵌套或无 sub_done_sender） =====
|
||||
// 9. 执行任务并处理结果
|
||||
let result = self.execute_task(agent, &session, &def, task.prompt.clone()).await;
|
||||
|
||||
match result {
|
||||
Ok(tool_result) => {
|
||||
let mut session = session;
|
||||
@ -1010,6 +1321,75 @@ impl SubAgentRuntime for DefaultSubAgentRuntime {
|
||||
fn available_subagent_names(&self) -> Vec<String> {
|
||||
self.subagent_runtime.available_names()
|
||||
}
|
||||
|
||||
async fn cancel_pending_for_topic(&self, topic_id: &str) -> usize {
|
||||
// 查询该 topic 下所有 running 的子代理
|
||||
let running = match self.store.list_pending_subagents(topic_id, Some("running")) {
|
||||
Ok(records) => records,
|
||||
Err(e) => {
|
||||
tracing::warn!(
|
||||
error = %e,
|
||||
topic_id = %topic_id,
|
||||
"Failed to list pending subagents for cancellation"
|
||||
);
|
||||
return 0;
|
||||
}
|
||||
};
|
||||
|
||||
let count = running.len();
|
||||
if count == 0 {
|
||||
return 0;
|
||||
}
|
||||
|
||||
tracing::info!(
|
||||
topic_id = %topic_id,
|
||||
count,
|
||||
"Cancelling pending subagents for topic"
|
||||
);
|
||||
|
||||
// 触发每个子代理的 CancellationToken
|
||||
let registry = self.cancel_registry.lock();
|
||||
for record in &running {
|
||||
if let Some(token) = registry.get(&record.task_id) {
|
||||
token.cancel();
|
||||
tracing::info!(
|
||||
task_id = %record.task_id,
|
||||
"Cancelled subagent token"
|
||||
);
|
||||
} else {
|
||||
// token 不在 registry 中（可能已完成但 DB 状态未更新，或进程重启后丢失）
|
||||
// 不变量 1：条件 UPDATE，仅在 status='running' 时转为 cancelled，
|
||||
// 避免 spawn 已完成的终态被覆盖（completed → cancelled 是非法转换）
|
||||
match self
|
||||
.store
|
||||
.try_update_pending_subagent_status(&record.task_id, "running", "cancelled")
|
||||
{
|
||||
Ok(true) => {
|
||||
tracing::info!(
|
||||
task_id = %record.task_id,
|
||||
"Marked subagent as cancelled in DB (token not in registry)"
|
||||
);
|
||||
}
|
||||
Ok(false) => {
|
||||
tracing::info!(
|
||||
task_id = %record.task_id,
|
||||
"Subagent status already updated by another path, skip cancel"
|
||||
);
|
||||
}
|
||||
Err(e) => {
|
||||
tracing::warn!(
|
||||
error = %e,
|
||||
task_id = %record.task_id,
|
||||
"Failed to mark subagent as cancelled in DB"
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
drop(registry);
|
||||
|
||||
count
|
||||
}
|
||||
}
|
||||
|
||||
/// 子代理定义目录
|
||||
|
||||
@ -17,6 +17,8 @@ pub enum TaskSessionState {
|
||||
Failed,
|
||||
/// 已超时
|
||||
Timeout,
|
||||
/// 已取消（用户 /stop 传播）
|
||||
Cancelled,
|
||||
/// 状态未知（如重启后从 DB 重建时无法可靠推断原状态）
|
||||
Unknown,
|
||||
}
|
||||
@ -27,6 +29,19 @@ impl Default for TaskSessionState {
|
||||
}
|
||||
}
|
||||
|
||||
impl TaskSessionState {
|
||||
pub fn as_str(&self) -> &'static str {
|
||||
match self {
|
||||
Self::Running => "running",
|
||||
Self::Completed => "completed",
|
||||
Self::Failed => "failed",
|
||||
Self::Timeout => "timeout",
|
||||
Self::Cancelled => "cancelled",
|
||||
Self::Unknown => "unknown",
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// 子代理来源
|
||||
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
|
||||
#[serde(rename_all = "snake_case")]
|
||||
@ -223,6 +238,13 @@ impl TaskSession {
|
||||
self.error = Some("Task execution timed out".to_string());
|
||||
self.updated_at = current_timestamp();
|
||||
}
|
||||
|
||||
/// 标记取消
|
||||
pub fn mark_cancelled(&mut self) {
|
||||
self.state = TaskSessionState::Cancelled;
|
||||
self.error = Some("Task cancelled by user".to_string());
|
||||
self.updated_at = current_timestamp();
|
||||
}
|
||||
}
|
||||
|
||||
/// 任务工具参数
|
||||
@ -271,7 +293,7 @@ pub struct TaskHandle {
|
||||
/// 任务执行结果
|
||||
#[derive(Debug, Clone, Serialize)]
|
||||
pub struct TaskToolResult {
|
||||
/// 状态: success/failed/timeout
|
||||
/// 状态: success/failed/timeout/running
|
||||
pub status: String,
|
||||
/// 任务完成总结
|
||||
pub summary: String,
|
||||
@ -280,3 +302,26 @@ pub struct TaskToolResult {
|
||||
/// 会话 ID（用于恢复）
|
||||
pub task_id: String,
|
||||
}
|
||||
|
||||
/// 异步子代理完成状态
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
|
||||
#[serde(rename_all = "lowercase")]
|
||||
pub enum SubagentStatus {
|
||||
Completed,
|
||||
Failed,
|
||||
Timeout,
|
||||
Cancelled,
|
||||
}
|
||||
|
||||
/// 异步子代理完成回调内容（通过 sub_done_q 传递给 wait 工具）
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct SubagentResult {
|
||||
/// 完成的子代理 task_id
|
||||
pub task_id: String,
|
||||
/// 完成状态
|
||||
pub status: SubagentStatus,
|
||||
/// 子代理输出（与 TaskToolResult.output 格式一致）
|
||||
pub output: String,
|
||||
/// 仍未完成的子代理 task_id 列表（供 LLM 判断全局进度）
|
||||
pub pending_task_ids: Vec<String>,
|
||||
}
|
||||
|
||||
@ -185,6 +185,7 @@ mod tests {
|
||||
tool_call_id: None,
|
||||
parent_capability: None,
|
||||
trace_id: None,
|
||||
..Default::default()
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@ -487,6 +487,7 @@ mod tests {
|
||||
tool_call_id: None,
|
||||
parent_capability: None,
|
||||
trace_id: None,
|
||||
..Default::default()
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@ -1,6 +1,11 @@
|
||||
use std::time::Duration;
|
||||
use std::sync::Arc;
|
||||
|
||||
use async_trait::async_trait;
|
||||
use tokio::sync::{mpsc, watch};
|
||||
|
||||
use crate::domain::CapabilityPolicy;
|
||||
use crate::tools::task::SubagentResult;
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct ToolResult {
|
||||
@ -9,7 +14,58 @@ pub struct ToolResult {
|
||||
pub error: Option<String>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Default)]
|
||||
/// wait_for_subagents 工具等待期间的事件。
|
||||
#[derive(Debug, Clone)]
|
||||
pub enum WaitEvent {
|
||||
/// 一个子代理完成，携带其结果
|
||||
SubagentResult(SubagentResult),
|
||||
/// wait 期间有新用户消息到达（已注入 history），携带新用户消息的内容列表
|
||||
UserMessage(Vec<String>),
|
||||
/// 等待超时
|
||||
Timeout,
|
||||
/// 收到取消信号（/stop），等待已优雅终止，状态已清理完毕
|
||||
Cancelled,
|
||||
}
|
||||
|
||||
/// wait_for_subagents 工具的协调器接口。
|
||||
///
|
||||
/// 封装了释放/重获取 serial_lock + select! 等待 + 管理等待状态的逻辑，
|
||||
/// 使 wait 工具不直接依赖 gateway 模块（避免循环依赖）。
|
||||
///
|
||||
/// 具体实现 `SessionWaitCoordinator` 在 gateway 模块中，持有 Session 引用。
|
||||
#[async_trait]
|
||||
pub trait WaitCoordinator: Send + Sync + 'static {
|
||||
/// 查询当前 topic 下仍处于 running 状态的子代理 task_id 列表。
|
||||
fn query_pending_task_ids(&self) -> Vec<String>;
|
||||
|
||||
/// 尝试排空 sub_done_q 中已缓冲的子代理结果（非阻塞 drain，但方法本身是 async
|
||||
/// 因为需要获取 Session 锁）。
|
||||
///
|
||||
/// 用于两种场景：
|
||||
/// 1. wait_for_subagents 入口处：即使 DB 中无 running 子代理，
|
||||
/// 队列可能仍缓冲了已完成子代理的结果（子代理完成 → send 到队列 →
|
||||
/// DB 更新为 completed，但 LLM 上轮未消费队列）。
|
||||
/// 2. wait() 进入 select! 前：多个子代理同时完成时，批量消费避免
|
||||
/// 每个结果各触发一次 LLM 调用。
|
||||
async fn try_drain_queued_results(&self) -> Vec<SubagentResult>;
|
||||
|
||||
/// 进入等待状态：释放 serial_lock → select! → 重获取 serial_lock。
|
||||
///
|
||||
/// 调用前提：serial_lock 已被执行路径获取，guard 存于 coordinator 内部。
|
||||
/// 返回后：serial_lock 已被重新获取，waiting 标志已清除。
|
||||
///
|
||||
/// `cancel_rx`：可选的取消信号接收端。当收到信号时（/stop 命令），
|
||||
/// select! 立即返回 `WaitEvent::Cancelled`，并完成完整的状态清理
|
||||
///（重获取锁、回填 guard、清除 is_waiting、归还 receiver）。
|
||||
/// 为 None 时退化为不检查取消（向后兼容，子代理场景）。
|
||||
async fn wait(
|
||||
&self,
|
||||
timeout: Duration,
|
||||
cancel_rx: Option<watch::Receiver<()>>,
|
||||
) -> WaitEvent;
|
||||
}
|
||||
|
||||
#[derive(Clone, Default)]
|
||||
pub struct ToolContext {
|
||||
pub channel_name: Option<String>,
|
||||
pub sender_id: Option<String>,
|
||||
@ -35,6 +91,46 @@ pub struct ToolContext {
|
||||
/// 端到端追踪 ID（从 InboundMessage 继承，用于 tool 执行路径的日志关联）。
|
||||
/// None 表示无追踪上下文（如子代理独立执行或测试环境）。
|
||||
pub trace_id: Option<String>,
|
||||
/// 异步子代理完成队列的 sender（按 topic 隔离）。
|
||||
/// 仅主 agent（nesting_depth=0）有值：agent_factory 构建时从 SessionHistory 注入。
|
||||
/// TaskTool spawn 异步子代理后，子代理完成时通过此 sender 发送 SubagentResult，
|
||||
/// 由 wait_for_subagents 工具的 receiver 端消费。
|
||||
/// 子代理自身（nesting_depth>0）为 None：嵌套层不支持异步，走同步路径。
|
||||
pub sub_done_sender: Option<mpsc::Sender<SubagentResult>>,
|
||||
/// wait_for_subagents 工具的协调器（仅主 agent 有值）。
|
||||
/// 封装了释放/重获取 serial_lock + select! 等待逻辑。
|
||||
/// wait 工具通过此接口实现真等待（释放锁让 process_one 注入用户消息）。
|
||||
pub wait_coordinator: Option<Arc<dyn WaitCoordinator>>,
|
||||
/// 取消信号接收端（仅主 agent 有值，由 agent_factory 从 cancel_token clone 注入）。
|
||||
/// wait_for_subagents 工具将其传给 coordinator.wait() 的 select!，
|
||||
/// 使 /stop 命令能立即中断等待并完成状态清理。
|
||||
/// watch::Receiver 可安全 clone，多个 receiver 共享同一 sender，
|
||||
/// 互不影响各自的 has_changed() / changed() 状态。
|
||||
pub cancel_rx: Option<watch::Receiver<()>>,
|
||||
}
|
||||
|
||||
impl std::fmt::Debug for ToolContext {
|
||||
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
||||
f.debug_struct("ToolContext")
|
||||
.field("channel_name", &self.channel_name)
|
||||
.field("sender_id", &self.sender_id)
|
||||
.field("chat_id", &self.chat_id)
|
||||
.field("session_id", &self.session_id)
|
||||
.field("topic_id", &self.topic_id)
|
||||
.field("message_id", &self.message_id)
|
||||
.field("message_seq", &self.message_seq)
|
||||
.field("subagent_description", &self.subagent_description)
|
||||
.field("nesting_depth", &self.nesting_depth)
|
||||
.field("task_id", &self.task_id)
|
||||
.field("parent_task_id", &self.parent_task_id)
|
||||
.field("tool_call_id", &self.tool_call_id)
|
||||
.field("parent_capability", &self.parent_capability)
|
||||
.field("trace_id", &self.trace_id)
|
||||
.field("sub_done_sender", &self.sub_done_sender)
|
||||
.field("wait_coordinator", &self.wait_coordinator.is_some())
|
||||
.field("cancel_rx", &self.cancel_rx.is_some())
|
||||
.finish()
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
|
||||
250
src/tools/wait_tool.rs
Normal file
250
src/tools/wait_tool.rs
Normal file
@ -0,0 +1,250 @@
|
||||
use std::time::Duration;
|
||||
|
||||
use async_trait::async_trait;
|
||||
use serde_json::json;
|
||||
|
||||
use crate::tools::{Tool, ToolContext, ToolResult, WaitEvent};
|
||||
|
||||
/// wait_for_subagents 工具 — 等待异步子代理完成或用户消息到达。
|
||||
///
|
||||
/// 调用后释放 serial_lock，进入 select! 等待：
|
||||
/// - 子代理完成 → 返回结果 + 未完成列表
|
||||
/// - 用户消息到达 → 返回 "有新用户消息"（消息已注入 history）
|
||||
/// - 超时 → 返回超时 + 未完成列表
|
||||
///
|
||||
/// 等待结束后重新获取 serial_lock，保证后续工具调用串行。
|
||||
pub struct WaitForSubagentsTool {
|
||||
/// 默认超时（秒），LLM 未指定时使用
|
||||
default_timeout_secs: u64,
|
||||
}
|
||||
|
||||
impl WaitForSubagentsTool {
|
||||
pub const TOOL_NAME: &'static str = "wait_for_subagents";
|
||||
|
||||
pub fn new(default_timeout_secs: u64) -> Self {
|
||||
Self {
|
||||
default_timeout_secs,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl Tool for WaitForSubagentsTool {
|
||||
fn name(&self) -> &str {
|
||||
Self::TOOL_NAME
|
||||
}
|
||||
|
||||
fn description(&self) -> &str {
|
||||
"Wait for asynchronous subagents to complete, or for new user messages to arrive. \
|
||||
Use this after launching subagents via the task tool to receive their results. \
|
||||
Returns the first completed subagent's result and a list of still-pending task IDs. \
|
||||
If pending_task_ids is non-empty, call this tool again to wait for the next one."
|
||||
}
|
||||
|
||||
fn parameters_schema(&self) -> serde_json::Value {
|
||||
json!({
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"timeout_secs": {
|
||||
"type": "integer",
|
||||
"description": "Maximum seconds to wait. Default 60. The tool returns early if a subagent completes or a user message arrives.",
|
||||
"default": 60
|
||||
}
|
||||
},
|
||||
"required": []
|
||||
})
|
||||
}
|
||||
|
||||
fn read_only(&self) -> bool {
|
||||
false
|
||||
}
|
||||
|
||||
fn exclusive(&self) -> bool {
|
||||
// wait 工具释放/重获取 serial_lock，不应与其他工具并发
|
||||
true
|
||||
}
|
||||
|
||||
async fn execute(&self, _args: serde_json::Value) -> anyhow::Result<ToolResult> {
|
||||
Ok(ToolResult {
|
||||
success: false,
|
||||
output: String::new(),
|
||||
error: Some(
|
||||
"wait_for_subagents requires tool context with wait_coordinator".to_string(),
|
||||
),
|
||||
})
|
||||
}
|
||||
|
||||
async fn execute_with_context(
|
||||
&self,
|
||||
context: &ToolContext,
|
||||
args: serde_json::Value,
|
||||
) -> anyhow::Result<ToolResult> {
|
||||
// 1. 获取 wait_coordinator（仅主 agent 有值）
|
||||
let coordinator = match &context.wait_coordinator {
|
||||
Some(c) => c.clone(),
|
||||
None => {
|
||||
return Ok(ToolResult {
|
||||
success: false,
|
||||
output: String::new(),
|
||||
error: Some(
|
||||
"wait_for_subagents is not available in this context (no wait_coordinator)"
|
||||
.to_string(),
|
||||
),
|
||||
});
|
||||
}
|
||||
};
|
||||
|
||||
// 2. 解析超时参数
|
||||
let timeout_secs = args
|
||||
.get("timeout_secs")
|
||||
.and_then(|v| v.as_u64())
|
||||
.unwrap_or(self.default_timeout_secs);
|
||||
let timeout = Duration::from_secs(timeout_secs);
|
||||
|
||||
// 3. 先尝试排空队列中已缓冲的结果
|
||||
// 场景：子代理已完成 → 结果 send 到队列 + DB 更新为 completed，
|
||||
// 但 LLM 上轮未调用 wait（或 wait 超时未消费）→ 结果缓冲在队列中。
|
||||
// 若不排空，query_pending_task_ids 返回空（DB 已非 running）→
|
||||
// 返回 "No pending" → 缓冲结果永远丢失。
|
||||
let drained = coordinator.try_drain_queued_results().await;
|
||||
if !drained.is_empty() {
|
||||
let pending_after = coordinator.query_pending_task_ids();
|
||||
let pending_str = if pending_after.is_empty() {
|
||||
"none".to_string()
|
||||
} else {
|
||||
pending_after.join(", ")
|
||||
};
|
||||
let formatted: Vec<String> = drained
|
||||
.iter()
|
||||
.map(|r| {
|
||||
format!(
|
||||
"Subagent {} completed (status: {:?}). Output: {}",
|
||||
r.task_id, r.status, r.output
|
||||
)
|
||||
})
|
||||
.collect();
|
||||
return Ok(ToolResult {
|
||||
success: true,
|
||||
output: format!(
|
||||
"Retrieved {} buffered subagent result(s):\n{}\nStill pending: [{}]",
|
||||
drained.len(),
|
||||
formatted.join("\n"),
|
||||
pending_str
|
||||
),
|
||||
error: None,
|
||||
});
|
||||
}
|
||||
|
||||
// 4. 查询 pending 子代理
|
||||
let pending = coordinator.query_pending_task_ids();
|
||||
if pending.is_empty() {
|
||||
return Ok(ToolResult {
|
||||
success: true,
|
||||
output: "No pending subagents to wait for.".to_string(),
|
||||
error: None,
|
||||
});
|
||||
}
|
||||
|
||||
tracing::info!(
|
||||
topic_id = ?context.topic_id,
|
||||
pending_count = pending.len(),
|
||||
timeout_secs,
|
||||
"wait_for_subagents: entering wait"
|
||||
);
|
||||
|
||||
// 4. 进入等待（coordinator 内部：释放锁 → select! → 重获取锁）
|
||||
// 传入 cancel_rx 使 /stop 命令能立即中断等待。
|
||||
// coordinator 在 select! 中以 biased 优先级处理：
|
||||
// 子代理结果 > 用户消息 > 取消信号 > 超时
|
||||
let event = coordinator
|
||||
.wait(timeout, context.cancel_rx.clone())
|
||||
.await;
|
||||
|
||||
// 5. 格式化返回结果
|
||||
let output = match event {
|
||||
WaitEvent::SubagentResult(result) => {
|
||||
let pending_str = if result.pending_task_ids.is_empty() {
|
||||
"none".to_string()
|
||||
} else {
|
||||
result.pending_task_ids.join(", ")
|
||||
};
|
||||
format!(
|
||||
"Subagent {} completed (status: {:?}). Output: {}\nStill pending: [{}]",
|
||||
result.task_id, result.status, result.output, pending_str
|
||||
)
|
||||
}
|
||||
WaitEvent::UserMessage(messages) => {
|
||||
let pending = coordinator.query_pending_task_ids();
|
||||
let pending_str = if pending.is_empty() {
|
||||
"none".to_string()
|
||||
} else {
|
||||
pending.join(", ")
|
||||
};
|
||||
if messages.is_empty() {
|
||||
format!(
|
||||
"A new user message arrived while waiting (content could not be retrieved). \
|
||||
Still pending subagents: [{}]",
|
||||
pending_str
|
||||
)
|
||||
} else {
|
||||
let formatted_msgs: Vec<String> = messages
|
||||
.iter()
|
||||
.enumerate()
|
||||
.map(|(i, msg)| format!(" [{}] {}", i + 1, msg))
|
||||
.collect();
|
||||
format!(
|
||||
"New user message(s) arrived while waiting:\n{}\n\
|
||||
These messages have been added to the conversation history. \
|
||||
Still pending subagents: [{}]",
|
||||
formatted_msgs.join("\n"),
|
||||
pending_str
|
||||
)
|
||||
}
|
||||
}
|
||||
WaitEvent::Timeout => {
|
||||
let pending = coordinator.query_pending_task_ids();
|
||||
format!(
|
||||
"Wait timed out after {}s. Still pending subagents: [{}]",
|
||||
timeout_secs,
|
||||
if pending.is_empty() {
|
||||
"none".to_string()
|
||||
} else {
|
||||
pending.join(", ")
|
||||
}
|
||||
)
|
||||
}
|
||||
WaitEvent::Cancelled => {
|
||||
// /stop 命令中断了等待。coordinator 已完成全部状态清理
|
||||
//（重获取 serial_lock、回填 guard、清除 is_waiting、归还 receiver）。
|
||||
// 返回提示性输出，Agent 下一轮迭代会检测到 cancel 并退出。
|
||||
let pending = coordinator.query_pending_task_ids();
|
||||
tracing::info!(
|
||||
topic_id = ?context.topic_id,
|
||||
pending_count = pending.len(),
|
||||
"wait_for_subagents: cancelled by /stop"
|
||||
);
|
||||
format!(
|
||||
"Wait was cancelled by /stop command. \
|
||||
Pending subagents (if any) have been cancelled separately. \
|
||||
Still pending in DB: [{}]",
|
||||
if pending.is_empty() {
|
||||
"none".to_string()
|
||||
} else {
|
||||
pending.join(", ")
|
||||
}
|
||||
)
|
||||
}
|
||||
};
|
||||
|
||||
tracing::info!(
|
||||
topic_id = ?context.topic_id,
|
||||
"wait_for_subagents: returning result"
|
||||
);
|
||||
|
||||
Ok(ToolResult {
|
||||
success: true,
|
||||
output,
|
||||
error: None,
|
||||
})
|
||||
}
|
||||
}
|
||||
Loading…
x
Reference in New Issue
Block a user