utim-cli 1.0.0__py3-none-any.whl

utim_cli/backup.py

@@ -0,0 +1,101 @@
+import os
+import shutil
+import json
+import sqlite3
+from utim_cli.logger import log_info, log_error, log_warning
+
+BACKUP_DIR = ".utim_backup"
+UTIM_DIR = ".utim"
+
+def is_json_valid(path: str) -> bool:
+    """Check if JSON file exists and is parseable."""
+    if not os.path.exists(path):
+        return True
+    try:
+        with open(path, "r", encoding="utf-8") as f:
+            json.load(f)
+        return True
+    except Exception:
+        return False
+
+def is_sqlite_valid(path: str) -> bool:
+    """Check if SQLite DB is valid via PRAGMA integrity_check."""
+    if not os.path.exists(path):
+        return True
+    try:
+        conn = sqlite3.connect(path)
+        c = conn.cursor()
+        c.execute("PRAGMA integrity_check(1)")
+        res = c.fetchone()
+        conn.close()
+        return res is not None and res[0] == "ok"
+    except Exception:
+        return False
+
+def is_file_valid(path: str) -> bool:
+    """Determine if a file is not corrupted based on its extension."""
+    if path.endswith(".json"):
+        return is_json_valid(path)
+    elif path.endswith(".db"):
+        return is_sqlite_valid(path)
+    return True
+
+def backup_state():
+    """Backup active configuration and database files to .utim_backup/."""
+    try:
+        if not os.path.exists(UTIM_DIR):
+            return
+            
+        os.makedirs(BACKUP_DIR, exist_ok=True)
+        
+        files_to_backup = [
+            "config.json", "mcp.json", "utim_local.db", 
+            "memory.json", "session_state.json"
+        ]
+        for f in files_to_backup:
+            src = os.path.join(UTIM_DIR, f)
+            dst = os.path.join(BACKUP_DIR, f)
+            if os.path.exists(src):
+                # Verify file is not corrupted before backing it up
+                if not is_file_valid(src):
+                    log_warning("backup", f"Skipping backup of corrupted file {f}")
+                    continue
+                
+                shutil.copy2(src, dst)
+                
+        log_info("backup", "State backup completed successfully.")
+    except Exception as e:
+        log_error("backup", f"State backup failed: {e}", e)
+
+def restore_state() -> bool:
+    """Restore configuration and database files from backup if corrupt or missing."""
+    try:
+        if not os.path.exists(BACKUP_DIR):
+            return False
+            
+        os.makedirs(UTIM_DIR, exist_ok=True)
+        restored = False
+        
+        files_to_restore = [
+            "config.json", "mcp.json", "utim_local.db", 
+            "memory.json", "session_state.json"
+        ]
+        for f in files_to_restore:
+            src = os.path.join(BACKUP_DIR, f)
+            dst = os.path.join(UTIM_DIR, f)
+            
+            # Check if destination is missing or corrupted
+            needs_restore = not os.path.exists(dst) or not is_file_valid(dst)
+            
+            if needs_restore and os.path.exists(src) and is_file_valid(src):
+                shutil.copy2(src, dst)
+                restored = True
+                log_info("backup", f"Restored {f} from backup.")
+                
+        if restored:
+            log_info("backup", "State backup restored successfully.")
+        return restored
+    except Exception as e:
+        log_error("backup", f"State restoration failed: {e}", e)
+        return False
+

utim_cli/billing.py

@@ -0,0 +1,40 @@
+import json
+import os
+
+class CreditManager:
+    def __init__(self, config_path=None):
+        if config_path is None:
+            config_path = os.path.join(".utim_tmp", "utim_state.json")
+        self.config_path = config_path
+        self.costs = {
+            "gemini": 1,      # 1 credit per 1000 tokens
+            "claude": 15,     # 15 credits per 1000 tokens (more expensive)
+            "codex": 10       # 10 credits per 1000 tokens
+        }
+        self.state = self._load_state()
+
+    def _load_state(self):
+        if os.path.exists(self.config_path):
+            with open(self.config_path, "r") as f:
+                return json.load(f)
+        return {"credits": 0}  # Start with 0 free credits
+
+    def save_state(self):
+        dir_name = os.path.dirname(self.config_path)
+        if dir_name:
+            os.makedirs(dir_name, exist_ok=True)
+        with open(self.config_path, "w") as f:
+            json.dump(self.state, f, indent=4)
+
+    def has_credits(self, agent_name, estimate_tokens=1000):
+        cost = (estimate_tokens / 1000) * self.costs.get(agent_name, 1)
+        return self.state["credits"] >= cost
+
+    def deduct_credits(self, agent_name, actual_tokens):
+        cost = (actual_tokens / 1000) * self.costs.get(agent_name, 1)
+        self.state["credits"] -= cost
+        self.save_state()
+        return cost
+
+    def get_balance(self):
+        return self.state["credits"]