fc3a95b152
后端核心实现： - 新增 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 崩溃恢复
382 lines
14 KiB
Rust
382 lines
14 KiB
Rust
//! Gateway Runtime - builds SessionManager with decoupled MCP integration
|
||
|
||
use std::collections::{HashMap, HashSet};
|
||
use std::path::PathBuf;
|
||
use std::sync::Arc;
|
||
|
||
use tokio::sync::RwLock;
|
||
|
||
use crate::agent::AgentError;
|
||
use crate::bus::MessageBus;
|
||
use crate::config::{
|
||
CompactionConfig, LLMProviderConfig, MemoryMaintenanceConfig, ModelResolver, SubagentsConfig,
|
||
TaskConfig,
|
||
};
|
||
use crate::gateway::model_selection::ModelSelectionStore;
|
||
use crate::gateway::tool_registry_factory::ToolRegistryFactory;
|
||
use crate::mcp::McpInitializer;
|
||
use crate::mcp::client::McpClientManager;
|
||
use crate::skills::SkillRuntime;
|
||
use crate::storage::{
|
||
ConversationRepository, MemoryRepository, PromptInjectionRepository, SchedulerJobRepository,
|
||
SessionStore, SkillEventRepository, TodoRepository,
|
||
};
|
||
use crate::tools::task::repository::TaskRepository;
|
||
use crate::tools::task::runtime::SubagentRuntime;
|
||
use crate::tools::todo_write::TodoItem;
|
||
use crate::tools::{
|
||
DefaultSubAgentRuntime, InMemoryTaskRepository, NoopSessionMessageSender, SessionMessageSender,
|
||
SubAgentRuntime, SubAgentRuntimeConfig, SubagentCatalog, TaskTool, ToolRegistry,
|
||
};
|
||
|
||
use super::agent_factory::AgentFactory;
|
||
use super::cli_session::CliSessionService;
|
||
use super::memory_maintenance_coordinator::MemoryMaintenanceCoordinator;
|
||
use super::provider_config_service::ProviderConfigService;
|
||
use super::scheduled_agent_task_service::ScheduledAgentTaskService;
|
||
use super::session::{SessionManager, SessionManagerServices};
|
||
use super::session_factory::SessionFactory;
|
||
use super::session_lifecycle::SessionLifecycleService;
|
||
use super::session_message_service::SessionMessageService;
|
||
|
||
/// Build SessionManager with optional MCP integration
|
||
///
|
||
/// MCP is initialized asynchronously in background, not blocking gateway startup.
|
||
pub(crate) fn build_session_manager(
|
||
agent_prompt_reinject_every: u64,
|
||
show_tool_results: bool,
|
||
default_timezone: String,
|
||
provider_config: LLMProviderConfig,
|
||
provider_configs: HashMap<String, LLMProviderConfig>,
|
||
skills: Arc<SkillRuntime>,
|
||
experts: Arc<crate::experts::ExpertRuntime>,
|
||
disabled_tools: HashSet<String>,
|
||
task_config: TaskConfig,
|
||
subagents_config: SubagentsConfig,
|
||
maintenance_config: MemoryMaintenanceConfig,
|
||
session_ttl_hours: Option<u64>,
|
||
mcp_config: crate::mcp::McpConfig,
|
||
mcp_tool_timeout_secs: u64,
|
||
bus: Option<Arc<MessageBus>>,
|
||
model_resolver: Arc<ModelResolver>,
|
||
compaction_config: CompactionConfig,
|
||
) -> Result<
|
||
(
|
||
SessionManager,
|
||
Arc<dyn TaskRepository>,
|
||
Option<Arc<McpClientManager>>,
|
||
Arc<SubagentRuntime>,
|
||
Arc<ModelSelectionStore>,
|
||
Option<Arc<dyn SubAgentRuntime>>,
|
||
),
|
||
AgentError,
|
||
> {
|
||
build_session_manager_with_sender(
|
||
agent_prompt_reinject_every,
|
||
show_tool_results,
|
||
default_timezone,
|
||
provider_config,
|
||
provider_configs,
|
||
skills,
|
||
experts,
|
||
Arc::new(NoopSessionMessageSender),
|
||
disabled_tools,
|
||
task_config,
|
||
subagents_config,
|
||
maintenance_config,
|
||
session_ttl_hours,
|
||
mcp_config,
|
||
mcp_tool_timeout_secs,
|
||
bus,
|
||
model_resolver,
|
||
compaction_config,
|
||
)
|
||
}
|
||
|
||
/// Build SessionManager with custom session message sender
|
||
pub(crate) fn build_session_manager_with_sender(
|
||
agent_prompt_reinject_every: u64,
|
||
show_tool_results: bool,
|
||
default_timezone: String,
|
||
provider_config: LLMProviderConfig,
|
||
provider_configs: HashMap<String, LLMProviderConfig>,
|
||
skills: Arc<SkillRuntime>,
|
||
experts: Arc<crate::experts::ExpertRuntime>,
|
||
session_message_sender: Arc<dyn SessionMessageSender>,
|
||
disabled_tools: HashSet<String>,
|
||
task_config: TaskConfig,
|
||
subagents_config: SubagentsConfig,
|
||
maintenance_config: MemoryMaintenanceConfig,
|
||
session_ttl_hours: Option<u64>,
|
||
mcp_config: crate::mcp::McpConfig,
|
||
mcp_tool_timeout_secs: u64,
|
||
bus: Option<Arc<MessageBus>>,
|
||
model_resolver: Arc<ModelResolver>,
|
||
compaction_config: CompactionConfig,
|
||
) -> Result<
|
||
(
|
||
SessionManager,
|
||
Arc<dyn TaskRepository>,
|
||
Option<Arc<McpClientManager>>,
|
||
Arc<SubagentRuntime>,
|
||
Arc<ModelSelectionStore>,
|
||
Option<Arc<dyn SubAgentRuntime>>,
|
||
),
|
||
AgentError,
|
||
> {
|
||
let store = Arc::new(
|
||
SessionStore::new()
|
||
.map_err(|err| AgentError::Other(format!("session store init error: {}", err)))?,
|
||
);
|
||
let known_agents = provider_configs.keys().cloned().collect::<HashSet<_>>();
|
||
let provider_configs = ProviderConfigService::new(
|
||
provider_config.clone(),
|
||
provider_configs,
|
||
maintenance_config,
|
||
);
|
||
|
||
if let Err(err) =
|
||
store.append_skill_event(None, "discovered", None, &skills.discovery_event_payload())
|
||
{
|
||
tracing::warn!(error = %err, "Failed to record skill discovery event");
|
||
}
|
||
|
||
let memories: Arc<dyn MemoryRepository> = store.clone();
|
||
let scheduler_jobs: Arc<dyn SchedulerJobRepository> = store.clone();
|
||
let skill_events: Arc<dyn SkillEventRepository> = store.clone();
|
||
let conversations: Arc<dyn ConversationRepository> = store.clone();
|
||
let todo_repository: Arc<dyn TodoRepository> = store.clone();
|
||
|
||
// Create ToolRegistryFactory
|
||
let factory = ToolRegistryFactory::new(
|
||
skills.clone(),
|
||
memories,
|
||
scheduler_jobs,
|
||
skill_events.clone(),
|
||
todo_repository,
|
||
session_message_sender.clone(),
|
||
known_agents,
|
||
default_timezone,
|
||
disabled_tools,
|
||
task_config.clone(),
|
||
);
|
||
|
||
// Create shared todo state for TodoWriteTool
|
||
let todo_state: Arc<RwLock<HashMap<String, Vec<TodoItem>>>> =
|
||
Arc::new(RwLock::new(HashMap::new()));
|
||
let factory = factory.with_todo_state(todo_state);
|
||
|
||
// Create MCP Initializer (async, non-blocking)
|
||
// MCP servers connect in background task
|
||
let mut mcp_initializer = McpInitializer::with_config(mcp_config);
|
||
|
||
// Add MCP manager to factory (if enabled)
|
||
let factory = if let Some(manager) = mcp_initializer.manager() {
|
||
factory.with_mcp_manager(manager)
|
||
} else {
|
||
factory
|
||
};
|
||
|
||
// Wait for MCP connections and collect MCP tools for subagents
|
||
// This needs to happen before building subagent tools
|
||
let mut mcp_tools_for_subagents: Vec<crate::mcp::tool_adapter::McpToolWrapper> = Vec::new();
|
||
if mcp_initializer.is_enabled() {
|
||
tokio::task::block_in_place(|| {
|
||
tokio::runtime::Handle::current().block_on(async {
|
||
// Wait for connections to complete
|
||
if let Err(e) = mcp_initializer.wait_for_connections().await {
|
||
tracing::error!(error = %e, "Failed to wait for MCP connections");
|
||
return;
|
||
}
|
||
|
||
// Collect MCP tools for subagents
|
||
if let Some(manager) = mcp_initializer.manager() {
|
||
let all_tools = manager.all_tools().await;
|
||
for (server_key, tool_info) in all_tools {
|
||
let wrapper = crate::mcp::tool_adapter::McpToolWrapper::new(
|
||
manager.clone(),
|
||
server_key.clone(),
|
||
tool_info,
|
||
mcp_tool_timeout_secs,
|
||
);
|
||
mcp_tools_for_subagents.push(wrapper);
|
||
}
|
||
tracing::info!(
|
||
tool_count = mcp_tools_for_subagents.len(),
|
||
"Collected MCP tools for subagents"
|
||
);
|
||
}
|
||
})
|
||
});
|
||
}
|
||
|
||
// Create SubAgentRuntime (if task tool is enabled)
|
||
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)
|
||
let subagent_tools = Arc::new(factory.build_subagent_tools(
|
||
if mcp_tools_for_subagents.is_empty() {
|
||
None
|
||
} else {
|
||
Some(mcp_tools_for_subagents.clone())
|
||
},
|
||
));
|
||
|
||
// Create subagent catalog with discovery, wrap in SubagentRuntime
|
||
let catalog = SubagentCatalog::discover(&subagents_config);
|
||
let subagent_runtime = Arc::new(SubagentRuntime::new(
|
||
subagents_config.clone(),
|
||
catalog,
|
||
std::env::current_dir().unwrap_or_else(|_| PathBuf::from(".")),
|
||
));
|
||
|
||
let runtime_config = SubAgentRuntimeConfig {
|
||
default_allowed_tools: task_config.allowed_tools.iter().cloned().collect(),
|
||
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(
|
||
runtime_config,
|
||
task_repository.clone(),
|
||
conversations.clone(),
|
||
subagent_tools.clone(),
|
||
provider_config.clone(),
|
||
model_resolver.clone(),
|
||
subagent_runtime.clone(),
|
||
bus.clone(),
|
||
store.clone(),
|
||
skills.clone(),
|
||
));
|
||
|
||
// 注册 task 工具到子代理工具集（需在 runtime 创建之后，打破循环依赖）
|
||
if factory.is_enabled("task") {
|
||
subagent_tools.register(TaskTool::new(
|
||
default_subagent_runtime.clone(),
|
||
Some(task_config.max_nesting_depth),
|
||
));
|
||
}
|
||
|
||
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 使用）
|
||
let subagent_runtime = Arc::new(SubagentRuntime::from_config(subagents_config.clone()));
|
||
(
|
||
factory,
|
||
Arc::new(InMemoryTaskRepository::new()),
|
||
subagent_runtime,
|
||
None,
|
||
)
|
||
};
|
||
|
||
// Build base tools
|
||
let tools = factory.build();
|
||
|
||
// 诊断日志：记录 MCP 初始化状态和基础工具数量
|
||
tracing::info!(
|
||
mcp_enabled = mcp_initializer.is_enabled(),
|
||
mcp_tools_collected = mcp_tools_for_subagents.len(),
|
||
base_tool_count = tools.tool_names().len(),
|
||
base_tools = ?tools.tool_names(),
|
||
"build_session_manager: base tools built (before MCP registration)"
|
||
);
|
||
|
||
// Register MCP tools to main agent (async)
|
||
// Note: MCP tools for subagents are already collected above
|
||
if mcp_initializer.is_enabled() {
|
||
tokio::task::block_in_place(|| {
|
||
tokio::runtime::Handle::current().block_on(async {
|
||
// Register pre-collected MCP tools
|
||
for tool in mcp_tools_for_subagents {
|
||
tools.register(tool);
|
||
}
|
||
tracing::info!("Registered MCP tools to main agent");
|
||
})
|
||
});
|
||
}
|
||
|
||
let tools = Arc::new(tools);
|
||
|
||
// 诊断日志：记录最终工具数量（含 MCP 工具，如果启用）
|
||
tracing::info!(
|
||
final_tool_count = tools.tool_names().len(),
|
||
final_tools = ?tools.tool_names(),
|
||
"build_session_manager: final tools registered to main agent"
|
||
);
|
||
|
||
let prompt_repository: Arc<dyn PromptInjectionRepository> = store.clone();
|
||
let model_selections = Arc::new(ModelSelectionStore::new());
|
||
let observer: Arc<dyn crate::observability::Observer> =
|
||
crate::observability::metrics::default_observer();
|
||
let agent_factory = AgentFactory::new(
|
||
tools.clone(),
|
||
skills.clone(),
|
||
experts.clone(),
|
||
subagent_runtime.clone(),
|
||
agent_prompt_reinject_every as usize,
|
||
prompt_repository.clone(),
|
||
model_resolver.clone(),
|
||
model_selections.clone(),
|
||
compaction_config,
|
||
Some(observer),
|
||
);
|
||
let session_factory = SessionFactory::new(
|
||
provider_config.clone(),
|
||
skills.clone(),
|
||
agent_factory,
|
||
conversations,
|
||
skill_events,
|
||
store.clone(),
|
||
);
|
||
let lifecycle = SessionLifecycleService::new(session_factory, session_ttl_hours);
|
||
let cli_sessions = CliSessionService::new(store.clone());
|
||
let messages = SessionMessageService::new(lifecycle.clone(), show_tool_results);
|
||
let scheduled_tasks = ScheduledAgentTaskService::new(
|
||
lifecycle.clone(),
|
||
provider_configs.clone(),
|
||
show_tool_results,
|
||
);
|
||
let memory_maintenance =
|
||
MemoryMaintenanceCoordinator::new(store.clone(), provider_configs.clone());
|
||
|
||
// Extract MCP manager for lifecycle management (e.g., disconnect on restart)
|
||
let mcp_manager = mcp_initializer.manager();
|
||
|
||
Ok((
|
||
SessionManager::from_services(SessionManagerServices {
|
||
tools: tools as Arc<ToolRegistry>,
|
||
skills,
|
||
experts,
|
||
subagent_runtime: subagent_runtime.clone(),
|
||
store,
|
||
show_tool_results,
|
||
lifecycle,
|
||
cli_sessions,
|
||
messages,
|
||
scheduled_tasks,
|
||
memory_maintenance,
|
||
task_repository: task_repository.clone(),
|
||
}),
|
||
task_repository,
|
||
mcp_manager,
|
||
subagent_runtime,
|
||
model_selections,
|
||
subagent_executor,
|
||
))
|
||
}
|