Coverage for apps / ai / services / quota.py: 91%

1"""AI quota checking and tracking service. 

2 

3Uses Django's database cache backend to enforce per-profile, per-feature 

4daily limits. Quotas only apply in passkey auth mode — home mode is unlimited. 

5""" 

6 

7from datetime import datetime, timedelta, UTC 

8 

9from django.conf import settings 

10from django.core.cache import cache 

11 

12from apps.core.models import AppSettings 

13 

14ALL_FEATURES = ( 

15 "remix", 

16 "remix_suggestions", 

17 "scale", 

18 "tips", 

19 "discover", 

20 "timer", 

21) 

22 

23FEATURE_LIMIT_FIELDS = { 

24 "remix": "daily_limit_remix", 

25 "remix_suggestions": "daily_limit_remix_suggestions", 

26 "scale": "daily_limit_scale", 

27 "tips": "daily_limit_tips", 

28 "discover": "daily_limit_discover", 

29 "timer": "daily_limit_timer", 

30} 

31 

32 

33def _cache_key(profile_id: int, feature: str) -> str: 

34 """Build the daily cache key for a profile/feature pair.""" 

35 today = datetime.now(UTC).date() 

36 return f"ai_quota:{profile_id}:{feature}:{today.isoformat()}" 

37 

38 

39def _seconds_until_midnight_utc() -> int: 

40 """Return seconds remaining until the next UTC midnight.""" 

41 now = datetime.now(UTC) 

42 tomorrow = datetime(now.year, now.month, now.day, tzinfo=UTC) + timedelta(days=1) 

43 return int((tomorrow - now).total_seconds()) 

44 

45 

46def _next_midnight_utc_iso() -> str: 

47 """Return next UTC midnight as an ISO 8601 string.""" 

48 now = datetime.now(UTC) 

49 tomorrow = datetime(now.year, now.month, now.day, tzinfo=UTC) + timedelta(days=1) 

50 return tomorrow.isoformat() 

51 

52 

53def check_quota(profile, feature: str) -> tuple[bool, dict]: 

54 """Check whether *profile* may use *feature* right now (read-only). 

55 

56 Returns (allowed, info_dict). Info is empty when allowed; 

57 contains remaining/limit/used/resets_at when denied. 

58 

59 NOTE: For atomic enforcement under concurrent load, use 

60 reserve_quota() instead of check_quota() + increment_quota(). 

61 """ 

62 if getattr(settings, "AUTH_MODE", "home") != "passkey": 

63 return (True, {}) 

64 

65 if profile.user and profile.user.is_staff: 

66 return (True, {}) 

67 

68 if profile.unlimited_ai: 

69 return (True, {}) 

70 

71 if feature not in FEATURE_LIMIT_FIELDS: 

72 raise ValueError(f"Unknown quota feature: {feature}") 

73 

74 app = AppSettings.get() 

75 limit_field = FEATURE_LIMIT_FIELDS[feature] 

76 limit = getattr(app, limit_field) 

77 

78 key = _cache_key(profile.pk, feature) 

79 used = cache.get(key, 0) 

80 

81 if used >= limit: 

82 return ( 

83 False, 

84 { 

85 "remaining": 0, 

86 "limit": limit, 

87 "used": used, 

88 "resets_at": _next_midnight_utc_iso(), 

89 }, 

90 ) 

91 

92 return (True, {}) 

93 

94 

95def reserve_quota(profile, feature: str) -> tuple[bool, dict]: 

96 """Atomically reserve a quota slot BEFORE executing the AI operation. 

97 

98 Uses cache.incr() which is atomic at the database level, preventing 

99 concurrent requests from bypassing the quota limit. 

100 

101 Returns (allowed, info_dict). If allowed, the counter is already 

102 incremented. On AI failure, call release_quota() to roll back. 

103 """ 

104 if getattr(settings, "AUTH_MODE", "home") != "passkey": 

105 return (True, {}) 

106 

107 if profile.user and profile.user.is_staff: 

108 return (True, {}) 

109 

110 if profile.unlimited_ai: 

111 return (True, {}) 

112 

113 if feature not in FEATURE_LIMIT_FIELDS: 

114 raise ValueError(f"Unknown quota feature: {feature}") 

115 

116 app = AppSettings.get() 

117 limit_field = FEATURE_LIMIT_FIELDS[feature] 

118 limit = getattr(app, limit_field) 

119 

120 key = _cache_key(profile.pk, feature) 

121 ttl = _seconds_until_midnight_utc() 

122 

123 # Atomic increment — if key doesn't exist, set to 1 

124 try: 

125 new_count = cache.incr(key) 

126 except ValueError: 

127 cache.set(key, 1, timeout=ttl) 

128 new_count = 1 

129 

130 if new_count > limit: 

131 # Over limit — roll back the increment we just did 

132 try: 

133 cache.decr(key) 

134 except ValueError: 

135 pass 

136 return ( 

137 False, 

138 { 

139 "remaining": 0, 

140 "limit": limit, 

141 "used": new_count - 1, 

142 "resets_at": _next_midnight_utc_iso(), 

143 }, 

144 ) 

145 

146 return (True, {}) 

147 

148 

149def release_quota(profile, feature: str) -> None: 

150 """Roll back a reserved quota slot on AI operation failure (FR-009).""" 

151 if getattr(settings, "AUTH_MODE", "home") != "passkey": 

152 return 

153 if profile.user and profile.user.is_staff: 

154 return 

155 if profile.unlimited_ai: 

156 return 

157 

158 key = _cache_key(profile.pk, feature) 

159 try: 

160 cache.decr(key) 

161 except ValueError: 

162 pass 

163 

164 

165def increment_quota(profile, feature: str) -> None: 

166 """Increment the daily counter after a successful AI operation. 

167 

168 DEPRECATED: Use reserve_quota() / release_quota() instead for 

169 atomic enforcement under concurrent load. 

170 """ 

171 key = _cache_key(profile.pk, feature) 

172 ttl = _seconds_until_midnight_utc() 

173 

174 try: 

175 cache.incr(key) 

176 except ValueError: 

177 cache.set(key, 1, timeout=ttl) 

178 

179 

180def get_usage(profile_id: int) -> dict: 

181 """Return {feature: count} for all features for today.""" 

182 return {feature: cache.get(_cache_key(profile_id, feature), 0) for feature in ALL_FEATURES}