Coverage for src/alprina_cli/api/middleware/usage_check.py: 0%

1""" 

2Usage Check Middleware 

3 

4Enforces usage limits for scans and API requests. 

5""" 

6 

7from fastapi import Request, HTTPException, Depends 

8from typing import Dict, Any, Optional 

9from loguru import logger 

10 

11from ..middleware.auth import get_current_user 

12from ..services.usage_service import usage_service 

13from ..services.neon_service import neon_service 

14 

15 

16async def check_usage_limits( 

17 request: Request, 

18 user: Dict[str, Any] = Depends(get_current_user) 

19) -> Dict[str, Any]: 

20 """ 

21 Check if user is within usage limits. 

22 

23 Args: 

24 request: FastAPI request object 

25 user: Current authenticated user 

26 

27 Returns: 

28 User object with usage info 

29 

30 Raises: 

31 HTTPException: If usage limits exceeded 

32 """ 

33 user_id = user["id"] 

34 tier = user["tier"] 

35 

36 # Check API rate limit 

37 within_limit, error = await usage_service.enforce_rate_limit( 

38 user_id, 

39 tier, 

40 neon_service 

41 ) 

42 

43 if not within_limit: 

44 logger.warning(f"Rate limit exceeded for user {user_id}") 

45 raise HTTPException( 

46 status_code=429, 

47 detail={ 

48 "error": "rate_limit_exceeded", 

49 "message": error, 

50 "tier": tier, 

51 "upgrade_url": "https://platform.alprina.ai/upgrade" 

52 } 

53 ) 

54 

55 # Add usage info to user object 

56 user["usage_checked"] = True 

57 return user 

58 

59 

60async def check_scan_permission( 

61 workflow_mode: str = "single", 

62 user: Dict[str, Any] = Depends(get_current_user) 

63) -> Dict[str, Any]: 

64 """ 

65 Check if user has permission to perform scan. 

66 

67 Args: 

68 workflow_mode: Workflow mode (single, parallel, sequential, coordinated) 

69 user: Current authenticated user 

70 

71 Returns: 

72 User object 

73 

74 Raises: 

75 HTTPException: If scan not permitted 

76 """ 

77 user_id = user["id"] 

78 tier = user["tier"] 

79 

80 # Check scan limit 

81 can_scan, error, usage_info = await usage_service.check_scan_limit( 

82 user_id, 

83 tier, 

84 neon_service 

85 ) 

86 

87 if not can_scan: 

88 logger.warning(f"Scan limit exceeded for user {user_id}") 

89 raise HTTPException( 

90 status_code=403, 

91 detail={ 

92 "error": "scan_limit_exceeded", 

93 "message": error, 

94 "usage": usage_info, 

95 "tier": tier, 

96 "upgrade_url": "https://platform.alprina.ai/upgrade" 

97 } 

98 ) 

99 

100 # Check workflow access 

101 has_access, error = await usage_service.check_workflow_access(tier, workflow_mode) 

102 

103 if not has_access: 

104 logger.warning( 

105 f"User {user_id} attempted {workflow_mode} workflow " 

106 f"without access (tier: {tier})" 

107 ) 

108 raise HTTPException( 

109 status_code=403, 

110 detail={ 

111 "error": "workflow_not_available", 

112 "message": error, 

113 "tier": tier, 

114 "required_tier": "pro", 

115 "upgrade_url": "https://platform.alprina.ai/upgrade" 

116 } 

117 ) 

118 

119 # Add usage info to user 

120 user["usage_info"] = usage_info 

121 return user 

122 

123 

124async def record_scan_usage( 

125 user_id: str, 

126 tier: str, 

127 workflow_mode: str, 

128 file_count: int, 

129 scan_data: Dict[str, Any] 

130): 

131 """ 

132 Record scan usage after successful scan. 

133 

134 Args: 

135 user_id: User ID 

136 tier: User tier 

137 workflow_mode: Workflow mode used 

138 file_count: Number of files scanned 

139 scan_data: Full scan data 

140 """ 

141 try: 

142 # Increment scan count 

143 await usage_service.increment_scan_count( 

144 user_id, 

145 tier, 

146 workflow_mode, 

147 file_count, 

148 neon_service 

149 ) 

150 

151 # Record in history 

152 await usage_service.record_scan( 

153 user_id, 

154 scan_data, 

155 neon_service 

156 ) 

157 

158 logger.info(f"Recorded scan usage for user {user_id}") 

159 

160 except Exception as e: 

161 logger.error(f"Failed to record scan usage: {e}") 

162 # Don't fail the scan if usage recording fails