mcpswitch-cli 0.1.0__py3-none-any.whl

mcpswitch/tier.py

@@ -0,0 +1,170 @@
+"""Tier management — Free, Pro (donation-ware), and Team (Lemon Squeezy).
+
+Free:  profile management, token estimation, auto-switch, analyze
+Pro:   waste detection, usage stats, weekly digest, community profiles, sync
+       -> always free; a periodic donation nudge is shown (non-blocking)
+Team:  team analytics, shared profiles, Slack alerts
+       -> requires a paid Team license (Lemon Squeezy)
+
+License keys are stored in ``~/.mcpswitch/license.json`` and validated
+offline via HMAC (see ``mcpswitch.billing``).
+"""
+
+import json
+import time
+from pathlib import Path
+
+from .billing import (
+    TEAM_PAYMENT_URL,
+    validate_license_key,
+)
+
+LICENSE_FILE = Path.home() / ".mcpswitch" / "license.json"
+NUDGE_FILE   = Path.home() / ".mcpswitch" / "nudge.json"
+
+# How often to show the donation nudge (seconds).  Default: once per 7 days.
+_NUDGE_INTERVAL = 7 * 24 * 3600
+
+
+
+# ── License persistence ───────────────────────────────────────────────────────
+
+def _load_license() -> dict:
+    if not LICENSE_FILE.exists():
+        return {}
+    try:
+        with open(LICENSE_FILE, "r", encoding="utf-8") as f:
+            return json.load(f)
+    except Exception:
+        return {}
+
+
+def _save_license(data: dict) -> None:
+    LICENSE_FILE.parent.mkdir(parents=True, exist_ok=True)
+    with open(LICENSE_FILE, "w", encoding="utf-8") as f:
+        json.dump(data, f, indent=2)
+
+
+# ── Activation ────────────────────────────────────────────────────────────────
+
+def activate_license(key: str) -> dict:
+    """Activate a Team license key.
+
+    Returns ``{"success": bool, "tier": str, "message": str}``.
+    """
+    key = key.strip()
+
+    result = validate_license_key(key)
+    if not result["valid"]:
+        return {
+            "success": False,
+            "tier":    "free",
+            "message": f"Invalid license key: {result['error']}",
+        }
+
+    tier  = result["tier"]   # always "team" for Lemon Squeezy-issued keys
+    seats = result.get("seats", 1)
+
+    _save_license({
+        "key":          key.upper(),
+        "tier":         tier,
+        "activated_at": time.time(),
+        "seats":        seats,
+        "email_hash":   result.get("email_hash", ""),
+    })
+
+    seat_str = f" ({seats} seat{'s' if seats != 1 else ''})"
+    return {
+        "success": True,
+        "tier":    tier,
+        "message": f"Activated {tier.upper()} tier{seat_str}.",
+    }
+
+
+# ── Tier queries ──────────────────────────────────────────────────────────────
+
+def get_tier() -> str:
+    """Return current tier: ``'free'``, ``'pro'``, or ``'team'``."""
+    return _load_license().get("tier", "free")
+
+
+def is_pro() -> bool:
+    # Pro features are free — everyone is considered "pro".
+    return True
+
+
+def is_team() -> bool:
+    return get_tier() == "team"
+
+
+def deactivate() -> None:
+    if LICENSE_FILE.exists():
+        LICENSE_FILE.unlink()
+
+
+# ── Feature gates ─────────────────────────────────────────────────────────────
+
+def require_pro(feature_name: str) -> bool:
+    """Always returns ``True`` — Pro features are free for everyone.
+
+    Call ``maybe_show_donation_nudge()`` separately to show an optional nudge.
+    """
+    return True
+
+
+def pro_gate_message(feature_name: str) -> str:
+    """Kept for backward-compat; now returns the donation nudge message."""
+    return donation_nudge_message()
+
+
+# ── Donation nudge (non-blocking) ─────────────────────────────────────────────
+
+def _load_nudge() -> dict:
+    if not NUDGE_FILE.exists():
+        return {}
+    try:
+        with open(NUDGE_FILE, "r", encoding="utf-8") as f:
+            return json.load(f)
+    except Exception:
+        return {}
+
+
+def _save_nudge(data: dict) -> None:
+    NUDGE_FILE.parent.mkdir(parents=True, exist_ok=True)
+    with open(NUDGE_FILE, "w", encoding="utf-8") as f:
+        json.dump(data, f, indent=2)
+
+
+def should_show_nudge() -> bool:
+    """Return True if enough time has passed since the last donation nudge."""
+    if is_team():
+        return False   # paid users don't see donation prompts
+    nudge = _load_nudge()
+    last  = nudge.get("last_shown", 0)
+    return (time.time() - last) >= _NUDGE_INTERVAL
+
+
+def record_nudge_shown() -> None:
+    """Mark that the nudge was shown now."""
+    _save_nudge({"last_shown": time.time()})
+
+
+def donation_nudge_message() -> str:
+    return (
+        "\n[dim]─────────────────────────────────────────────────────[/dim]\n"
+        "[bold yellow]MCPSwitch is free and open-source.[/bold yellow] "
+        "If it saves you context, consider supporting it:\n\n"
+        "  [bold cyan]mcpswitch donate[/bold cyan]   "
+        "→ one-time tip, any amount, no account needed\n\n"
+        "Want team features (shared profiles, Slack alerts)?\n"
+        "  [bold cyan]mcpswitch upgrade[/bold cyan]  "
+        f"→ {TEAM_PAYMENT_URL}\n"
+        "[dim]─────────────────────────────────────────────────────[/dim]\n"
+    )
+
+
+def maybe_show_donation_nudge(console) -> None:
+    """Print the donation nudge if the weekly timer has elapsed."""
+    if should_show_nudge():
+        console.print(donation_nudge_message())
+        record_nudge_shown()

mcpswitch/tokens.py

@@ -0,0 +1,426 @@
+"""Estimate token cost of MCP server tool schemas.
+
+Strategy (in order):
+1. Cache hit  — return stored exact count from previous live query
+2. Known list — return hardcoded count for popular servers
+3. Live query — start the server process, call tools/list, measure real schemas
+4. Fallback   — assume 5 tools if server won't start
+"""
+
+import json
+import sqlite3
+import subprocess
+import time
+from pathlib import Path
+from typing import Optional
+
+import tiktoken
+
+# ── Constants ────────────────────────────────────────────────────────────────
+
+TOKENS_PER_TOOL_AVG = 180       # avg tokens per tool definition (name+desc+schema)
+BASE_OVERHEAD_PER_SERVER = 50   # MCP server metadata / transport overhead
+CONTEXT_WINDOW = 200_000
+CACHE_DB = Path.home() / ".mcpswitch" / "token_cache.db"
+CACHE_TTL_DAYS = 7              # re-query after 7 days
+
+# Hardcoded counts for the most popular servers (fast path, no process spawn)
+KNOWN_TOOL_COUNTS = {
+    "github": 30,
+    "context7": 3,
+    "brave-search": 2,
+    "playwright": 25,
+    "sequential-thinking": 1,
+    "filesystem": 11,
+    "postgres": 6,
+    "sqlite": 8,
+    "memory": 5,
+    "fetch": 2,
+    "slack": 12,
+    "notion": 8,
+    "obsidian": 4,
+    "prompts-local": 2,
+    "linear": 20,
+    "jira": 18,
+    "figma": 10,
+    "google-drive": 8,
+    "gmail": 6,
+    "calendar": 5,
+    "stripe": 14,
+    "supabase": 12,
+    "vercel": 8,
+    "cloudflare": 10,
+}
+
+_enc: Optional[tiktoken.Encoding] = None
+
+
+def _get_encoder() -> tiktoken.Encoding:
+    global _enc
+    if _enc is None:
+        _enc = tiktoken.get_encoding("cl100k_base")
+    return _enc
+
+
+def _count_tokens(text: str) -> int:
+    return len(_get_encoder().encode(text))
+
+
+# ── Cache (SQLite) ────────────────────────────────────────────────────────────
+
+def _get_cache_db() -> sqlite3.Connection:
+    CACHE_DB.parent.mkdir(parents=True, exist_ok=True)
+    conn = sqlite3.connect(str(CACHE_DB))
+    conn.execute("""
+        CREATE TABLE IF NOT EXISTS tool_cache (
+            server_name  TEXT PRIMARY KEY,
+            token_count  INTEGER NOT NULL,
+            tool_count   INTEGER NOT NULL,
+            queried_at   REAL NOT NULL
+        )
+    """)
+    conn.commit()
+    return conn
+
+
+def _cache_get(server_name: str) -> Optional[int]:
+    """Return cached token count if fresh, else None."""
+    try:
+        conn = _get_cache_db()
+        row = conn.execute(
+            "SELECT token_count, queried_at FROM tool_cache WHERE server_name = ?",
+            (server_name,)
+        ).fetchone()
+        conn.close()
+        if row:
+            token_count, queried_at = row
+            age_days = (time.time() - queried_at) / 86400
+            if age_days < CACHE_TTL_DAYS:
+                return token_count
+    except Exception:
+        pass
+    return None
+
+
+def _cache_set(server_name: str, token_count: int, tool_count: int) -> None:
+    try:
+        conn = _get_cache_db()
+        conn.execute(
+            """INSERT OR REPLACE INTO tool_cache
+               (server_name, token_count, tool_count, queried_at)
+               VALUES (?, ?, ?, ?)""",
+            (server_name, token_count, tool_count, time.time())
+        )
+        conn.commit()
+        conn.close()
+    except Exception:
+        pass
+
+
+def _cache_set_bulk(entries: list[tuple[str, int, int]]) -> None:
+    """Batch insert many (server_name, token_count, tool_count) entries at once.
+    Much faster than calling _cache_set() in a loop for 100+ entries.
+    """
+    if not entries:
+        return
+    try:
+        now = time.time()
+        conn = _get_cache_db()
+        conn.executemany(
+            """INSERT OR REPLACE INTO tool_cache
+               (server_name, token_count, tool_count, queried_at)
+               VALUES (?, ?, ?, ?)""",
+            [(name, tok, tools, now) for name, tok, tools in entries],
+        )
+        conn.commit()
+        conn.close()
+    except Exception:
+        pass
+
+
+def clear_cache(server_name: Optional[str] = None) -> int:
+    """Clear cache. Pass server_name to clear one entry, or None to clear all.
+    Returns number of rows deleted."""
+    try:
+        conn = _get_cache_db()
+        if server_name:
+            cur = conn.execute("DELETE FROM tool_cache WHERE server_name = ?", (server_name,))
+        else:
+            cur = conn.execute("DELETE FROM tool_cache")
+        conn.commit()
+        count = cur.rowcount
+        conn.close()
+        return count
+    except Exception:
+        return 0
+
+
+def list_cache() -> list[dict]:
+    """Return all cached entries."""
+    try:
+        conn = _get_cache_db()
+        rows = conn.execute(
+            "SELECT server_name, token_count, tool_count, queried_at FROM tool_cache ORDER BY token_count DESC"
+        ).fetchall()
+        conn.close()
+        return [
+            {
+                "server": r[0],
+                "tokens": r[1],
+                "tools": r[2],
+                "age_hours": round((time.time() - r[3]) / 3600, 1),
+            }
+            for r in rows
+        ]
+    except Exception:
+        return []
+
+
+# ── Live Query via MCP Protocol ───────────────────────────────────────────────
+
+_MCP_INIT_MSG = json.dumps({
+    "jsonrpc": "2.0",
+    "id": 1,
+    "method": "initialize",
+    "params": {
+        "protocolVersion": "2024-11-05",
+        "capabilities": {},
+        "clientInfo": {"name": "mcpswitch", "version": "0.1.0"},
+    },
+}) + "\n"
+
+_MCP_TOOLS_MSG = json.dumps({
+    "jsonrpc": "2.0",
+    "id": 2,
+    "method": "tools/list",
+    "params": {},
+}) + "\n"
+
+
+def _live_query_tools(server_config: dict, timeout: float = 8.0) -> Optional[list[dict]]:
+    """Start the MCP server process and call tools/list.
+
+    Returns list of tool dicts, or None if the server can't be reached.
+    Each tool dict has 'name', 'description', 'inputSchema'.
+    """
+    command = server_config.get("command")
+    args = server_config.get("args", [])
+    env_extra = server_config.get("env", {})
+
+    if not command:
+        return None
+
+    import os
+    env = os.environ.copy()
+    env.update(env_extra)
+
+    try:
+        proc = subprocess.Popen(
+            [command] + args,
+            stdin=subprocess.PIPE,
+            stdout=subprocess.PIPE,
+            stderr=subprocess.DEVNULL,
+            env=env,
+            text=True,
+        )
+
+        # Send initialize
+        proc.stdin.write(_MCP_INIT_MSG)
+        proc.stdin.flush()
+
+        # Read initialize response (with timeout via poll)
+        deadline = time.time() + timeout
+        init_resp = None
+        while time.time() < deadline:
+            line = proc.stdout.readline()
+            if not line:
+                time.sleep(0.05)
+                continue
+            try:
+                msg = json.loads(line.strip())
+                if msg.get("id") == 1:
+                    init_resp = msg
+                    break
+            except json.JSONDecodeError:
+                continue
+
+        if init_resp is None:
+            proc.terminate()
+            return None
+
+        # Send initialized notification
+        proc.stdin.write(json.dumps({
+            "jsonrpc": "2.0",
+            "method": "notifications/initialized",
+            "params": {},
+        }) + "\n")
+        proc.stdin.flush()
+
+        # Send tools/list
+        proc.stdin.write(_MCP_TOOLS_MSG)
+        proc.stdin.flush()
+
+        # Read tools/list response
+        tools_resp = None
+        deadline = time.time() + timeout
+        while time.time() < deadline:
+            line = proc.stdout.readline()
+            if not line:
+                time.sleep(0.05)
+                continue
+            try:
+                msg = json.loads(line.strip())
+                if msg.get("id") == 2:
+                    tools_resp = msg
+                    break
+            except json.JSONDecodeError:
+                continue
+
+        proc.terminate()
+
+        if tools_resp and "result" in tools_resp:
+            return tools_resp["result"].get("tools", [])
+
+    except Exception:
+        pass
+
+    return None
+
+
+def _tokens_from_tool_list(tools: list[dict]) -> int:
+    """Count exact tokens from real tool schemas."""
+    total = BASE_OVERHEAD_PER_SERVER
+    for tool in tools:
+        schema_str = json.dumps(tool)
+        total += _count_tokens(schema_str)
+    return total
+
+
+# ── Public API ────────────────────────────────────────────────────────────────
+
+def get_server_tokens(
+    server_name: str,
+    server_config: dict,
+    live: bool = False,
+) -> tuple[int, str]:
+    """Return (token_count, source) for a server.
+
+    source is one of: 'cache', 'known', 'live', 'fallback'
+
+    Args:
+        server_name:   name from mcpServers config
+        server_config: full server config dict
+        live:          if True, force a live query even if cache/known exists
+    """
+    name_lower = server_name.lower()
+
+    # 1. Cache
+    if not live:
+        cached = _cache_get(server_name)
+        if cached is not None:
+            return cached, "cache"
+
+    # 2. Known list (fast path for popular servers)
+    if not live:
+        for known, count in KNOWN_TOOL_COUNTS.items():
+            if known in name_lower:
+                tokens = BASE_OVERHEAD_PER_SERVER + (count * TOKENS_PER_TOOL_AVG)
+                return tokens, "known"
+
+    # 3. Live query — start the server, call tools/list
+    tools = _live_query_tools(server_config)
+    if tools is not None:
+        tokens = _tokens_from_tool_list(tools)
+        _cache_set(server_name, tokens, len(tools))
+        return tokens, "live"
+
+    # 4. Fallback — assume 5 tools
+    tokens = BASE_OVERHEAD_PER_SERVER + (5 * TOKENS_PER_TOOL_AVG)
+    return tokens, "fallback"
+
+
+def _cache_get_bulk(names: list[str]) -> dict[str, int]:
+    """Fetch multiple cache entries in a single SQLite query.
+    Returns {server_name: token_count} for cache-fresh entries only.
+    """
+    if not names:
+        return {}
+    try:
+        conn = _get_cache_db()
+        placeholders = ",".join("?" * len(names))
+        rows = conn.execute(
+            f"SELECT server_name, token_count, queried_at FROM tool_cache WHERE server_name IN ({placeholders})",
+            names,
+        ).fetchall()
+        conn.close()
+        now = time.time()
+        return {
+            r[0]: r[1]
+            for r in rows
+            if (now - r[2]) / 86400 < CACHE_TTL_DAYS
+        }
+    except Exception:
+        return {}
+
+
+def estimate_total_tokens(servers: dict, live: bool = False) -> dict:
+    """Return token estimates for a set of MCP servers.
+
+    Uses a single bulk cache query for all servers — fast even at 1000+ servers.
+
+    Returns:
+        {
+            "total": int,
+            "per_server": { "name": {"tokens": int, "source": str} },
+            "context_pct": float
+        }
+    """
+    # Bulk cache lookup — one SQLite query for all servers
+    cache_hits = {} if live else _cache_get_bulk(list(servers.keys()))
+
+    per_server = {}
+    for name, cfg in servers.items():
+        # 1. Cache hit
+        if name in cache_hits:
+            per_server[name] = {"tokens": cache_hits[name], "source": "cache"}
+            continue
+
+        # 2. Known list
+        if not live:
+            name_lower = name.lower()
+            matched = False
+            for known, count in KNOWN_TOOL_COUNTS.items():
+                if known in name_lower:
+                    tokens = BASE_OVERHEAD_PER_SERVER + (count * TOKENS_PER_TOOL_AVG)
+                    per_server[name] = {"tokens": tokens, "source": "known"}
+                    matched = True
+                    break
+            if matched:
+                continue
+
+        # 3. Live query
+        tools = _live_query_tools(cfg)
+        if tools is not None:
+            tokens = _tokens_from_tool_list(tools)
+            _cache_set(name, tokens, len(tools))
+            per_server[name] = {"tokens": tokens, "source": "live"}
+            continue
+
+        # 4. Fallback
+        tokens = BASE_OVERHEAD_PER_SERVER + (5 * TOKENS_PER_TOOL_AVG)
+        per_server[name] = {"tokens": tokens, "source": "fallback"}
+
+    total = sum(v["tokens"] for v in per_server.values())
+    pct = round((total / CONTEXT_WINDOW) * 100, 1)
+
+    return {
+        "total": total,
+        "per_server": per_server,
+        "context_pct": pct,
+    }
+
+
+def format_token_count(n: int) -> str:
+    if n >= 1000:
+        return f"{n:,} (~{n/1000:.1f}k)"
+    return str(n)