Vapora/crates/vapora-agents/src/bin/server.rs

//! VAPORA Agent Server Binary
//! Provides HTTP server for agent coordination and health checks

use anyhow::Result;
use axum::{extract::State, routing::get, Json, Router};
use serde_json::json;
use std::sync::Arc;
use tokio::net::TcpListener;
use tracing::{error, info};
use vapora_agents::{config::AgentConfig, coordinator::AgentCoordinator, registry::AgentRegistry};
use vapora_llm_router::{BudgetConfig, BudgetManager};

#[derive(Clone)]
struct AppState {
    coordinator: Arc<AgentCoordinator>,
    #[allow(dead_code)]
    budget_manager: Option<Arc<BudgetManager>>,
}

#[tokio::main]
async fn main() -> Result<()> {
    // Initialize tracing
    tracing_subscriber::fmt()
        .with_env_filter(
            tracing_subscriber::EnvFilter::from_default_env()
                .add_directive("vapora_agents=debug".parse()?),
        )
        .init();

    info!("Starting VAPORA Agent Server");

    // Load configuration
    let config = AgentConfig::from_env()?;
    info!("Loaded configuration from environment");

    // Load budget configuration
    let budget_config_path = std::env::var("BUDGET_CONFIG_PATH")
        .unwrap_or_else(|_| "config/agent-budgets.toml".to_string());
    let budget_manager = match BudgetConfig::load_or_default(&budget_config_path) {
        Ok(budget_config) => {
            if budget_config.budgets.is_empty() {
                info!(
                    "No budget configuration found at {}, running without budget enforcement",
                    budget_config_path
                );
                None
            } else {
                let manager = Arc::new(BudgetManager::new(budget_config.budgets));
                info!(
                    "Loaded budget configuration for {} roles",
                    manager.list_budgets().await.len()
                );
                Some(manager)
            }
        }
        Err(e) => {
            error!("Failed to load budget configuration: {}", e);
            return Err(e.into());
        }
    };

    // Initialize agent registry and coordinator
    // Max 10 agents per role (can be configured via environment)
    let max_agents_per_role = std::env::var("MAX_AGENTS_PER_ROLE")
        .ok()
        .and_then(|v| v.parse().ok())
        .unwrap_or(10);
    let registry = Arc::new(AgentRegistry::new(max_agents_per_role));
    let mut coordinator = AgentCoordinator::new(config, registry).await?;

    // Attach budget manager to coordinator if available
    if let Some(ref bm) = budget_manager {
        coordinator = coordinator.with_budget_manager(bm.clone());
        info!("Budget enforcement enabled for agent coordinator");
    }

    let coordinator = Arc::new(coordinator);

    // Start coordinator
    let _coordinator_handle = {
        let coordinator = coordinator.clone();
        tokio::spawn(async move {
            if let Err(e) = coordinator.start().await {
                error!("Coordinator error: {}", e);
            }
        })
    };

    // Build application state
    let state = AppState {
        coordinator,
        budget_manager,
    };

    // Build HTTP router
    let app = Router::new()
        .route("/health", get(health_handler))
        .route("/ready", get(readiness_handler))
        .with_state(state);

    // Start HTTP server
    let addr = std::env::var("BIND_ADDR").unwrap_or_else(|_| "0.0.0.0:9000".to_string());
    info!("Agent server listening on {}", addr);

    let listener = TcpListener::bind(&addr).await?;

    axum::serve(listener, app).await?;

    // Note: coordinator_handle would be awaited here if needed,
    // but axum::serve blocks until server shutdown
    Ok(())
}

/// Health check endpoint
async fn health_handler() -> Json<serde_json::Value> {
    Json(json!({
        "status": "healthy",
        "service": "vapora-agents",
        "version": env!("CARGO_PKG_VERSION")
    }))
}

/// Readiness check endpoint
async fn readiness_handler(State(state): State<AppState>) -> Json<serde_json::Value> {
    let is_ready = state.coordinator.is_ready().await;

    Json(json!({
        "ready": is_ready,
        "service": "vapora-agents",
        "agents": state.coordinator.get_agent_count().await
    }))
}