PicoBot/src/gateway/runtime.rs

//! 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>,
        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>,
        Arc<ModelSelectionStore>,
        Option<Arc<dyn SubAgentRuntime>>,
    ),
    AgentError,
> {
    let store = Arc::new(
        SessionStore::new()
            .map_err(|err| AgentError::Other(format!("session store init error: {}", err)))?,
    );
    // 模型选择内存缓存：session 级 + topic 级（topic 级启动时从 topics 表预热）
    let model_selections = Arc::new(ModelSelectionStore::new());
    let topic_model_selections = Arc::new(ModelSelectionStore::new());
    match store.list_topic_model_selections() {
        Ok(entries) => {
            for (topic_id, provider, model) in entries {
                topic_model_selections.set(&topic_id, provider, model);
            }
        }
        Err(err) => {
            tracing::warn!(
                error = %err,
                "build_session_manager: failed to preload topic model selections"
            );
        }
    }
    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(),
            Some(model_selections.clone()),
            Some(topic_model_selections.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 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(),
        topic_model_selections.clone(),
        store.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,
        topic_model_selections,
        subagent_executor,
    ))
}