From 6ef2860b268d031f8efcf7ffebc93a5b86fe869f Mon Sep 17 00:00:00 2001 From: Nurfog Date: Mon, 23 Mar 2026 17:18:04 -0300 Subject: [PATCH] feat: Enforce Automatic Token Limits in AI Handlers CMS Service: - generate_quiz: Check 2000 tokens limit - generate_course: Check 5000 tokens limit - generate_hotspots: Check 2000 tokens limit - generate_role_play: Check 2500 tokens limit - summarize_lesson: Check 1500 tokens limit LMS Service: - chat_with_tutor: Check 1000 tokens limit - evaluate_audio_response: Check 1500 tokens limit Behavior: - Returns HTTP 429 Too Many Requests when limit exceeded - Error message indicates monthly limit exceeded - Users must contact admin for limit increase - Works with automatic alert system (Phase 3) Total Functions Protected: 7 Co-authored-by: Qwen-Coder --- services/cms-service/src/handlers.rs | 33 ++++++++++++++++++++++++---- services/lms-service/src/handlers.rs | 12 +++++++++- 2 files changed, 40 insertions(+), 5 deletions(-) diff --git a/services/cms-service/src/handlers.rs b/services/cms-service/src/handlers.rs index f6f7591..9ea78a8 100644 --- a/services/cms-service/src/handlers.rs +++ b/services/cms-service/src/handlers.rs @@ -1074,6 +1074,11 @@ pub async fn summarize_lesson( Path(id): Path, ) -> Result, StatusCode> { tracing::info!("Received summarization request for lesson: {}", id); + + // Check token limit before proceeding (estimate 1500 tokens for summary) + if let Err(_) = common::token_limits::check_ai_token_limit(&pool, claims.sub, 1500).await { + return Err(StatusCode::TOO_MANY_REQUESTS); + } // 1. Fetch lesson let lesson = sqlx::query_as::<_, Lesson>("SELECT * FROM lessons WHERE id = $1 AND organization_id = $2") @@ -1213,6 +1218,11 @@ pub async fn generate_quiz( Json(quiz_req): Json, ) -> Result, StatusCode> { tracing::info!("Received quiz generation request for lesson: {}", id); + + // Check token limit before proceeding (estimate 2000 tokens for quiz) + if let Err(_) = common::token_limits::check_ai_token_limit(&pool, claims.sub, 2000).await { + return Err(StatusCode::TOO_MANY_REQUESTS); + } // 1. Fetch lesson let lesson = sqlx::query_as::<_, Lesson>("SELECT * FROM lessons WHERE id = $1 AND organization_id = $2") @@ -2086,11 +2096,16 @@ pub struct GenerateHotspotsPayload { pub async fn generate_hotspots( Org(org_ctx): Org, - _claims: common::auth::Claims, + claims: common::auth::Claims, State(pool): State, Path(lesson_id): Path, Json(payload): Json, ) -> Result, StatusCode> { + // Check token limit before proceeding (estimate 2000 tokens for hotspots) + if let Err(_) = common::token_limits::check_ai_token_limit(&pool, claims.sub, 2000).await { + return Err(StatusCode::TOO_MANY_REQUESTS); + } + // 1. Resolve image path // imageUrl in frontend is like "/assets/filename.ext" // We need to map it to "uploads/filename.ext" @@ -2267,7 +2282,7 @@ pub async fn generate_hotspots( let total_tokens = input_tokens + output_tokens; let _ = sqlx::query("SELECT log_ai_usage($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11)") - .bind(_claims.sub) + .bind(claims.sub) .bind(org_ctx.id) .bind(total_tokens) .bind(input_tokens) @@ -2294,11 +2309,16 @@ pub struct GenerateRolePlayPayload { pub async fn generate_role_play( Org(org_ctx): Org, - _claims: common::auth::Claims, + claims: common::auth::Claims, State(pool): State, Path(lesson_id): Path, Json(payload): Json, ) -> Result, StatusCode> { + // Check token limit before proceeding (estimate 2500 tokens for role-play) + if let Err(_) = common::token_limits::check_ai_token_limit(&pool, claims.sub, 2500).await { + return Err(StatusCode::TOO_MANY_REQUESTS); + } + // 1. Fetch lesson context let lesson = sqlx::query_as::<_, Lesson>("SELECT * FROM lessons WHERE id = $1 AND organization_id = $2") .bind(lesson_id) @@ -2396,7 +2416,7 @@ pub async fn generate_role_play( let total_tokens = input_tokens + output_tokens; let _ = sqlx::query("SELECT log_ai_usage($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11)") - .bind(_claims.sub) + .bind(claims.sub) .bind(org_ctx.id) .bind(total_tokens) .bind(input_tokens) @@ -4032,6 +4052,11 @@ pub async fn generate_course( "Starting AI course generation for prompt: {}", payload.prompt ); + + // Check token limit before proceeding (estimate 5000 tokens for course generation) + if let Err(_) = common::token_limits::check_ai_token_limit(&pool, claims.sub, 5000).await { + return Err(StatusCode::TOO_MANY_REQUESTS); + } // 1. Determine target org let target_org_id = payload.target_organization_id.unwrap_or(org_ctx.id); diff --git a/services/lms-service/src/handlers.rs b/services/lms-service/src/handlers.rs index 251f381..cd1a035 100644 --- a/services/lms-service/src/handlers.rs +++ b/services/lms-service/src/handlers.rs @@ -2102,9 +2102,14 @@ pub async fn get_recommendations( pub async fn evaluate_audio_response( Org(_org_ctx): Org, - _claims: Claims, + claims: Claims, Json(payload): Json, ) -> Result, (StatusCode, String)> { + // Check token limit before proceeding (estimate 1500 tokens for audio evaluation) + if let Err(_) = common::token_limits::check_ai_token_limit(&pool, claims.sub, 1500).await { + return Err((StatusCode::TOO_MANY_REQUESTS, "Token limit exceeded".to_string())); + } + let client = reqwest::Client::new(); let provider = env::var("AI_PROVIDER").unwrap_or_else(|_| "openai".to_string()); @@ -2496,6 +2501,11 @@ pub async fn chat_with_tutor( Path(lesson_id): Path, Json(payload): Json, ) -> Result, (StatusCode, String)> { + // Check token limit before proceeding (estimate 1000 tokens for chat) + if let Err(_) = common::token_limits::check_ai_token_limit(&pool, claims.sub, 1000).await { + return Err((StatusCode::TOO_MANY_REQUESTS, "Monthly AI token limit exceeded. Please contact your administrator.".to_string())); + } + // 1. Fetch lesson context with access check (matches get_lesson_content) let is_preview = claims.token_type.as_deref() == Some("preview");