token-tracker 0.1.0__py3-none-any.whl

src/adapters/claude.py

@@ -0,0 +1,161 @@
+import json
+import os
+from datetime import datetime, timezone
+from pathlib import Path
+
+from .types import AgentInfo, UsageEntry
+
+CLAUDE_DIRS = [
+    os.path.expanduser("~/.claude/projects"),
+    os.path.expanduser("~/.config/claude/projects"),
+]
+
+
+def detect() -> AgentInfo | None:
+    for d in _get_claude_dirs():
+        if Path(d).is_dir():
+            return AgentInfo(
+                id="claude-code",
+                name="Claude Code",
+                data_dir=d,
+                installed=True,
+            )
+    return None
+
+
+def load_entries(hours_back: int = 0) -> list[UsageEntry]:
+    entries: list[UsageEntry] = []
+    seen: set[str] = set()
+    cutoff = None
+    if hours_back > 0:
+        from datetime import timedelta
+        cutoff = datetime.now(timezone.utc) - timedelta(hours=hours_back)
+
+    for base_dir in _get_claude_dirs():
+        base = Path(base_dir)
+        if not base.is_dir():
+            continue
+        for jsonl_path in base.rglob("*.jsonl"):
+            fallback_project = _extract_project_from_dir(jsonl_path, base)
+            _parse_jsonl(jsonl_path, fallback_project, entries, seen, cutoff)
+
+    entries.sort(key=lambda e: e.timestamp)
+    return entries
+
+
+def _get_claude_dirs() -> list[str]:
+    dirs = list(CLAUDE_DIRS)
+    env = os.environ.get("CLAUDE_CONFIG_DIR")
+    if env:
+        for p in env.split(","):
+            projects_dir = os.path.join(p.strip(), "projects")
+            if projects_dir not in dirs:
+                dirs.insert(0, projects_dir)
+    return dirs
+
+
+def _project_from_cwd(cwd: str) -> str:
+    home = os.path.expanduser("~")
+    if cwd.startswith(home):
+        rel = cwd[len(home):].strip(os.sep)
+    else:
+        rel = cwd.strip(os.sep)
+    parts = rel.split(os.sep)
+    return parts[-1] if parts and parts[-1] else rel or "unknown"
+
+
+def _extract_project_from_dir(jsonl_path: Path, base: Path) -> str:
+    rel = jsonl_path.relative_to(base)
+    project_dir = str(rel.parts[0]) if rel.parts else "unknown"
+    decoded = project_dir.replace("-", os.sep).strip(os.sep)
+    home = os.path.expanduser("~").strip(os.sep)
+    if decoded.startswith(home):
+        decoded = decoded[len(home):].strip(os.sep)
+    parts = decoded.split(os.sep)
+    return parts[-1] if parts else "unknown"
+
+
+def _parse_jsonl(
+    path: Path,
+    project: str,
+    entries: list[UsageEntry],
+    seen: set[str],
+    cutoff: datetime | None,
+) -> None:
+    try:
+        with open(path, "r", encoding="utf-8") as f:
+            for line in f:
+                line = line.strip()
+                if not line:
+                    continue
+                try:
+                    data = json.loads(line)
+                except json.JSONDecodeError:
+                    continue
+
+                if data.get("type") != "assistant":
+                    continue
+
+                entry = _parse_assistant_entry(data, project)
+                if entry is None:
+                    continue
+
+                if cutoff and entry.timestamp < cutoff:
+                    continue
+
+                if entry.dedup_key in seen:
+                    continue
+                seen.add(entry.dedup_key)
+
+                entries.append(entry)
+    except (OSError, PermissionError):
+        pass
+
+
+def _parse_assistant_entry(data: dict, project: str) -> UsageEntry | None:
+    message = data.get("message")
+    if not message or not isinstance(message, dict):
+        return None
+
+    usage = message.get("usage")
+    if not usage or not isinstance(usage, dict):
+        return None
+
+    input_tokens = usage.get("input_tokens") or 0
+    output_tokens = usage.get("output_tokens") or 0
+    cache_creation = usage.get("cache_creation_input_tokens") or 0
+    cache_read = usage.get("cache_read_input_tokens") or 0
+
+    if input_tokens == 0 and output_tokens == 0 and cache_creation == 0 and cache_read == 0:
+        return None
+
+    timestamp_str = data.get("timestamp", "")
+    try:
+        ts = datetime.fromisoformat(timestamp_str.replace("Z", "+00:00"))
+    except (ValueError, AttributeError):
+        return None
+
+    message_id = message.get("id", "")
+    request_id = data.get("requestId") or ""
+    model = message.get("model", "unknown")
+    session_id = data.get("sessionId", "")
+    cost_usd = data.get("costUSD")
+
+    cwd = data.get("cwd", "")
+    if cwd:
+        project = _project_from_cwd(cwd)
+
+    return UsageEntry(
+        timestamp=ts,
+        session_id=session_id,
+        message_id=message_id,
+        request_id=request_id,
+        model=model,
+        input_tokens=input_tokens,
+        output_tokens=output_tokens,
+        cache_creation_tokens=cache_creation,
+        cache_read_tokens=cache_read,
+        cost_usd=cost_usd,
+        project=project,
+        agent_id="claude-code",
+    )

src/adapters/codex.py

@@ -0,0 +1,226 @@
+import json
+import os
+import sqlite3
+from datetime import datetime, timedelta, timezone
+from pathlib import Path
+
+from .types import AgentInfo, RateLimits, UsageEntry
+
+CODEX_DIR = os.path.expanduser("~/.codex")
+SESSIONS_DIR = os.path.join(CODEX_DIR, "sessions")
+STATE_DB = os.path.join(CODEX_DIR, "state_5.sqlite")
+
+
+def detect() -> AgentInfo | None:
+    if Path(SESSIONS_DIR).is_dir():
+        return AgentInfo(
+            id="codex",
+            name="Codex",
+            data_dir=SESSIONS_DIR,
+            installed=True,
+        )
+    return None
+
+
+def load_entries(hours_back: int = 0) -> list[UsageEntry]:
+    entries: list[UsageEntry] = []
+    seen: set[str] = set()
+    cutoff = None
+    if hours_back > 0:
+        cutoff = datetime.now(timezone.utc) - timedelta(hours=hours_back)
+
+    models = _load_thread_models()
+
+    sessions_path = Path(SESSIONS_DIR)
+    if not sessions_path.is_dir():
+        return entries
+
+    for jsonl_path in sessions_path.rglob("*.jsonl"):
+        _parse_jsonl(jsonl_path, models, entries, seen, cutoff)
+
+    entries.sort(key=lambda e: e.timestamp)
+    return entries
+
+
+def _load_thread_models() -> dict[str, str]:
+    if not os.path.exists(STATE_DB):
+        return {}
+    try:
+        conn = sqlite3.connect(f"file:{STATE_DB}?mode=ro", uri=True)
+        rows = conn.execute("SELECT id, model FROM threads WHERE model IS NOT NULL").fetchall()
+        conn.close()
+        return {row[0]: row[1] for row in rows}
+    except (sqlite3.Error, OSError):
+        return {}
+
+
+def load_rate_limits() -> RateLimits | None:
+    sessions_path = Path(SESSIONS_DIR)
+    if not sessions_path.is_dir():
+        return None
+
+    jsonl_files = sorted(sessions_path.rglob("*.jsonl"), key=lambda p: p.stat().st_mtime, reverse=True)
+    models = _load_thread_models()
+
+    for path in jsonl_files[:5]:
+        rl = _extract_rate_limits(path, models)
+        if rl:
+            return rl
+    return None
+
+
+def _extract_rate_limits(path: Path, models: dict[str, str]) -> RateLimits | None:
+    session_id = ""
+    last_rl = None
+    try:
+        with open(path, "r", encoding="utf-8") as f:
+            for line in f:
+                line = line.strip()
+                if not line:
+                    continue
+                try:
+                    data = json.loads(line)
+                except json.JSONDecodeError:
+                    continue
+                if data.get("type") == "session_meta":
+                    session_id = data.get("payload", {}).get("id", "")
+                if data.get("type") != "event_msg":
+                    continue
+                payload = data.get("payload", {})
+                if payload.get("type") != "token_count":
+                    continue
+                rl = payload.get("rate_limits")
+                if rl:
+                    last_rl = (rl, data.get("timestamp", ""), session_id)
+    except (OSError, PermissionError):
+        return None
+
+    if not last_rl:
+        return None
+
+    rl, ts, sid = last_rl
+    primary = rl.get("primary") or {}
+    secondary = rl.get("secondary") or {}
+
+    five_pct = primary.get("used_percent")
+    five_reset = primary.get("resets_at")
+    seven_pct = secondary.get("used_percent")
+    seven_reset = secondary.get("resets_at")
+
+    now_ts = datetime.now(timezone.utc).timestamp()
+    if five_reset and five_reset < now_ts:
+        five_pct = 0.0
+    if seven_reset and seven_reset < now_ts:
+        seven_pct = 0.0
+
+    if five_pct is None and seven_pct is None:
+        return None
+
+    model_name = models.get(sid, "")
+
+    return RateLimits(
+        five_hour_pct=five_pct,
+        five_hour_resets_at=five_reset,
+        seven_day_pct=seven_pct,
+        seven_day_resets_at=seven_reset,
+        model=model_name,
+        updated_at=ts,
+    )
+
+
+def _project_from_cwd(cwd: str) -> str:
+    home = os.path.expanduser("~")
+    if cwd.startswith(home):
+        rel = cwd[len(home):].strip(os.sep)
+    else:
+        rel = cwd.strip(os.sep)
+    parts = rel.split(os.sep)
+    return parts[-1] if parts and parts[-1] else rel or "unknown"
+
+
+def _parse_jsonl(
+    path: Path,
+    models: dict[str, str],
+    entries: list[UsageEntry],
+    seen: set[str],
+    cutoff: datetime | None,
+) -> None:
+    session_id = ""
+    session_ts = ""
+    project = "unknown"
+    model = "unknown"
+    last_usage = None
+    msg_count = 0
+
+    try:
+        with open(path, "r", encoding="utf-8") as f:
+            for line in f:
+                line = line.strip()
+                if not line:
+                    continue
+                try:
+                    data = json.loads(line)
+                except json.JSONDecodeError:
+                    continue
+
+                row_type = data.get("type")
+
+                if row_type == "session_meta":
+                    payload = data.get("payload", {})
+                    session_id = payload.get("id", "")
+                    session_ts = payload.get("timestamp", "")
+                    cwd = payload.get("cwd", "")
+                    if cwd:
+                        project = _project_from_cwd(cwd)
+                    model = models.get(session_id, "unknown")
+                    continue
+
+                if row_type != "event_msg":
+                    continue
+
+                payload = data.get("payload", {})
+                if payload.get("type") == "token_count":
+                    info = payload.get("info")
+                    if info and info.get("total_token_usage"):
+                        last_usage = info["total_token_usage"]
+                        msg_count += 1
+    except (OSError, PermissionError):
+        return
+
+    if not last_usage or not session_id:
+        return
+
+    cached = last_usage.get("cached_input_tokens", 0)
+    input_tokens = last_usage.get("input_tokens", 0) - cached
+    output_tokens = last_usage.get("output_tokens", 0) + last_usage.get("reasoning_output_tokens", 0)
+
+    if input_tokens == 0 and output_tokens == 0:
+        return
+
+    try:
+        ts = datetime.fromisoformat(session_ts.replace("Z", "+00:00"))
+    except (ValueError, AttributeError):
+        return
+
+    if cutoff and ts < cutoff:
+        return
+
+    if session_id in seen:
+        return
+    seen.add(session_id)
+
+    entries.append(UsageEntry(
+        timestamp=ts,
+        session_id=session_id,
+        message_id=session_id,
+        request_id="",
+        model=model,
+        input_tokens=input_tokens,
+        output_tokens=output_tokens,
+        cache_creation_tokens=0,
+        cache_read_tokens=cached,
+        cost_usd=None,
+        project=project,
+        agent_id="codex",
+        message_count=msg_count,
+    ))

src/adapters/rate_limits.py

@@ -0,0 +1,48 @@
+import json
+import os
+from datetime import datetime, timezone
+
+from .types import RateLimits
+
+STATUS_FILE = os.path.expanduser("~/.claude/tt-status.json")
+
+
+def load_rate_limits() -> RateLimits | None:
+    if not os.path.exists(STATUS_FILE):
+        return None
+
+    try:
+        with open(STATUS_FILE, "r", encoding="utf-8") as f:
+            data = json.load(f)
+    except (json.JSONDecodeError, OSError):
+        return None
+
+    rl = data.get("rate_limits") or {}
+    five = rl.get("five_hour") or {}
+    seven = rl.get("seven_day") or {}
+
+    now_ts = datetime.now(timezone.utc).timestamp()
+    five_pct = five.get("used_percentage")
+    five_reset = five.get("resets_at")
+    if five_reset and five_reset < now_ts:
+        five_pct = 0.0
+
+    seven_pct = seven.get("used_percentage")
+    seven_reset = seven.get("resets_at")
+    if seven_reset and seven_reset < now_ts:
+        seven_pct = 0.0
+
+    model_info = data.get("model") or {}
+    model_name = model_info.get("display_name") or model_info.get("id") or ""
+
+    if five_pct is None and seven_pct is None and not model_name:
+        return None
+
+    return RateLimits(
+        five_hour_pct=five_pct,
+        five_hour_resets_at=five_reset,
+        seven_day_pct=seven_pct,
+        seven_day_resets_at=seven_reset,
+        model=model_name,
+        updated_at=data.get("_received_at", ""),
+    )

src/adapters/registry.py

@@ -0,0 +1,11 @@
+from .types import AgentInfo
+from . import claude, codex
+
+
+def detect_agents() -> list[AgentInfo]:
+    agents: list[AgentInfo] = []
+    for detector in [claude.detect, codex.detect]:
+        info = detector()
+        if info:
+            agents.append(info)
+    return agents

src/adapters/types.py

@@ -0,0 +1,129 @@
+from dataclasses import dataclass, field
+from datetime import datetime
+
+
+@dataclass
+class UsageEntry:
+    timestamp: datetime
+    session_id: str
+    message_id: str
+    request_id: str
+    model: str
+    input_tokens: int
+    output_tokens: int
+    cache_creation_tokens: int
+    cache_read_tokens: int
+    cost_usd: float | None
+    project: str
+    agent_id: str
+    message_count: int = 1
+
+    @property
+    def total_tokens(self) -> int:
+        return self.input_tokens + self.output_tokens + self.cache_creation_tokens + self.cache_read_tokens
+
+    @property
+    def dedup_key(self) -> str:
+        return f"{self.message_id}:{self.request_id}"
+
+
+@dataclass
+class AgentInfo:
+    id: str
+    name: str
+    data_dir: str
+    installed: bool
+
+
+@dataclass
+class DailyStats:
+    date: str
+    input_tokens: int = 0
+    output_tokens: int = 0
+    cache_creation_tokens: int = 0
+    cache_read_tokens: int = 0
+    total_tokens: int = 0
+    cost_usd: float = 0.0
+    session_count: int = 0
+    message_count: int = 0
+    models: dict[str, int] = field(default_factory=dict)
+
+
+@dataclass
+class WeeklyStats:
+    week: str
+    week_start: str = ""
+    week_end: str = ""
+    input_tokens: int = 0
+    output_tokens: int = 0
+    cache_creation_tokens: int = 0
+    cache_read_tokens: int = 0
+    total_tokens: int = 0
+    cost_usd: float = 0.0
+    session_count: int = 0
+    message_count: int = 0
+    models: dict[str, int] = field(default_factory=dict)
+
+
+@dataclass
+class SessionStats:
+    session_id: str
+    project: str
+    model: str
+    start_time: datetime
+    end_time: datetime
+    duration_minutes: float
+    input_tokens: int = 0
+    output_tokens: int = 0
+    cache_creation_tokens: int = 0
+    cache_read_tokens: int = 0
+    total_tokens: int = 0
+    cost_usd: float = 0.0
+    message_count: int = 0
+
+
+@dataclass
+class MonthlyStats:
+    month: str
+    input_tokens: int = 0
+    output_tokens: int = 0
+    cache_creation_tokens: int = 0
+    cache_read_tokens: int = 0
+    total_tokens: int = 0
+    cost_usd: float = 0.0
+    session_count: int = 0
+    message_count: int = 0
+    models: dict[str, int] = field(default_factory=dict)
+
+
+@dataclass
+class RateLimits:
+    five_hour_pct: float | None = None
+    five_hour_resets_at: int | None = None
+    seven_day_pct: float | None = None
+    seven_day_resets_at: int | None = None
+    model: str = ""
+    updated_at: str = ""
+
+
+@dataclass
+class P90Limits:
+    token_limit: int = 0
+    cost_limit: float = 0.0
+    message_limit: int = 0
+
+
+@dataclass
+class SessionBlock:
+    start_time: datetime
+    end_time: datetime
+    entries: list[UsageEntry] = field(default_factory=list)
+    input_tokens: int = 0
+    output_tokens: int = 0
+    cache_creation_tokens: int = 0
+    cache_read_tokens: int = 0
+    total_tokens: int = 0
+    cost_usd: float = 0.0
+    is_active: bool = False
+    burn_rate: float = 0.0
+    is_gap: bool = False