usage-lens 0.0.1__tar.gz

usage_lens-0.0.1/PKG-INFO

@@ -0,0 +1,13 @@
+Metadata-Version: 2.4
+Name: usage-lens
+Version: 0.0.1
+Summary: AI usage tracker.
+Author: Karthikeya
+Requires-Python: >=3.9
+Requires-Dist: requests
+Requires-Dist: websocket-client
+Requires-Dist: keyring
+Requires-Dist: pywin32
+Requires-Dist: pystray
+Requires-Dist: Pillow
+Requires-Dist: pyyaml

usage_lens-0.0.1/pyproject.toml

@@ -0,0 +1,26 @@
+[build-system]
+requires = ["setuptools>=61.0"]
+build-backend = "setuptools.build_meta"
+
+[project]
+name = "usage-lens"
+version = "0.0.1"
+description = "AI usage tracker."
+authors = [{name = "Karthikeya"}]
+requires-python = ">=3.9"
+
+dependencies = [
+    "requests",
+    "websocket-client",
+    "keyring",
+    "pywin32",
+    "pystray",
+    "Pillow",
+    "pyyaml",
+]
+
+[project.scripts]
+usage-lens = "ai_tokenkeeper.cli:main"
+
+[tool.setuptools.packages.find]
+where = ["src"]

usage_lens-0.0.1/setup.cfg

@@ -0,0 +1,4 @@
+[egg_info]
+tag_build = 
+tag_date = 0
+

usage_lens-0.0.1/src/ai_tokenkeeper/__init__.py

@@ -0,0 +1,4 @@
+# src/ai_tokenkeeper/__init__.py
+from .core.engine import TokenKeeperEngine
+
+__version__ = "0.1.0"

usage_lens-0.0.1/src/ai_tokenkeeper/__main__.py

@@ -0,0 +1,4 @@
+from ai_tokenkeeper.cli import main
+
+if __name__ == "__main__":
+    main()

usage_lens-0.0.1/src/ai_tokenkeeper/cli.py

@@ -0,0 +1,84 @@
+import sys
+import argparse
+import tkinter as tk
+from ai_tokenkeeper.core.engine import TokenKeeperEngine
+from ai_tokenkeeper.ui.dashboard import TokenKeeperDashboard
+from ai_tokenkeeper.platform.tray import SystemTrayManager
+from ai_tokenkeeper.platform.autostart import AutoStartManager
+
+def main():
+    """Main terminal entry point for managing the TokenKeeper application."""
+    parser = argparse.ArgumentParser(description="AI TokenKeeper: Cross-platform LLM session recovery and monitoring utility.")
+    subparsers = parser.add_subparsers(dest="command", help="Available utility operations")
+
+    # Command: start
+    start_parser = subparsers.add_parser("start", help="Launch the primary monitoring daemon and desktop dashboard")
+    start_parser.add_argument("--headless", action="store_true", help="Launch directly into the system tray without showing the UI window")
+
+    # Command: autostart
+    autostart_parser = subparsers.add_parser("autostart", help="Configure system boot integrations")
+    autostart_parser.add_argument("action", choices=["enable", "disable"], help="Enable or disable launch on system startup")
+
+    args = parser.parse_args()
+
+    # Default to 'start' action if no argument is provided
+    if not args.command:
+        args.command = "start"
+        args.headless = False
+
+    if args.command == "autostart":
+        if args.action == "enable":
+            success = AutoStartManager.enable()
+            print("[CLI] System startup integration successfully enabled.") if success else print("[CLI] Failed to write startup configuration.")
+        elif args.action == "disable":
+            success = AutoStartManager.disable()
+            print("[CLI] System startup integration disabled.") if success else print("[CLI] Failed to clear startup configuration.")
+        sys.exit(0)
+
+    if args.command == "start":
+        # Initialize the underlying orchestration logic engine
+        engine = TokenKeeperEngine()
+        engine.start_background_loop()
+
+        root = tk.Tk()
+        app = TokenKeeperDashboard(root, engine)
+
+        def open_dashboard_from_tray(icon, item):
+            root.after(0, lambda: root.deiconify())
+
+        def force_refresh_from_tray(icon, item):
+            for provider_id in engine.providers:
+                engine.update_provider_state(provider_id)
+
+        def complete_shutdown():
+            engine.stop_background_loop()
+            if tray_manager.icon:
+                tray_manager.icon.stop()
+            root.destroy()
+            sys.exit(0)
+
+        # Wire up the cross-platform system tray hooks
+        tray_manager = SystemTrayManager(
+            on_open_dashboard=open_dashboard_from_tray,
+            on_force_refresh=force_refresh_from_tray,
+            on_exit=complete_shutdown
+        )
+
+        # Override default window exit protocol to intercept and minimize cleanly to the system tray
+        def hide_window_to_tray():
+            root.withdraw()
+
+        root.protocol("WM_DELETE_WINDOW", hide_window_to_tray)
+
+        import threading
+        tray_thread = threading.Thread(target=tray_manager.run_tray_icon, daemon=True)
+        tray_thread.start()
+
+        if args.headless:
+            root.withdraw()
+
+        # Run the primary thread lifecycle loop
+        root.mainloop()
+
+if __name__ == "__main__":
+    main()

usage_lens-0.0.1/src/ai_tokenkeeper/config/__init__.py

@@ -0,0 +1,4 @@
+from ai_tokenkeeper.config.settings import AppConfig
+
+# Bridge the new premium AppConfig to the class name expected by the core engine
+ConfigManager = AppConfig

usage_lens-0.0.1/src/ai_tokenkeeper/config/settings.py

@@ -0,0 +1,50 @@
+import json
+import os
+from pathlib import Path
+from typing import Any, Dict
+
+class AppConfig:
+    """
+    Handles local JSON configuration persistence for premium widget modes, 
+    provider selections, and pinning choices.
+    """
+    def __init__(self):
+        self.config_dir = Path(os.path.expanduser("~/.config/ai_tokenkeeper"))
+        self.config_file = self.config_dir / "config.json"
+        
+        # Default application states
+        self.defaults = {
+            "selected_provider": "claude",
+            "ui_mode": "dashboard",        # Options: dashboard, widget, minibar
+            "always_on_top": True,
+            "threshold_alert_pct": 80
+        }
+        self.data = self._load_config()
+
+    def _load_config(self) -> Dict[str, Any]:
+        try:
+            self.config_dir.mkdir(parents=True, exist_ok=True)
+            if self.config_file.exists():
+                with open(self.config_file, "r") as f:
+                    user_data = json.load(f)
+                    # Merge user flags over defaults to maintain compatibility
+                    return {**self.defaults, **user_data}
+        except Exception:
+            pass
+        return self.defaults.copy()
+
+    def save(self) -> None:
+        """Saves current memory runtime states safely to disk."""
+        try:
+            with open(self.config_file, "w") as f:
+                json.dump(self.data, f, indent=4)
+        except Exception as e:
+            print(f"[Config Engine] Save error: {e}")
+
+    def get(self, key: str, default: Any = None) -> Any:
+        """Retrieves a configuration metric or returns the default fallback value."""
+        return self.data.get(key, default if default is not None else self.defaults.get(key))
+
+    def set(self, key: str, value: Any) -> None:
+        self.data[key] = value
+        self.save()

usage_lens-0.0.1/src/ai_tokenkeeper/config.py

@@ -0,0 +1,46 @@
+import os
+import json
+from typing import Dict, Any
+
+class ConfigManager:
+    """
+    Manages non-sensitive application configurations and state caching.
+    Saves metadata to the user's local application data directory.
+    """
+    def __init__(self):
+        self.app_data_dir = os.path.join(
+            os.environ.get("LOCALAPPDATA", os.path.expanduser("~")), 
+            "AITokenKeeper"
+        )
+        os.makedirs(self.app_data_dir, exist_ok=True)
+        self.config_path = os.path.join(self.app_data_dir, "config.json")
+        self.defaults = {
+            "polling_interval_seconds": 60,
+            "auto_start": False,
+            "enabled_providers": ["claude"]
+        }
+        self.settings = self._load_config()
+
+    def _load_config(self) -> Dict[str, Any]:
+        if os.path.exists(self.config_path):
+            try:
+                with open(self.config_path, "r") as f:
+                    loaded = json.load(f)
+                    # Merge defaults for any missing keys
+                    return {**self.defaults, **loaded}
+            except Exception as e:
+                print(f"[Config] Error reading configuration file: {e}")
+        return self.defaults.copy()
+
+    def get(self, key: str, default: Any = None) -> Any:
+        """Retrieves a configuration metric or returns default value."""
+        return self.settings.get(key, default if default is not None else self.defaults.get(key))
+
+    def set(self, key: str, value: Any) -> None:
+        """Updates a configuration metric and commits the change to disk."""
+        self.settings[key] = value
+        try:
+            with open(self.config_path, "w") as f:
+                json.dump(self.settings, f, indent=4)
+        except Exception as e:
+            print(f"[Config] Failed to persist configuration settings: {e}")

usage_lens-0.0.1/src/ai_tokenkeeper/core/engine.py

@@ -0,0 +1,102 @@
+import threading
+import time
+from typing import Dict, Any, Optional
+from ai_tokenkeeper.config import ConfigManager
+from ai_tokenkeeper.secure_storage import SecureStorage
+from ai_tokenkeeper.core.plugin_loader import PluginLoader
+
+class TokenKeeperEngine:
+    """
+    Central orchestration runner that maintains background execution states,
+    coordinates tracking cycles, and provides isolated data streams to the view layer.
+    """
+    def __init__(self):
+        self.config = ConfigManager()
+        self.storage = SecureStorage()
+        self.providers: Dict[str, Any] = {}
+        self.shared_state: Dict[str, Any] = {}
+        self.is_running = False
+        self._loop_thread: Optional[threading.Thread] = None
+
+        self._initialize_plugins()
+
+    def _initialize_plugins(self) -> None:
+        """Discovers internal/external provider components and boots them."""
+        discovered_classes = PluginLoader.discover_providers()
+        
+        for p_id, provider_class in discovered_classes.items():
+            # Instantiate each discovered provider with core resources
+            instance = provider_class(config_manager=self.config, secure_store=self.storage)
+            self.providers[p_id] = instance
+            
+            # Populate initial safe data mapping state
+            self.shared_state[p_id] = {
+                "display_name": instance.display_name,
+                "is_authenticated": False,
+                "metrics": {"status": "Initializing", "input_tokens": 0, "output_tokens": 0}
+            }
+            
+            # Run quick non-blocking validation sweep to catch existing cookies
+            instance.refresh_session()
+
+    def start_background_loop(self) -> None:
+        """Fires up the non-blocking background orchestration interval worker."""
+        if self.is_running:
+            return
+        self.is_running = True
+        self._loop_thread = threading.Thread(target=self._execution_loop, daemon=True)
+        self._loop_thread.start()
+
+    def stop_background_loop(self) -> None:
+        """Triggers clean teardown flags for the background operations thread."""
+        self.is_running = False
+
+    def trigger_provider_auth(self, provider_id: str) -> bool:
+        """Executes targeted runtime interaction sequence for a specific engine."""
+        if provider_id not in self.providers:
+            return False
+            
+        provider = self.providers[provider_id]
+        success = provider.authenticate()
+        
+        # Instantly refresh localized engine metrics on tracking maps
+        self.update_provider_state(provider_id)
+        return success
+
+    def update_provider_state(self, provider_id: str) -> None:
+        """Queries explicit metric fields to populate current view data boundaries."""
+        provider = self.providers[provider_id]
+        is_auth = provider.refresh_session()
+        
+        metrics = {"status": "Unauthenticated", "input_tokens": 0, "output_tokens": 0}
+        if is_auth:
+            metrics = provider.get_usage()
+            
+        self.shared_state[provider_id] = {
+            "display_name": provider.display_name,
+            "is_authenticated": is_auth,
+            "metrics": metrics
+        }
+
+    def _execution_loop(self) -> None:
+        """Continuous low-priority background thread execution lifecycle loop."""
+        while self.is_running:
+            enabled_targets = self.config.get("enabled_providers", [])
+            
+            for p_id in enabled_targets:
+                if p_id in self.providers:
+                    try:
+                        self.update_provider_state(p_id)
+                    except Exception as loop_err:
+                        print(f"[Engine] Exception during active polling cycle for {p_id}: {loop_err}")
+                        
+            interval = self.config.get("polling_interval_seconds", 60)
+            # Sleep in tiny blocks to allow fast application exit steps
+            for _ in range(int(interval)):
+                if not self.is_running:
+                    break
+                time.sleep(1)
+
+    def get_latest_state(self) -> Dict[str, Any]:
+        """Provides the presentation interface with thread-safe current execution states."""
+        return self.shared_state.copy()

usage_lens-0.0.1/src/ai_tokenkeeper/core/plugin_loader.py

@@ -0,0 +1,50 @@
+import os
+import inspect
+import importlib
+from typing import Dict, Type
+from ai_tokenkeeper.providers.base import BaseProvider
+
+class PluginLoader:
+    """
+    Discovers, inspects, and registers provider plugins within the package environment.
+    """
+    
+    @staticmethod
+    def discover_providers() -> Dict[str, Type[BaseProvider]]:
+        """
+        Scans the providers directory, looks for valid subclasses of BaseProvider,
+        and builds a registration dictionary mapped by their provider_id.
+        """
+        discovered: Dict[str, Type[BaseProvider]] = {}
+        
+        # Locate the absolute directory paths relative to this runtime file
+        current_dir = os.path.dirname(os.path.abspath(__file__))
+        providers_dir = os.path.normpath(os.path.join(current_dir, "..", "providers"))
+        
+        if not os.path.exists(providers_dir):
+            return discovered
+
+        # Scan folder files
+        for filename in os.listdir(providers_dir):
+            if filename.endswith(".py") and filename != "base.py" and not filename.startswith("__"):
+                module_name = f"ai_tokenkeeper.providers.{filename[:-3]}"
+                
+                try:
+                    module = importlib.import_module(module_name)
+                    
+                    # Inspect every class declared inside the found python file
+                    for _, obj in inspect.getmembers(module, inspect.isclass):
+                        # Ensure the class inherits from BaseProvider but isn't BaseProvider itself
+                        if issubclass(obj, BaseProvider) and obj is not BaseProvider:
+                            # Safely read properties to register the identifier string
+                            try:
+                                provider_instance = obj.__new__(obj)
+                                p_id = provider_instance.provider_id
+                                discovered[p_id] = obj
+                            except Exception as instantiation_err:
+                                print(f"[PluginLoader] Failed tracking attributes for class {obj.__name__}: {instantiation_err}")
+                                
+                except Exception as import_err:
+                    print(f"[PluginLoader] Could not load provider module {module_name}: {import_err}")
+                    
+        return discovered

usage_lens-0.0.1/src/ai_tokenkeeper/database/history.py

@@ -0,0 +1,54 @@
+import json
+import os
+from datetime import datetime
+from pathlib import Path
+from typing import Any, Dict, List
+
+class UsageHistory:
+    """
+    Lightweight append-only time-series data store mapping historical usage percentages 
+    to drive UI sparkline visualization components.
+    """
+    def __init__(self):
+        self.data_dir = Path(os.path.expanduser("~/.config/ai_tokenkeeper"))
+        self.history_file = self.data_dir / "usage_history.json"
+        self._init_store()
+
+    def _init_store(self) -> None:
+        self.data_dir.mkdir(parents=True, exist_ok=True)
+        if not self.history_file.exists():
+            with open(self.history_file, "w") as f:
+                json.dump([], f)
+
+    def log_entry(self, provider_id: str, session_pct: int, weekly_pct: int) -> None:
+        """Appends a snapshot metric into the local time-series repository."""
+        try:
+            with open(self.history_file, "r") as f:
+                history = json.load(f)
+            
+            entry = {
+                "timestamp": datetime.now().isoformat(),
+                "provider_id": provider_id,
+                "session_pct": session_pct,
+                "weekly_pct": weekly_pct
+            }
+            history.append(entry)
+            
+            # Keep log bounded to last 200 data tracks to prevent infinite disk leaks
+            if len(history) > 200:
+                history = history[-200:]
+                
+            with open(self.history_file, "w") as f:
+                json.dump(history, f, indent=2)
+        except Exception as e:
+            print(f"[History Store] Log write failure: {e}")
+
+    def get_records(self, provider_id: str, limit: int = 20) -> List[Dict[str, Any]]:
+        """Retrieves filtered chronological records for a designated core provider."""
+        try:
+            with open(self.history_file, "r") as f:
+                history = json.load(f)
+            filtered = [r for r in history if r.get("provider_id") == provider_id]
+            return filtered[-limit:]
+        except Exception:
+            return []

usage_lens-0.0.1/src/ai_tokenkeeper/platform/autostart.py

@@ -0,0 +1,108 @@
+import os
+import sys
+
+class AutoStartManager:
+    """
+    Manages OS-specific hooks to register or deregister the utility from system startup lists.
+    """
+    APP_NAME = "AITokenKeeper"
+
+    @classmethod
+    def enable(cls) -> bool:
+        """Registers the application executable into the target platform startup lifecycle."""
+        platform = sys.platform
+        # Determine the accurate target running script handle
+        exec_command = f'"{sys.executable}" -m ai_tokenkeeper'
+
+        try:
+            if platform == "win32":
+                import winreg
+                key = winreg.OpenKey(
+                    winreg.HKEY_CURRENT_USER,
+                    r"Software\Microsoft\Windows\CurrentVersion\Run",
+                    0, winreg.KEY_SET_VALUE
+                )
+                winreg.SetValueEx(key, cls.APP_NAME, 0, winreg.REG_SZ, exec_command)
+                winreg.CloseKey(key)
+                return True
+
+            elif platform == "darwin":  # macOS Launchd Agent
+                plist_dir = os.path.expanduser("~/Library/LaunchAgents")
+                os.makedirs(plist_dir, exist_ok=True)
+                plist_path = os.path.join(plist_dir, f"com.tokenkeeper.app.plist")
+                
+                plist_content = f"""<?xml version="1.0" encoding="UTF-8"?>
+<!DOCTYPE plist PUBLIC "-//Apple//DTD PLIST 1.0//EN" "http://www.apple.com/DTDs/PropertyList-1.0.dtd">
+<plist version="1.0">
+<dict>
+    <key>Label</key>
+    <string>com.tokenkeeper.app</string>
+    <key>ProgramArguments</key>
+    <array>
+        <string>{sys.executable}</string>
+        <string>-m</string>
+        <string>ai_tokenkeeper</string>
+    </array>
+    <key>RunAtLoad</key>
+    <true/>
+</dict>
+</plist>"""
+                with open(plist_path, "w") as f:
+                    f.write(plist_content)
+                return True
+
+            elif platform.startswith("linux"):  # XDG Desktop Autostart standard
+                autostart_dir = os.path.expanduser("~/.config/autostart")
+                os.makedirs(autostart_dir, exist_ok=True)
+                desktop_path = os.path.join(autostart_dir, "ai_tokenkeeper.desktop")
+                
+                desktop_content = f"""[Desktop Entry]
+Type=Application
+Name=AI TokenKeeper
+Exec={sys.executable} -m ai_tokenkeeper
+Hidden=false
+NoDisplay=false
+X-GNOME-Autostart-enabled=true
+Comment=LLM Session Recovery Agent
+"""
+                with open(desktop_path, "w") as f:
+                    f.write(desktop_content)
+                return True
+
+        except Exception as e:
+            print(f"[AutoStart] Failed to enable startup registration: {e}")
+        return False
+
+    @classmethod
+    def disable(cls) -> bool:
+        """Removes the application registration from the system startup sequences."""
+        platform = sys.platform
+        try:
+            if platform == "win32":
+                import winreg
+                key = winreg.OpenKey(
+                    winreg.HKEY_CURRENT_USER,
+                    r"Software\Microsoft\Windows\CurrentVersion\Run",
+                    0, winreg.KEY_SET_VALUE
+                )
+                try:
+                    winreg.DeleteValue(key, cls.APP_NAME)
+                except FileNotFoundError:
+                    pass
+                winreg.CloseKey(key)
+                return True
+
+            elif platform == "darwin":
+                plist_path = os.path.expanduser("~/Library/LaunchAgents/com.tokenkeeper.app.plist")
+                if os.path.exists(plist_path):
+                    os.remove(plist_path)
+                return True
+
+            elif platform.startswith("linux"):
+                desktop_path = os.path.expanduser("~/.config/autostart/ai_tokenkeeper.desktop")
+                if os.path.exists(desktop_path):
+                    os.remove(desktop_path)
+                return True
+        except Exception as e:
+            print(f"[AutoStart] Failed to tear down autostart configs: {e}")
+        return False

usage_lens-0.0.1/src/ai_tokenkeeper/platform/browser_detector.py

@@ -0,0 +1,45 @@
+import os
+import sys
+
+class BrowserDetector:
+    """
+    Locates compatible chromium browser engines across diverse operating systems
+    to enable targeted stealth automation loops.
+    """
+    
+    @staticmethod
+    def get_edge_path() -> str:
+        """
+        Scans predictable system directory frames to return the binary location of Edge.
+        Raises FileNotFoundError if no match is discovered.
+        """
+        platform = sys.platform
+        paths = []
+
+        if platform == "win32":
+            paths = [
+                r"C:\Program Files (x86)\Microsoft\Edge\Application\msedge.exe",
+                r"C:\Program Files\Microsoft\Edge\Application\msedge.exe",
+                os.path.join(os.environ.get("LOCALAPPDATA", ""), "Microsoft", "Edge", "Application", "msedge.exe")
+            ]
+        elif platform == "darwin":  # macOS
+            paths = [
+                "/Applications/Microsoft Edge.app/Contents/MacOS/Microsoft Edge",
+                os.path.expanduser("~/Applications/Microsoft Edge.app/Contents/MacOS/Microsoft Edge")
+            ]
+        elif platform.startswith("linux"):
+            paths = [
+                "/usr/bin/microsoft-edge",
+                "/usr/bin/microsoft-edge-stable",
+                "/usr/local/bin/microsoft-edge",
+                "/snap/bin/microsoft-edge"
+            ]
+
+        for path in paths:
+            if os.path.exists(path):
+                return path
+
+        raise FileNotFoundError(
+            f"Microsoft Edge binary could not be automatically detected on platform: {platform}. "
+            "Please ensure it is installed or specify the path manually."
+        )

usage_lens-0.0.1/src/ai_tokenkeeper/platform/tray.py

@@ -0,0 +1,49 @@
+import threading
+from typing import Callable
+import pystray
+from PIL import Image, ImageDraw
+
+class SystemTrayManager:
+    """
+    Controls the application presence in the OS taskbar/tray area.
+    Keeps the engine processing tokens even when the dashboard window is hidden.
+    """
+    def __init__(self, on_open_dashboard: Callable, on_force_refresh: Callable, on_exit: Callable):
+        self.on_open_dashboard = on_open_dashboard
+        self.on_force_refresh = on_force_refresh
+        self.on_exit = on_exit
+        self.icon = None
+
+    def _create_fallback_icon(self) -> Image.Image:
+        """Generates a simple geometric placeholder image if a dedicated PNG asset is missing."""
+        image = Image.new("RGBA", (64, 64), (0, 0, 0, 0))
+        dc = ImageDraw.Draw(image)
+        # Draw a clean circular indicator frame
+        dc.ellipse([8, 8, 56, 56], fill=(30, 41, 59), outline=(99, 102, 241), width=4)
+        dc.text((24, 22), "TK", fill=(255, 255, 255))
+        return image
+
+    def run_tray_icon(self) -> None:
+        """Spawns the system tray instance inside an unblocked running environment loop."""
+        menu = pystray.Menu(
+            pystray.MenuItem("Open Dashboard", self.on_open_dashboard, default=True),
+            pystray.MenuItem("Force Refresh Sync", self.on_force_refresh),
+            pystray.Menu.SEPARATOR,
+            pystray.MenuItem("Exit Utility", self.on_quit_triggered)
+        )
+
+        self.icon = pystray.Icon(
+            "AITokenKeeper",
+            icon=self._create_fallback_icon(),
+            title="AI TokenKeeper (Running)",
+            menu=menu
+        )
+        
+        # Runs the internal platform monitoring process
+        self.icon.run()
+
+    def on_quit_triggered(self, icon, item) -> None:
+        """Internal callback to stop the icon instance before handing execution off to global hooks."""
+        if self.icon:
+            self.icon.stop()
+        self.on_exit()