SurfSense/_bmad-output/implementation-artifacts/5-4-usage-tracking-rate-limit-enforcement.md

Story 5.4: Hệ thống Khóa Tác vụ dựa trên Hạn Mức (Usage Tracking & Rate Limit Enforcement)

Status: done

Story

As a Kỹ sư Hệ thống, I want những người dùng hết quota (vượt quá file upload hoặc số lượng tin nhắn) bị từ chối dịch vụ cho đến khi nâng cấp, so that mô hình kinh doanh không bị lỗ do chi phí LLM và Storage, áp dụng theo FR13.

Acceptance Criteria

1. Endpoint Upload (Document) kiểm tra pages_used vs pages_limit (dựa trên subscription tier) → từ chối nếu vượt.
2. Endpoint Chat (RAG AI) kiểm tra tokens_used_this_month vs monthly_token_limit → từ chối nếu vượt (đã cover ở Story 3.5 Task 4).
3. API trả lỗi 402 Payment Required khi quota exceeded.
4. UI hiển thị Toast Error hoặc Modal up-sell "Upgrade to Pro".
5. Dashboard hiển thị quota indicator (pages used / limit, tokens used / limit).

As-Is (Code hiện tại)

Component	Hiện trạng	File
Page Quota Service	Đã implement đầy đủ — check_page_limit(), update_page_usage(), get_page_usage(), estimate_pages_*()	surfsense_backend/app/services/page_limit_service.py
Enforcement trong Celery	Đã có — check quota trước khi process document	document_tasks.py
Enforcement trong Connectors	Đã có — Google Drive, OneDrive, Dropbox indexers check remaining_quota	connector_indexers/*.py
Enforcement tại HTTP API	Chưa có — document upload route không check quota trước khi enqueue task	routes/
Token Quota Service	Chưa có — sẽ tạo ở Story 3.5
Frontend Quota UI	Chưa có
Frontend Error Handling	Chưa có — không bắt 402 từ upload/chat

Gap chính:

1. Thiếu quota pre-check tại HTTP layer (upload route) — user phải đợi Celery task fail mới biết bị reject
2. Thiếu frontend feedback khi quota exceeded
3. Thiếu quota indicator trong dashboard

Tasks / Subtasks

• Task 1: Thêm Page Quota Pre-check vào Document Upload Route (Backend)

  • Subtask 1.1: Tìm document upload route — surfsense_backend/app/routes/documents_routes.py
  • Subtask 1.2: Inject PageLimitService, gọi estimate_pages_from_metadata(filename, file_size) rồi check_page_limit(user_id, estimated_pages).
  • Subtask 1.3: PageLimitExceededError → HTTPException(402) với message mô tả quota.
  • Subtask 1.4: Giữ nguyên enforcement trong Celery tasks (double-check layer).
  • Subtask 1.5: Thêm pre-check cho create_documents endpoint (extension/YouTube connector).

• Task 2: Plan-based Limits (Backend)

  • Subtask 2.1: PLAN_LIMITS config đã có trong config/__init__.py:314.
  • Subtask 2.2: Webhook Story 5.3 + admin approval đã update limits khi subscription change.

• Task 3: Frontend — Bắt lỗi Quota Exceeded

  • Subtask 3.1: Document upload — DocumentUploadTab.tsx bắt HTTP 402 (check error.status) cho cả file và folder upload.
  • Subtask 3.2: Toast error với "Upgrade" action button → navigate to /pricing.
  • Subtask 3.3: Chat SSE — đã có QuotaExceededError pattern từ Story 3.5 (page.tsx).

• Task 4: Frontend — Quota Indicator (Dashboard)

  • Subtask 4.1: Extend UserRead schema + Zod type với monthly_token_limit, tokens_used_this_month, plan_id, subscription_status.
  • Subtask 4.2: Extend PageUsageDisplay — hiển thị 2 progress bars (Pages, Tokens) trong sidebar.
  • Subtask 4.3: Warning state (amber) khi usage > 80%, critical state (red) khi > 95%.

• Task 5: Anti-spam Rate Limiting (Optional — skipped per spec)

Dev Notes

PageLimitService — đã sẵn sàng, không cần viết lại

page_limit_service.py đã có đầy đủ:

• check_page_limit(user_id, estimated_pages) → raises PageLimitExceededError
• update_page_usage(user_id, pages_to_add)
• estimate_pages_before_processing(file_path)
• estimate_pages_from_metadata(filename, file_size) (pure function, dùng cho HTTP layer)

Enforcement Layers

HTTP Upload Route → PageLimitService.check_page_limit() [FAST — pre-check]
       ↓ (nếu pass)
Celery Task → PageLimitService.check_page_limit() [ACCURATE — sau khi parse xong]
       ↓ (nếu pass)
PageLimitService.update_page_usage() [COMMIT — tăng pages_used]

References

• surfsense_backend/app/services/page_limit_service.py — page quota (đọc, ít sửa)
• surfsense_backend/app/tasks/celery_tasks/document_tasks.py — enforcement pattern hiện tại
• surfsense_backend/app/routes/documents_routes.py — document upload route
• surfsense_web/components/sources/DocumentUploadTab.tsx — frontend upload component

Dev Agent Record

Implementation (2026-04-15)

Backend:

• documents_routes.py: Added page quota pre-check to both create_documents and create_documents_file_upload endpoints. Uses PageLimitService.estimate_pages_from_metadata() for fast estimation before any I/O, raises HTTP 402 on exceeded quota.
• schemas/users.py: Extended UserRead with monthly_token_limit, tokens_used_this_month, plan_id, subscription_status — exposed via /api/v1/users/me.

Frontend:

• DocumentUploadTab.tsx: 402 handling in both file upload (onError) and folder upload (catch) paths — shows toast with "Upgrade" action linking to /pricing.
• user.types.ts: Zod schema extended with token quota fields.
• PageUsageDisplay.tsx: Extended with token usage bar, color-coded warning (amber >80%, red >95%).
• layout.types.ts: PageUsage interface extended with tokensUsed/tokensLimit.
• LayoutDataProvider.tsx: Passes token data to sidebar.

File List

• surfsense_backend/app/routes/documents_routes.py — page quota pre-check in upload endpoints
• surfsense_backend/app/schemas/users.py — extended UserRead
• surfsense_web/components/sources/DocumentUploadTab.tsx — 402 error handling
• surfsense_web/contracts/types/user.types.ts — quota fields in Zod schema
• surfsense_web/components/layout/ui/sidebar/PageUsageDisplay.tsx — token usage bar + color states
• surfsense_web/components/layout/types/layout.types.ts — extended PageUsage interface
• surfsense_web/components/layout/providers/LayoutDataProvider.tsx — pass token data
• surfsense_web/components/layout/ui/sidebar/Sidebar.tsx — pass token props to PageUsageDisplay

Review Findings (2026-04-15)

• [Review][Patch] PAST_DUE users retain pro-tier limits and can continue consuming resources [page_limit_service.py, token_quota_service.py] — Fixed: Both services now check subscription_status and apply free-tier limits when PAST_DUE