| from __future__ import annotations |
|
|
| import json |
| import sqlite3 |
| import tempfile |
| import unittest |
| from datetime import datetime |
| from pathlib import Path |
| from zoneinfo import ZoneInfo |
|
|
| from codex_usage_profiler.core import build_report |
| from codex_usage_profiler.render import render_html, render_markdown |
|
|
|
|
| class UsageProfilerTests(unittest.TestCase): |
| def test_aggregates_threads_edges_and_rollout_breakdowns(self) -> None: |
| with tempfile.TemporaryDirectory() as tmp: |
| root = Path(tmp) |
| state = root / "state_5.sqlite" |
| sessions_dir = root / "sessions" |
| sessions_dir.mkdir() |
| main_rollout = sessions_dir / "rollout-main.jsonl" |
| child_rollout = sessions_dir / "rollout-child.jsonl" |
|
|
| _write_rollout( |
| main_rollout, |
| [ |
| {"type": "turn_context", "payload": {"model": "gpt-5.4", "cwd": "/home/test/project"}}, |
| {"type": "event_msg", "payload": {"type": "task_started"}}, |
| { |
| "type": "event_msg", |
| "payload": { |
| "type": "token_count", |
| "info": { |
| "total_token_usage": { |
| "input_tokens": 70, |
| "cached_input_tokens": 20, |
| "output_tokens": 25, |
| "reasoning_output_tokens": 5, |
| "total_tokens": 100, |
| } |
| }, |
| }, |
| }, |
| ], |
| malformed=True, |
| ) |
| _write_rollout( |
| child_rollout, |
| [ |
| {"type": "turn_context", "payload": {"model": "gpt-5.4-mini", "cwd": "/home/test/project"}}, |
| { |
| "type": "event_msg", |
| "payload": { |
| "type": "token_count", |
| "info": { |
| "total_token_usage": { |
| "input_tokens": 35, |
| "cached_input_tokens": 10, |
| "output_tokens": 12, |
| "reasoning_output_tokens": 3, |
| "total_tokens": 50, |
| } |
| }, |
| }, |
| }, |
| ], |
| ) |
| _create_state( |
| state, |
| [ |
| { |
| "id": "main", |
| "rollout_path": str(main_rollout), |
| "created_at_ms": _epoch_ms("2026-04-20T22:30:00+00:00"), |
| "updated_at_ms": _epoch_ms("2026-04-20T23:00:00+00:00"), |
| "model_provider": "openai", |
| "cwd": "/home/test/project", |
| "title": "Main work", |
| "tokens_used": 100, |
| "agent_nickname": "", |
| "agent_role": "", |
| "model": None, |
| }, |
| { |
| "id": "child", |
| "rollout_path": str(child_rollout), |
| "created_at_ms": _epoch_ms("2026-04-21T10:00:00+00:00"), |
| "updated_at_ms": _epoch_ms("2026-04-21T10:05:00+00:00"), |
| "model_provider": "openai", |
| "cwd": "/home/test/project", |
| "title": "Explore", |
| "tokens_used": 50, |
| "agent_nickname": "Noether", |
| "agent_role": "explorer", |
| "model": "gpt-5.4-mini", |
| }, |
| ], |
| [("main", "child")], |
| ) |
|
|
| report = build_report( |
| state_path=state, |
| sessions_dir=sessions_dir, |
| tz_name="Europe/Berlin", |
| now=datetime(2026, 4, 22, 12, tzinfo=ZoneInfo("Europe/Berlin")), |
| ) |
|
|
| self.assertEqual(report.total_tokens, 150) |
| self.assertEqual(report.daily[0].label, "2026-04-21") |
| self.assertEqual(report.daily[0].tokens, 150) |
| self.assertEqual(report.weekly[0].label, "2026-W17") |
| self.assertEqual({bucket.label: bucket.tokens for bucket in report.by_kind}, {"main": 100, "subagent": 50}) |
| self.assertEqual({bucket.label: bucket.tokens for bucket in report.by_role}, {"main": 100, "explorer": 50}) |
| self.assertEqual({bucket.label: bucket.tokens for bucket in report.by_model}, {"gpt-5.4": 100, "gpt-5.4-mini": 50}) |
| self.assertEqual(report.rollout_totals.total_tokens, 150) |
| self.assertEqual(report.rollout_totals.cached_input_tokens, 30) |
| self.assertEqual(len(report.warnings), 1) |
|
|
| def test_days_filter_uses_local_calendar_days(self) -> None: |
| with tempfile.TemporaryDirectory() as tmp: |
| root = Path(tmp) |
| state = root / "state_5.sqlite" |
| sessions_dir = root / "sessions" |
| sessions_dir.mkdir() |
| _create_state( |
| state, |
| [ |
| { |
| "id": "old", |
| "rollout_path": "", |
| "created_at_ms": _epoch_ms("2026-04-18T12:00:00+00:00"), |
| "updated_at_ms": _epoch_ms("2026-04-18T12:00:00+00:00"), |
| "model_provider": "openai", |
| "cwd": "/tmp/a", |
| "title": "Old", |
| "tokens_used": 99, |
| "agent_nickname": "", |
| "agent_role": "", |
| "model": "gpt-5.4", |
| }, |
| { |
| "id": "new", |
| "rollout_path": "", |
| "created_at_ms": _epoch_ms("2026-04-21T12:00:00+00:00"), |
| "updated_at_ms": _epoch_ms("2026-04-21T12:00:00+00:00"), |
| "model_provider": "openai", |
| "cwd": "/tmp/a", |
| "title": "New", |
| "tokens_used": 10, |
| "agent_nickname": "", |
| "agent_role": "", |
| "model": "gpt-5.4", |
| }, |
| ], |
| [], |
| ) |
|
|
| report = build_report( |
| state_path=state, |
| sessions_dir=sessions_dir, |
| tz_name="Europe/Berlin", |
| days=2, |
| now=datetime(2026, 4, 22, 12, tzinfo=ZoneInfo("Europe/Berlin")), |
| ) |
|
|
| self.assertEqual(len(report.sessions), 1) |
| self.assertEqual(report.sessions[0].thread_id, "new") |
| self.assertEqual(report.total_tokens, 10) |
|
|
| def test_renderers_include_expected_sections_and_escape_html(self) -> None: |
| with tempfile.TemporaryDirectory() as tmp: |
| root = Path(tmp) |
| state = root / "state_5.sqlite" |
| sessions_dir = root / "sessions" |
| sessions_dir.mkdir() |
| _create_state( |
| state, |
| [ |
| { |
| "id": "s1", |
| "rollout_path": "", |
| "created_at_ms": _epoch_ms("2026-04-21T12:00:00+00:00"), |
| "updated_at_ms": _epoch_ms("2026-04-21T12:00:00+00:00"), |
| "model_provider": "openai", |
| "cwd": "/tmp/a", |
| "title": "<unsafe>", |
| "tokens_used": 10, |
| "agent_nickname": "", |
| "agent_role": "", |
| "model": "gpt-5.4", |
| } |
| ], |
| [], |
| ) |
| report = build_report( |
| state_path=state, |
| sessions_dir=sessions_dir, |
| tz_name="Europe/Berlin", |
| now=datetime(2026, 4, 22, 12, tzinfo=ZoneInfo("Europe/Berlin")), |
| ) |
|
|
| markdown = render_markdown(report) |
| html = render_html(report) |
| self.assertIn("## Daily Tokens", markdown) |
| self.assertIn("## Top Sessions", markdown) |
| self.assertNotIn("<unsafe>", html) |
| self.assertNotIn("<unsafe>", html) |
|
|
|
|
| def _create_state(state: Path, rows: list[dict[str, object]], edges: list[tuple[str, str]]) -> None: |
| with sqlite3.connect(state) as conn: |
| conn.execute( |
| """ |
| CREATE TABLE threads ( |
| id TEXT PRIMARY KEY, |
| rollout_path TEXT NOT NULL, |
| created_at INTEGER NOT NULL, |
| updated_at INTEGER NOT NULL, |
| created_at_ms INTEGER, |
| updated_at_ms INTEGER, |
| model_provider TEXT NOT NULL, |
| cwd TEXT NOT NULL, |
| title TEXT NOT NULL, |
| tokens_used INTEGER NOT NULL, |
| agent_nickname TEXT, |
| agent_role TEXT, |
| model TEXT |
| ) |
| """ |
| ) |
| conn.execute( |
| """ |
| CREATE TABLE thread_spawn_edges ( |
| parent_thread_id TEXT NOT NULL, |
| child_thread_id TEXT NOT NULL PRIMARY KEY |
| ) |
| """ |
| ) |
| for row in rows: |
| conn.execute( |
| """ |
| INSERT INTO threads ( |
| id, |
| rollout_path, |
| created_at, |
| updated_at, |
| created_at_ms, |
| updated_at_ms, |
| model_provider, |
| cwd, |
| title, |
| tokens_used, |
| agent_nickname, |
| agent_role, |
| model |
| ) |
| VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) |
| """, |
| ( |
| row["id"], |
| row["rollout_path"], |
| int(row["created_at_ms"]) // 1000, |
| int(row["updated_at_ms"]) // 1000, |
| row["created_at_ms"], |
| row["updated_at_ms"], |
| row["model_provider"], |
| row["cwd"], |
| row["title"], |
| row["tokens_used"], |
| row["agent_nickname"], |
| row["agent_role"], |
| row["model"], |
| ), |
| ) |
| for parent, child in edges: |
| conn.execute( |
| "INSERT INTO thread_spawn_edges (parent_thread_id, child_thread_id) VALUES (?, ?)", |
| (parent, child), |
| ) |
|
|
|
|
| def _write_rollout(path: Path, events: list[dict[str, object]], malformed: bool = False) -> None: |
| with path.open("w", encoding="utf-8") as handle: |
| for event in events: |
| handle.write(json.dumps(event) + "\n") |
| if malformed: |
| handle.write("{not-json\n") |
|
|
|
|
| def _epoch_ms(value: str) -> int: |
| return int(datetime.fromisoformat(value).timestamp() * 1000) |
|
|
|
|
| if __name__ == "__main__": |
| unittest.main() |
|
|
