Track token usage and cost of every LLM call

.gitignore

@@ -90,3 +90,4 @@ Thumbs.db
 !tests/test_updater_change_report.py
 !tests/test_updater_write_guard.py
 !tests/test_updater_orchestrator.py
+!tests/test_usage_ledger.py

codewiki/cli/adapters/doc_generator.py

@@ -68,6 +68,7 @@ def __init__(
         self.commit_id = commit_id
         self.progress_tracker = ProgressTracker(total_stages=5, verbose=verbose)
         self.job = DocumentationJob()
+        self._usage_ledger = None
 
         # Setup job metadata
         self.job.repository_path = str(repo_path)
@@ -168,6 +169,7 @@ def generate(self) -> DocumentationJob:
 
             # Run backend documentation generation
             asyncio.run(self._run_backend_generation(backend_config))
+            self._record_usage_statistics()
 
             # Stage 4: HTML Generation (optional)
             if self.generate_html:
@@ -198,6 +200,7 @@ async def _run_backend_generation(self, backend_config: BackendConfig):
 
         # Create documentation generator
         doc_generator = DocumentationGenerator(backend_config, commit_id=self.commit_id)
+        self._usage_ledger = doc_generator.usage_ledger
 
         # Incremental update (component-level): keep the previous graph before
         # the builder overwrites it, so the updater can diff old against new.
@@ -309,18 +312,21 @@ async def _run_backend_generation(self, backend_config: BackendConfig):
                             "Clustering modules with LLM...",
                         )
                 cluster_model = backend_config.cluster_model or None
-                module_tree = cluster_modules(
-                    leaf_nodes,
-                    components,
-                    backend_config,
-                    completer=lambda p: doc_generator.backend.complete(p, model=cluster_model),
-                )
-                if module_tree:
-                    module_tree = super_group_modules(
-                        module_tree,
+                with doc_generator.usage_ledger.in_phase("clustering"):
+                    module_tree = cluster_modules(
+                        leaf_nodes,
+                        components,
                         backend_config,
                         completer=lambda p: doc_generator.backend.complete(p, model=cluster_model),
                     )
+                    if module_tree:
+                        module_tree = super_group_modules(
+                            module_tree,
+                            backend_config,
+                            completer=lambda p: doc_generator.backend.complete(
+                                p, model=cluster_model
+                            ),
+                        )
                 # Artifact nodes the clustering LLM dropped get a fixed module
                 # so build/CI/config coverage does not depend on the LLM.
                 if getattr(backend_config, "artifacts_enabled", True):
@@ -440,7 +446,8 @@ async def _run_incremental_update(
         updater = IncrementalUpdater(
             backend_config, doc_generator.backend, doc_generator, update_opts
         )
-        record = await updater.run(prev_graph_path, components, leaf_nodes, revision)
+        with doc_generator.usage_ledger.in_phase("update"):
+            record = await updater.run(prev_graph_path, components, leaf_nodes, revision)
         self._last_update_record = record
         summary = record.summary()
         if self.verbose:
@@ -547,6 +554,29 @@ def _run_html_generation(self):
 
         self.progress_tracker.complete_stage()
 
+    def _record_usage_statistics(self):
+        """Copy the run's token and cost totals from the usage ledger into the job."""
+        if self._usage_ledger is None:
+            return
+        summary = self._usage_ledger.summary()
+        totals = summary["totals"]
+        stats = self.job.statistics
+        stats.usage_totals = totals
+        stats.total_tokens_used = int(
+            sum(
+                totals.get(k, 0)
+                for k in (
+                    "input_tokens",
+                    "output_tokens",
+                    "cache_read_tokens",
+                    "cache_write_tokens",
+                )
+            )
+            or totals.get("total_tokens", 0)
+        )
+        stats.total_cost_usd = totals.get("cost_usd")
+        stats.unpriced_models = summary["unpriced_models"]
+
     def _finalize_job(self):
         """Finalize the job (metadata already created by backend)."""
         # Just verify metadata exists

codewiki/cli/commands/generate.py

@@ -851,6 +851,10 @@ def generate_command(
                 "total_files_analyzed": job.statistics.total_files_analyzed,
                 "generation_time": generation_time,
                 "total_tokens_used": job.statistics.total_tokens_used,
+                "usage_totals": job.statistics.usage_totals,
+                "total_cost_usd": job.statistics.total_cost_usd,
+                "unpriced_models": job.statistics.unpriced_models,
+                "provider": getattr(config, "provider", "openai-compatible"),
             },
         )
 

codewiki/cli/models/job.py

@@ -37,6 +37,11 @@ class JobStatistics:
     leaf_nodes: int = 0
     max_depth: int = 0
     total_tokens_used: int = 0
+    # Summed provider usage of this run (input/output/cache tokens, cost_usd)
+    usage_totals: Dict[str, float] = field(default_factory=dict)
+    total_cost_usd: Optional[float] = None
+    # Models whose calls reported tokens but no price
+    unpriced_models: List[str] = field(default_factory=list)
 
 
 @dataclass

codewiki/cli/utils/instructions.py

@@ -45,6 +45,52 @@ def get_pr_creation_url(repo_url: str, branch_name: str) -> str:
     return f"{base_url}/compare/{branch_name}"
 
 
+def _compact(n: float) -> str:
+    """1234567 -> '1.2M', 8900 -> '8.9k'."""
+    for div, suffix in ((1_000_000, "M"), (1_000, "k")):
+        if n >= div:
+            return f"{n / div:.1f}{suffix}"
+    return f"{int(n)}"
+
+
+def _echo_usage(statistics: dict) -> None:
+    """Print the run's token usage and estimated cost, if any was recorded."""
+    tokens = statistics.get("total_tokens_used") or 0
+    if not tokens:
+        return
+    totals = statistics.get("usage_totals") or {}
+    parts = [
+        f"{label} {_compact(totals[key])}"
+        for label, key in (
+            ("in", "input_tokens"),
+            ("out", "output_tokens"),
+            ("cache-read", "cache_read_tokens"),
+            ("cache-write", "cache_write_tokens"),
+        )
+        if totals.get(key)
+    ]
+    detail = f" ({' / '.join(parts)})" if parts else ""
+    click.echo(f"  Tokens used:       {tokens:,}{detail}")
+
+    cost = statistics.get("total_cost_usd")
+    provider = statistics.get("provider") or ""
+    if cost:
+        note = (
+            " (API-equivalent; billed to your subscription)"
+            if provider in ("claude-code", "codex")
+            else ""
+        )
+        click.echo(f"  Est. cost:         ${cost:,.2f}{note}")
+    unpriced = statistics.get("unpriced_models") or []
+    if unpriced:
+        click.secho(f"  Cost unknown for:  {', '.join(unpriced)}", fg="yellow")
+        if provider == "codex":
+            # Claude Code reports its own cost; codex cost comes from caw's table
+            click.echo(
+                "    set a price: caw pricing set codex <model> --input <usd/1M> --output <usd/1M>"
+            )
+
+
 def display_post_generation_instructions(
     output_dir: Path,
     repo_name: str,
@@ -95,9 +141,7 @@ def display_post_generation_instructions(
             minutes = int(statistics["generation_time"] // 60)
             seconds = int(statistics["generation_time"] % 60)
             click.echo(f"  Generation time:   {minutes} minutes {seconds} seconds")
-        # if 'total_tokens_used' in statistics:
-        #     tokens = statistics['total_tokens_used']
-        #     click.echo(f"  Tokens used:       ~{tokens:,}")
+        _echo_usage(statistics)
         click.echo()
 
     # Next steps

codewiki/src/be/backend.py

@@ -18,6 +18,9 @@
 from __future__ import annotations
 
 import abc
+import contextlib
+import threading
+from collections.abc import Iterator
 from dataclasses import dataclass, field
 from typing import TYPE_CHECKING, Any, Dict, List
 
@@ -65,6 +68,87 @@ def usage_to_dict(usage: Any) -> dict[str, Any] | None:
     return out or None
 
 
+_TOKEN_KEYS = ("input_tokens", "output_tokens", "prompt_tokens", "completion_tokens")
+
+
+@dataclass
+class UsageLedger:
+    """Every LLM call of one CodeWiki run, with its tokens and cost.
+
+    Calls are recorded once per CLI process / agent run, so nested caw
+    sub-module agents (separate ``claude`` / ``codex`` processes) add up
+    without double counting.  ``add`` can be called from caw's MCP worker
+    threads, hence the lock.  ``phase`` labels the stage the calls belong to
+    (clustering, modules, overview, ...) and ``target`` the page a
+    ``complete()`` call is for; stages run sequentially.
+    """
+
+    phase: str = "generate"
+    target: str = ""
+    entries: list[dict[str, Any]] = field(default_factory=list)
+    _lock: threading.Lock = field(default_factory=threading.Lock, repr=False, compare=False)
+
+    @contextlib.contextmanager
+    def in_phase(self, name: str, target: str = "") -> Iterator[None]:
+        previous = (self.phase, self.target)
+        self.phase, self.target = name, target
+        try:
+            yield
+        finally:
+            self.phase, self.target = previous
+
+    def add(
+        self,
+        kind: str,
+        target: str,
+        usage: dict[str, Any] | None,
+        *,
+        seconds: float = 0.0,
+        model: str | None = None,
+    ) -> None:
+        entry: dict[str, Any] = {
+            "phase": self.phase,
+            "kind": kind,
+            "target": target or self.target,
+            "model": model or "",
+            "seconds": round(seconds, 2),
+            **(usage or {}),
+        }
+        with self._lock:
+            self.entries.append(entry)
+
+    def summary(self) -> dict[str, Any]:
+        """Totals overall and per phase, plus the per-call list.
+
+        A call with tokens but no ``cost_usd`` is *unpriced*: its model is
+        listed under ``unpriced_models`` so a $0 total is never mistaken for
+        a free run.
+        """
+        with self._lock:
+            entries = [dict(e) for e in self.entries]
+        totals: dict[str, float] = {}
+        by_phase: dict[str, dict[str, float]] = {}
+        unpriced: set[str] = set()
+        for e in entries:
+            phase_totals = by_phase.setdefault(e["phase"], {"calls": 0})
+            phase_totals["calls"] += 1
+            for k, v in e.items():
+                if k == "seconds" or not isinstance(v, (int, float)) or isinstance(v, bool):
+                    continue
+                totals[k] = totals.get(k, 0) + v
+                phase_totals[k] = phase_totals.get(k, 0) + v
+            if not e.get("cost_usd") and any(e.get(k) for k in _TOKEN_KEYS):
+                unpriced.add(e["model"] or "<default>")
+        return {
+            "scope": "this run only (modules resumed from disk cost nothing)",
+            "n_calls": len(entries),
+            "totals": totals,
+            "by_phase": by_phase,
+            "unpriced_models": sorted(unpriced),
+            "calls": entries,
+        }
+
+
 CAW_PROVIDERS = frozenset({"claude-code", "codex"})
 
 
@@ -79,9 +163,32 @@ class LLMBackend(abc.ABC):
     ``last_usage`` holds the token usage of the most recent ``complete`` /
     ``run_module_agent`` / ``run_update_agent`` call when the provider exposes
     it (``None`` otherwise). The incremental updater reads it for its record.
+    ``usage_ledger`` accumulates every call of the run for the cost report.
     """
 
     last_usage: dict[str, Any] | None = None
+    _usage_ledger: UsageLedger | None = None
+
+    @property
+    def usage_ledger(self) -> UsageLedger:
+        if self._usage_ledger is None:
+            self._usage_ledger = UsageLedger()
+        return self._usage_ledger
+
+    def _record_usage(
+        self,
+        kind: str,
+        target: str,
+        usage: Any,
+        *,
+        seconds: float = 0.0,
+        model: str | None = None,
+    ) -> dict[str, Any] | None:
+        """Set ``last_usage`` from a provider usage object and log it in the ledger."""
+        usage_dict = usage_to_dict(usage)
+        self.last_usage = usage_dict
+        self.usage_ledger.add(kind, target, usage_dict, seconds=seconds, model=model)
+        return usage_dict
 
     @abc.abstractmethod
     def complete(

codewiki/src/be/caw_backend.py

@@ -30,7 +30,7 @@
 from caw import ToolGroup
 
 from codewiki.src.be.agent_tools.deps import CodeWikiDeps
-from codewiki.src.be.backend import AgentReply, LLMBackend, usage_to_dict
+from codewiki.src.be.backend import AgentReply, LLMBackend
 from codewiki.src.be.cluster_modules import format_potential_core_components
 from codewiki.src.be.dependency_analyzer.models.core import Node
 from codewiki.src.be.prompt_template import (
@@ -260,8 +260,15 @@ def complete(
             tools=ToolGroup.READER,
             system_prompt=system_prompt,
         )
+        started = time.time()
         traj = agent.completion(prompt)
-        self.last_usage = usage_to_dict(getattr(traj, "total_usage", None))
+        self._record_usage(
+            "complete",
+            "",
+            getattr(traj, "total_usage", None),
+            seconds=time.time() - started,
+            model=getattr(traj, "model", None) or effective_model,
+        )
         return traj.result
 
     # ------------------------------------------------------------------
@@ -305,12 +312,18 @@ def _run_update_agent_sync(
         except Exception as e:
             logger.error("Update agent for %s failed via caw: %s", deps.current_module_name, e)
             raise
-        usage = usage_to_dict(getattr(traj, "total_usage", None))
-        self.last_usage = usage
+        seconds = time.time() - started
+        usage = self._record_usage(
+            "update_agent",
+            deps.current_module_name,
+            getattr(traj, "total_usage", None),
+            seconds=seconds,
+            model=getattr(traj, "model", None) or self._model,
+        )
         return AgentReply(
             text=traj.result or "",
             usage=usage,
-            seconds=time.time() - started,
+            seconds=seconds,
             meta={"turns": traj.num_turns, "tool_calls": traj.total_tool_calls},
         )
 
@@ -485,19 +498,29 @@ def _run_module_agent_sync(
             run_cwd = working_dir
         else:
             run_cwd = self._repo_root
+        started = time.time()
         try:
             os.chdir(run_cwd)
             try:
                 traj = agent.completion(user_prompt)
             finally:
                 os.chdir(original_cwd)
+            # Nested sub-module agents are separate CLI processes that record
+            # their own usage, so this covers only this module's own run.
+            usage = self._record_usage(
+                "module_agent",
+                module_name,
+                getattr(traj, "total_usage", None),
+                seconds=time.time() - started,
+                model=getattr(traj, "model", None) or self._model,
+            )
             logger.info(
-                "Module %s completed via caw (turns=%d, tool_calls=%d)",
+                "Module %s completed via caw (turns=%d, tool_calls=%d, cost=$%.4f)",
                 module_name,
                 traj.num_turns,
                 traj.total_tool_calls,
+                (usage or {}).get("cost_usd", 0.0),
             )
-            self.last_usage = usage_to_dict(getattr(traj, "total_usage", None))
             file_manager.save_json(deps.module_tree, module_tree_path)
             return deps.module_tree
         except Exception as e: