| from __future__ import annotations |
|
|
| import json |
| import sqlite3 |
| from collections import defaultdict |
| from dataclasses import dataclass, field |
| from datetime import datetime, timedelta, timezone |
| from pathlib import Path |
| from typing import Any, Callable, Iterable |
| from urllib.parse import quote |
| from zoneinfo import ZoneInfo |
|
|
|
|
| @dataclass(slots=True) |
| class TokenBreakdown: |
| input_tokens: int = 0 |
| cached_input_tokens: int = 0 |
| output_tokens: int = 0 |
| reasoning_output_tokens: int = 0 |
| total_tokens: int = 0 |
|
|
| @classmethod |
| def from_usage(cls, usage: dict[str, Any]) -> "TokenBreakdown": |
| return cls( |
| input_tokens=_int_value(usage.get("input_tokens")), |
| cached_input_tokens=_int_value(usage.get("cached_input_tokens")), |
| output_tokens=_int_value(usage.get("output_tokens")), |
| reasoning_output_tokens=_int_value(usage.get("reasoning_output_tokens")), |
| total_tokens=_int_value(usage.get("total_tokens")), |
| ) |
|
|
| def add(self, other: "TokenBreakdown | None") -> None: |
| if other is None: |
| return |
| self.input_tokens += other.input_tokens |
| self.cached_input_tokens += other.cached_input_tokens |
| self.output_tokens += other.output_tokens |
| self.reasoning_output_tokens += other.reasoning_output_tokens |
| self.total_tokens += other.total_tokens |
|
|
|
|
| @dataclass(slots=True) |
| class RolloutStats: |
| path: Path | None |
| found: bool |
| token_breakdown: TokenBreakdown | None = None |
| model: str | None = None |
| cwd: str | None = None |
| turn_count: int = 0 |
| malformed_lines: int = 0 |
| warnings: list[str] = field(default_factory=list) |
|
|
|
|
| @dataclass(slots=True) |
| class SessionUsage: |
| thread_id: str |
| created_at: datetime |
| updated_at: datetime |
| cwd: str |
| model: str |
| model_provider: str |
| role: str |
| agent_role: str |
| agent_nickname: str |
| kind: str |
| parent_thread_id: str | None |
| title: str |
| tokens: int |
| rollout_path: Path | None |
| rollout_found: bool |
| rollout_breakdown: TokenBreakdown | None |
| turn_count: int |
|
|
| @property |
| def week_label(self) -> str: |
| year, week, _ = self.created_at.isocalendar() |
| return f"{year}-W{week:02d}" |
|
|
| @property |
| def day_label(self) -> str: |
| return self.created_at.date().isoformat() |
|
|
|
|
| @dataclass(slots=True) |
| class Bucket: |
| label: str |
| sessions: int = 0 |
| tokens: int = 0 |
| rollout: TokenBreakdown = field(default_factory=TokenBreakdown) |
|
|
| def add(self, session: SessionUsage) -> None: |
| self.sessions += 1 |
| self.tokens += session.tokens |
| self.rollout.add(session.rollout_breakdown) |
|
|
|
|
| @dataclass(slots=True) |
| class UsageReport: |
| generated_at: datetime |
| timezone_name: str |
| state_path: Path |
| sessions_dir: Path |
| sessions: list[SessionUsage] |
| daily: list[Bucket] |
| weekly: list[Bucket] |
| by_cwd: list[Bucket] |
| by_model: list[Bucket] |
| by_role: list[Bucket] |
| by_kind: list[Bucket] |
| top_sessions: list[SessionUsage] |
| rollout_totals: TokenBreakdown |
| warnings: list[str] |
| days_filter: int | None = None |
|
|
| @property |
| def total_tokens(self) -> int: |
| return sum(session.tokens for session in self.sessions) |
|
|
| @property |
| def parsed_rollout_count(self) -> int: |
| return sum(1 for session in self.sessions if session.rollout_breakdown is not None) |
|
|
| @property |
| def missing_rollout_count(self) -> int: |
| return sum(1 for session in self.sessions if not session.rollout_found) |
|
|
| @property |
| def date_range_label(self) -> str: |
| if not self.sessions: |
| return "n/a" |
| first = min(session.created_at for session in self.sessions) |
| last = max(session.created_at for session in self.sessions) |
| return f"{first.date().isoformat()} to {last.date().isoformat()}" |
|
|
|
|
| def build_report( |
| *, |
| state_path: Path, |
| sessions_dir: Path, |
| tz_name: str, |
| days: int | None = None, |
| top: int = 20, |
| now: datetime | None = None, |
| ) -> UsageReport: |
| tz = ZoneInfo(tz_name) |
| generated_at = (now or datetime.now(tz)).astimezone(tz) |
| rows, child_to_parent = _load_threads(state_path) |
| sessions: list[SessionUsage] = [] |
| warnings: list[str] = [] |
|
|
| if days is not None and days <= 0: |
| raise ValueError("--days must be greater than zero") |
| cutoff_date = None |
| if days is not None: |
| cutoff_date = generated_at.date() - timedelta(days=days - 1) |
|
|
| for row in rows: |
| created_at = _local_datetime(row.get("created_at_ms"), row.get("created_at"), tz) |
| if cutoff_date is not None and created_at.date() < cutoff_date: |
| continue |
|
|
| updated_at = _local_datetime(row.get("updated_at_ms"), row.get("updated_at"), tz) |
| thread_id = str(row["id"]) |
| rollout_path = _resolve_rollout_path(row.get("rollout_path"), sessions_dir, thread_id) |
| rollout = parse_rollout(rollout_path) |
| warnings.extend(rollout.warnings) |
|
|
| parent_thread_id = child_to_parent.get(thread_id) |
| kind = "subagent" if parent_thread_id else "main" |
| agent_role = str(row.get("agent_role") or "") |
| role = agent_role or ("main" if kind == "main" else "unknown") |
| model_provider = str(row.get("model_provider") or "unknown") |
| model = str(row.get("model") or rollout.model or "unknown") |
| cwd = str(row.get("cwd") or rollout.cwd or "unknown") |
|
|
| sessions.append( |
| SessionUsage( |
| thread_id=thread_id, |
| created_at=created_at, |
| updated_at=updated_at, |
| cwd=cwd, |
| model=model, |
| model_provider=model_provider, |
| role=role, |
| agent_role=agent_role, |
| agent_nickname=str(row.get("agent_nickname") or ""), |
| kind=kind, |
| parent_thread_id=parent_thread_id, |
| title=str(row.get("title") or ""), |
| tokens=_int_value(row.get("tokens_used")), |
| rollout_path=rollout.path, |
| rollout_found=rollout.found, |
| rollout_breakdown=rollout.token_breakdown, |
| turn_count=rollout.turn_count, |
| ) |
| ) |
|
|
| sessions.sort(key=lambda session: (session.created_at, session.thread_id)) |
| rollout_totals = TokenBreakdown() |
| for session in sessions: |
| rollout_totals.add(session.rollout_breakdown) |
|
|
| return UsageReport( |
| generated_at=generated_at, |
| timezone_name=tz_name, |
| state_path=Path(state_path).expanduser(), |
| sessions_dir=Path(sessions_dir).expanduser(), |
| sessions=sessions, |
| daily=_group_sessions(sessions, lambda session: session.day_label, sort_by_label=True), |
| weekly=_group_sessions(sessions, lambda session: session.week_label, sort_by_label=True), |
| by_cwd=_group_sessions(sessions, lambda session: _compact_path(session.cwd)), |
| by_model=_group_sessions(sessions, lambda session: session.model or "unknown"), |
| by_role=_group_sessions(sessions, lambda session: session.role or "unknown"), |
| by_kind=_group_sessions(sessions, lambda session: session.kind), |
| top_sessions=sorted(sessions, key=lambda session: session.tokens, reverse=True)[:top], |
| rollout_totals=rollout_totals, |
| warnings=warnings, |
| days_filter=days, |
| ) |
|
|
|
|
| def parse_rollout(path: Path | None) -> RolloutStats: |
| if path is None: |
| return RolloutStats(path=None, found=False, warnings=["No rollout path for a session"]) |
|
|
| expanded = path.expanduser() |
| if not expanded.exists(): |
| return RolloutStats( |
| path=expanded, |
| found=False, |
| warnings=[f"Missing rollout file: {expanded}"], |
| ) |
|
|
| stats = RolloutStats(path=expanded, found=True) |
| try: |
| with expanded.open("r", encoding="utf-8") as handle: |
| for line_number, line in enumerate(handle, start=1): |
| if not line.strip(): |
| continue |
| try: |
| event = json.loads(line) |
| except json.JSONDecodeError: |
| stats.malformed_lines += 1 |
| continue |
| _read_rollout_event(stats, event) |
| except OSError as exc: |
| stats.found = False |
| stats.warnings.append(f"Could not read rollout file {expanded}: {exc}") |
| return stats |
|
|
| if stats.malformed_lines: |
| stats.warnings.append( |
| f"Malformed JSONL lines in {expanded}: {stats.malformed_lines}" |
| ) |
| return stats |
|
|
|
|
| def _read_rollout_event(stats: RolloutStats, event: dict[str, Any]) -> None: |
| event_type = event.get("type") |
| payload = event.get("payload") |
| if not isinstance(payload, dict): |
| payload = {} |
|
|
| if event_type == "session_meta": |
| stats.cwd = stats.cwd or _string_or_none(payload.get("cwd")) |
| return |
|
|
| if event_type == "turn_context": |
| stats.model = stats.model or _string_or_none(payload.get("model")) |
| stats.cwd = stats.cwd or _string_or_none(payload.get("cwd")) |
| return |
|
|
| if event_type != "event_msg": |
| return |
|
|
| payload_type = payload.get("type") |
| if payload_type == "task_started": |
| stats.turn_count += 1 |
| return |
| if payload_type != "token_count": |
| return |
|
|
| info = payload.get("info") |
| if not isinstance(info, dict): |
| return |
| usage = info.get("total_token_usage") |
| if isinstance(usage, dict): |
| stats.token_breakdown = TokenBreakdown.from_usage(usage) |
|
|
|
|
| def _load_threads(state_path: Path) -> tuple[list[dict[str, Any]], dict[str, str]]: |
| expanded = Path(state_path).expanduser() |
| if not expanded.exists(): |
| raise FileNotFoundError(f"State database not found: {expanded}") |
|
|
| uri = f"file:{quote(str(expanded), safe='/')}?mode=ro" |
| with sqlite3.connect(uri, uri=True) as conn: |
| conn.row_factory = sqlite3.Row |
| conn.execute("PRAGMA query_only = ON") |
| rows = [ |
| dict(row) |
| for row in conn.execute( |
| """ |
| SELECT |
| id, |
| rollout_path, |
| created_at, |
| updated_at, |
| created_at_ms, |
| updated_at_ms, |
| model_provider, |
| cwd, |
| title, |
| tokens_used, |
| agent_nickname, |
| agent_role, |
| model |
| FROM threads |
| ORDER BY created_at_ms ASC, id ASC |
| """ |
| ) |
| ] |
| edges = [ |
| dict(edge) |
| for edge in conn.execute( |
| """ |
| SELECT parent_thread_id, child_thread_id |
| FROM thread_spawn_edges |
| """ |
| ) |
| ] |
| return rows, {str(edge["child_thread_id"]): str(edge["parent_thread_id"]) for edge in edges} |
|
|
|
|
| def _resolve_rollout_path( |
| rollout_path: Any, sessions_dir: Path, thread_id: str |
| ) -> Path | None: |
| sessions_root = Path(sessions_dir).expanduser() |
| if rollout_path: |
| candidate = Path(str(rollout_path)).expanduser() |
| if not candidate.is_absolute(): |
| candidate = sessions_root / candidate |
| if candidate.exists(): |
| return candidate |
|
|
| if sessions_root.exists(): |
| matches = sorted(sessions_root.rglob(f"*{thread_id}.jsonl")) |
| if matches: |
| return matches[0] |
|
|
| if rollout_path: |
| return candidate |
| return None |
|
|
|
|
| def _group_sessions( |
| sessions: Iterable[SessionUsage], |
| key: Callable[[SessionUsage], str], |
| *, |
| sort_by_label: bool = False, |
| ) -> list[Bucket]: |
| buckets: dict[str, Bucket] = {} |
| for session in sessions: |
| label = key(session) |
| bucket = buckets.setdefault(label, Bucket(label=label)) |
| bucket.add(session) |
|
|
| values = list(buckets.values()) |
| if sort_by_label: |
| values.sort(key=lambda bucket: bucket.label) |
| else: |
| values.sort(key=lambda bucket: bucket.tokens, reverse=True) |
| return values |
|
|
|
|
| def _local_datetime(ms_value: Any, seconds_value: Any, tz: ZoneInfo) -> datetime: |
| if ms_value is not None: |
| timestamp = _int_value(ms_value) / 1000 |
| else: |
| timestamp = _int_value(seconds_value) |
| return datetime.fromtimestamp(timestamp, timezone.utc).astimezone(tz) |
|
|
|
|
| def _compact_path(value: str) -> str: |
| if not value: |
| return "unknown" |
| try: |
| path = Path(value).expanduser() |
| home = Path.home() |
| if path.is_absolute(): |
| relative = path.relative_to(home) |
| return "~" if str(relative) == "." else "~/" + str(relative) |
| return str(path) |
| except ValueError: |
| return value |
|
|
|
|
| def _int_value(value: Any) -> int: |
| if value is None: |
| return 0 |
| try: |
| return int(value) |
| except (TypeError, ValueError): |
| return 0 |
|
|
|
|
| def _string_or_none(value: Any) -> str | None: |
| if value is None: |
| return None |
| text = str(value) |
| return text if text else None |
|
|
