gitdataai/apps/adminrpc/src/main.rs

use std::net::SocketAddr;
use actix_web::{web, App as ActixApp, HttpResponse, HttpServer};
use anyhow::Context as _;
use clap::Parser;
use config::AppConfig;
use deadpool_redis::{cluster, Runtime};
use session_manager::{SessionManager, SessionStorage};
use rpc::admin::server::{serve, DEFAULT_GRPC_PORT};
use uuid::Uuid;

mod args;
use args::Args;

#[tokio::main]
async fn main() -> anyhow::Result<()> {
    let cfg = AppConfig::load();
    let log_level = cfg.log_level().unwrap_or_else(|_| "info".to_string());
    observability::init_tracing_subscriber(&log_level);

    let args = Args::parse();
    let grpc_addr: SocketAddr = args
        .bind
        .map(|s| s.parse())
        .unwrap_or_else(|| format!("0.0.0.0:{}", DEFAULT_GRPC_PORT).parse())
        .context("invalid grpc bind address")?;

    // Admin HTTP port is gRPC port + 1 (e.g., 9091)
    let admin_port: u16 = args.http_port.unwrap_or(grpc_addr.port() + 1);
    let admin_addr: SocketAddr = format!("0.0.0.0:{}", admin_port).parse().unwrap();

    tracing::info!(
        app_name = %cfg.app_name().unwrap_or_default(),
        grpc_addr = %grpc_addr,
        admin_addr = %admin_addr,
        "Starting admin RPC server"
    );

    // ── OTLP tracing ─────────────────────────────────────────────────────────
    let _otel_guard = if cfg.otel_enabled().unwrap_or(false) {
        let endpoint = cfg.otel_endpoint().unwrap_or_else(|_| "http://localhost:4317".to_string());
        let service_name = cfg.otel_service_name().unwrap_or_else(|_| "adminrpc".to_string());
        let service_version = cfg.otel_service_version().unwrap_or_else(|_| "0.1.0".to_string());
        tracing::info!(endpoint = %endpoint, service = %service_name, "OTLP tracing enabled");
        let guard = observability::init_otlp(&endpoint, &service_name, &service_version, &log_level)
            .map_err(|e| anyhow::anyhow!("OTLP init failed: {}", e))?;
        guard
    } else {
        None
    };

    // Redis connection pool
    let redis_url = cfg.redis_url()?;
    tracing::info!(redis_url = %redis_url, "Connecting to Redis");
    let manager = cluster::Manager::new(vec![redis_url.clone()], false)
        .map_err(|e| anyhow::anyhow!("failed to create redis cluster manager: {}", e))?;
    let pool: cluster::Pool = cluster::Pool::builder(manager)
        .max_size(16)
        .runtime(Runtime::Tokio1)
        .build()
        .map_err(|e| anyhow::anyhow!("failed to build redis pool: {}", e))?;

    let _conn = pool.get().await.context("redis pool connection failed")?;
    tracing::info!("Redis connected");

    let storage = SessionStorage::new(pool.clone());
    let session_manager = SessionManager::new(storage);

    // Spawn gRPC server in background
    let sm_for_grpc = session_manager.clone();
    let grpc_handle = tokio::spawn(async move {
        if let Err(e) = serve(grpc_addr, sm_for_grpc).await {
            tracing::error!(error = %e, "Admin gRPC server error");
        }
    });

    // Start HTTP REST server
    let http_handle = tokio::spawn(async move {
        let pool_for_http = pool.clone();
        let sm_for_http = session_manager.clone();

        let result = HttpServer::new(move || {
            ActixApp::new()
                .app_data(web::Data::new(pool_for_http.clone()))
                .app_data(web::Data::new(sm_for_http.clone()))
                .route("/health", web::get().to(health))
                .route("/admin/metrics/export", web::get().to(metrics_export))
                .service(
                    web::scope("/api/admin")
                        // Sessions
                        .route("/sessions/workspace/{workspace_id}", web::get().to(list_workspace_sessions))
                        .route("/sessions/user/{user_id}", web::get().to(list_user_sessions))
                        .route("/sessions/user/{user_id}/status", web::get().to(get_user_status))
                        .route("/sessions/user/{user_id}/info", web::get().to(get_user_info))
                        .route("/sessions/workspace/{workspace_id}/online-users", web::get().to(get_workspace_online_users))
                        .route("/sessions/user/{user_id}/online", web::get().to(is_user_online))
                        .route("/sessions/kick", web::post().to(kick_user))
                        .route("/sessions/kick-workspace", web::post().to(kick_user_from_workspace))
                        // Metrics
                        .route("/metrics", web::get().to(get_metrics))
                        .route("/metrics/export", web::get().to(metrics_export)),
                )
        })
        .bind(admin_addr)
        .expect("failed to bind HTTP server")
        .workers(2)
        .run()
        .await;

        if let Err(e) = result {
            tracing::error!(error = %e, "Admin HTTP server error");
        }
    });

    tracing::info!(addr = %grpc_addr, "Admin gRPC server listening");
    tracing::info!(addr = %admin_addr, "Admin HTTP server listening");

    grpc_handle.await?;
    http_handle.abort();
    tracing::info!("Admin RPC server stopped");
    Ok(())
}

// ─── HTTP Handlers ────────────────────────────────────────────────────────────

async fn health() -> HttpResponse {
    HttpResponse::Ok().json(serde_json::json!({ "ok": true }))
}

async fn metrics_export(pool: web::Data<cluster::Pool>) -> HttpResponse {
    match observability::export_all_metrics_csv(pool.get_ref(), "").await {
        Ok(csv) => HttpResponse::Ok()
            .content_type("text/csv; charset=utf-8")
            .body(csv),
        Err(e) => HttpResponse::InternalServerError().json(serde_json::json!({ "error": e.to_string() })),
    }
}

async fn get_metrics(pool: web::Data<cluster::Pool>) -> HttpResponse {
    match observability::query_all_instance_metrics(pool.get_ref(), "", 100).await {
        Ok(instances) => HttpResponse::Ok().json(instances),
        Err(e) => HttpResponse::InternalServerError().json(serde_json::json!({ "error": e.to_string() })),
    }
}

fn parse_uuid(s: &str) -> Option<Uuid> {
    Uuid::parse_str(s).ok()
}

async fn list_workspace_sessions(
    sm: web::Data<SessionManager>,
    path: web::Path<String>,
) -> HttpResponse {
    let workspace_id = match parse_uuid(&path) {
        Some(id) => id,
        None => return HttpResponse::BadRequest().json(serde_json::json!({ "error": "invalid workspace_id" })),
    };
    match sm.get_workspace_sessions(&workspace_id).await {
        Ok(sessions) => HttpResponse::Ok().json(sessions),
        Err(e) => HttpResponse::InternalServerError().json(serde_json::json!({ "error": e.to_string() })),
    }
}

async fn list_user_sessions(
    sm: web::Data<SessionManager>,
    path: web::Path<String>,
) -> HttpResponse {
    let user_id = match parse_uuid(&path) {
        Some(id) => id,
        None => return HttpResponse::BadRequest().json(serde_json::json!({ "error": "invalid user_id" })),
    };
    match sm.get_user_sessions(&user_id).await {
        Ok(sessions) => HttpResponse::Ok().json(sessions),
        Err(e) => HttpResponse::InternalServerError().json(serde_json::json!({ "error": e.to_string() })),
    }
}

async fn get_user_status(
    sm: web::Data<SessionManager>,
    path: web::Path<String>,
) -> HttpResponse {
    let user_id = match parse_uuid(&path) {
        Some(id) => id,
        None => return HttpResponse::BadRequest().json(serde_json::json!({ "error": "invalid user_id" })),
    };
    match sm.get_user_status(&user_id).await {
        Ok(status) => HttpResponse::Ok().json(serde_json::json!({ "status": format!("{:?}", status) })),
        Err(e) => HttpResponse::InternalServerError().json(serde_json::json!({ "error": e.to_string() })),
    }
}

async fn get_user_info(
    sm: web::Data<SessionManager>,
    path: web::Path<String>,
) -> HttpResponse {
    let user_id = match parse_uuid(&path) {
        Some(id) => id,
        None => return HttpResponse::BadRequest().json(serde_json::json!({ "error": "invalid user_id" })),
    };
    match sm.get_user_info(&user_id).await {
        Ok(info) => HttpResponse::Ok().json(info),
        Err(e) => HttpResponse::InternalServerError().json(serde_json::json!({ "error": e.to_string() })),
    }
}

async fn get_workspace_online_users(
    sm: web::Data<SessionManager>,
    path: web::Path<String>,
) -> HttpResponse {
    let workspace_id = match parse_uuid(&path) {
        Some(id) => id,
        None => return HttpResponse::BadRequest().json(serde_json::json!({ "error": "invalid workspace_id" })),
    };
    match sm.get_workspace_online_users(&workspace_id).await {
        Ok(user_ids) => HttpResponse::Ok().json(serde_json::json!({ "user_ids": user_ids })),
        Err(e) => HttpResponse::InternalServerError().json(serde_json::json!({ "error": e.to_string() })),
    }
}

async fn is_user_online(
    sm: web::Data<SessionManager>,
    path: web::Path<String>,
) -> HttpResponse {
    let user_id = match parse_uuid(&path) {
        Some(id) => id,
        None => return HttpResponse::BadRequest().json(serde_json::json!({ "error": "invalid user_id" })),
    };
    match sm.is_user_online(&user_id).await {
        Ok(online) => HttpResponse::Ok().json(serde_json::json!({ "online": online })),
        Err(e) => HttpResponse::InternalServerError().json(serde_json::json!({ "error": e.to_string() })),
    }
}

#[derive(serde::Deserialize)]
struct KickUserPayload {
    user_id: String,
}

async fn kick_user(
    sm: web::Data<SessionManager>,
    body: web::Json<KickUserPayload>,
) -> HttpResponse {
    let user_id = match parse_uuid(&body.user_id) {
        Some(id) => id,
        None => return HttpResponse::BadRequest().json(serde_json::json!({ "error": "invalid user_id" })),
    };
    match sm.kick_user(&user_id).await {
        Ok(count) => HttpResponse::Ok().json(serde_json::json!({ "kicked_count": count })),
        Err(e) => HttpResponse::InternalServerError().json(serde_json::json!({ "error": e.to_string() })),
    }
}

#[derive(serde::Deserialize)]
struct KickWorkspacePayload {
    user_id: String,
    workspace_id: String,
}

async fn kick_user_from_workspace(
    sm: web::Data<SessionManager>,
    body: web::Json<KickWorkspacePayload>,
) -> HttpResponse {
    let user_id = match parse_uuid(&body.user_id) {
        Some(id) => id,
        None => return HttpResponse::BadRequest().json(serde_json::json!({ "error": "invalid user_id" })),
    };
    let workspace_id = match parse_uuid(&body.workspace_id) {
        Some(id) => id,
        None => return HttpResponse::BadRequest().json(serde_json::json!({ "error": "invalid workspace_id" })),
    };
    match sm.kick_user_from_workspace(&user_id, &workspace_id).await {
        Ok(count) => HttpResponse::Ok().json(serde_json::json!({ "kicked_count": count })),
        Err(e) => HttpResponse::InternalServerError().json(serde_json::json!({ "error": e.to_string() })),
    }
}