1818from __future__ import annotations
1919
2020import abc
21+ import contextlib
22+ import threading
23+ from collections .abc import Iterator
2124from dataclasses import dataclass , field
2225from typing import TYPE_CHECKING , Any , Dict , List
2326
@@ -65,6 +68,87 @@ def usage_to_dict(usage: Any) -> dict[str, Any] | None:
6568 return out or None
6669
6770
71+ _TOKEN_KEYS = ("input_tokens" , "output_tokens" , "prompt_tokens" , "completion_tokens" )
72+
73+
74+ @dataclass
75+ class UsageLedger :
76+ """Every LLM call of one CodeWiki run, with its tokens and cost.
77+
78+ Calls are recorded once per CLI process / agent run, so nested caw
79+ sub-module agents (separate ``claude`` / ``codex`` processes) add up
80+ without double counting. ``add`` can be called from caw's MCP worker
81+ threads, hence the lock. ``phase`` labels the stage the calls belong to
82+ (clustering, modules, overview, ...) and ``target`` the page a
83+ ``complete()`` call is for; stages run sequentially.
84+ """
85+
86+ phase : str = "generate"
87+ target : str = ""
88+ entries : list [dict [str , Any ]] = field (default_factory = list )
89+ _lock : threading .Lock = field (default_factory = threading .Lock , repr = False , compare = False )
90+
91+ @contextlib .contextmanager
92+ def in_phase (self , name : str , target : str = "" ) -> Iterator [None ]:
93+ previous = (self .phase , self .target )
94+ self .phase , self .target = name , target
95+ try :
96+ yield
97+ finally :
98+ self .phase , self .target = previous
99+
100+ def add (
101+ self ,
102+ kind : str ,
103+ target : str ,
104+ usage : dict [str , Any ] | None ,
105+ * ,
106+ seconds : float = 0.0 ,
107+ model : str | None = None ,
108+ ) -> None :
109+ entry : dict [str , Any ] = {
110+ "phase" : self .phase ,
111+ "kind" : kind ,
112+ "target" : target or self .target ,
113+ "model" : model or "" ,
114+ "seconds" : round (seconds , 2 ),
115+ ** (usage or {}),
116+ }
117+ with self ._lock :
118+ self .entries .append (entry )
119+
120+ def summary (self ) -> dict [str , Any ]:
121+ """Totals overall and per phase, plus the per-call list.
122+
123+ A call with tokens but no ``cost_usd`` is *unpriced*: its model is
124+ listed under ``unpriced_models`` so a $0 total is never mistaken for
125+ a free run.
126+ """
127+ with self ._lock :
128+ entries = [dict (e ) for e in self .entries ]
129+ totals : dict [str , float ] = {}
130+ by_phase : dict [str , dict [str , float ]] = {}
131+ unpriced : set [str ] = set ()
132+ for e in entries :
133+ phase_totals = by_phase .setdefault (e ["phase" ], {"calls" : 0 })
134+ phase_totals ["calls" ] += 1
135+ for k , v in e .items ():
136+ if k == "seconds" or not isinstance (v , (int , float )) or isinstance (v , bool ):
137+ continue
138+ totals [k ] = totals .get (k , 0 ) + v
139+ phase_totals [k ] = phase_totals .get (k , 0 ) + v
140+ if not e .get ("cost_usd" ) and any (e .get (k ) for k in _TOKEN_KEYS ):
141+ unpriced .add (e ["model" ] or "<default>" )
142+ return {
143+ "scope" : "this run only (modules resumed from disk cost nothing)" ,
144+ "n_calls" : len (entries ),
145+ "totals" : totals ,
146+ "by_phase" : by_phase ,
147+ "unpriced_models" : sorted (unpriced ),
148+ "calls" : entries ,
149+ }
150+
151+
68152CAW_PROVIDERS = frozenset ({"claude-code" , "codex" })
69153
70154
@@ -79,9 +163,32 @@ class LLMBackend(abc.ABC):
79163 ``last_usage`` holds the token usage of the most recent ``complete`` /
80164 ``run_module_agent`` / ``run_update_agent`` call when the provider exposes
81165 it (``None`` otherwise). The incremental updater reads it for its record.
166+ ``usage_ledger`` accumulates every call of the run for the cost report.
82167 """
83168
84169 last_usage : dict [str , Any ] | None = None
170+ _usage_ledger : UsageLedger | None = None
171+
172+ @property
173+ def usage_ledger (self ) -> UsageLedger :
174+ if self ._usage_ledger is None :
175+ self ._usage_ledger = UsageLedger ()
176+ return self ._usage_ledger
177+
178+ def _record_usage (
179+ self ,
180+ kind : str ,
181+ target : str ,
182+ usage : Any ,
183+ * ,
184+ seconds : float = 0.0 ,
185+ model : str | None = None ,
186+ ) -> dict [str , Any ] | None :
187+ """Set ``last_usage`` from a provider usage object and log it in the ledger."""
188+ usage_dict = usage_to_dict (usage )
189+ self .last_usage = usage_dict
190+ self .usage_ledger .add (kind , target , usage_dict , seconds = seconds , model = model )
191+ return usage_dict
85192
86193 @abc .abstractmethod
87194 def complete (
0 commit comments