feat: Monthly token limits per user

Database: - Add monthly_token_limit and token_limit_reset_day to users table - Create ai_usage_monthly view for current month usage - Add check_token_limit() function to verify available tokens - Add get_user_usage_stats() function for historical usage API Endpoints: - PUT /admin/users/{user_id}/token-limit - Set monthly limit - GET /admin/users/{user_id}/token-usage - Get user's current usage - GET /admin/users/{user_id}/token-limit/check - Check if user has tokens Features: - Default limit: 100,000 tokens/month - Reset day: 1st of month (configurable 1-28) - Limit of 0 = unlimited tokens - Enforce limits at API level (check before AI requests) Co-authored-by: Qwen-Coder <qwen-coder@alibabacloud.com>
2026-03-23 16:10:27 -03:00
parent ef213a61a0
commit 9e57756528
3 changed files with 265 additions and 1 deletions
@@ -0,0 +1,138 @@
+-- Monthly Token Limits per User
+-- Allows setting and enforcing monthly AI token usage limits
+
+-- Add monthly token limit column to users table
+ALTER TABLE users 
+ADD COLUMN IF NOT EXISTS monthly_token_limit INTEGER DEFAULT 100000,
+ADD COLUMN IF NOT EXISTS token_limit_reset_day INTEGER DEFAULT 1; -- Day of month to reset (1-28)
+
+COMMENT ON COLUMN users.monthly_token_limit IS 'Maximum AI tokens user can consume per month (0 = unlimited)';
+COMMENT ON COLUMN users.token_limit_reset_day IS 'Day of month when token counter resets (1-28, to avoid month-end issues)';
+
+-- Create view for current month usage per user
+CREATE OR REPLACE VIEW ai_usage_monthly AS
+SELECT 
+    au.user_id,
+    au.organization_id,
+    DATE_TRUNC('month', au.created_at + (u.token_limit_reset_day - 1) * INTERVAL '1 day') AS usage_month,
+    COUNT(*) AS total_requests,
+    SUM(au.tokens_used) AS total_tokens,
+    SUM(au.input_tokens) AS input_tokens,
+    SUM(au.output_tokens) AS output_tokens,
+    SUM(au.estimated_cost_usd) AS total_cost_usd
+FROM ai_usage_logs au
+JOIN users u ON au.user_id = u.id
+WHERE au.created_at >= DATE_TRUNC('month', NOW() + (u.token_limit_reset_day - 1) * INTERVAL '1 day') - (u.token_limit_reset_day - 1) * INTERVAL '1 day'
+GROUP BY au.user_id, au.organization_id, usage_month;
+
+COMMENT ON VIEW ai_usage_monthly IS 'Current month AI usage per user with custom reset day';
+
+-- Function to check if user has exceeded token limit
+CREATE OR REPLACE FUNCTION check_token_limit(
+    p_user_id UUID,
+    p_additional_tokens INTEGER DEFAULT 0
+)
+RETURNS TABLE (
+    has_available_tokens BOOLEAN,
+    monthly_limit INTEGER,
+    used_tokens BIGINT,
+    remaining_tokens BIGINT,
+    reset_date TIMESTAMPTZ
+) AS $$
+DECLARE
+    v_limit INTEGER;
+    v_reset_day INTEGER;
+    v_used BIGINT;
+    v_month_start TIMESTAMPTZ;
+    v_next_reset TIMESTAMPTZ;
+BEGIN
+    -- Get user's limit settings
+    SELECT monthly_token_limit, token_limit_reset_day
+    INTO v_limit, v_reset_day
+    FROM users
+    WHERE id = p_user_id;
+    
+    -- Default values if not set
+    v_limit := COALESCE(v_limit, 100000);
+    v_reset_day := COALESCE(v_reset_day, 1);
+    
+    -- If limit is 0, unlimited
+    IF v_limit = 0 THEN
+        RETURN QUERY SELECT 
+            TRUE,
+            v_limit,
+            0::BIGINT,
+            0::BIGINT,
+            NOW() + INTERVAL '1 month';
+        RETURN;
+    END IF;
+    
+    -- Calculate current month start based on reset day
+    v_month_start := DATE_TRUNC('month', NOW() + (v_reset_day - 1) * INTERVAL '1 day') 
+                     - (v_reset_day - 1) * INTERVAL '1 day';
+    
+    -- Calculate next reset date
+    v_next_reset := v_month_start + INTERVAL '1 month';
+    
+    -- Get current usage
+    SELECT COALESCE(SUM(tokens_used), 0)
+    INTO v_used
+    FROM ai_usage_logs
+    WHERE user_id = p_user_id
+      AND created_at >= v_month_start;
+    
+    -- Return results
+    RETURN QUERY SELECT
+        (v_used + p_additional_tokens) < v_limit,
+        v_limit,
+        v_used,
+        (v_limit - v_used)::BIGINT,
+        v_next_reset;
+END;
+$$ LANGUAGE plpgsql;
+
+COMMENT ON FUNCTION check_token_limit IS 'Check if user has available tokens for current period';
+
+-- Function to get usage statistics for a user
+CREATE OR REPLACE FUNCTION get_user_usage_stats(
+    p_user_id UUID,
+    p_months INTEGER DEFAULT 3
+)
+RETURNS TABLE (
+    month_start DATE,
+    total_requests BIGINT,
+    total_tokens BIGINT,
+    input_tokens BIGINT,
+    output_tokens BIGINT,
+    total_cost_usd NUMERIC,
+    monthly_limit INTEGER,
+    percentage_used NUMERIC
+) AS $$
+BEGIN
+    RETURN QUERY
+    SELECT 
+        DATE_TRUNC('month', au.created_at)::DATE AS month_start,
+        COUNT(*) AS total_requests,
+        SUM(au.tokens_used) AS total_tokens,
+        SUM(au.input_tokens) AS input_tokens,
+        SUM(au.output_tokens) AS output_tokens,
+        SUM(au.estimated_cost_usd) AS total_cost_usd,
+        u.monthly_token_limit,
+        CASE 
+            WHEN u.monthly_token_limit = 0 THEN 0
+            ELSE (SUM(au.tokens_used)::NUMERIC / u.monthly_token_limit * 100)::NUMERIC(5,2)
+        END AS percentage_used
+    FROM ai_usage_logs au
+    JOIN users u ON au.user_id = u.id
+    WHERE au.user_id = p_user_id
+      AND au.created_at >= NOW() - (p_months || ' months')::INTERVAL
+    GROUP BY DATE_TRUNC('month', au.created_at), u.monthly_token_limit
+    ORDER BY month_start DESC;
+END;
+$$ LANGUAGE plpgsql;
+
+COMMENT ON FUNCTION get_user_usage_stats IS 'Get AI usage statistics for user over last N months';
+
+-- Add index for faster monthly queries
+CREATE INDEX IF NOT EXISTS idx_ai_usage_user_month 
+ON ai_usage_logs(user_id, created_at);
@@ -1,12 +1,13 @@
 use axum::{
    Json,
-    extract::{Query, State},
+    extract::{Path, Query, State},
    http::StatusCode,
 };
 use common::{auth::Claims, middleware::Org};
 use serde::{Deserialize, Serialize};
 use sqlx::PgPool;
 use chrono::{DateTime, Utc};
+use uuid::Uuid;

 // ==================== Token Usage Tracking ====================

@@ -717,3 +718,116 @@ pub struct GlobalAiUsageResponse {
    pub top_users: Vec<TopUserUsage>,
    pub student_chat_usage: Vec<StudentChatUsage>,
 }
+
+// ==================== User Token Limits ====================
+
+#[derive(Debug, Deserialize)]
+pub struct SetTokenLimitPayload {
+    pub monthly_token_limit: i32,
+    pub token_limit_reset_day: Option<i32>,
+}
+
+/// PUT /admin/users/{user_id}/token-limit - Set monthly token limit for user
+pub async fn set_user_token_limit(
+    _org: Org,
+    _claims: Claims,
+    State(pool): State<PgPool>,
+    Path(user_id): Path<Uuid>,
+    Json(payload): Json<SetTokenLimitPayload>,
+) -> Result<StatusCode, (StatusCode, String)> {
+    // Validate reset day (1-28 to avoid month-end issues)
+    let reset_day = payload.token_limit_reset_day.unwrap_or(1).clamp(1, 28);
+    
+    sqlx::query(
+        r#"
+        UPDATE users
+        SET monthly_token_limit = $1,
+            token_limit_reset_day = $2,
+            updated_at = NOW()
+        WHERE id = $3
+        "#
+    )
+    .bind(payload.monthly_token_limit)
+    .bind(reset_day)
+    .bind(user_id)
+    .execute(&pool)
+    .await
+    .map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, format!("Failed to set limit: {}", e)))?;
+    
+    Ok(StatusCode::OK)
+}
+
+/// GET /admin/users/{user_id}/token-usage - Get token usage for specific user
+pub async fn get_user_token_usage(
+    _org: Org,
+    _claims: Claims,
+    State(pool): State<PgPool>,
+    Path(user_id): Path<Uuid>,
+) -> Result<Json<UserTokenUsageDetail>, (StatusCode, String)> {
+    // Get current month usage
+    let usage: UserTokenUsageDetail = sqlx::query_as(
+        r#"
+        SELECT
+            u.id as user_id,
+            u.email,
+            u.full_name,
+            u.monthly_token_limit,
+            u.token_limit_reset_day,
+            COALESCE(SUM(au.tokens_used), 0) as used_tokens,
+            COUNT(au.id) as total_requests,
+            COALESCE(SUM(au.estimated_cost_usd), 0) as total_cost_usd,
+            MAX(au.created_at) as last_used
+        FROM users u
+        LEFT JOIN ai_usage_logs au ON u.id = au.user_id
+            AND au.created_at >= DATE_TRUNC('month', NOW())
+        WHERE u.id = $1
+        GROUP BY u.id, u.email, u.full_name, u.monthly_token_limit, u.token_limit_reset_day
+        "#
+    )
+    .bind(user_id)
+    .fetch_one(&pool)
+    .await
+    .map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, format!("Failed to fetch usage: {}", e)))?;
+    
+    Ok(Json(usage))
+}
+
+/// GET /admin/users/{user_id}/token-limit/check - Check if user has available tokens
+pub async fn check_user_token_limit(
+    _org: Org,
+    _claims: Claims,
+    State(pool): State<PgPool>,
+    Path(user_id): Path<Uuid>,
+) -> Result<Json<TokenLimitCheckResponse>, (StatusCode, String)> {
+    let result: TokenLimitCheckResponse = sqlx::query_as(
+        "SELECT * FROM check_token_limit($1, 0)"
+    )
+    .bind(user_id)
+    .fetch_one(&pool)
+    .await
+    .map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, format!("Failed to check limit: {}", e)))?;
+    
+    Ok(Json(result))
+}
+
+#[derive(Debug, Serialize, sqlx::FromRow)]
+pub struct UserTokenUsageDetail {
+    pub user_id: Uuid,
+    pub email: String,
+    pub full_name: String,
+    pub monthly_token_limit: i32,
+    pub token_limit_reset_day: i32,
+    pub used_tokens: i64,
+    pub total_requests: i64,
+    pub total_cost_usd: f64,
+    pub last_used: Option<DateTime<Utc>>,
+}
+
+#[derive(Debug, Serialize, sqlx::FromRow)]
+pub struct TokenLimitCheckResponse {
+    pub has_available_tokens: bool,
+    pub monthly_limit: i32,
+    pub used_tokens: i64,
+    pub remaining_tokens: i64,
+    pub reset_date: DateTime<Utc>,
+}
@@ -401,6 +401,18 @@ async fn main() {
            "/admin/ai-usage/global",
            get(handlers_admin::get_ai_usage_global),
        )
+        .route(
+            "/admin/users/{user_id}/token-limit",
+            put(handlers_admin::set_user_token_limit),
+        )
+        .route(
+            "/admin/users/{user_id}/token-usage",
+            get(handlers_admin::get_user_token_usage),
+        )
+        .route(
+            "/admin/users/{user_id}/token-limit/check",
+            get(handlers_admin::check_user_token_limit),
+        )
        .route_layer(middleware::from_fn(
            common::middleware::org_extractor_middleware,
        ));