feat: 增强MCP集成,支持配置路径和连接重试机制
This commit is contained in:
parent
7d11ab8067
commit
4fabbe47e9
@ -873,7 +873,9 @@ impl Config {
|
|||||||
}
|
}
|
||||||
|
|
||||||
pub fn load_default() -> Result<Self, Box<dyn std::error::Error>> {
|
pub fn load_default() -> Result<Self, Box<dyn std::error::Error>> {
|
||||||
let path = get_default_config_path();
|
let path = std::env::var("CONFIG_PATH")
|
||||||
|
.map(std::path::PathBuf::from)
|
||||||
|
.unwrap_or_else(|_| get_default_config_path());
|
||||||
Self::load_from(&path)
|
Self::load_from(&path)
|
||||||
}
|
}
|
||||||
|
|
||||||
@ -907,6 +909,8 @@ impl Config {
|
|||||||
mcp_servers = config.mcp_servers.len(),
|
mcp_servers = config.mcp_servers.len(),
|
||||||
"MCP servers loaded from config"
|
"MCP servers loaded from config"
|
||||||
);
|
);
|
||||||
|
} else {
|
||||||
|
tracing::info!("No mcpServers found in config");
|
||||||
}
|
}
|
||||||
|
|
||||||
Ok(config)
|
Ok(config)
|
||||||
|
|||||||
@ -62,6 +62,7 @@ pub struct GatewayState {
|
|||||||
pub task_repository: Arc<dyn TaskRepository>,
|
pub task_repository: Arc<dyn TaskRepository>,
|
||||||
pub cancel_manager: CancelManager,
|
pub cancel_manager: CancelManager,
|
||||||
pub restart_tx: watch::Sender<bool>,
|
pub restart_tx: watch::Sender<bool>,
|
||||||
|
pub mcp_manager: Option<Arc<crate::mcp::client::McpClientManager>>,
|
||||||
}
|
}
|
||||||
|
|
||||||
impl GatewayState {
|
impl GatewayState {
|
||||||
@ -86,7 +87,7 @@ impl GatewayState {
|
|||||||
mcp_servers: config.mcp_servers.clone(),
|
mcp_servers: config.mcp_servers.clone(),
|
||||||
};
|
};
|
||||||
|
|
||||||
let (session_manager, task_repository) = build_session_manager_with_sender(
|
let (session_manager, task_repository, mcp_manager) = build_session_manager_with_sender(
|
||||||
agent_prompt_reinject_every,
|
agent_prompt_reinject_every,
|
||||||
show_tool_results,
|
show_tool_results,
|
||||||
config.time.timezone.clone(),
|
config.time.timezone.clone(),
|
||||||
@ -113,6 +114,7 @@ impl GatewayState {
|
|||||||
task_repository,
|
task_repository,
|
||||||
cancel_manager,
|
cancel_manager,
|
||||||
restart_tx,
|
restart_tx,
|
||||||
|
mcp_manager,
|
||||||
})
|
})
|
||||||
}
|
}
|
||||||
|
|
||||||
@ -241,6 +243,7 @@ pub async fn run(
|
|||||||
let (result_tx, result_rx) = tokio::sync::oneshot::channel::<bool>();
|
let (result_tx, result_rx) = tokio::sync::oneshot::channel::<bool>();
|
||||||
let channel_manager = state.channel_manager.clone();
|
let channel_manager = state.channel_manager.clone();
|
||||||
let cancel_manager = state.cancel_manager.clone();
|
let cancel_manager = state.cancel_manager.clone();
|
||||||
|
let mcp_manager = state.mcp_manager.clone();
|
||||||
|
|
||||||
// Spawn ctrl_c / restart handler
|
// Spawn ctrl_c / restart handler
|
||||||
tokio::spawn(async move {
|
tokio::spawn(async move {
|
||||||
@ -249,6 +252,10 @@ pub async fn run(
|
|||||||
tracing::info!("Shutdown signal received");
|
tracing::info!("Shutdown signal received");
|
||||||
cancel_manager.cancel_all().await;
|
cancel_manager.cancel_all().await;
|
||||||
let _ = scheduler_shutdown_tx.send(true);
|
let _ = scheduler_shutdown_tx.send(true);
|
||||||
|
if let Some(ref mgr) = mcp_manager {
|
||||||
|
tracing::info!("Disconnecting MCP servers before shutdown");
|
||||||
|
let _ = mgr.disconnect_all().await;
|
||||||
|
}
|
||||||
let _ = channel_manager.stop_all().await;
|
let _ = channel_manager.stop_all().await;
|
||||||
let _ = result_tx.send(false);
|
let _ = result_tx.send(false);
|
||||||
let _ = shutdown_tx.send(());
|
let _ = shutdown_tx.send(());
|
||||||
@ -258,6 +265,10 @@ pub async fn run(
|
|||||||
tracing::info!("Restart signal received");
|
tracing::info!("Restart signal received");
|
||||||
cancel_manager.cancel_all().await;
|
cancel_manager.cancel_all().await;
|
||||||
let _ = scheduler_shutdown_tx.send(true);
|
let _ = scheduler_shutdown_tx.send(true);
|
||||||
|
if let Some(ref mgr) = mcp_manager {
|
||||||
|
tracing::info!("Disconnecting MCP servers before restart");
|
||||||
|
let _ = mgr.disconnect_all().await;
|
||||||
|
}
|
||||||
let _ = channel_manager.stop_all().await;
|
let _ = channel_manager.stop_all().await;
|
||||||
let _ = result_tx.send(true);
|
let _ = result_tx.send(true);
|
||||||
let _ = shutdown_tx.send(());
|
let _ = shutdown_tx.send(());
|
||||||
|
|||||||
@ -10,6 +10,7 @@ use crate::bus::MessageBus;
|
|||||||
use crate::config::{LLMProviderConfig, MemoryMaintenanceConfig, SubagentsConfig, TaskConfig};
|
use crate::config::{LLMProviderConfig, MemoryMaintenanceConfig, SubagentsConfig, TaskConfig};
|
||||||
use crate::gateway::tool_registry_factory::ToolRegistryFactory;
|
use crate::gateway::tool_registry_factory::ToolRegistryFactory;
|
||||||
use crate::mcp::McpInitializer;
|
use crate::mcp::McpInitializer;
|
||||||
|
use crate::mcp::client::McpClientManager;
|
||||||
use crate::skills::SkillRuntime;
|
use crate::skills::SkillRuntime;
|
||||||
use crate::storage::{
|
use crate::storage::{
|
||||||
ConversationRepository, MemoryRepository, PromptInjectionRepository, SchedulerJobRepository,
|
ConversationRepository, MemoryRepository, PromptInjectionRepository, SchedulerJobRepository,
|
||||||
@ -49,7 +50,7 @@ pub(crate) fn build_session_manager(
|
|||||||
session_ttl_hours: Option<u64>,
|
session_ttl_hours: Option<u64>,
|
||||||
mcp_config: crate::mcp::McpConfig,
|
mcp_config: crate::mcp::McpConfig,
|
||||||
bus: Option<Arc<MessageBus>>,
|
bus: Option<Arc<MessageBus>>,
|
||||||
) -> Result<(SessionManager, Arc<dyn TaskRepository>), AgentError> {
|
) -> Result<(SessionManager, Arc<dyn TaskRepository>, Option<Arc<McpClientManager>>), AgentError> {
|
||||||
build_session_manager_with_sender(
|
build_session_manager_with_sender(
|
||||||
agent_prompt_reinject_every,
|
agent_prompt_reinject_every,
|
||||||
show_tool_results,
|
show_tool_results,
|
||||||
@ -84,7 +85,7 @@ pub(crate) fn build_session_manager_with_sender(
|
|||||||
session_ttl_hours: Option<u64>,
|
session_ttl_hours: Option<u64>,
|
||||||
mcp_config: crate::mcp::McpConfig,
|
mcp_config: crate::mcp::McpConfig,
|
||||||
bus: Option<Arc<MessageBus>>,
|
bus: Option<Arc<MessageBus>>,
|
||||||
) -> Result<(SessionManager, Arc<dyn TaskRepository>), AgentError> {
|
) -> Result<(SessionManager, Arc<dyn TaskRepository>, Option<Arc<McpClientManager>>), AgentError> {
|
||||||
let store = Arc::new(
|
let store = Arc::new(
|
||||||
SessionStore::new()
|
SessionStore::new()
|
||||||
.map_err(|err| AgentError::Other(format!("session store init error: {}", err)))?,
|
.map_err(|err| AgentError::Other(format!("session store init error: {}", err)))?,
|
||||||
@ -265,6 +266,9 @@ pub(crate) fn build_session_manager_with_sender(
|
|||||||
let memory_maintenance =
|
let memory_maintenance =
|
||||||
MemoryMaintenanceCoordinator::new(store.clone(), provider_configs.clone());
|
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 {
|
Ok((SessionManager::from_services(SessionManagerServices {
|
||||||
tools: tools as Arc<ToolRegistry>,
|
tools: tools as Arc<ToolRegistry>,
|
||||||
skills,
|
skills,
|
||||||
@ -276,5 +280,5 @@ pub(crate) fn build_session_manager_with_sender(
|
|||||||
scheduled_tasks,
|
scheduled_tasks,
|
||||||
memory_maintenance,
|
memory_maintenance,
|
||||||
task_repository: task_repository.clone(),
|
task_repository: task_repository.clone(),
|
||||||
}), task_repository))
|
}), task_repository, mcp_manager))
|
||||||
}
|
}
|
||||||
|
|||||||
@ -680,7 +680,7 @@ impl SessionManager {
|
|||||||
mcp_config,
|
mcp_config,
|
||||||
None,
|
None,
|
||||||
)
|
)
|
||||||
.map(|(session_manager, _)| session_manager)
|
.map(|(session_manager, _, _)| session_manager)
|
||||||
}
|
}
|
||||||
|
|
||||||
pub fn tools(&self) -> Arc<ToolRegistry> {
|
pub fn tools(&self) -> Arc<ToolRegistry> {
|
||||||
|
|||||||
@ -77,33 +77,90 @@ impl McpClientManager {
|
|||||||
/// This method is designed to be called asynchronously without
|
/// This method is designed to be called asynchronously without
|
||||||
/// blocking the main gateway startup flow.
|
/// blocking the main gateway startup flow.
|
||||||
/// Takes a list of (key, config) pairs from the mcpServers map.
|
/// Takes a list of (key, config) pairs from the mcpServers map.
|
||||||
|
///
|
||||||
|
/// Each server connection is retried up to 3 times with exponential backoff
|
||||||
|
/// (1s, 2s, 4s) to handle transient failures during gateway restart.
|
||||||
pub async fn connect_all(&self, servers: Vec<(String, McpServerConfig)>) -> anyhow::Result<()> {
|
pub async fn connect_all(&self, servers: Vec<(String, McpServerConfig)>) -> anyhow::Result<()> {
|
||||||
|
let total = servers.iter().filter(|(_, c)| c.is_active).count();
|
||||||
|
if total == 0 {
|
||||||
|
tracing::info!("No active MCP servers to connect");
|
||||||
|
return Ok(());
|
||||||
|
}
|
||||||
|
|
||||||
|
tracing::info!(server_count = total, "Connecting to MCP servers");
|
||||||
|
|
||||||
|
let mut connected = 0usize;
|
||||||
|
let mut failed = 0usize;
|
||||||
|
const MAX_RETRIES: u32 = 3;
|
||||||
|
|
||||||
for (key, config) in servers {
|
for (key, config) in servers {
|
||||||
if !config.is_active {
|
if !config.is_active {
|
||||||
tracing::info!(key = %key, "Skipping inactive MCP server");
|
tracing::info!(key = %key, "Skipping inactive MCP server");
|
||||||
continue;
|
continue;
|
||||||
}
|
}
|
||||||
|
|
||||||
// Each server connection is independent
|
let mut last_error = None;
|
||||||
match self.connect_server(&key, &config).await {
|
for attempt in 1..=MAX_RETRIES {
|
||||||
Ok(info) => {
|
match self.connect_server(&key, &config).await {
|
||||||
tracing::info!(
|
Ok(info) => {
|
||||||
key = %key,
|
tracing::info!(
|
||||||
name = %info.name,
|
key = %key,
|
||||||
tools_count = info.tools.len(),
|
name = %info.name,
|
||||||
"Connected to MCP server"
|
tools_count = info.tools.len(),
|
||||||
);
|
attempt,
|
||||||
}
|
"Connected to MCP server"
|
||||||
Err(e) => {
|
);
|
||||||
// Log error but continue with other servers
|
connected += 1;
|
||||||
tracing::error!(
|
last_error = None;
|
||||||
key = %key,
|
break;
|
||||||
error = %e,
|
}
|
||||||
"Failed to connect to MCP server"
|
Err(e) => {
|
||||||
);
|
last_error = Some(e);
|
||||||
|
if attempt < MAX_RETRIES {
|
||||||
|
let delay_secs = 1u64 << (attempt - 1); // 1s, 2s, 4s
|
||||||
|
tracing::warn!(
|
||||||
|
key = %key,
|
||||||
|
attempt,
|
||||||
|
max_retries = MAX_RETRIES,
|
||||||
|
retry_delay_secs = delay_secs,
|
||||||
|
error = %last_error.as_ref().unwrap(),
|
||||||
|
"MCP connection failed, retrying"
|
||||||
|
);
|
||||||
|
tokio::time::sleep(std::time::Duration::from_secs(delay_secs)).await;
|
||||||
|
}
|
||||||
|
}
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
if let Some(e) = last_error {
|
||||||
|
tracing::error!(
|
||||||
|
key = %key,
|
||||||
|
error = %e,
|
||||||
|
attempts = MAX_RETRIES,
|
||||||
|
"Failed to connect to MCP server after all retries"
|
||||||
|
);
|
||||||
|
failed += 1;
|
||||||
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
tracing::info!(
|
||||||
|
total,
|
||||||
|
connected,
|
||||||
|
failed,
|
||||||
|
"MCP connection summary: {}/{} connected, {} failed",
|
||||||
|
connected,
|
||||||
|
total,
|
||||||
|
failed
|
||||||
|
);
|
||||||
|
|
||||||
|
// Only return error if all connections failed
|
||||||
|
if connected == 0 && failed > 0 {
|
||||||
|
return Err(anyhow::anyhow!(
|
||||||
|
"All {} MCP server connection(s) failed",
|
||||||
|
failed
|
||||||
|
));
|
||||||
|
}
|
||||||
|
|
||||||
Ok(())
|
Ok(())
|
||||||
}
|
}
|
||||||
|
|
||||||
@ -343,10 +400,31 @@ impl McpInitializer {
|
|||||||
/// This spawns a background task to connect to MCP servers,
|
/// This spawns a background task to connect to MCP servers,
|
||||||
/// allowing the gateway to start immediately.
|
/// allowing the gateway to start immediately.
|
||||||
pub fn with_config(config: crate::mcp::McpConfig) -> Self {
|
pub fn with_config(config: crate::mcp::McpConfig) -> Self {
|
||||||
|
let server_count = config.mcp_servers.len();
|
||||||
|
let active_count = config.active_servers().len();
|
||||||
|
|
||||||
if !config.has_active_servers() {
|
if !config.has_active_servers() {
|
||||||
|
if server_count > 0 {
|
||||||
|
tracing::info!(
|
||||||
|
server_count,
|
||||||
|
active_count,
|
||||||
|
"MCP disabled: {} server(s) configured but none active",
|
||||||
|
server_count
|
||||||
|
);
|
||||||
|
} else {
|
||||||
|
tracing::info!("MCP disabled: no mcpServers configured in config");
|
||||||
|
}
|
||||||
return Self::disabled();
|
return Self::disabled();
|
||||||
}
|
}
|
||||||
|
|
||||||
|
tracing::info!(
|
||||||
|
server_count,
|
||||||
|
active_count,
|
||||||
|
"MCP enabled: {} active server(s) out of {} configured",
|
||||||
|
active_count,
|
||||||
|
server_count
|
||||||
|
);
|
||||||
|
|
||||||
let manager = Arc::new(McpClientManager::new());
|
let manager = Arc::new(McpClientManager::new());
|
||||||
let servers = config.active_servers();
|
let servers = config.active_servers();
|
||||||
|
|
||||||
|
|||||||
Loading…
x
Reference in New Issue
Block a user