tokenmizer 0.2.4__py3-none-any.whl

tokenmizer/api/rate_limiter.py

@@ -0,0 +1,110 @@
+"""
+Simple token-bucket rate limiter for the proxy endpoint.
+No external deps — pure stdlib.
+"""
+from __future__ import annotations
+
+import asyncio
+import logging
+import time
+from collections import defaultdict
+from dataclasses import dataclass, field
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass
+class _Bucket:
+    tokens: float
+    last_refill: float = field(default_factory=time.monotonic)
+
+
+class RateLimiter:
+    """
+    Per-client token-bucket rate limiter.
+    Default: 60 requests/minute per API key (or IP if no key).
+
+    Usage:
+        limiter = RateLimiter(rate=60, per_seconds=60, burst=10)
+        allowed, retry_after = limiter.check("client-id")
+    """
+
+    def __init__(
+        self,
+        rate: int = 60,
+        per_seconds: int = 60,
+        burst: int = 10,
+        max_clients: int = 50_000,
+    ):
+        self.rate = rate                  # tokens per window
+        self.per_seconds = per_seconds    # window length
+        self.burst = burst                # max burst above window rate
+        self.max_clients = max_clients    # hard cap — prevents unbounded growth
+        self.capacity = rate + burst
+        self.refill_rate = rate / per_seconds   # tokens per second
+        self._buckets: dict[str, _Bucket] = defaultdict(
+            lambda: _Bucket(tokens=float(self.capacity))
+        )
+        self._lock = asyncio.Lock()
+        # Cleanup: evict stale buckets every 5 minutes
+        self._last_cleanup = time.monotonic()
+        self._cleanup_interval = 300
+
+    async def check(self, client_id: str) -> tuple[bool, float]:
+        """
+        Returns (allowed, retry_after_seconds).
+        retry_after is 0.0 if allowed.
+        """
+        async with self._lock:
+            now = time.monotonic()
+            bucket = self._buckets[client_id]
+
+            # Refill
+            elapsed = now - bucket.last_refill
+            bucket.tokens = min(
+                self.capacity,
+                bucket.tokens + elapsed * self.refill_rate
+            )
+            bucket.last_refill = now
+
+            # Evict stale buckets periodically
+            if now - self._last_cleanup > self._cleanup_interval:
+                self._evict_stale(now)
+
+            # Hard cap: if still over limit after cleanup, evict oldest entries
+            if len(self._buckets) >= self.max_clients:
+                # Remove ~10% oldest to avoid thrashing
+                evict_count = max(1, self.max_clients // 10)
+                oldest = sorted(self._buckets.items(), key=lambda x: x[1].last_refill)[:evict_count]
+                for k, _ in oldest:
+                    del self._buckets[k]
+                logger.warning(f"Rate limiter hard cap hit — evicted {evict_count} oldest buckets")
+
+            if bucket.tokens >= 1.0:
+                bucket.tokens -= 1.0
+                return True, 0.0
+            else:
+                # How long until 1 token refills
+                retry_after = (1.0 - bucket.tokens) / self.refill_rate
+                logger.warning(f"Rate limit hit for client '{client_id}'")
+                return False, retry_after
+
+    def _evict_stale(self, now: float, stale_after: float = 600.0) -> None:
+        """Remove buckets inactive for >10 minutes to prevent memory leak."""
+        stale = [k for k, b in self._buckets.items()
+                 if (now - b.last_refill) > stale_after]
+        for k in stale:
+            del self._buckets[k]
+        if stale:
+            logger.debug(f"Rate limiter: evicted {len(stale)} stale buckets")
+        self._last_cleanup = now
+
+
+# Singleton
+_limiter: RateLimiter | None = None
+
+def get_rate_limiter(rate: int = 60, per_seconds: int = 60, burst: int = 10) -> RateLimiter:
+    global _limiter
+    if _limiter is None:
+        _limiter = RateLimiter(rate=rate, per_seconds=per_seconds, burst=burst)
+    return _limiter

tokenmizer/checkpoints/manager.py

@@ -0,0 +1,383 @@
+"""
+Checkpoint Manager — creates and restores session checkpoints.
+
+Key fixes over V3:
+- Extracts from FULL message history, not just last 10
+- Tiered resume blocks (critical / standard / full)
+- Graph diff between checkpoints
+- Accurate token counting via tiktoken
+- SQLite persistence
+"""
+from __future__ import annotations
+
+import json
+import logging
+import time
+import uuid
+from dataclasses import dataclass
+from pathlib import Path
+from typing import TYPE_CHECKING, Optional
+
+if TYPE_CHECKING:
+    import sqlite3
+
+from tokenmizer.core.errors import CheckpointPersistError
+from tokenmizer.core.tokenizer import count_tokens
+from tokenmizer.graph_memory.graph import GraphMemory
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass
+class Checkpoint:
+    checkpoint_id: str
+    session_id: str
+    created_at: float
+    context_pct: float
+    trigger: str                   # "auto_threshold" | "manual" | "provider_switch"
+    message_count: int
+    graph_snapshot: dict           # full graph state at checkpoint time
+    graph_diff: dict               # diff from previous checkpoint
+    resume_critical: str           # ~100 tokens — must-know facts
+    resume_standard: str           # ~300 tokens — normal resume
+    resume_full: str               # ~600 tokens — deep resume
+    model: str = ""
+    next_action: str = ""
+
+    @property
+    def resume_tokens(self) -> int:
+        return count_tokens(self.resume_standard)
+
+
+class CheckpointManager:
+    """
+    Manages checkpoints for all sessions.
+    Stored in SQLite — survives restarts.
+    """
+
+    def __init__(self, storage_dir: str = "./checkpoints"):
+        self._dir = Path(storage_dir)
+        self._dir.mkdir(parents=True, exist_ok=True)
+        self._db_path = self._dir / "checkpoints.db"
+        self._safe_init_db()
+        self._prev_snapshots: dict[str, dict] = {}  # session_id → last snapshot
+
+    def _safe_init_db(self) -> None:
+        """Initialize DB, deleting corrupt file if necessary."""
+        try:
+            self._init_db()
+        except Exception:
+            logger.warning(f"Checkpoint DB corrupt or unreadable — recreating: {self._db_path}")
+            try:
+                self._db_path.unlink(missing_ok=True)
+            except Exception as del_err:
+                logger.error(f"Could not delete corrupt checkpoint DB: {del_err}")
+            try:
+                self._init_db()
+            except Exception as e:
+                logger.error(f"Cannot initialize checkpoint DB after cleanup: {e}")
+
+    def _db_connect(self) -> sqlite3.Connection:
+        """SQLite connection with WAL mode and timeout for concurrent safety."""
+        import sqlite3 as _sqlite3
+        conn = _sqlite3.connect(str(self._db_path), timeout=5.0, check_same_thread=False)
+        conn.execute("PRAGMA journal_mode=WAL")
+        conn.execute("PRAGMA synchronous=NORMAL")
+        return conn
+
+    def _init_db(self) -> None:
+        conn = self._db_connect()
+        try:
+            conn.execute("""
+                CREATE TABLE IF NOT EXISTS checkpoints (
+                    checkpoint_id TEXT PRIMARY KEY,
+                    session_id TEXT NOT NULL,
+                    created_at REAL NOT NULL,
+                    context_pct REAL,
+                    trigger TEXT,
+                    message_count INTEGER,
+                    data_json TEXT NOT NULL
+                )
+            """)
+            conn.execute("CREATE INDEX IF NOT EXISTS idx_ckpt_session ON checkpoints(session_id)")
+            conn.commit()
+        finally:
+            conn.close()
+
+    def create(
+        self,
+        session_id: str,
+        messages: list[dict],      # FULL message history — not just recent
+        graph: GraphMemory,
+        context_pct: float,
+        trigger: str = "auto_threshold",
+        model: str = "",
+    ) -> Checkpoint:
+        """Create a checkpoint from the current session state."""
+        checkpoint_id = f"ckpt_{uuid.uuid4().hex[:12]}"
+
+        # Force full extraction from ALL messages before snapshotting
+        graph.extract_from_messages(messages, incremental=True)
+
+        graph_snapshot = {
+            "nodes": [
+                {
+                    "id": n.id,
+                    "type": n.type.value,
+                    "label": n.label,
+                    "status": n.status.value,
+                    "summary": n.summary,
+                    "importance": n.importance,
+                }
+                for n in graph._nodes.values()
+                if not n._evicted
+            ],
+            "edge_count": len(graph._edges),
+        }
+
+        # Compute diff from previous checkpoint
+        prev = self._prev_snapshots.get(session_id, {"nodes": []})
+        graph_diff = self._compute_diff(prev, graph_snapshot)
+        self._prev_snapshots[session_id] = graph_snapshot
+
+        # Get last user message as next_action hint
+        next_action = ""
+        for msg in reversed(messages):
+            if msg.get("role") == "user":
+                next_action = msg.get("content", "")[:200]
+                break
+
+        # Build tiered resume blocks
+        resume_critical = self._build_critical(graph, next_action)
+        resume_standard = self._build_standard(graph, next_action)
+        resume_full = self._build_full(graph, messages, next_action)
+
+        ckpt = Checkpoint(
+            checkpoint_id=checkpoint_id,
+            session_id=session_id,
+            created_at=time.time(),
+            context_pct=context_pct,
+            trigger=trigger,
+            message_count=len(messages),
+            graph_snapshot=graph_snapshot,
+            graph_diff=graph_diff,
+            resume_critical=resume_critical,
+            resume_standard=resume_standard,
+            resume_full=resume_full,
+            model=model,
+            next_action=next_action,
+        )
+
+        self._save(ckpt)
+        logger.info(
+            f"Checkpoint {checkpoint_id}: session={session_id} "
+            f"msgs={len(messages)} nodes={len(graph._nodes)} "
+            f"context={context_pct:.0%} resume_tokens={ckpt.resume_tokens}"
+        )
+        return ckpt
+
+    def _build_critical(self, graph: GraphMemory, next_action: str) -> str:
+        """~100 tokens. Only open blockers + critical decisions."""
+        from tokenmizer.graph_memory.graph import NodeStatus, NodeType
+        lines = []
+
+        open_errors = [n for n in graph._nodes.values()
+                       if n.type == NodeType.ERROR and n.status == NodeStatus.FAILED]
+        if open_errors:
+            lines.append("OPEN BUGS: " + " | ".join(e.label for e in open_errors[:3]))
+
+        high_priority_tasks = [n for n in graph._nodes.values()
+                               if n.type == NodeType.TASK
+                               and n.status == NodeStatus.IN_PROGRESS
+                               and n.importance >= 0.8]
+        if high_priority_tasks:
+            lines.append("CRITICAL WIP: " + " | ".join(t.label for t in high_priority_tasks[:3]))
+
+        decisions = sorted(
+            [n for n in graph._nodes.values() if n.type == NodeType.DECISION],
+            key=lambda x: x.importance, reverse=True
+        )
+        if decisions:
+            lines.append("KEY DECISIONS: " + " | ".join(d.label for d in decisions[:3]))
+
+        if next_action:
+            lines.append(f"LAST REQUEST: {next_action[:100]}")
+
+        return "\n".join(lines)
+
+    def _build_standard(self, graph: GraphMemory, next_action: str) -> str:
+        """~300 tokens. Normal resume — goals, tasks, decisions, files."""
+        block = graph.to_context_block(token_budget=300)
+        if next_action:
+            block += f"\nContinue from: {next_action[:150]}"
+        return block
+
+    def _build_full(self, graph: GraphMemory, messages: list[dict], next_action: str) -> str:
+        """~600 tokens. Deep resume with environment, schemas, dependencies."""
+        from tokenmizer.graph_memory.graph import NodeStatus, NodeType
+        parts = [self._build_standard(graph, "")]
+
+        env_nodes = [n for n in graph._nodes.values() if n.type == NodeType.ENVIRONMENT]
+        if env_nodes:
+            parts.append("Environment: " + ", ".join(e.label for e in env_nodes[:8]))
+
+        dep_nodes = [n for n in graph._nodes.values() if n.type == NodeType.DEPENDENCY]
+        if dep_nodes:
+            parts.append("Dependencies: " + ", ".join(d.label for d in dep_nodes[:10]))
+
+        schema_nodes = [n for n in graph._nodes.values() if n.type == NodeType.SCHEMA]
+        if schema_nodes:
+            parts.append("Schemas: " + " | ".join(s.label for s in schema_nodes[:4]))
+
+        endpoint_nodes = [n for n in graph._nodes.values() if n.type == NodeType.ENDPOINT]
+        if endpoint_nodes:
+            parts.append("Endpoints: " + ", ".join(e.label for e in endpoint_nodes[:8]))
+
+        done_tasks = [n for n in graph._nodes.values()
+                      if n.type == NodeType.TASK and n.status == NodeStatus.COMPLETED]
+        done_tasks.sort(key=lambda x: x.updated_at, reverse=True)
+        if done_tasks:
+            parts.append("Recently completed: " + " | ".join(t.label for t in done_tasks[:6]))
+
+        if next_action:
+            parts.append(f"Continue from: {next_action[:150]}")
+
+        return "\n".join(parts)
+
+    def _compute_diff(self, prev: dict, current: dict) -> dict:
+        prev_nodes = {n["id"]: n for n in prev.get("nodes", [])}
+        curr_nodes = {n["id"]: n for n in current.get("nodes", [])}
+        return {
+            "added": [n for nid, n in curr_nodes.items() if nid not in prev_nodes],
+            "removed": [n for nid, n in prev_nodes.items() if nid not in curr_nodes],
+            "status_changed": [
+                {"id": nid, "from": prev_nodes[nid]["status"], "to": curr_nodes[nid]["status"]}
+                for nid in curr_nodes
+                if nid in prev_nodes
+                and curr_nodes[nid]["status"] != prev_nodes[nid]["status"]
+            ],
+        }
+
+    def _save(self, ckpt: Checkpoint) -> None:
+        """
+        Persist a checkpoint to SQLite.
+
+        FIXED: previously this caught Exception, logged it, and returned
+        None — silently. The caller (create()) had no way to know the
+        save failed, so callers (including the auto-checkpoint trigger and
+        the manual /api/checkpoint endpoint) would report a checkpoint as
+        successfully created when nothing was actually written to disk.
+        For a tool whose entire pitch is "never lose context," silently
+        losing the checkpoint on save failure is the worst possible
+        failure mode — the user trusts the safety net fired and finds out
+        otherwise only when they try to resume and there's nothing there.
+
+        Now raises CheckpointPersistError so callers can decide how to
+        handle it (the API layer already wraps checkpoint creation in
+        try/except and returns a proper 500 — this just makes that path
+        reachable instead of dead code).
+        """
+        try:
+            conn = self._db_connect()
+            try:
+                conn.execute(
+                    """INSERT OR REPLACE INTO checkpoints
+                       (checkpoint_id, session_id, created_at, context_pct, trigger, message_count, data_json)
+                       VALUES (?, ?, ?, ?, ?, ?, ?)""",
+                    (
+                        ckpt.checkpoint_id,
+                        ckpt.session_id,
+                        ckpt.created_at,
+                        ckpt.context_pct,
+                        ckpt.trigger,
+                        ckpt.message_count,
+                        json.dumps({
+                            "graph_snapshot": ckpt.graph_snapshot,
+                            "graph_diff": ckpt.graph_diff,
+                            "resume_critical": ckpt.resume_critical,
+                            "resume_standard": ckpt.resume_standard,
+                            "resume_full": ckpt.resume_full,
+                            "model": ckpt.model,
+                            "next_action": ckpt.next_action,
+                        }),
+                    ),
+                )
+                conn.commit()
+            finally:
+                conn.close()
+        except Exception as e:
+            logger.error(f"Checkpoint save failed for {ckpt.checkpoint_id}: {e}")
+            raise CheckpointPersistError(
+                f"Failed to persist checkpoint {ckpt.checkpoint_id} for "
+                f"session {ckpt.session_id}: {e}"
+            ) from e
+
+    def get_latest(self, session_id: str) -> Optional[Checkpoint]:
+        try:
+            conn = self._db_connect()
+            try:
+                row = conn.execute(
+                    """SELECT checkpoint_id, session_id, created_at, context_pct,
+                              trigger, message_count, data_json
+                       FROM checkpoints WHERE session_id=?
+                       ORDER BY created_at DESC LIMIT 1""",
+                    (session_id,),
+                ).fetchone()
+            finally:
+                conn.close()
+            if not row:
+                return None
+            return self._row_to_checkpoint(row)
+        except Exception as e:
+            logger.error(f"Checkpoint load failed: {e}")
+            return None
+
+    def list_checkpoints(self, session_id: str) -> list[dict]:
+        """
+        Returns checkpoint metadata for a session, newest first.
+
+        FIXED: previously a DB read failure here was indistinguishable from
+        "this session genuinely has zero checkpoints" — both returned `[]`
+        with zero logging. A caller (e.g. the /api/checkpoints/{session_id}
+        endpoint) would show an empty list to the user with no way to tell
+        whether checkpointing is broken or just hasn't run yet. We still
+        return [] on failure (changing the return type here would break the
+        API contract), but now we log it at error level so it's actually
+        visible in production instead of invisible by design.
+        """
+        try:
+            conn = self._db_connect()
+            try:
+                rows = conn.execute(
+                    """SELECT checkpoint_id, created_at, context_pct, trigger, message_count
+                       FROM checkpoints WHERE session_id=? ORDER BY created_at DESC""",
+                    (session_id,),
+                ).fetchall()
+            finally:
+                conn.close()
+            return [
+                {"checkpoint_id": r[0], "created_at": r[1], "context_pct": r[2],
+                 "trigger": r[3], "message_count": r[4]}
+                for r in rows
+            ]
+        except Exception as e:
+            logger.error(f"Checkpoint list query failed for session {session_id}: {e}")
+            return []
+
+    def _row_to_checkpoint(self, row) -> Checkpoint:
+        data = json.loads(row[6])
+        return Checkpoint(
+            checkpoint_id=row[0],
+            session_id=row[1],
+            created_at=row[2],
+            context_pct=row[3] or 0.0,
+            trigger=row[4] or "unknown",
+            message_count=row[5] or 0,
+            graph_snapshot=data.get("graph_snapshot", {}),
+            graph_diff=data.get("graph_diff", {}),
+            resume_critical=data.get("resume_critical", ""),
+            resume_standard=data.get("resume_standard", ""),
+            resume_full=data.get("resume_full", ""),
+            model=data.get("model", ""),
+            next_action=data.get("next_action", ""),
+        )

tokenmizer/cli.py

@@ -0,0 +1,153 @@
+"""TokenMizer CLI"""
+from __future__ import annotations
+
+from typing import Optional
+
+import typer
+from rich.console import Console
+from rich.panel import Panel
+
+app = typer.Typer(
+    name="tokenmizer",
+    help="🧠 TokenMizer — Never lose your AI context again.",
+    add_completion=False,
+)
+console = Console()
+
+
+@app.command()
+def serve(
+    host: str = typer.Option("0.0.0.0", help="Bind host"),
+    port: int = typer.Option(8000, help="Bind port"),
+    config: Optional[str] = typer.Option(None, help="Path to tokenmizer.yaml"),
+    reload: bool = typer.Option(False, help="Auto-reload on code changes (dev only)"),
+    workers: int = typer.Option(1, help="Number of worker processes"),
+):
+    """Start TokenMizer proxy + dashboard."""
+    import os
+
+    import uvicorn
+
+    if config:
+        os.environ["TOKENMIZER_CONFIG"] = config
+
+    console.print(Panel.fit(
+        "[bold green]🧠 TokenMizer[/bold green]\n"
+        f"[dim]Proxy:     http://{host}:{port}/v1/chat/completions[/dim]\n"
+        f"[dim]Dashboard: http://{host}:{port}[/dim]\n"
+        f"[dim]API Docs:  http://{host}:{port}/docs[/dim]\n"
+        f"[dim]Health:    http://{host}:{port}/health[/dim]",
+        border_style="green",
+    ))
+
+    uvicorn.run(
+        "tokenmizer.api.app:app",
+        host=host,
+        port=port,
+        reload=reload,
+        workers=workers if not reload else 1,
+        log_level="info",
+    )
+
+
+@app.command()
+def stats(
+    session_id: Optional[str] = typer.Argument(None, help="Session ID (optional)"),
+    server: str = typer.Option("http://localhost:8000", help="TokenMizer server URL"),
+    api_key: Optional[str] = typer.Option(None, envvar="TOKENMIZER_API_KEY"),
+):
+    """Print session/global analytics."""
+    import httpx
+
+    headers = {}
+    if api_key:
+        headers["Authorization"] = f"Bearer {api_key}"
+
+    url = f"{server}/api/stats"
+    if session_id:
+        url += f"?session_id={session_id}"
+
+    try:
+        data = httpx.get(url, headers=headers, timeout=5).json()
+    except Exception as e:
+        console.print(f"[red]Cannot reach server: {e}[/red]")
+        raise typer.Exit(1)
+
+    d = data.get("daily", {})
+    console.print(Panel.fit(
+        f"[bold]Daily Stats[/bold]\n"
+        f"[green]Requests:     {d.get('requests', 0):,}[/green]\n"
+        f"[green]Tokens saved: {d.get('tokens_saved', 0):,} ({d.get('savings_pct', 0):.1f}%)[/green]\n"
+        f"[yellow]Cost saved:   ${d.get('cost_saved_usd', 0):.4f}[/yellow]",
+        border_style="green",
+    ))
+
+
+@app.command()
+def checkpoint(
+    session_id: str = typer.Argument(..., help="Session ID to checkpoint"),
+    server: str = typer.Option("http://localhost:8000"),
+    api_key: Optional[str] = typer.Option(None, envvar="TOKENMIZER_API_KEY"),
+    level: str = typer.Option("standard", help="Resume level: critical | standard | full"),
+):
+    """Create a manual checkpoint and show resume context."""
+    import httpx
+
+    headers = {}
+    if api_key:
+        headers["Authorization"] = f"Bearer {api_key}"
+
+    r = httpx.post(
+        f"{server}/api/checkpoint?session_id={session_id}",
+        headers=headers,
+        timeout=30,
+    )
+    if r.status_code != 200:
+        console.print(f"[red]Error: {r.text}[/red]")
+        raise typer.Exit(1)
+
+    data = r.json()
+    console.print(Panel.fit(
+        f"[green]✅ Checkpoint created[/green]\n"
+        f"[dim]ID:            {data['checkpoint_id']}[/dim]\n"
+        f"[dim]Nodes:         {data.get('node_count', 0)}[/dim]\n"
+        f"[dim]Resume tokens: {data['resume_tokens']}[/dim]\n\n"
+        f"[bold]Resume context ({level}):[/bold]\n"
+        f"{data.get('resume_standard', '')}",
+        border_style="green",
+    ))
+
+
+@app.command()
+def resume(
+    session_id: str = typer.Argument(..., help="Session ID to resume"),
+    server: str = typer.Option("http://localhost:8000"),
+    level: str = typer.Option("standard", help="critical | standard | full"),
+    api_key: Optional[str] = typer.Option(None, envvar="TOKENMIZER_API_KEY"),
+):
+    """Get the resume context for a session checkpoint."""
+    import httpx
+
+    headers = {}
+    if api_key:
+        headers["Authorization"] = f"Bearer {api_key}"
+
+    r = httpx.get(
+        f"{server}/api/resume/{session_id}?level={level}",
+        headers=headers,
+        timeout=10,
+    )
+    if r.status_code == 404:
+        console.print(f"[yellow]No checkpoint found for session: {session_id}[/yellow]")
+        raise typer.Exit(1)
+
+    data = r.json()
+    console.print(Panel(
+        data["resume_context"],
+        title=f"[green]Resume — {session_id[:16]}... ({data['token_count']} tokens)[/green]",
+        border_style="green",
+    ))
+
+
+if __name__ == "__main__":
+    app()