Update app.py

app.py

@@ -225,91 +225,203 @@ class ModelCacheManager:
 # --- TOKEN MANAGER ---
 class TokenManager:
     def __init__(self):
-        self.user_tokens = {}  # {session_id: {"balance": float, "start_time": float, "purchases": {}}}
-        self.active_sessions = {}
+        self.user_tokens = {}  # {username: {"balance": float, "start_time": float, "purchases": {}}}
+        self.owner_username = "turtle170"  # Owner gets infinite tokens
     
-    def get_session_id(self) -> str:
-        """Generate or retrieve session ID from Gradio request"""
-        import hashlib
-        import time
-        # Simple session ID based on timestamp (in production, use gr.Request)
-        return hashlib.md5(str(time.time()).encode()).hexdigest()[:8]
+    def is_owner(self, username: str) -> bool:
+        """Check if user is the owner"""
+        if not username:
+            return False
+        return username.lower() == self.owner_username.lower()
     
-    def initialize_user(self, session_id: str):
-        """Initialize new user with monthly credits"""
-        if session_id not in self.user_tokens:
-            self.user_tokens[session_id] = {
-                "balance": MONTHLY_TOKEN_CREDITS,
-                "start_time": time.time(),
-                "purchases": {"batch_multiplier": 1, "token_limit": 2048},
-                "total_spent": 0.0
-            }
-            logger.info(f"[TOKEN] New user {session_id}: {MONTHLY_TOKEN_CREDITS} tokens")
+    def initialize_user(self, username: str):
+        """Initialize new user with monthly credits (or infinite for owner)"""
+        if not username:
+            username = "anonymous"
+        
+        if username not in self.user_tokens:
+            # Owner gets infinite tokens
+            if self.is_owner(username):
+                self.user_tokens[username] = {
+                    "balance": float('inf'),
+                    "start_time": time.time(),
+                    "purchases": {"batch_multiplier": 1, "token_limit": 2048},
+                    "total_spent": 0.0,
+                    "is_owner": True,
+                    "username": username
+                }
+                logger.info(f"[TOKEN] 👑 OWNER {username} initialized with INFINITE tokens!")
+            else:
+                self.user_tokens[username] = {
+                    "balance": MONTHLY_TOKEN_CREDITS,
+                    "start_time": time.time(),
+                    "purchases": {"batch_multiplier": 1, "token_limit": 2048},
+                    "total_spent": 0.0,
+                    "is_owner": False,
+                    "username": username,
+                    "last_reset": time.time()
+                }
+                logger.info(f"[TOKEN] New user {username}: {MONTHLY_TOKEN_CREDITS} tokens")
+    
+    def check_monthly_reset(self, username: str):
+        """Reset tokens if a month has passed"""
+        if not username or username not in self.user_tokens:
+            return
+        
+        if self.user_tokens[username].get("is_owner", False):
+            return  # Owner never needs reset
+        
+        last_reset = self.user_tokens[username].get("last_reset", time.time())
+        month_in_seconds = 30 * 24 * 60 * 60  # 30 days
+        
+        if time.time() - last_reset > month_in_seconds:
+            self.user_tokens[username]["balance"] = MONTHLY_TOKEN_CREDITS
+            self.user_tokens[username]["last_reset"] = time.time()
+            self.user_tokens[username]["total_spent"] = 0.0
+            logger.info(f"[TOKEN] Monthly reset for {username}: {MONTHLY_TOKEN_CREDITS} tokens")
     
-    def charge_usage(self, session_id: str, duration_ms: float) -> bool:
-        """Charge user for inference time. Returns True if successful"""
-        self.initialize_user(session_id)
+    def charge_usage(self, username: str, duration_ms: float) -> bool:
+        """Charge user for inference time. Returns True if successful. Owner never charged."""
+        if not username:
+            username = "anonymous"
+        
+        self.initialize_user(username)
+        self.check_monthly_reset(username)
+        
+        # Owner never gets charged
+        if self.user_tokens[username].get("is_owner", False):
+            return True
         
         cost = (duration_ms / 100.0) * TOKEN_COST_PER_100MS
         
-        if self.user_tokens[session_id]["balance"] >= cost:
-            self.user_tokens[session_id]["balance"] -= cost
-            self.user_tokens[session_id]["total_spent"] += cost
-            logger.info(f"[TOKEN] Charged {cost:.4f} tokens ({duration_ms:.0f}ms) | Remaining: {self.user_tokens[session_id]['balance']:.2f}")
+        # Check if user has enough balance
+        if self.user_tokens[username]["balance"] <= 0:
+            logger.warning(f"[TOKEN] ❌ {username} has 0 tokens! Access denied.")
+            return False
+        
+        if self.user_tokens[username]["balance"] >= cost:
+            self.user_tokens[username]["balance"] -= cost
+            self.user_tokens[username]["balance"] = max(0, self.user_tokens[username]["balance"])  # Never go below 0
+            self.user_tokens[username]["total_spent"] += cost
+            logger.info(f"[TOKEN] Charged {cost:.4f} tokens ({duration_ms:.0f}ms) | Remaining: {self.user_tokens[username]['balance']:.2f}")
             return True
         else:
-            logger.warning(f"[TOKEN] Insufficient balance! Need {cost:.4f}, have {self.user_tokens[session_id]['balance']:.2f}")
+            # Insufficient balance - set to 0 and deny
+            self.user_tokens[username]["balance"] = 0
+            logger.warning(f"[TOKEN] ❌ Insufficient balance! {username} now at 0 tokens.")
             return False
     
-    def purchase_batch_upgrade(self, session_id: str) -> tuple:
-        """Purchase batch size upgrade (exponential cost)"""
-        self.initialize_user(session_id)
+    def can_use_engine(self, username: str) -> tuple:
+        """Check if user can use the engine. Returns (bool, message)"""
+        if not username:
+            username = "anonymous"
+        
+        self.initialize_user(username)
+        self.check_monthly_reset(username)
+        
+        if self.user_tokens[username].get("is_owner", False):
+            return True, "👑 Owner access granted"
+        
+        balance = self.user_tokens[username]["balance"]
+        
+        if balance <= 0:
+            last_reset = self.user_tokens[username].get("last_reset", time.time())
+            time_until_reset = 30 * 24 * 60 * 60 - (time.time() - last_reset)
+            days_left = int(time_until_reset / (24 * 60 * 60))
+            return False, f"❌ Out of tokens! Resets in {days_left} days. Current balance: 0.00"
+        
+        return True, f"✅ Access granted. Balance: {balance:.2f} tokens"
+    
+    def purchase_batch_upgrade(self, username: str) -> tuple:
+        """Purchase batch size upgrade (exponential cost). Free for owner."""
+        if not username:
+            return False, "❌ Please login first"
+        
+        self.initialize_user(username)
         
-        current_mult = self.user_tokens[session_id]["purchases"]["batch_multiplier"]
+        # Owner gets free upgrades
+        if self.user_tokens[username].get("is_owner", False):
+            current_mult = self.user_tokens[username]["purchases"]["batch_multiplier"]
+            self.user_tokens[username]["purchases"]["batch_multiplier"] = current_mult * 2
+            new_mult = current_mult * 2
+            logger.info(f"[TOKEN] 👑 OWNER free batch upgrade: {current_mult}x → {new_mult}x")
+            return True, f"👑 FREE UPGRADE! Batch now {new_mult}x!"
+        
+        current_mult = self.user_tokens[username]["purchases"]["batch_multiplier"]
         upgrade_level = int(math.log2(current_mult)) if current_mult > 1 else 0
         cost = BATCH_UPGRADE_BASE_COST * (2 ** upgrade_level)
         
-        if self.user_tokens[session_id]["balance"] >= cost:
-            self.user_tokens[session_id]["balance"] -= cost
-            self.user_tokens[session_id]["purchases"]["batch_multiplier"] = current_mult * 2
+        if self.user_tokens[username]["balance"] >= cost:
+            self.user_tokens[username]["balance"] -= cost
+            self.user_tokens[username]["purchases"]["batch_multiplier"] = current_mult * 2
             new_mult = current_mult * 2
             logger.info(f"[TOKEN] Batch upgrade: {current_mult}x → {new_mult}x | Cost: {cost:.5f}")
             return True, f"✅ Batch upgraded to {new_mult}x! (-{cost:.5f} tokens)"
         else:
-            return False, f"❌ Insufficient tokens! Need {cost:.5f}, have {self.user_tokens[session_id]['balance']:.2f}"
+            return False, f"❌ Insufficient tokens! Need {cost:.5f}, have {self.user_tokens[username]['balance']:.2f}"
     
-    def purchase_token_upgrade(self, session_id: str, extra_tokens: int = 1000) -> tuple:
-        """Purchase extra response token length"""
-        self.initialize_user(session_id)
+    def purchase_token_upgrade(self, username: str, extra_tokens: int = 1000) -> tuple:
+        """Purchase extra response token length. Free for owner."""
+        if not username:
+            return False, "❌ Please login first"
+        
+        self.initialize_user(username)
+        
+        # Owner gets free upgrades
+        if self.user_tokens[username].get("is_owner", False):
+            self.user_tokens[username]["purchases"]["token_limit"] += extra_tokens
+            new_limit = self.user_tokens[username]["purchases"]["token_limit"]
+            logger.info(f"[TOKEN] 👑 OWNER free token upgrade: +{extra_tokens} tokens")
+            return True, f"👑 FREE UPGRADE! Token limit now {new_limit}!"
         
         cost = (extra_tokens / 1000) * TOKEN_UPGRADE_COST_PER_1K
         
-        if self.user_tokens[session_id]["balance"] >= cost:
-            self.user_tokens[session_id]["balance"] -= cost
-            self.user_tokens[session_id]["purchases"]["token_limit"] += extra_tokens
-            new_limit = self.user_tokens[session_id]["purchases"]["token_limit"]
+        if self.user_tokens[username]["balance"] >= cost:
+            self.user_tokens[username]["balance"] -= cost
+            self.user_tokens[username]["purchases"]["token_limit"] += extra_tokens
+            new_limit = self.user_tokens[username]["purchases"]["token_limit"]
             logger.info(f"[TOKEN] Token limit upgrade: +{extra_tokens} tokens | Cost: {cost:.5f}")
             return True, f"✅ Token limit now {new_limit}! (-{cost:.5f} tokens)"
         else:
-            return False, f"❌ Insufficient tokens! Need {cost:.5f}, have {self.user_tokens[session_id]['balance']:.2f}"
+            return False, f"❌ Insufficient tokens! Need {cost:.5f}, have {self.user_tokens[username]['balance']:.2f}"
     
-    def get_balance(self, session_id: str) -> float:
+    def get_balance(self, username: str) -> float:
         """Get user's current token balance"""
-        self.initialize_user(session_id)
-        return round(self.user_tokens[session_id]["balance"], 2)
+        if not username:
+            username = "anonymous"
+        
+        self.initialize_user(username)
+        self.check_monthly_reset(username)
+        
+        balance = self.user_tokens[username]["balance"]
+        
+        # Show ∞ for owner
+        if balance == float('inf'):
+            return balance
+        
+        return round(max(0, balance), 2)  # Never show negative
     
-    def get_purchases(self, session_id: str) -> dict:
+    def get_purchases(self, username: str) -> dict:
         """Get user's current purchases"""
-        self.initialize_user(session_id)
-        return self.user_tokens[session_id]["purchases"]
+        if not username:
+            username = "anonymous"
+        
+        self.initialize_user(username)
+        return self.user_tokens[username]["purchases"]
     
-    def end_session(self, session_id: str):
+    def end_session(self, username: str):
         """End user session and log stats"""
-        if session_id in self.user_tokens:
-            stats = self.user_tokens[session_id]
+        if not username:
+            return "No active session found."
+        
+        if username in self.user_tokens:
+            stats = self.user_tokens[username]
+            
+            if stats.get("is_owner", False):
+                return f"👑 Owner session ended. Welcome back anytime, {stats['username']}!"
+            
             logger.info(f"[TOKEN] Session ended: Spent {stats['total_spent']:.2f}, Remaining {stats['balance']:.2f}")
-            # Don't delete - keep for monthly tracking
-            return f"Session ended. You spent {stats['total_spent']:.2f} tokens this session."
+            return f"Session ended. You spent {stats['total_spent']:.2f} tokens this session. Balance: {stats['balance']:.2f}"
         return "No active session found."
 
 # Global token manager
@@ -830,7 +942,7 @@ class ZeroEngine:
         threading.Thread(target=_bg_eval, daemon=True).start()
         return "⚡ Primed"
 
-    def inference_generator(self, prompt: str, history: List[Dict], ghost_context: str, repo: str, quant: str) -> Generator:
+    def inference_generator(self, prompt: str, history: List[Dict], ghost_context: str, repo: str, quant: str, username: str) -> Generator:
         # Update activity timestamp
         self.update_activity()
         
@@ -885,8 +997,8 @@ class ZeroEngine:
         try:
             # Get max tokens from user purchases
             max_tokens = 2048
-            if session_id:
-                max_tokens = token_manager.get_purchases(session_id)["token_limit"]
+            if username:
+                max_tokens = token_manager.get_purchases(username)["token_limit"]
             
             # HYPER-OPTIMIZED CPU INFERENCE SETTINGS
             stream = self.llm(
@@ -925,17 +1037,17 @@ class ZeroEngine:
                     self.perf_stats["peak_tps"] = tps
                 
                 # Charge tokens every second
-                if int(elapsed * 1000) % 1000 < 100 and session_id:  # Every ~1 second
-                    token_manager.charge_usage(session_id, elapsed * 1000)
+                if int(elapsed * 1000) % 1000 < 100 and username:  # Every ~1 second
+                    token_manager.charge_usage(username, elapsed * 1000)
                 
                 # Update history with streaming content + performance metrics
-                balance = token_manager.get_balance(session_id) if session_id else 0
+                balance = token_manager.get_balance(username) if username else 0
                 history[-1]["content"] = f"{response_text}\n\n`⚡ {tps} t/s | 🎯 Peak: {self.perf_stats['peak_tps']:.1f} t/s | 💰 {balance:.2f} tokens`"
                 yield history
             
             # Final token charge for remaining time
-            if session_id:
-                token_manager.charge_usage(session_id, elapsed * 1000)
+            if username:
+                token_manager.charge_usage(username, elapsed * 1000)
             
             # Update global performance stats
             self.perf_stats["total_tokens"] += tokens_count
@@ -1061,7 +1173,7 @@ h1, h2, h3, h4, h5, h6 {
 kernel = ZeroEngine()
 
 # Session ID for token tracking
-session_id = token_manager.get_session_id()
+username = token_manager.get_username()
 
 with gr.Blocks(title="ZeroEngine V0.2", css=CUSTOM_CSS) as demo:
     # Header with Token Display