bamboo_server/app_state/

init.rs

//! Infrastructure initialization helpers.
//!
//! These functions create domain-level services (storage, skill manager, provider handles,
//! MCP servers, metrics, permissions, schedules) that are shared between `AppState`
//! and `AgentRuntime`.

use std::collections::HashMap;
use std::path::{Path, PathBuf};
use std::sync::Arc;
use std::time::Duration;

use chrono::Utc;
use tokio::sync::{broadcast, RwLock};

use bamboo_agent_core::storage::Storage;
use bamboo_agent_core::AgentEvent;
use bamboo_engine::Agent;
use bamboo_llm::Config;
use bamboo_llm::LLMProvider;
use bamboo_mcp::manager::McpServerManager;
use bamboo_skills::{SkillManager, SkillStoreConfig};
use bamboo_storage::LockedSessionStore;
use bamboo_storage::SessionStoreV2;

use crate::error::AppError;
use crate::schedule_app::manager::{build_schedule_context, ScheduleContext};
use crate::schedule_app::ScheduleManager;
use crate::schedule_app::ScheduleStore;
use bamboo_engine::execution::spawn::{SpawnContext, SpawnScheduler};
use bamboo_metrics::metrics_service::MetricsService;

use super::{AgentRunner, AgentStatus};

/// Type alias for the permission checker trait object.
pub(super) type PermissionChecker = dyn bamboo_tools::permission::PermissionChecker;

/// Type alias for the provider lock + handle pair returned by [`build_provider_handles`].
type ProviderHandles = (Arc<RwLock<Arc<dyn LLMProvider>>>, Arc<dyn LLMProvider>);

/// Initialize storage components (session store + storage trait object).
///
/// Spawns a background task to rebuild the search index on startup.
pub async fn init_storage(
    data_dir: &Path,
) -> Result<(Arc<SessionStoreV2>, Arc<dyn Storage>), AppError> {
    tracing::info!("Initializing session store V2 at: {:?}", data_dir);
    let session_store = Arc::new(SessionStoreV2::new(data_dir.to_path_buf()).await.map_err(
        |error| {
            tracing::error!(
                "Failed to initialize SessionStoreV2 at {:?}: {}",
                data_dir,
                error
            );
            AppError::StorageError(error)
        },
    )?);

    // One-shot, idempotent migration: split each legacy session's runtime
    // control-plane into a `runtime.json` sidecar so the O(1) runtime-save path
    // (used heavily during sub-agent spawn) is active immediately. Run inline
    // before the server starts serving so it cannot race concurrent saves; it is
    // a no-op on already-migrated stores (marker file).
    match session_store.migrate_runtime_sidecars().await {
        Ok(0) => {}
        Ok(count) => tracing::info!("Runtime sidecar migration created {count} sidecar(s)"),
        Err(error) => {
            // Non-fatal: loading tolerates a missing sidecar. Log and continue.
            tracing::warn!("Runtime sidecar migration failed (continuing): {error}");
        }
    }

    let session_store_for_rebuild = session_store.clone();
    tokio::spawn(async move {
        let purged_rows = match session_store_for_rebuild
            .search_index()
            .prune_stale_sessions()
            .await
        {
            Ok(count) => count,
            Err(error) => {
                tracing::warn!("Background session search index prune failed: {}", error);
                0
            }
        };

        if let Err(error) = session_store_for_rebuild.rebuild_search_index().await {
            tracing::warn!("Background session search index rebuild failed: {}", error);
        } else {
            tracing::info!("Background session search index rebuild completed");
        }

        match session_store_for_rebuild
            .search_index()
            .maybe_vacuum_if_needed(purged_rows)
            .await
        {
            Ok(true) => tracing::info!("Background session search index vacuum completed"),
            Ok(false) => {}
            Err(error) => {
                tracing::warn!("Background session search index vacuum failed: {}", error)
            }
        }
    });

    let storage: Arc<dyn Storage> = session_store.clone();
    tracing::info!(
        "Session store V2 initialized (index: {:?}, sessions: {:?})",
        session_store.index_path(),
        session_store.sessions_root_dir()
    );
    Ok((session_store, storage))
}

/// Initialize the skill manager with workspace directory and active mode from environment.
pub async fn init_skill_manager(data_dir: &Path) -> Arc<SkillManager> {
    let project_dir = std::env::var_os("BAMBOO_WORKSPACE_DIR")
        .map(PathBuf::from)
        .or_else(|| std::env::current_dir().ok());
    let active_mode = std::env::var("BAMBOO_SKILL_MODE")
        .ok()
        .map(|value| value.trim().to_string())
        .filter(|value| !value.is_empty());

    let skill_manager = Arc::new(SkillManager::with_config(SkillStoreConfig {
        skills_dir: data_dir.join("skills"),
        project_dir,
        active_mode,
    }));
    if let Err(error) = skill_manager.initialize().await {
        tracing::warn!("Failed to initialize skill manager: {}", error);
    }
    skill_manager
}

/// Build reloadable provider handles.
///
/// Returns `(provider_lock, provider_handle)` where `provider_lock` is the
/// hot-reloadable `RwLock` and `provider_handle` is a stable `ReloadableProvider`
/// that always delegates to the latest value in the lock.
pub fn build_provider_handles(provider: Arc<dyn LLMProvider>) -> ProviderHandles {
    let provider_lock: Arc<RwLock<Arc<dyn LLMProvider>>> = Arc::new(RwLock::new(provider));
    let provider_handle: Arc<dyn LLMProvider> = Arc::new(
        crate::reloadable_provider::ReloadableProvider::new(provider_lock.clone()),
    );
    (provider_lock, provider_handle)
}

/// Load permission configuration from disk.
///
/// Falls back to disabled permissions if no config exists or loading fails.
pub async fn load_permission_checker(bamboo_home_dir: &Path) -> Arc<PermissionChecker> {
    use bamboo_tools::permission::{PermissionConfig, RiskLevel};

    let storage = bamboo_tools::permission::storage::PermissionStorage::new(bamboo_home_dir);
    let permission_config = match storage.load().await {
        Ok(Some(config)) => config,
        Ok(None) => {
            // First run, no saved config: default posture is "ask on high-risk".
            // Checks are enabled (PermissionConfig::new defaults to enabled +
            // Default mode); only high-risk operations (execute command, delete,
            // git write, terminal session) require confirmation. Lower-risk ops
            // (file writes, HTTP) auto-allow.
            let cfg = PermissionConfig::new();
            cfg.set_confirm_threshold(RiskLevel::High);
            cfg
        }
        Err(error) => {
            tracing::warn!(
                "Failed to load permission config; defaulting to ask-on-high-risk: {error}"
            );
            let cfg = PermissionConfig::new();
            cfg.set_confirm_threshold(RiskLevel::High);
            cfg
        }
    };
    permission_config.cleanup_expired_grants();

    // Wrap the config checker in a mode-aware checker so the active PermissionMode
    // (Default / Plan / AcceptEdits / DontAsk / BypassPermissions) takes effect at
    // the checker level. Both share the same Arc<PermissionConfig>, so session
    // grants recorded on approval (see the respond handler) are visible here.
    let config = Arc::new(permission_config);
    let inner: Arc<dyn bamboo_tools::permission::PermissionChecker> = Arc::new(
        bamboo_tools::permission::ConfigPermissionChecker::new(config.clone()),
    );
    Arc::new(bamboo_tools::permission::ModeAwarePermissionChecker::new(
        inner, config,
    ))
}

/// Initialize MCP server manager with background server initialization.
pub fn init_mcp_manager(config: Arc<RwLock<Config>>) -> Arc<McpServerManager> {
    let mcp_manager = Arc::new(McpServerManager::new_with_config(config.clone()));

    // Initialize MCP servers in the background so the HTTP API is responsive quickly.
    {
        let mcp_manager = mcp_manager.clone();
        let config = config.clone();
        tokio::spawn(async move {
            let mcp_config = config.read().await.mcp.clone();
            mcp_manager.initialize_from_config(&mcp_config).await;
        });
    }

    mcp_manager
}

/// Initialize metrics service with SQLite backend.
pub async fn init_metrics_service(data_dir: &Path) -> Result<Arc<MetricsService>, AppError> {
    let service = MetricsService::new(data_dir.join("metrics.db"))
        .await
        .map_err(|error| {
            tracing::error!("Failed to initialize metrics storage: {}", error);
            AppError::InternalError(anyhow::anyhow!(
                "Failed to initialize metrics storage: {error}"
            ))
        })?;
    Ok(Arc::new(service))
}

/// Idle-eviction TTL for terminal runners / session senders, in seconds.
///
/// A completed runner is retained for at least this long after `completed_at`
/// so a client that subscribes shortly after the run finishes still gets the
/// cached `last_critical_events` / `last_budget_event` replay.
pub(crate) const SESSION_MAP_IDLE_TTL_SECS: i64 = 300;

/// Single idle-eviction pass over the per-session runner + event-sender maps.
///
/// Drops a `(agent_runners, session_event_senders)` pair together when the
/// runner is terminal, older than `ttl_secs`, and its broadcast channel has **no
/// live receivers**. Returns the number of runners evicted.
///
/// # Why `receiver_count() == 0` is required
///
/// A runner's `event_sender` is a clone of the session's long-lived sender (see
/// `reserve_runner` / `try_reserve_runner`), so `receiver_count()` reflects
/// every SSE/WS subscriber on the session stream. Evicting while a receiver is
/// live would (a) yank the replay cache out from under a client that just
/// subscribed after `Complete`, and (b) silently drop a terminal parent whose
/// still-running children forward events into this stream (that parent keeps a
/// live subscription). So we only reclaim genuinely-idle terminal sessions —
/// exactly the ephemeral sub-agent / guardian / scheduled runs that never had a
/// UI subscription and whose count therefore reaches zero.
///
/// Removing the map's `Sender` (the last clone once the terminal runner's own
/// clone is dropped by the `retain`) closes the channel, which also lets any
/// per-session notification relay task exit on `RecvError::Closed`.
///
/// This runs synchronously under both write guards held by the caller (no
/// `.await` between the check and the removals), so a concurrent `reserve_runner`
/// / `get_or_create_event_sender` (which need those same locks) cannot interleave.
pub(crate) fn evict_idle_session_entries(
    runners: &mut HashMap<String, AgentRunner>,
    senders: &mut HashMap<String, broadcast::Sender<AgentEvent>>,
    ttl_secs: i64,
    now: chrono::DateTime<Utc>,
    log_prefix: Option<&'static str>,
) -> usize {
    let mut evicted: Vec<String> = Vec::new();

    runners.retain(|session_id, runner| {
        let keep = match &runner.status {
            AgentStatus::Running => true,
            _ => {
                let age =
                    now.signed_duration_since(runner.completed_at.unwrap_or(runner.started_at));
                let expired = age.num_seconds() >= ttl_secs;
                let has_receivers = runner.event_sender.receiver_count() > 0;
                // Keep unless it is BOTH past the TTL AND has no live receiver.
                !expired || has_receivers
            }
        };
        if !keep {
            evicted.push(session_id.clone());
            if let Some(prefix) = log_prefix {
                tracing::debug!("[{}:{}] Evicting idle terminal runner", prefix, session_id);
            } else {
                tracing::debug!("[{}] Evicting idle terminal runner", session_id);
            }
        }
        keep
    });

    // Drop the paired long-lived sender for each evicted runner, but only if it
    // too has no live receiver (re-checked here; it is the same channel as the
    // runner's `event_sender`, so this is normally a formality — a session sender
    // can, however, outlive its runner, and we must never close a channel a
    // client is actively reading).
    for session_id in &evicted {
        if senders
            .get(session_id)
            .is_some_and(|sender| sender.receiver_count() == 0)
        {
            senders.remove(session_id);
        }
    }

    evicted.len()
}

/// Spawn a background task that idle-evicts completed agent runners **and their
/// paired session event senders** after [`SESSION_MAP_IDLE_TTL_SECS`].
///
/// Fire-and-forget (runs for the process lifetime), matching the other
/// background maintenance tickers. See [`evict_idle_session_entries`] for the
/// eviction predicate and its race-freedom argument.
pub fn spawn_session_map_cleanup_task(
    runners: Arc<RwLock<HashMap<String, AgentRunner>>>,
    senders: Arc<RwLock<HashMap<String, broadcast::Sender<AgentEvent>>>>,
    log_prefix: Option<&'static str>,
) {
    tokio::spawn(async move {
        loop {
            tokio::time::sleep(Duration::from_secs(60)).await;

            // Lock order: runners ⊃ senders. This matches `reserve_runner_core`
            // (which nests senders inside the runners write lock to re-assert the
            // sender) and is never inverted anywhere, so nesting is deadlock-free.
            //
            // Both write locks are intentionally held across the whole O(n) pass
            // rather than snapshot-then-remove: the atomicity is load-bearing.
            // `reserve_runner_core` inserts a `Running` runner AND re-asserts its
            // sender under the runners lock; holding both here means a concurrent
            // reservation cannot interleave between our runner-eviction and our
            // sender-removal, so we can never drop a sender that a just-reserved
            // run re-asserted. The pass is cheap (a receiver_count atomic load +
            // timestamp math per entry) at the 60s cadence, so the widened window
            // is acceptable.
            let mut runners_guard = runners.write().await;
            let mut senders_guard = senders.write().await;
            let now = Utc::now();
            let evicted = evict_idle_session_entries(
                &mut runners_guard,
                &mut senders_guard,
                SESSION_MAP_IDLE_TTL_SECS,
                now,
                log_prefix,
            );
            drop(senders_guard);
            drop(runners_guard);

            if evicted > 0 {
                tracing::debug!("Idle-evicted {evicted} completed session runner(s)/sender(s)");
            }
        }
    });
}

/// Initialize schedule store for timed tasks.
pub async fn init_schedule_store(data_dir: &PathBuf) -> Result<Arc<ScheduleStore>, AppError> {
    let store = ScheduleStore::new(data_dir.clone())
        .await
        .map_err(|error| {
            tracing::error!(
                "Failed to initialize ScheduleStore at {:?}: {}",
                data_dir,
                error
            );
            AppError::StorageError(error)
        })?;
    Ok(Arc::new(store))
}

/// Build sub-session spawn scheduler.
#[allow(clippy::too_many_arguments)]
pub fn build_spawn_scheduler(
    agent: Arc<Agent>,
    child_tools: Arc<dyn bamboo_agent_core::tools::ToolExecutor>,
    sessions: bamboo_engine::SessionCache,
    agent_runners: Arc<RwLock<HashMap<String, AgentRunner>>>,
    session_event_senders: Arc<RwLock<HashMap<String, broadcast::Sender<AgentEvent>>>>,
    external_child_runner: Arc<dyn bamboo_engine::runtime::execution::ExternalChildRunner>,
    provider_router: Option<Arc<bamboo_llm::ProviderModelRouter>>,
    completion_handler: Option<Arc<dyn bamboo_engine::execution::ChildCompletionHandler>>,
    app_data_dir: Option<std::path::PathBuf>,
    account_feed_inbox: Option<bamboo_engine::execution::AccountFeedInbox>,
) -> Arc<SpawnScheduler> {
    Arc::new(SpawnScheduler::new(SpawnContext {
        agent,
        tools: child_tools,
        sessions_cache: sessions,
        agent_runners,
        session_event_senders,
        external_child_runner,
        provider_router,
        app_data_dir,
        completion_handler,
        account_feed_inbox,
    }))
}

/// Build schedule manager with minimal tool surface for background automation.
#[allow(clippy::too_many_arguments)]
pub fn build_schedule_manager(
    schedule_store: Arc<ScheduleStore>,
    agent: Arc<Agent>,
    tools_for_schedules: Arc<dyn bamboo_agent_core::tools::ToolExecutor>,
    sessions: bamboo_engine::SessionCache,
    agent_runners: Arc<RwLock<HashMap<String, AgentRunner>>>,
    session_event_senders: Arc<RwLock<HashMap<String, broadcast::Sender<AgentEvent>>>>,
    persistence: Arc<LockedSessionStore>,
    config: Arc<RwLock<Config>>,
    provider_registry: Arc<bamboo_llm::ProviderRegistry>,
    app_data_dir: Option<std::path::PathBuf>,
    account_feed_inbox: Option<bamboo_engine::execution::AccountFeedInbox>,
) -> Arc<ScheduleManager> {
    let base_ctx = ScheduleContext {
        schedule_store,
        agent,
        tools: tools_for_schedules,
        sessions_cache: sessions,
        agent_runners,
        session_event_senders,
        account_feed_inbox,
        persistence,
        app_data_dir,
        trigger_engine: crate::schedule_app::default_trigger_engine(),
        resolve_run_config: Arc::new(|_| unimplemented!("replaced by build_schedule_context")),
    };
    Arc::new(ScheduleManager::new(build_schedule_context(
        base_ctx,
        config,
        provider_registry,
    )))
}

#[cfg(test)]
mod eviction_tests {
    use super::*;
    use chrono::Duration as ChronoDuration;

    const TTL: i64 = SESSION_MAP_IDLE_TTL_SECS;

    /// A terminal (Completed) runner whose `event_sender` shares `tx`'s channel,
    /// exactly as production sets it in `reserve_runner`.
    fn terminal_runner(
        tx: &broadcast::Sender<AgentEvent>,
        completed_secs_ago: i64,
        now: chrono::DateTime<Utc>,
    ) -> AgentRunner {
        let mut runner = AgentRunner::new();
        runner.status = AgentStatus::Completed;
        runner.completed_at = Some(now - ChronoDuration::seconds(completed_secs_ago));
        runner.event_sender = tx.clone();
        runner
    }

    fn insert_pair(
        runners: &mut HashMap<String, AgentRunner>,
        senders: &mut HashMap<String, broadcast::Sender<AgentEvent>>,
        id: &str,
        runner: AgentRunner,
        tx: broadcast::Sender<AgentEvent>,
    ) {
        runners.insert(id.to_string(), runner);
        senders.insert(id.to_string(), tx);
    }

    #[test]
    fn evicts_idle_terminal_pair_past_ttl_with_no_receivers() {
        let now = Utc::now();
        let (tx, _) = broadcast::channel::<AgentEvent>(16); // receiver dropped -> count 0
        let mut runners = HashMap::new();
        let mut senders = HashMap::new();
        insert_pair(
            &mut runners,
            &mut senders,
            "s1",
            terminal_runner(&tx, TTL + 100, now),
            tx.clone(),
        );

        let evicted = evict_idle_session_entries(&mut runners, &mut senders, TTL, now, None);

        assert_eq!(evicted, 1);
        assert!(runners.is_empty(), "terminal idle runner must be dropped");
        assert!(
            senders.is_empty(),
            "paired session sender must be dropped together with the runner"
        );
    }

    #[test]
    fn retains_terminal_runner_with_live_receiver() {
        let now = Utc::now();
        let (tx, _) = broadcast::channel::<AgentEvent>(16);
        // A client subscribed just after completion — its receiver is live.
        let _live_rx = tx.subscribe();
        assert_eq!(tx.receiver_count(), 1);

        let mut runners = HashMap::new();
        let mut senders = HashMap::new();
        insert_pair(
            &mut runners,
            &mut senders,
            "s1",
            terminal_runner(&tx, TTL + 100, now),
            tx.clone(),
        );

        let evicted = evict_idle_session_entries(&mut runners, &mut senders, TTL, now, None);

        assert_eq!(evicted, 0, "must not evict while a receiver is live");
        assert!(runners.contains_key("s1"));
        assert!(senders.contains_key("s1"));
    }

    #[test]
    fn retains_young_terminal_runner() {
        let now = Utc::now();
        let (tx, _) = broadcast::channel::<AgentEvent>(16);
        let mut runners = HashMap::new();
        let mut senders = HashMap::new();
        // Completed only 60s ago — inside the replay window.
        insert_pair(
            &mut runners,
            &mut senders,
            "s1",
            terminal_runner(&tx, 60, now),
            tx.clone(),
        );

        let evicted = evict_idle_session_entries(&mut runners, &mut senders, TTL, now, None);

        assert_eq!(
            evicted, 0,
            "must preserve the post-completion replay window"
        );
        assert!(runners.contains_key("s1"));
        assert!(senders.contains_key("s1"));
    }

    #[test]
    fn retains_running_runner_regardless_of_age() {
        let now = Utc::now();
        let (tx, _) = broadcast::channel::<AgentEvent>(16);
        let mut runner = AgentRunner::new();
        runner.status = AgentStatus::Running;
        runner.started_at = now - ChronoDuration::seconds(TTL + 100_000);
        runner.completed_at = None;
        runner.event_sender = tx.clone();

        let mut runners = HashMap::new();
        let mut senders = HashMap::new();
        insert_pair(&mut runners, &mut senders, "s1", runner, tx.clone());

        let evicted = evict_idle_session_entries(&mut runners, &mut senders, TTL, now, None);

        assert_eq!(evicted, 0, "a Running runner is never evicted");
        assert!(runners.contains_key("s1"));
        assert!(senders.contains_key("s1"));
    }

    #[test]
    fn keeps_sender_if_it_has_receivers_even_when_runner_evicted() {
        // Defensive edge: the runner qualifies for eviction, but the session
        // sender independently has a live receiver (e.g. a subscriber that
        // reused the channel). The sender must survive so we never close a
        // channel a client is reading.
        let now = Utc::now();
        let (runner_tx, _) = broadcast::channel::<AgentEvent>(16);
        // Distinct channel for the map entry, with a live receiver.
        let (sender_tx, _) = broadcast::channel::<AgentEvent>(16);
        let _live = sender_tx.subscribe();

        let mut runners = HashMap::new();
        let mut senders = HashMap::new();
        runners.insert(
            "s1".to_string(),
            terminal_runner(&runner_tx, TTL + 100, now),
        );
        senders.insert("s1".to_string(), sender_tx.clone());

        let evicted = evict_idle_session_entries(&mut runners, &mut senders, TTL, now, None);

        assert_eq!(evicted, 1, "runner (no receivers) is evicted");
        assert!(runners.is_empty());
        assert!(
            senders.contains_key("s1"),
            "sender with a live receiver must be retained even when the runner is dropped"
        );
    }
}