feat: token pipeline and device fingerprinting

cost_tracker.py

@@ -52,8 +52,13 @@ def get_log_file_path():
 LOG_FILE = get_log_file_path()
 LOG_DIR = os.path.dirname(LOG_FILE)
 
+import contextvars
+
+# Stores the total tokens used in the current solve request
+current_request_tokens = contextvars.ContextVar('current_request_tokens', default=0)
+
 def log_api_usage(usage_metadata, source="unknown"):
-    """Log API usage to proper location"""
+    """Log API usage to proper location and add to current ContextVar"""
     if not usage_metadata:
         return
 
@@ -73,6 +78,12 @@ def log_api_usage(usage_metadata, source="unknown"):
         with open(LOG_FILE, "a", encoding="utf-8") as f:
             f.write(json.dumps(entry) + "\n")
             
+        # Update the ContextVar for the current request
+        try:
+            current_request_tokens.set(current_request_tokens.get() + entry["total_tokens"])
+        except Exception:
+            pass
+
         # Also print to stdout for immediate visibility in console logs
         print(f"💰 [USAGE] {source}: In={entry['input_tokens']}, Out={entry['output_tokens']}")
         

main.py

@@ -246,9 +246,10 @@ async def solve_stream_v2(
     uid = decoded_token.get('uid')
     final_student_name = student_name or user or "תלמיד"
     print(f"🚀 🟢 [V2] Received request from UID: {uid} ({final_student_name}). Grade: {grade}")
+    device_id = request.headers.get('Device-ID')
     
     # V2 Quota Check (Firestore)
-    is_allowed, msg, current_usage, limit = quota_manager_v2.check_limit(uid)
+    is_allowed, msg, current_usage, limit = quota_manager_v2.check_limit(uid, device_id=device_id)
     if not is_allowed:
         response_content = build_standard_response(
             final_answer=f"הגעת למכסה היומית ({limit} שאלות)",
@@ -289,6 +290,8 @@ async def solve_stream_v2(
         print("🚀 [TRACE-V2] Initiating streaming orchestrator.solve_problem with multiple images...")
 
         async def event_generator():
+            import cost_tracker
+            cost_tracker.current_request_tokens.set(0)
             try:
                 # orchestrator handles the rest using student_name for pedagogical stuff, but quota is already handled.
                 async for event in orchestrator.solve_problem(
@@ -313,6 +316,11 @@ async def solve_stream_v2(
                     "event": "error",
                     "data": json.dumps({"error": str(e)})
                 }
+            finally:
+                total_tokens = cost_tracker.current_request_tokens.get()
+                if total_tokens > 0:
+                    quota_manager_v2.increment_usage(uid, increment_questions=0, tokens_used=total_tokens)
+                    print(f"🪙 [V2-QUOTA] Deducted {total_tokens} tokens for UID: {uid}")
 
         return EventSourceResponse(event_generator())
 
@@ -380,7 +388,7 @@ async def upgrade_success(req: UpgradeRequest):
         user_ref.set({
             "tier": "parent_premium",
             "parent_email": req.parent_email,
-            "quota_limit": 999
+            "monthly_token_budget": 2800000
         }, merge=True)
         
         logger.info(f"🎉 UPGRADED USER {req.uid} to parent_premium")

quota_system_v2.py

@@ -22,8 +22,11 @@ class QuotaManagerV2:
 
     def get_today_key(self):
         return datetime.date.today().isoformat()
+        
+    def get_this_month_key(self):
+        return datetime.date.today().strftime("%Y-%m")
 
-    def check_limit(self, uid: str):
+    def check_limit(self, uid: str, device_id: str = None):
         """
         Returns:
             (allowed: bool, message: str, current_usage: int, limit: int)
@@ -54,16 +57,35 @@ class QuotaManagerV2:
             if data.get("role") == "admin" or data.get("status") == "admin":
                 return True, "Unlimited (Admin)", 0, -1
 
+            # --- Device ID Enforcement ---
+            if device_id:
+                tier = data.get("tier", "student_basic")
+                allowed_devices = data.get("allowed_devices", [])
+                
+                max_devices = 2 if tier == "parent_premium" else 1
+                
+                if device_id not in allowed_devices:
+                    if len(allowed_devices) >= max_devices:
+                        return False, f"Device Limit Exceeded ({max_devices})", 0, 0
+                    else:
+                        # Add device atomically
+                        doc_ref.update({"allowed_devices": firestore.ArrayUnion([device_id])})
+            # -----------------------------
+
             limit = data.get("quota_limit", self.default_limit)
+            monthly_token_budget = data.get("monthly_token_budget")
             
             # Special limits
             if limit < 0:
                 return True, "Unlimited", 0, -1
 
             today = self.get_today_key()
+            this_month = self.get_this_month_key()
+            
             last_usage_date = data.get("last_usage_date", "")
+            last_usage_month = data.get("last_usage_month", "")
             
-            # If it's a new day, usage is effectively 0
+            # Daily Question Quota
             if last_usage_date != today:
                 current_count = 0
             else:
@@ -71,6 +93,16 @@ class QuotaManagerV2:
 
             if current_count >= limit:
                 return False, f"Daily limit reached ({limit})", current_count, limit
+                
+            # Monthly Token Quota (if defined)
+            if monthly_token_budget is not None:
+                if last_usage_month != this_month:
+                    used_tokens = 0
+                else:
+                    used_tokens = data.get("used_tokens_this_month", 0)
+                    
+                if used_tokens >= monthly_token_budget:
+                    return False, f"Monthly token budget reached", current_count, limit
 
             return True, "Allowed", current_count, limit
 
@@ -79,7 +111,7 @@ class QuotaManagerV2:
             # Fail closed or open? Let's fail open temporarily so we don't block users if DB hiccups
             return True, "Error - Fallback Allow", 0, self.default_limit
 
-    def increment_usage(self, uid: str):
+    def increment_usage(self, uid: str, increment_questions: int = 1, tokens_used: int = 0):
         if not uid:
             return
 
@@ -88,6 +120,7 @@ class QuotaManagerV2:
             return
 
         today = self.get_today_key()
+        this_month = self.get_this_month_key()
         doc_ref = db.collection(self.collection_name).document(uid)
 
         try:
@@ -95,32 +128,43 @@ class QuotaManagerV2:
             def update_in_transaction(transaction, doc_ref):
                 snapshot = doc_ref.get(transaction=transaction)
                 if not snapshot.exists:
-                    # Create document if it doesn't exist (though migration should handle this)
                     transaction.set(doc_ref, {
                         "quota_limit": self.default_limit,
-                        "used_today": 1,
+                        "used_today": increment_questions,
+                        "used_tokens_this_month": tokens_used,
                         "last_usage_date": today,
+                        "last_usage_month": this_month,
                         "last_seen": firestore.SERVER_TIMESTAMP
                     })
                     return
 
                 data = snapshot.to_dict()
                 last_date = data.get("last_usage_date", "")
+                last_month = data.get("last_usage_month", "")
                 
+                # Daily logic
                 if last_date != today:
-                    new_usage = 1
+                    new_usage = increment_questions
+                else:
+                    new_usage = data.get("used_today", 0) + increment_questions
+                    
+                # Monthly Token logic
+                if last_month != this_month:
+                    new_tokens = tokens_used
                 else:
-                    new_usage = data.get("used_today", 0) + 1
+                    new_tokens = data.get("used_tokens_this_month", 0) + tokens_used
 
                 transaction.update(doc_ref, {
                     "used_today": new_usage,
+                    "used_tokens_this_month": new_tokens,
                     "last_usage_date": today,
+                    "last_usage_month": this_month,
                     "last_seen": firestore.SERVER_TIMESTAMP
                 })
 
             transaction = db.transaction()
             update_in_transaction(transaction, doc_ref)
-            logger.info(f"📊 [QUOTA_V2] Incremented usage for {uid}. Date: {today}")
+            logger.info(f"📊 [QUOTA_V2] Incremented usage for {uid}. Date: {today}, Tokens: {tokens_used}")
 
         except Exception as e:
             logger.error(f"❌ [QUOTA_V2] Failed to increment usage for {uid}: {e}")