From 9e57756528a7e8f7cd61ae2f435101bf8b43de12 Mon Sep 17 00:00:00 2001 From: Nurfog Date: Mon, 23 Mar 2026 16:10:27 -0300 Subject: [PATCH] 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 --- .../20260323000000_monthly_token_limits.sql | 138 ++++++++++++++++++ services/cms-service/src/handlers_admin.rs | 116 ++++++++++++++- services/cms-service/src/main.rs | 12 ++ 3 files changed, 265 insertions(+), 1 deletion(-) create mode 100644 services/cms-service/migrations/20260323000000_monthly_token_limits.sql diff --git a/services/cms-service/migrations/20260323000000_monthly_token_limits.sql b/services/cms-service/migrations/20260323000000_monthly_token_limits.sql new file mode 100644 index 0000000..72efb6b --- /dev/null +++ b/services/cms-service/migrations/20260323000000_monthly_token_limits.sql @@ -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); diff --git a/services/cms-service/src/handlers_admin.rs b/services/cms-service/src/handlers_admin.rs index 2b64bf9..80ee042 100644 --- a/services/cms-service/src/handlers_admin.rs +++ b/services/cms-service/src/handlers_admin.rs @@ -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, pub student_chat_usage: Vec, } + +// ==================== User Token Limits ==================== + +#[derive(Debug, Deserialize)] +pub struct SetTokenLimitPayload { + pub monthly_token_limit: i32, + pub token_limit_reset_day: Option, +} + +/// 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, + Path(user_id): Path, + Json(payload): Json, +) -> Result { + // 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, + Path(user_id): Path, +) -> Result, (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, + Path(user_id): Path, +) -> Result, (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>, +} + +#[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, +} diff --git a/services/cms-service/src/main.rs b/services/cms-service/src/main.rs index cc292f6..40922f7 100644 --- a/services/cms-service/src/main.rs +++ b/services/cms-service/src/main.rs @@ -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, ));