diff --git a/.gitignore b/.gitignore index 1595ab90..9402bacf 100644 --- a/.gitignore +++ b/.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 diff --git a/codewiki/cli/adapters/doc_generator.py b/codewiki/cli/adapters/doc_generator.py index 472cba4e..072bbd88 100644 --- a/codewiki/cli/adapters/doc_generator.py +++ b/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 diff --git a/codewiki/cli/commands/generate.py b/codewiki/cli/commands/generate.py index daa1a1f8..d5d057f6 100644 --- a/codewiki/cli/commands/generate.py +++ b/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"), }, ) diff --git a/codewiki/cli/models/job.py b/codewiki/cli/models/job.py index 19bde6f6..11db6e8e 100644 --- a/codewiki/cli/models/job.py +++ b/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 diff --git a/codewiki/cli/utils/instructions.py b/codewiki/cli/utils/instructions.py index f20e6e72..d292d7fe 100644 --- a/codewiki/cli/utils/instructions.py +++ b/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 --input --output " + ) + + 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 diff --git a/codewiki/src/be/backend.py b/codewiki/src/be/backend.py index d31733f9..8d1f7864 100644 --- a/codewiki/src/be/backend.py +++ b/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 "") + 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( diff --git a/codewiki/src/be/caw_backend.py b/codewiki/src/be/caw_backend.py index 60e76403..6da530f5 100644 --- a/codewiki/src/be/caw_backend.py +++ b/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: diff --git a/codewiki/src/be/documentation_generator.py b/codewiki/src/be/documentation_generator.py index eddaaecd..e645cc8b 100644 --- a/codewiki/src/be/documentation_generator.py +++ b/codewiki/src/be/documentation_generator.py @@ -6,7 +6,7 @@ from typing import Any from codewiki import __version__ -from codewiki.src.be.backend import LLMBackend, get_backend +from codewiki.src.be.backend import LLMBackend, UsageLedger, get_backend from codewiki.src.be.cluster_modules import ( cluster_modules, ensure_artifact_module, @@ -67,6 +67,14 @@ def __init__( self.graph_builder = DependencyGraphBuilder(config) self.backend: LLMBackend = backend or get_backend(config) + @property + def usage_ledger(self) -> UsageLedger: + """The backend's usage ledger (a throwaway one for backends without it).""" + ledger = getattr(getattr(self, "backend", None), "usage_ledger", None) + if not isinstance(ledger, UsageLedger): + ledger = self.__dict__.setdefault("_fallback_ledger", UsageLedger()) + return ledger + def create_documentation_metadata( self, working_dir: str, components: dict[str, Any], num_leaf_nodes: int ): @@ -76,6 +84,7 @@ def create_documentation_metadata( metadata = { "generation_info": { "timestamp": datetime.now(UTC).isoformat(), + "provider": getattr(self.config, "provider", None), "main_model": self.config.main_model, "generator_version": __version__, "repo_path": self.config.repo_path, @@ -91,6 +100,8 @@ def create_documentation_metadata( "max_depth": self.config.max_depth, }, "files_generated": ["overview.md", "module_tree.json", "first_module_tree.json"], + # Tokens and cost of every LLM call in this run (see UsageLedger) + "usage": self.usage_ledger.summary(), } # Add generated markdown files (docs-relative paths) to the metadata @@ -240,7 +251,20 @@ async def generate_module_documentation( # regenerated overview.md is rebuilt from the module docs via # generate_parent_module_docs([]) rather than by the whole-repo agent. processing_order = self.get_processing_order(module_tree) + with self.usage_ledger.in_phase("modules"): + return await self._generate_modules_in_order( + components, leaf_nodes, working_dir, module_tree_path, module_tree, processing_order + ) + async def _generate_modules_in_order( + self, + components: dict[str, Any], + leaf_nodes: list[str], + working_dir: str, + module_tree_path: str, + module_tree: dict[str, Any], + processing_order: list[tuple[list[str], str]], + ) -> str: # Process modules in dependency order final_module_tree = module_tree processed_modules = set() @@ -377,11 +401,13 @@ async def generate_parent_module_docs( ) logger.debug(f"Overview prompt for {module_name}: {len(prompt)} chars") + phase = "overview" if not module_path else "parent_modules" try: - parent_docs = self.backend.complete( - prompt, - system_prompt=format_overview_system_prompt(self.config.get_prompt_addition()), - ) + with self.usage_ledger.in_phase(phase, target=module_name): + parent_docs = self.backend.complete( + prompt, + system_prompt=format_overview_system_prompt(self.config.get_prompt_addition()), + ) if not parent_docs: raise RuntimeError( f"LLM returned empty content for {module_name} overview " @@ -451,18 +477,19 @@ async def run(self) -> None: # one is set. Caw mode's cluster_model is typically empty — # complete() falls back to its own _model in that case. cluster_model = self.config.cluster_model or None - module_tree = cluster_modules( - leaf_nodes, - components, - self.config, - completer=lambda p: self.backend.complete(p, model=cluster_model), - ) - if module_tree: - module_tree = super_group_modules( - module_tree, + with self.usage_ledger.in_phase("clustering"): + module_tree = cluster_modules( + leaf_nodes, + components, self.config, completer=lambda p: self.backend.complete(p, model=cluster_model), ) + if module_tree: + module_tree = super_group_modules( + module_tree, + self.config, + completer=lambda p: self.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(self.config, "artifacts_enabled", True): diff --git a/codewiki/src/be/pydantic_ai_backend.py b/codewiki/src/be/pydantic_ai_backend.py index 9b7e6a99..0aabd7f4 100644 --- a/codewiki/src/be/pydantic_ai_backend.py +++ b/codewiki/src/be/pydantic_ai_backend.py @@ -68,8 +68,15 @@ def complete( system_prompt: str | None = None, ) -> str: pop_last_usage() + started = time.time() result = call_llm(prompt, self._config, model=model, system_prompt=system_prompt) - self.last_usage = pop_last_usage() + self._record_usage( + "complete", + "", + pop_last_usage(), + seconds=time.time() - started, + model=model or self._config.main_model, + ) return result async def run_update_agent( @@ -93,8 +100,13 @@ async def run_update_agent( usage_limits=UsageLimits(request_limit=self._config.request_limit), ) seconds = time.time() - started - usage = _run_usage(result) - self.last_usage = usage + usage = self._record_usage( + "update_agent", + deps.current_module_name, + _run_usage(result), + seconds=seconds, + model=self._config.main_model, + ) text = result.output if isinstance(result.output, str) else str(result.output) return AgentReply(text=text, usage=usage, seconds=seconds) @@ -167,6 +179,7 @@ async def run_module_agent( custom_instructions=self._custom_instructions, ) + started = time.time() try: result = await agent.run( format_user_prompt( @@ -180,7 +193,15 @@ async def run_module_agent( deps=deps, usage_limits=UsageLimits(request_limit=config.request_limit), ) - self.last_usage = _run_usage(result) + # Covers this module's own run; pydantic-ai sub-module agents run + # with their own usage and are not counted here. + self._record_usage( + "module_agent", + module_name, + _run_usage(result), + seconds=time.time() - started, + model=config.main_model, + ) file_manager.save_json(deps.module_tree, module_tree_path) return deps.module_tree except Exception as e: diff --git a/guides/providers.md b/guides/providers.md index e702a401..aa6d3913 100644 --- a/guides/providers.md +++ b/guides/providers.md @@ -106,6 +106,27 @@ Things to know: - Large repositories can push a module prompt past the CLI's input limit. CodeWiki trims the module tree in that case and retries. +### Cost tracking + +At the end of `codewiki generate`, CodeWiki prints the tokens used and an +estimated cost for the run. `metadata.json` gets a `usage` block with totals, +a per-phase breakdown (`clustering`, `modules`, `parent_modules`, `overview`, +`update`) and one entry per LLM call. Each nested sub-module agent is counted +once. Modules resumed from an earlier run cost nothing and don't appear. + +The cost is what the same tokens would cost on the API. Your subscription +is billed as usual. + +- **Claude Code** reports the cost itself. +- **Codex** is priced from caw's rate table using the model Codex actually + runs (`--main-model`, or `model` in `~/.codex/config.toml`). When a model + has no rate, CodeWiki lists it as "cost unknown" instead of showing $0. Add + a rate with: + + ```bash + caw pricing set codex --input --cached --output + ``` + ## Which model to use - The **main model** writes the pages. It needs strong code understanding diff --git a/pyproject.toml b/pyproject.toml index 1a3867c7..e101fbb2 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -61,7 +61,7 @@ dependencies = [ "python-multipart>=0.0.20", "colorama>=0.4.6", "logfire>=4.1.0", - "coding-agent-wrapper @ git+https://github.com/anhnh2002/caw@fix/codex-exec-robustness", + "coding-agent-wrapper @ git+https://github.com/anhnh2002/caw@main", "mcp>=1.0.0" ] diff --git a/tests/test_usage_ledger.py b/tests/test_usage_ledger.py new file mode 100644 index 00000000..afe809c1 --- /dev/null +++ b/tests/test_usage_ledger.py @@ -0,0 +1,188 @@ +"""Usage ledger: every LLM call of a run, summed into tokens and cost.""" + +import json +import threading +from types import SimpleNamespace + +import pytest + +from codewiki.cli.utils.instructions import _echo_usage +from codewiki.src.be.backend import LLMBackend, UsageLedger +from codewiki.src.be.documentation_generator import DocumentationGenerator + + +def _usage(input_tokens=0, output_tokens=0, cache_read_tokens=0, cost_usd=0.0): + """Stand-in for caw's UsageStats (CI does not install caw).""" + return SimpleNamespace( + input_tokens=input_tokens, + output_tokens=output_tokens, + cache_read_tokens=cache_read_tokens, + cache_write_tokens=0, + cost_usd=cost_usd, + ) + + +def test_summary_totals_and_phases(): + ledger = UsageLedger() + with ledger.in_phase("clustering"): + ledger.add("complete", "", {"input_tokens": 100, "output_tokens": 10, "cost_usd": 0.5}) + with ledger.in_phase("modules"): + ledger.add( + "module_agent", "auth", {"input_tokens": 200, "output_tokens": 20, "cost_usd": 1.0} + ) + ledger.add( + "module_agent", "db", {"input_tokens": 300, "output_tokens": 30, "cost_usd": 1.5} + ) + + summary = ledger.summary() + assert summary["n_calls"] == 3 + assert summary["totals"] == {"input_tokens": 600, "output_tokens": 60, "cost_usd": 3.0} + assert summary["by_phase"]["clustering"]["calls"] == 1 + assert summary["by_phase"]["modules"] == { + "calls": 2, + "input_tokens": 500, + "output_tokens": 50, + "cost_usd": 2.5, + } + assert summary["unpriced_models"] == [] + assert [c["target"] for c in summary["calls"]] == ["", "auth", "db"] + + +def test_in_phase_restores_and_labels_target(): + ledger = UsageLedger() + with ledger.in_phase("overview", target="repo"): + ledger.add("complete", "", None) + ledger.add("complete", "", None) + calls = ledger.summary()["calls"] + assert (calls[0]["phase"], calls[0]["target"]) == ("overview", "repo") + assert (calls[1]["phase"], calls[1]["target"]) == ("generate", "") + + +def test_tokens_without_cost_are_unpriced(): + ledger = UsageLedger() + ledger.add("module_agent", "a", {"input_tokens": 10, "cost_usd": 0.0}, model="gpt-x") + ledger.add("module_agent", "b", {"input_tokens": 10}, model="") + ledger.add("module_agent", "c", None, model="ignored") # no usage reported at all + assert ledger.summary()["unpriced_models"] == ["", "gpt-x"] + + +def test_add_is_thread_safe(): + ledger = UsageLedger() + + def worker(): + for _ in range(500): + ledger.add("module_agent", "m", {"input_tokens": 1, "cost_usd": 0.001}) + + threads = [threading.Thread(target=worker) for _ in range(8)] + for t in threads: + t.start() + for t in threads: + t.join() + summary = ledger.summary() + assert summary["n_calls"] == 4000 + assert summary["totals"]["input_tokens"] == 4000 + + +def _caw_backend(monkeypatch, usage): + pytest.importorskip("caw") + from codewiki.src.be import caw_backend + from codewiki.src.be.caw_backend import CawBackend + + class FakeCawAgent: + def __init__(self, **kwargs): + self.model = kwargs.get("model") + + def completion(self, prompt): + return SimpleNamespace( + result="ok", + total_usage=usage, + model="opus", + num_turns=1, + total_tool_calls=0, + ) + + monkeypatch.setattr(caw_backend, "CawAgent", FakeCawAgent) + backend = object.__new__(CawBackend) # skip the CLI-on-PATH check + backend._caw_provider = "claude_code" + backend._model = None + return backend + + +def test_caw_complete_records_usage(monkeypatch): + usage = _usage(input_tokens=1000, output_tokens=200, cache_read_tokens=50, cost_usd=0.12) + backend = _caw_backend(monkeypatch, usage) + + with backend.usage_ledger.in_phase("overview", target="repo"): + assert backend.complete("prompt") == "ok" + + assert backend.last_usage["cost_usd"] == 0.12 + (call,) = backend.usage_ledger.summary()["calls"] + assert call["phase"] == "overview" + assert call["target"] == "repo" + assert call["kind"] == "complete" + assert call["model"] == "opus" + assert call["input_tokens"] == 1000 + assert call["cost_usd"] == 0.12 + + +class _StubBackend(LLMBackend): + def complete(self, prompt, *, model=None, system_prompt=None): + self._record_usage("complete", "", _usage(input_tokens=10, cost_usd=0.25)) + return "ok" + + async def run_module_agent(self, *args, **kwargs): + return {} + + +def test_metadata_includes_usage(tmp_path): + backend = _StubBackend() + backend.complete("prompt") + + gen = object.__new__(DocumentationGenerator) + gen.config = SimpleNamespace( + main_model="opus", provider="claude-code", repo_path=".", max_depth=2, layout="flat" + ) + gen.commit_id = None + gen.backend = backend + gen.create_documentation_metadata(str(tmp_path), {}, 0) + + metadata = json.loads((tmp_path / "metadata.json").read_text()) + assert metadata["generation_info"]["provider"] == "claude-code" + assert metadata["usage"]["n_calls"] == 1 + assert metadata["usage"]["totals"]["cost_usd"] == 0.25 + + +def test_cli_prints_tokens_and_cost(capsys): + _echo_usage( + { + "total_tokens_used": 1_250_000, + "usage_totals": {"input_tokens": 1_000_000, "output_tokens": 250_000}, + "total_cost_usd": 4.2, + "unpriced_models": [], + "provider": "claude-code", + } + ) + out = capsys.readouterr().out + assert "1,250,000 (in 1.0M / out 250.0k)" in out + assert "$4.20 (API-equivalent; billed to your subscription)" in out + + +def test_cli_flags_unpriced_codex_models(capsys): + _echo_usage( + { + "total_tokens_used": 10, + "usage_totals": {"input_tokens": 10}, + "total_cost_usd": 0.0, + "unpriced_models": ["gpt-x"], + "provider": "codex", + } + ) + out = capsys.readouterr().out + assert "Est. cost" not in out + assert "Cost unknown for: gpt-x" in out + assert "caw pricing set codex" in out + + +def test_cli_silent_without_usage(capsys): + _echo_usage({"total_tokens_used": 0}) + assert capsys.readouterr().out == ""