Skip to content

Commit 54042e2

Browse files
committed
feat(backend): Add user tier/limits system (#96)
- Create UserLimitsService for centralized tier management - Define tier limits (free/pro/enterprise) - Add methods: check_repo_count, check_repo_size, get_usage_summary - Add Supabase migration for user_profiles table - Initialize service in dependencies.py Tiers: - Free: 3 repos, 500 files, 2000 functions - Pro: 20 repos, 5000 files, 20000 functions - Enterprise: Unlimited Used by #93, #94, #95 for enforcing limits.
1 parent bf40c59 commit 54042e2

4 files changed

Lines changed: 477 additions & 0 deletions

File tree

backend/dependencies.py

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,7 @@
1717
from services.rate_limiter import RateLimiter, APIKeyManager
1818
from services.supabase_service import get_supabase_service
1919
from services.input_validator import InputValidator, CostController
20+
from services.user_limits import init_user_limits_service, get_user_limits_service
2021

2122
# Service instances (singleton pattern)
2223
indexer = OptimizedCodeIndexer()
@@ -31,6 +32,12 @@
3132
api_key_manager = APIKeyManager(get_supabase_service().client)
3233
cost_controller = CostController(get_supabase_service().client)
3334

35+
# User tier and limits management
36+
user_limits = init_user_limits_service(
37+
supabase_client=get_supabase_service().client,
38+
redis_client=cache.redis if cache.redis else None
39+
)
40+
3441

3542
def get_repo_or_404(repo_id: str, user_id: str) -> dict:
3643
"""

backend/services/user_limits.py

Lines changed: 338 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,338 @@
1+
"""
2+
User Tier & Limits Service
3+
Centralized system for managing user tiers and resource limits.
4+
5+
Tiers:
6+
- free: Default tier for new users
7+
- pro: Paid tier with higher limits
8+
- enterprise: Custom limits for large organizations
9+
10+
Used by:
11+
- #93: Playground rate limiting
12+
- #94: Repo size limits
13+
- #95: Repo count limits
14+
"""
15+
from dataclasses import dataclass
16+
from typing import Optional, Dict, Any
17+
from enum import Enum
18+
19+
from services.observability import logger, metrics
20+
21+
22+
class UserTier(str, Enum):
23+
"""User subscription tiers"""
24+
FREE = "free"
25+
PRO = "pro"
26+
ENTERPRISE = "enterprise"
27+
28+
29+
@dataclass(frozen=True)
30+
class TierLimits:
31+
"""Resource limits for a tier"""
32+
# Repository limits
33+
max_repos: Optional[int] # None = unlimited
34+
max_files_per_repo: int
35+
max_functions_per_repo: int
36+
37+
# Playground limits (anti-abuse, not business gate)
38+
playground_searches_per_day: Optional[int] # None = unlimited
39+
40+
# Future limits (placeholders)
41+
max_team_members: Optional[int] = None
42+
priority_indexing: bool = False
43+
mcp_access: bool = True
44+
45+
46+
# Tier definitions
47+
TIER_LIMITS: Dict[UserTier, TierLimits] = {
48+
UserTier.FREE: TierLimits(
49+
max_repos=3,
50+
max_files_per_repo=500,
51+
max_functions_per_repo=2000,
52+
playground_searches_per_day=50, # Generous, anti-abuse only
53+
max_team_members=1,
54+
priority_indexing=False,
55+
mcp_access=True,
56+
),
57+
UserTier.PRO: TierLimits(
58+
max_repos=20,
59+
max_files_per_repo=5000,
60+
max_functions_per_repo=20000,
61+
playground_searches_per_day=None, # Unlimited
62+
max_team_members=10,
63+
priority_indexing=True,
64+
mcp_access=True,
65+
),
66+
UserTier.ENTERPRISE: TierLimits(
67+
max_repos=None, # Unlimited
68+
max_files_per_repo=50000,
69+
max_functions_per_repo=200000,
70+
playground_searches_per_day=None,
71+
max_team_members=None,
72+
priority_indexing=True,
73+
mcp_access=True,
74+
),
75+
}
76+
77+
78+
@dataclass
79+
class LimitCheckResult:
80+
"""Result of a limit check"""
81+
allowed: bool
82+
current: int
83+
limit: Optional[int]
84+
message: str
85+
86+
@property
87+
def limit_display(self) -> str:
88+
"""Display limit as string (handles unlimited)"""
89+
return str(self.limit) if self.limit is not None else "∞"
90+
91+
def to_dict(self) -> Dict[str, Any]:
92+
return {
93+
"allowed": self.allowed,
94+
"current": self.current,
95+
"limit": self.limit,
96+
"limit_display": self.limit_display,
97+
"message": self.message,
98+
}
99+
100+
101+
class UserLimitsService:
102+
"""
103+
Service for checking and enforcing user tier limits.
104+
105+
Usage:
106+
limits = UserLimitsService(supabase_client, redis_client)
107+
108+
# Check if user can add another repo
109+
result = await limits.check_repo_count(user_id)
110+
if not result.allowed:
111+
raise HTTPException(403, result.message)
112+
113+
# Check if repo size is within limits
114+
result = await limits.check_repo_size(user_id, file_count, function_count)
115+
if not result.allowed:
116+
raise HTTPException(400, result.message)
117+
"""
118+
119+
def __init__(self, supabase_client, redis_client=None):
120+
self.supabase = supabase_client
121+
self.redis = redis_client
122+
self._tier_cache_ttl = 300 # Cache tier for 5 minutes
123+
124+
# ===== TIER MANAGEMENT =====
125+
126+
async def get_user_tier(self, user_id: str) -> UserTier:
127+
"""
128+
Get user's current tier.
129+
130+
Checks Redis cache first, then Supabase.
131+
Defaults to FREE if not found.
132+
"""
133+
# Try cache first
134+
if self.redis:
135+
cache_key = f"user:tier:{user_id}"
136+
cached = self.redis.get(cache_key)
137+
if cached:
138+
try:
139+
return UserTier(cached.decode() if isinstance(cached, bytes) else cached)
140+
except ValueError:
141+
pass
142+
143+
# Query Supabase
144+
tier = await self._get_tier_from_db(user_id)
145+
146+
# Cache the result
147+
if self.redis:
148+
cache_key = f"user:tier:{user_id}"
149+
self.redis.setex(cache_key, self._tier_cache_ttl, tier.value)
150+
151+
return tier
152+
153+
async def _get_tier_from_db(self, user_id: str) -> UserTier:
154+
"""Get tier from Supabase user_profiles table"""
155+
try:
156+
result = self.supabase.table("user_profiles").select("tier").eq("user_id", user_id).execute()
157+
158+
if result.data and result.data[0].get("tier"):
159+
tier_value = result.data[0]["tier"]
160+
return UserTier(tier_value)
161+
except Exception as e:
162+
logger.warning("Failed to get user tier from DB", user_id=user_id, error=str(e))
163+
164+
return UserTier.FREE
165+
166+
def get_limits(self, tier: UserTier) -> TierLimits:
167+
"""Get limits for a tier"""
168+
return TIER_LIMITS.get(tier, TIER_LIMITS[UserTier.FREE])
169+
170+
async def get_user_limits(self, user_id: str) -> TierLimits:
171+
"""Get limits for a specific user"""
172+
tier = await self.get_user_tier(user_id)
173+
return self.get_limits(tier)
174+
175+
# ===== REPO COUNT LIMITS (#95) =====
176+
177+
async def get_user_repo_count(self, user_id: str) -> int:
178+
"""Get current repo count for user"""
179+
try:
180+
result = self.supabase.table("repositories").select("id", count="exact").eq("user_id", user_id).execute()
181+
return result.count or 0
182+
except Exception as e:
183+
logger.error("Failed to get repo count", user_id=user_id, error=str(e))
184+
return 0
185+
186+
async def check_repo_count(self, user_id: str) -> LimitCheckResult:
187+
"""
188+
Check if user can add another repository.
189+
190+
Returns:
191+
LimitCheckResult with allowed=True if under limit
192+
"""
193+
tier = await self.get_user_tier(user_id)
194+
limits = self.get_limits(tier)
195+
current_count = await self.get_user_repo_count(user_id)
196+
197+
# Unlimited repos
198+
if limits.max_repos is None:
199+
return LimitCheckResult(
200+
allowed=True,
201+
current=current_count,
202+
limit=None,
203+
message=f"OK ({current_count}/∞ repos)"
204+
)
205+
206+
# Check limit
207+
if current_count >= limits.max_repos:
208+
metrics.increment("user_limit_exceeded", tags={"limit": "repo_count", "tier": tier.value})
209+
logger.info("Repo count limit reached", user_id=user_id, current=current_count, limit=limits.max_repos)
210+
return LimitCheckResult(
211+
allowed=False,
212+
current=current_count,
213+
limit=limits.max_repos,
214+
message=f"Repository limit reached ({current_count}/{limits.max_repos}). Upgrade for more repos."
215+
)
216+
217+
return LimitCheckResult(
218+
allowed=True,
219+
current=current_count,
220+
limit=limits.max_repos,
221+
message=f"OK ({current_count}/{limits.max_repos} repos)"
222+
)
223+
224+
# ===== REPO SIZE LIMITS (#94) =====
225+
226+
async def check_repo_size(
227+
self,
228+
user_id: str,
229+
file_count: int,
230+
function_count: int
231+
) -> LimitCheckResult:
232+
"""
233+
Check if repo size is within user's tier limits.
234+
235+
Args:
236+
user_id: The user attempting to index
237+
file_count: Number of code files in repo
238+
function_count: Number of functions/classes detected
239+
240+
Returns:
241+
LimitCheckResult with allowed=True if within limits
242+
"""
243+
tier = await self.get_user_tier(user_id)
244+
limits = self.get_limits(tier)
245+
246+
# Check file count
247+
if file_count > limits.max_files_per_repo:
248+
metrics.increment("user_limit_exceeded", tags={"limit": "file_count", "tier": tier.value})
249+
logger.info(
250+
"Repo file count exceeds limit",
251+
user_id=user_id,
252+
file_count=file_count,
253+
limit=limits.max_files_per_repo
254+
)
255+
return LimitCheckResult(
256+
allowed=False,
257+
current=file_count,
258+
limit=limits.max_files_per_repo,
259+
message=f"Repository too large ({file_count:,} files). {tier.value.title()} tier allows up to {limits.max_files_per_repo:,} files."
260+
)
261+
262+
# Check function count
263+
if function_count > limits.max_functions_per_repo:
264+
metrics.increment("user_limit_exceeded", tags={"limit": "function_count", "tier": tier.value})
265+
logger.info(
266+
"Repo function count exceeds limit",
267+
user_id=user_id,
268+
function_count=function_count,
269+
limit=limits.max_functions_per_repo
270+
)
271+
return LimitCheckResult(
272+
allowed=False,
273+
current=function_count,
274+
limit=limits.max_functions_per_repo,
275+
message=f"Repository has too many functions ({function_count:,}). {tier.value.title()} tier allows up to {limits.max_functions_per_repo:,} functions."
276+
)
277+
278+
return LimitCheckResult(
279+
allowed=True,
280+
current=file_count,
281+
limit=limits.max_files_per_repo,
282+
message=f"OK ({file_count:,} files, {function_count:,} functions)"
283+
)
284+
285+
# ===== PLAYGROUND RATE LIMITS (#93) =====
286+
287+
def get_playground_limit(self, tier: UserTier = UserTier.FREE) -> Optional[int]:
288+
"""Get playground search limit for tier"""
289+
return self.get_limits(tier).playground_searches_per_day
290+
291+
# ===== USAGE SUMMARY =====
292+
293+
async def get_usage_summary(self, user_id: str) -> Dict[str, Any]:
294+
"""
295+
Get complete usage summary for user.
296+
Useful for dashboard display.
297+
"""
298+
tier = await self.get_user_tier(user_id)
299+
limits = self.get_limits(tier)
300+
repo_count = await self.get_user_repo_count(user_id)
301+
302+
return {
303+
"tier": tier.value,
304+
"repositories": {
305+
"current": repo_count,
306+
"limit": limits.max_repos,
307+
"display": f"{repo_count}/{limits.max_repos if limits.max_repos else '∞'}"
308+
},
309+
"limits": {
310+
"max_files_per_repo": limits.max_files_per_repo,
311+
"max_functions_per_repo": limits.max_functions_per_repo,
312+
"playground_searches_per_day": limits.playground_searches_per_day,
313+
},
314+
"features": {
315+
"priority_indexing": limits.priority_indexing,
316+
"mcp_access": limits.mcp_access,
317+
}
318+
}
319+
320+
321+
# Singleton instance (initialized in dependencies.py)
322+
_user_limits_service: Optional[UserLimitsService] = None
323+
324+
325+
def get_user_limits_service() -> UserLimitsService:
326+
"""Get or create UserLimitsService instance"""
327+
global _user_limits_service
328+
if _user_limits_service is None:
329+
raise RuntimeError("UserLimitsService not initialized. Call init_user_limits_service first.")
330+
return _user_limits_service
331+
332+
333+
def init_user_limits_service(supabase_client, redis_client=None) -> UserLimitsService:
334+
"""Initialize the UserLimitsService singleton"""
335+
global _user_limits_service
336+
_user_limits_service = UserLimitsService(supabase_client, redis_client)
337+
logger.info("UserLimitsService initialized")
338+
return _user_limits_service

0 commit comments

Comments
 (0)