clawmates/crates/cm-api/src/quota.rs

//! Per-workspace resource quotas keyed off the workspace plan. Phase 1 of the
//! multi-tenant scaling work: keep one tenant from exhausting the shared pool.

use axum::extract::State;
use axum::Json;
use cm_domain::WorkspaceId;
use serde::Serialize;
use sqlx::Row;

use crate::{ApiError, AppState, Authed};

/// Caps for a plan tier.
pub struct Quota {
    pub max_agents: i64,
    pub max_live_containers: i64,
    /// Ceiling on `queued` + `running` topology runs at once. Prevents one
    /// workspace flooding the shared queue (a single team run also spawns a
    /// tier-tree of children, so the practical cap grows with the topology).
    pub max_active_runs: i64,
}

/// Per-plan limits. Unknown plans fall back to the free tier.
pub fn plan_quota(plan: &str) -> Quota {
    match plan {
        "team" => Quota {
            max_agents: 50,
            max_live_containers: 50,
            max_active_runs: 100,
        },
        "pro" => Quota {
            max_agents: 20,
            max_live_containers: 20,
            max_active_runs: 25,
        },
        _ => Quota {
            max_agents: 3,
            max_live_containers: 3,
            max_active_runs: 5,
        },
    }
}

/// The workspace's plan name (defaults to "free").
async fn plan_of(state: &AppState, workspace_id: WorkspaceId) -> Result<String, ApiError> {
    let row = sqlx::query("SELECT plan FROM workspaces WHERE id = $1")
        .bind(workspace_id.as_uuid())
        .fetch_optional(&state.pool)
        .await?;
    Ok(row
        .map(|r| r.get::<String, _>("plan"))
        .unwrap_or_else(|| "free".to_string()))
}

/// Reject creating another agent if the workspace is at its plan cap.
pub async fn enforce_new_agent(
    state: &AppState,
    workspace_id: WorkspaceId,
) -> Result<(), ApiError> {
    let plan = plan_of(state, workspace_id).await?;
    let quota = plan_quota(&plan);
    let used = cm_db::repo::agents::count_active(&state.pool, workspace_id).await?;
    if used >= quota.max_agents {
        return Err(ApiError::Quota(format!(
            "agent limit reached ({} on the {plan} plan) — upgrade your plan or remove an agent",
            quota.max_agents
        )));
    }
    Ok(())
}

/// Reject enqueueing another topology run if the workspace is at its plan cap.
pub async fn enforce_new_run(state: &AppState, workspace_id: WorkspaceId) -> Result<(), ApiError> {
    let plan = plan_of(state, workspace_id).await?;
    let quota = plan_quota(&plan);
    let used = cm_db::repo::topology_runs::count_active(&state.pool, workspace_id).await?;
    if used >= quota.max_active_runs {
        return Err(ApiError::Quota(format!(
            "active-run limit reached ({} on the {plan} plan) — wait for a run to finish or upgrade",
            quota.max_active_runs
        )));
    }
    Ok(())
}

/// Reject spinning up another container if the workspace is at its plan cap.
pub async fn enforce_new_container(
    state: &AppState,
    workspace_id: WorkspaceId,
) -> Result<(), ApiError> {
    let plan = plan_of(state, workspace_id).await?;
    let quota = plan_quota(&plan);
    let used =
        cm_db::repo::agent_containers::count_for_workspace(&state.pool, workspace_id).await?;
    if used >= quota.max_live_containers {
        return Err(ApiError::Quota(format!(
            "live container limit reached ({} on the {plan} plan) — close a terminal/agent or upgrade",
            quota.max_live_containers
        )));
    }
    Ok(())
}

#[derive(Serialize)]
pub struct QuotaUsage {
    plan: String,
    agents_used: i64,
    max_agents: i64,
    containers_used: i64,
    max_live_containers: i64,
    active_runs: i64,
    max_active_runs: i64,
}

/// `GET /api/quota` — the caller's workspace usage + limits (for the UI).
pub async fn get_quota(
    State(state): State<AppState>,
    Authed(user): Authed,
) -> Result<Json<QuotaUsage>, ApiError> {
    let plan = plan_of(&state, user.workspace_id).await?;
    let quota = plan_quota(&plan);
    let agents_used = cm_db::repo::agents::count_active(&state.pool, user.workspace_id).await?;
    let containers_used =
        cm_db::repo::agent_containers::count_for_workspace(&state.pool, user.workspace_id).await?;
    let active_runs =
        cm_db::repo::topology_runs::count_active(&state.pool, user.workspace_id).await?;
    Ok(Json(QuotaUsage {
        plan,
        agents_used,
        max_agents: quota.max_agents,
        containers_used,
        max_live_containers: quota.max_live_containers,
        active_runs,
        max_active_runs: quota.max_active_runs,
    }))
}