diff --git a/codeframe/cli/app.py b/codeframe/cli/app.py index 9783e079..7f0509e0 100644 --- a/codeframe/cli/app.py +++ b/codeframe/cli/app.py @@ -28,6 +28,7 @@ from codeframe.cli.env_commands import env_app from codeframe.cli.engines_commands import engines_app from codeframe.cli.hooks_commands import hooks_app +from codeframe.cli.stats_commands import stats_app # Load environment variables from .env files # Priority: workspace .env > home .env @@ -4870,6 +4871,7 @@ def templates_apply( app.add_typer(engines_app, name="engines") app.add_typer(hooks_app, name="hooks") +app.add_typer(stats_app, name="stats") # ============================================================================= diff --git a/codeframe/cli/stats_commands.py b/codeframe/cli/stats_commands.py new file mode 100644 index 00000000..f4c5026f --- /dev/null +++ b/codeframe/cli/stats_commands.py @@ -0,0 +1,311 @@ +"""CLI stats commands for headless token/cost tracking. + +This module provides commands for viewing token usage and cost statistics +directly from the local workspace database (no server required): + +- tokens: View workspace token usage summary +- costs: View cost report with optional period filtering +- export: Export usage data to CSV or JSON + +Usage: + cf stats tokens # Workspace token summary + cf stats tokens --task # Per-task breakdown + cf stats costs # All-time costs + cf stats costs --period month # Last 30 days + cf stats export --format csv --output tokens.csv +""" + +import logging +from datetime import datetime, timedelta, timezone +from pathlib import Path +from typing import Optional + +import typer +from rich.table import Table + +from codeframe.cli.helpers import console + +logger = logging.getLogger(__name__) + +stats_app = typer.Typer( + name="stats", + help="Token usage and cost statistics", + no_args_is_help=True, +) + + +def _get_db(): + """Get database from current workspace. + + Looks for .codeframe/state.db relative to the current directory. + + Returns: + Initialized Database instance. + + Raises: + typer.Exit: If no workspace is found. + """ + from codeframe.persistence.database import Database + + db_path = Path(".codeframe/state.db") + if not db_path.exists(): + console.print("[red]Error:[/red] No workspace found. Run 'cf init' first.") + raise typer.Exit(1) + db = Database(db_path) + db.initialize() + return db + + +def _get_tracker(db): + """Create a MetricsTracker from a database instance. + + Args: + db: Initialized Database instance. + + Returns: + MetricsTracker instance. + """ + from codeframe.lib.metrics_tracker import MetricsTracker + + return MetricsTracker(db=db) + + +def _format_number(n: int) -> str: + """Format number with thousands separator.""" + return f"{n:,}" + + +@stats_app.command() +def tokens( + task: Optional[int] = typer.Option( + None, "--task", "-t", help="Filter by task ID for per-task breakdown" + ), +): + """Show workspace token usage summary. + + Displays total tokens used across all tasks, with input/output breakdown + and per-model statistics. Use --task to filter to a specific task. + + Examples: + cf stats tokens # Workspace summary + cf stats tokens --task 1 # Task 1 breakdown + """ + db = _get_db() + try: + tracker = _get_tracker(db) + + if task is not None: + # Per-task summary + summary = tracker.get_task_token_summary(task) + + console.print(f"\n[bold]Token Usage for Task {task}[/bold]\n") + + table = Table(show_header=True, title=None) + table.add_column("Metric", style="cyan") + table.add_column("Value", justify="right") + + table.add_row("Total Tokens", _format_number(summary["total_tokens"])) + table.add_row("Input Tokens", _format_number(summary["total_input_tokens"])) + table.add_row("Output Tokens", _format_number(summary["total_output_tokens"])) + table.add_row("Total Cost", f"${summary['total_cost_usd']:.4f}") + table.add_row("LLM Calls", str(summary["call_count"])) + + console.print(table) + else: + # Workspace-wide summary + records = db.get_workspace_token_usage() + + total_input = 0 + total_output = 0 + total_cost = 0.0 + model_stats: dict[str, dict] = {} + + for record in records: + total_input += record["input_tokens"] + total_output += record["output_tokens"] + total_cost += record["estimated_cost_usd"] + + model = record["model_name"] + if model not in model_stats: + model_stats[model] = { + "input_tokens": 0, + "output_tokens": 0, + "cost_usd": 0.0, + "calls": 0, + } + model_stats[model]["input_tokens"] += record["input_tokens"] + model_stats[model]["output_tokens"] += record["output_tokens"] + model_stats[model]["cost_usd"] += record["estimated_cost_usd"] + model_stats[model]["calls"] += 1 + + total_tokens = total_input + total_output + + console.print("\n[bold]Workspace Token Usage Summary[/bold]\n") + + summary_table = Table(show_header=True) + summary_table.add_column("Metric", style="cyan") + summary_table.add_column("Value", justify="right") + + summary_table.add_row("Total Tokens", _format_number(total_tokens)) + summary_table.add_row("Input Tokens", _format_number(total_input)) + summary_table.add_row("Output Tokens", _format_number(total_output)) + summary_table.add_row("Total Cost", f"${total_cost:.4f}") + summary_table.add_row("LLM Calls", str(len(records))) + + console.print(summary_table) + + if model_stats: + console.print("\n[bold]By Model:[/bold]") + model_table = Table(show_header=True) + model_table.add_column("Model", style="cyan") + model_table.add_column("Tokens", justify="right") + model_table.add_column("Cost", justify="right") + model_table.add_column("Calls", justify="right") + + for model_name, stats in model_stats.items(): + model_table.add_row( + model_name, + _format_number(stats["input_tokens"] + stats["output_tokens"]), + f"${stats['cost_usd']:.4f}", + str(stats["calls"]), + ) + + console.print(model_table) + finally: + db.close() + + +@stats_app.command() +def costs( + period: Optional[str] = typer.Option( + None, + "--period", + "-p", + help="Time period: 'day' (24h), 'week' (7d), 'month' (30d)", + ), +): + """Show cost report. + + Displays total costs and per-model breakdown. Use --period to filter + to a recent time window. + + Examples: + cf stats costs # All-time costs + cf stats costs --period month # Last 30 days + cf stats costs --period week # Last 7 days + cf stats costs --period day # Last 24 hours + """ + db = _get_db() + try: + # Calculate date range from period + start_date = None + end_date = None + now = datetime.now(timezone.utc) + + if period == "day": + start_date = now - timedelta(days=1) + elif period == "week": + start_date = now - timedelta(weeks=1) + elif period == "month": + start_date = now - timedelta(days=30) + elif period is not None: + console.print( + f"[red]Error:[/red] Unknown period '{period}'. Use 'day', 'week', or 'month'." + ) + raise typer.Exit(1) + + # Single fetch: get raw records and compute summary + per-model breakdown in one pass + records = db.get_workspace_token_usage(start_date=start_date, end_date=end_date) + + total_cost = 0.0 + total_tokens = 0 + model_costs: dict[str, dict] = {} + for record in records: + cost = record["estimated_cost_usd"] + tokens = record["input_tokens"] + record["output_tokens"] + total_cost += cost + total_tokens += tokens + + model = record["model_name"] + if model not in model_costs: + model_costs[model] = {"cost_usd": 0.0, "tokens": 0, "calls": 0} + model_costs[model]["cost_usd"] += cost + model_costs[model]["tokens"] += tokens + model_costs[model]["calls"] += 1 + + period_label = f" ({period})" if period else " (all time)" + console.print(f"\n[bold]Cost Report{period_label}[/bold]\n") + + table = Table(show_header=True) + table.add_column("Metric", style="cyan") + table.add_column("Value", justify="right") + + table.add_row("Total Cost", f"${total_cost:.4f}") + table.add_row("Total Tokens", _format_number(total_tokens)) + table.add_row("LLM Calls", str(len(records))) + + console.print(table) + + if model_costs: + console.print("\n[bold]By Model:[/bold]") + model_table = Table(show_header=True) + model_table.add_column("Model", style="cyan") + model_table.add_column("Cost", justify="right") + model_table.add_column("Tokens", justify="right") + model_table.add_column("Calls", justify="right") + + for model_name, stats in model_costs.items(): + model_table.add_row( + model_name, + f"${stats['cost_usd']:.4f}", + _format_number(stats["tokens"]), + str(stats["calls"]), + ) + + console.print(model_table) + finally: + db.close() + + +@stats_app.command("export") +def export_data( + format: str = typer.Option( + "csv", "--format", "-f", help="Output format: csv or json" + ), + output: str = typer.Option( + ..., "--output", "-o", help="Output file path" + ), + task: Optional[int] = typer.Option( + None, "--task", "-t", help="Filter by task ID" + ), +): + """Export usage data to CSV or JSON. + + Exports raw token usage records to a file for external analysis. + Use --task to export records for a single task only. + + Examples: + cf stats export --format csv --output tokens.csv + cf stats export --format json --output tokens.json + cf stats export --format csv --output task1.csv --task 1 + """ + from codeframe.lib.metrics_tracker import MetricsTracker + + db = _get_db() + try: + if task is not None: + records = db.get_batch_token_usage(task_ids=[task]) + else: + records = db.get_workspace_token_usage() + + if format == "csv": + MetricsTracker.export_to_csv(records, output) + elif format == "json": + MetricsTracker.export_to_json(records, output) + else: + console.print(f"[red]Error:[/red] Unknown format '{format}'. Use 'csv' or 'json'.") + raise typer.Exit(1) + + console.print(f"Exported {len(records)} records to {output}") + finally: + db.close() diff --git a/codeframe/core/react_agent.py b/codeframe/core/react_agent.py index b675ebce..7a3129b6 100644 --- a/codeframe/core/react_agent.py +++ b/codeframe/core/react_agent.py @@ -142,6 +142,9 @@ def __init__( self.fix_tracker = FixAttemptTracker() self.blocker_id: Optional[str] = None + # Token usage tracking: accumulate per-call records across the run. + self._token_records: list[dict] = [] + # Stall detection self._stall_triggered = threading.Event() self._stall_event: Optional[StallEvent] = None @@ -254,6 +257,14 @@ def run(self, task_id: str) -> AgentStatus: return AgentStatus.FAILED finally: self._stall_monitor.stop() + try: + self._persist_token_usage(task_id) + except Exception: + logger.debug( + "Failed to persist token usage for task %s", + task_id, + exc_info=True, + ) except StallDetectedError: raise # Monitor stopped by finally above; let runtime handle retry except Exception: @@ -265,6 +276,100 @@ def run(self, task_id: str) -> AgentStatus: self._emit_stream_error(task_id, "exception") return AgentStatus.FAILED + def get_token_usage(self) -> list[dict]: + """Return the accumulated per-call token usage records. + + Each record is a dict with keys: input_tokens, output_tokens, model, + call_type, iteration. + """ + return list(self._token_records) + + def get_total_tokens(self) -> dict: + """Return aggregated token totals and estimated cost. + + Returns: + Dict with input_tokens, output_tokens, total_tokens, + and estimated_cost_usd. + """ + total_in = sum(r["input_tokens"] for r in self._token_records) + total_out = sum(r["output_tokens"] for r in self._token_records) + return { + "input_tokens": total_in, + "output_tokens": total_out, + "total_tokens": total_in + total_out, + "estimated_cost_usd": self._estimate_total_cost(), + } + + # ------------------------------------------------------------------ + # Token persistence + # ------------------------------------------------------------------ + + def _estimate_total_cost(self) -> float: + """Estimate total USD cost from accumulated token records. + + Delegates per-model pricing to MetricsTracker.calculate_cost to keep + pricing logic in a single location. Falls back to 0.0 on any error. + """ + try: + from codeframe.lib.metrics_tracker import MetricsTracker + + total = 0.0 + for record in self._token_records: + total += MetricsTracker.calculate_cost( + record["model"], + record["input_tokens"], + record["output_tokens"], + ) + return round(total, 6) + except Exception: + return 0.0 + + def _persist_token_usage(self, task_id: str) -> None: + """Persist accumulated token records to the workspace database. + + Uses MetricsTracker.record_token_usage_sync for synchronous writes. + Failures are logged but never propagated to the caller. + The database connection is always closed via try/finally. + """ + if not self._token_records: + return + + db = None + try: + from codeframe.lib.metrics_tracker import MetricsTracker + from codeframe.persistence.database import Database + + db = Database(str(self.workspace.db_path)) + db.initialize() + tracker = MetricsTracker(db=db) + + # Cast task_id to int for the persistence layer (core uses str, DB uses int). + try: + task_id_int: int | None = int(task_id) + except (ValueError, TypeError): + task_id_int = None + + for record in self._token_records: + tracker.record_token_usage_sync( + task_id=task_id_int, + agent_id="react-agent", + project_id=0, + model_name=record["model"], + input_tokens=record["input_tokens"], + output_tokens=record["output_tokens"], + call_type=record["call_type"], + ) + except Exception: + logger.debug( + "Token usage persistence failed for task %s", task_id, exc_info=True, + ) + finally: + if db is not None: + try: + db.close() + except Exception: + pass + # ------------------------------------------------------------------ # ReAct loop # ------------------------------------------------------------------ @@ -336,6 +441,15 @@ def _react_loop(self, system_prompt: str) -> AgentStatus: ) iterations += 1 + # Record token usage for this LLM call. + self._token_records.append({ + "input_tokens": response.input_tokens, + "output_tokens": response.output_tokens, + "model": response.model, + "call_type": "task_execution", + "iteration": iterations, + }) + if not response.has_tool_calls: # Text-only response — agent thinks it's done. # Check for blocker patterns before accepting completion. @@ -565,6 +679,15 @@ def _run_final_verification( system=system_prompt, ) + # Record token usage for verification fix calls. + self._token_records.append({ + "input_tokens": response.input_tokens, + "output_tokens": response.output_tokens, + "model": response.model, + "call_type": "verification_fix", + "iteration": attempt, + }) + if not response.has_tool_calls: break # Agent done fixing → re-run gates diff --git a/codeframe/lib/metrics_tracker.py b/codeframe/lib/metrics_tracker.py index fe5294ab..2ee97268 100644 --- a/codeframe/lib/metrics_tracker.py +++ b/codeframe/lib/metrics_tracker.py @@ -3,10 +3,11 @@ This module provides token usage tracking and cost estimation for LLM calls across agents and projects. It supports: -- Recording token usage per LLM call +- Recording token usage per LLM call (async and sync) - Cost calculation for Claude models (Sonnet 4.5, Opus 4, Haiku 4) -- Cost aggregation by project, agent, model, and call type +- Cost aggregation by project, agent, model, task, and workspace - Timeline-based token usage statistics +- Export to CSV and JSON Example: >>> from codeframe.lib.metrics_tracker import MetricsTracker @@ -17,8 +18,8 @@ >>> db.initialize() >>> tracker = MetricsTracker(db=db) >>> - >>> # Record token usage after LLM call - >>> usage_id = await tracker.record_token_usage( + >>> # Record token usage after LLM call (sync) + >>> usage_id = tracker.record_token_usage_sync( ... task_id=27, ... agent_id="backend-001", ... project_id=1, @@ -34,9 +35,12 @@ Total: $0.01 """ +import csv +import json import logging +import re from datetime import datetime, timedelta, timezone -from typing import Dict, Any, Optional +from typing import Dict, Any, List, Optional from codeframe.core.models import CallType, TokenUsage from codeframe.persistence.database import Database @@ -50,6 +54,36 @@ "claude-haiku-4": {"input": 0.80, "output": 4.00}, } +# Regex to strip -YYYYMMDD date suffixes from Anthropic API model names +# (e.g., "claude-sonnet-4-5-20250514" → "claude-sonnet-4-5") +_DATE_SUFFIX_RE = re.compile(r"-\d{8}$") + + +def normalize_model_name(raw_model: str) -> str: + """Normalize a model name by stripping date suffixes. + + The Anthropic API returns model names like 'claude-sonnet-4-5-20250514' + but our pricing dict uses 'claude-sonnet-4-5'. This function strips + the date suffix and returns the canonical name. + + Args: + raw_model: Raw model name from the API (e.g., 'claude-sonnet-4-5-20250514') + + Returns: + Normalized model name (e.g., 'claude-sonnet-4-5') + """ + # If it already matches a known model, return as-is + if raw_model in MODEL_PRICING: + return raw_model + + # Try stripping date suffix (8 digits at the end) + stripped = _DATE_SUFFIX_RE.sub("", raw_model) + if stripped in MODEL_PRICING: + return stripped + + # Unknown model - return as-is + return raw_model + class MetricsTracker: """Tracks token usage and costs for LLM API calls. @@ -89,16 +123,17 @@ def calculate_cost(model_name: str, input_tokens: int, output_tokens: int) -> fl - Claude Opus 4: $15.00 input / $75.00 output per MTok - Claude Haiku 4: $0.80 input / $4.00 output per MTok + Handles model names with date suffixes (e.g., 'claude-sonnet-4-5-20250514') + by normalizing them first. Unknown models return $0.00 cost instead of + raising, to avoid crashing the agent during recording. + Args: - model_name: Model identifier (e.g., "claude-sonnet-4-5") + model_name: Model identifier (e.g., "claude-sonnet-4-5" or "claude-sonnet-4-5-20250514") input_tokens: Number of input tokens output_tokens: Number of output tokens Returns: - Estimated cost in USD (rounded to 6 decimal places) - - Raises: - ValueError: If model_name is not recognized + Estimated cost in USD (rounded to 6 decimal places), or 0.0 for unknown models Example: >>> cost = MetricsTracker.calculate_cost( @@ -107,13 +142,16 @@ def calculate_cost(model_name: str, input_tokens: int, output_tokens: int) -> fl >>> print(f"${cost:.4f}") $0.0105 """ - if model_name not in MODEL_PRICING: - raise ValueError( - f"Unknown model: {model_name}. " - f"Supported models: {', '.join(MODEL_PRICING.keys())}" + normalized = normalize_model_name(model_name) + + if normalized not in MODEL_PRICING: + logger.warning( + f"Unknown model '{model_name}' (normalized: '{normalized}'). " + f"Returning $0.00 cost. Supported: {', '.join(MODEL_PRICING.keys())}" ) + return 0.0 - prices = MODEL_PRICING[model_name] + prices = MODEL_PRICING[normalized] # Calculate cost: (tokens * price_per_mtok) / 1,000,000 input_cost = (input_tokens * prices["input"]) / 1_000_000 @@ -170,12 +208,8 @@ async def record_token_usage( if input_tokens < 0 or output_tokens < 0: raise ValueError("Token counts cannot be negative") - # Calculate cost - try: - estimated_cost = self.calculate_cost(model_name, input_tokens, output_tokens) - except ValueError as e: - logger.error(f"Cost calculation failed: {e}") - raise + # Calculate cost (returns 0.0 for unknown models) + estimated_cost = self.calculate_cost(model_name, input_tokens, output_tokens) # Create TokenUsage model token_usage = TokenUsage( @@ -202,6 +236,170 @@ async def record_token_usage( return usage_id + def record_token_usage_sync( + self, + task_id: Optional[int], + agent_id: str, + project_id: int, + model_name: str, + input_tokens: int, + output_tokens: int, + call_type: CallType = CallType.OTHER, + session_id: Optional[str] = None, + ) -> int: + """Record token usage for an LLM call (synchronous version). + + Identical to record_token_usage but synchronous, for use from + synchronous code paths like the ReactAgent. + + Args: + task_id: Task ID if this call is related to a task (None for non-task calls) + agent_id: ID of the agent making the call + project_id: Project ID + model_name: Model identifier (e.g., "claude-sonnet-4-5") + input_tokens: Number of input tokens + output_tokens: Number of output tokens + call_type: Type of call (TASK_EXECUTION, CODE_REVIEW, COORDINATION, OTHER) + session_id: Optional SDK session ID for conversation tracking + + Returns: + Database ID of the created token usage record + + Raises: + ValueError: If token counts are negative + """ + if input_tokens < 0 or output_tokens < 0: + raise ValueError("Token counts cannot be negative") + + estimated_cost = self.calculate_cost(model_name, input_tokens, output_tokens) + + token_usage = TokenUsage( + task_id=task_id, + actual_cost_usd=None, + agent_id=agent_id, + project_id=project_id, + model_name=model_name, + input_tokens=input_tokens, + output_tokens=output_tokens, + estimated_cost_usd=estimated_cost, + call_type=call_type, + session_id=session_id, + timestamp=datetime.now(timezone.utc), + ) + + usage_id = self.db.save_token_usage(token_usage) + + logger.info( + f"Recorded token usage (sync): agent={agent_id}, model={model_name}, " + f"tokens={input_tokens + output_tokens}, cost=${estimated_cost:.6f}" + ) + + return usage_id + + def get_task_token_summary(self, task_id: int) -> Dict[str, Any]: + """Get aggregated token usage summary for a single task. + + Args: + task_id: Task ID to summarize + + Returns: + Dictionary with aggregated token data: + { + "task_id": int, + "total_input_tokens": int, + "total_output_tokens": int, + "total_tokens": int, + "total_cost_usd": float, + "call_count": int, + } + """ + return self.db.get_task_token_summary(task_id) + + def get_workspace_costs( + self, + start_date: Optional[datetime] = None, + end_date: Optional[datetime] = None, + ) -> Dict[str, Any]: + """Get aggregated costs across all tasks in the workspace. + + Args: + start_date: Optional start of date range (inclusive) + end_date: Optional end of date range (inclusive) + + Returns: + Dictionary with cost breakdown: + { + "total_cost_usd": float, + "total_tokens": int, + "total_calls": int, + } + """ + records = self.db.get_workspace_token_usage( + start_date=start_date, end_date=end_date + ) + + result: Dict[str, Any] = { + "total_cost_usd": 0.0, + "total_tokens": 0, + "total_calls": len(records), + } + + for record in records: + result["total_cost_usd"] += record["estimated_cost_usd"] + result["total_tokens"] += record["input_tokens"] + record["output_tokens"] + + result["total_cost_usd"] = round(result["total_cost_usd"], 6) + return result + + @staticmethod + def export_to_csv(records: List[Dict[str, Any]], output_path: str) -> None: + """Export token usage records to a CSV file. + + Args: + records: List of token usage record dictionaries + output_path: Path to write the CSV file + """ + fieldnames = [ + "id", "task_id", "agent_id", "project_id", "model_name", + "input_tokens", "output_tokens", "estimated_cost_usd", + "actual_cost_usd", "call_type", "session_id", "timestamp", + ] + + with open(output_path, "w", newline="") as f: + writer = csv.DictWriter(f, fieldnames=fieldnames, extrasaction="ignore") + writer.writeheader() + for record in records: + writer.writerow(record) + + @staticmethod + def export_to_json(records: List[Dict[str, Any]], output_path: str) -> None: + """Export token usage records to a JSON file with metadata. + + Args: + records: List of token usage record dictionaries + output_path: Path to write the JSON file + """ + # Convert sqlite3.Row objects to plain dicts if needed + serializable_records = [] + for record in records: + row = dict(record) + # Ensure all values are JSON-serializable + for key, value in row.items(): + if isinstance(value, datetime): + row[key] = value.isoformat() + serializable_records.append(row) + + data = { + "metadata": { + "exported_at": datetime.now(timezone.utc).isoformat(), + "record_count": len(serializable_records), + }, + "records": serializable_records, + } + + with open(output_path, "w") as f: + json.dump(data, f, indent=2, default=str) + async def get_project_costs( self, project_id: int, diff --git a/codeframe/persistence/database.py b/codeframe/persistence/database.py index 07360c2d..5674d279 100644 --- a/codeframe/persistence/database.py +++ b/codeframe/persistence/database.py @@ -756,6 +756,18 @@ def get_project_costs_aggregate(self, *args, **kwargs): """Delegate to token_usage.get_project_costs_aggregate().""" return self.token_usage.get_project_costs_aggregate(*args, **kwargs) + def get_task_token_summary(self, *args, **kwargs): + """Delegate to token_usage.get_task_token_summary().""" + return self.token_usage.get_task_token_summary(*args, **kwargs) + + def get_batch_token_usage(self, *args, **kwargs): + """Delegate to token_usage.get_batch_token_usage().""" + return self.token_usage.get_batch_token_usage(*args, **kwargs) + + def get_workspace_token_usage(self, *args, **kwargs): + """Delegate to token_usage.get_workspace_token_usage().""" + return self.token_usage.get_workspace_token_usage(*args, **kwargs) + def create_correction_attempt(self, *args, **kwargs): """Delegate to correction_attempts.create_correction_attempt().""" return self.correction_attempts.create_correction_attempt(*args, **kwargs) diff --git a/codeframe/persistence/repositories/token_repository.py b/codeframe/persistence/repositories/token_repository.py index 036c0696..67111ec0 100644 --- a/codeframe/persistence/repositories/token_repository.py +++ b/codeframe/persistence/repositories/token_repository.py @@ -164,6 +164,116 @@ def get_token_usage( + def get_task_token_summary(self, task_id: int) -> Dict[str, Any]: + """Get aggregated token usage summary for a single task. + + Args: + task_id: Task ID to summarize + + Returns: + Dictionary with aggregated token data: + { + "task_id": int, + "total_input_tokens": int, + "total_output_tokens": int, + "total_tokens": int, + "total_cost_usd": float, + "call_count": int, + } + """ + cursor = self.conn.cursor() + cursor.execute( + """ + SELECT + COALESCE(SUM(input_tokens), 0) as total_input_tokens, + COALESCE(SUM(output_tokens), 0) as total_output_tokens, + COALESCE(SUM(input_tokens + output_tokens), 0) as total_tokens, + COALESCE(SUM(estimated_cost_usd), 0.0) as total_cost_usd, + COUNT(*) as call_count + FROM token_usage + WHERE task_id = ? + """, + (task_id,), + ) + row = cursor.fetchone() + + return { + "task_id": task_id, + "total_input_tokens": row["total_input_tokens"], + "total_output_tokens": row["total_output_tokens"], + "total_tokens": row["total_tokens"], + "total_cost_usd": row["total_cost_usd"], + "call_count": row["call_count"], + } + + def get_batch_token_usage( + self, + task_ids: List[int], + start_date: Optional[datetime] = None, + end_date: Optional[datetime] = None, + ) -> List[Dict[str, Any]]: + """Get token usage records filtered by a list of task IDs. + + Args: + task_ids: List of task IDs to filter by + start_date: Optional start of date range (inclusive) + end_date: Optional end of date range (inclusive) + + Returns: + List of token usage records as dictionaries + """ + if not task_ids: + return [] + + cursor = self.conn.cursor() + placeholders = ",".join("?" for _ in task_ids) + query = f"SELECT * FROM token_usage WHERE task_id IN ({placeholders})" + params: list = list(task_ids) + + if start_date is not None: + query += " AND timestamp >= ?" + params.append(start_date.isoformat()) + + if end_date is not None: + query += " AND timestamp <= ?" + params.append(end_date.isoformat()) + + query += " ORDER BY timestamp DESC" + + cursor.execute(query, params) + return [dict(row) for row in cursor.fetchall()] + + def get_workspace_token_usage( + self, + start_date: Optional[datetime] = None, + end_date: Optional[datetime] = None, + ) -> List[Dict[str, Any]]: + """Get all token usage records across the workspace. + + Args: + start_date: Optional start of date range (inclusive) + end_date: Optional end of date range (inclusive) + + Returns: + List of token usage records as dictionaries + """ + cursor = self.conn.cursor() + query = "SELECT * FROM token_usage WHERE 1=1" + params: list = [] + + if start_date is not None: + query += " AND timestamp >= ?" + params.append(start_date.isoformat()) + + if end_date is not None: + query += " AND timestamp <= ?" + params.append(end_date.isoformat()) + + query += " ORDER BY timestamp DESC" + + cursor.execute(query, params) + return [dict(row) for row in cursor.fetchall()] + def get_project_costs_aggregate(self, project_id: int) -> Dict[str, Any]: """Get aggregated cost statistics for a project. diff --git a/tests/cli/test_stats_commands.py b/tests/cli/test_stats_commands.py new file mode 100644 index 00000000..2da220e9 --- /dev/null +++ b/tests/cli/test_stats_commands.py @@ -0,0 +1,251 @@ +"""Tests for CLI stats commands (headless token/cost tracking). + +TDD approach: Write tests first, then implement. +Tests the `cf stats tokens`, `cf stats costs`, and `cf stats export` commands. +""" + +import csv +import json +import os +from datetime import datetime, timezone + +import pytest +from typer.testing import CliRunner + +from codeframe.cli.stats_commands import stats_app +from codeframe.core.models import CallType, TokenUsage +from codeframe.persistence.database import Database + +pytestmark = pytest.mark.v2 + +runner = CliRunner() + + +def _seed_project_and_tasks(db): + """Create a project and tasks to satisfy FK constraints.""" + cursor = db.conn.cursor() + cursor.execute( + "INSERT INTO projects (name, description, source_type, source_branch, workspace_path) " + "VALUES (?, ?, ?, ?, ?)", + ("test-project", "Test project", "empty", "main", "/tmp/test"), + ) + cursor.execute( + "INSERT INTO tasks (project_id, title, description, status, priority) " + "VALUES (?, ?, ?, ?, ?)", + (1, "Task 1", "First task", "in_progress", 0), + ) + cursor.execute( + "INSERT INTO tasks (project_id, title, description, status, priority) " + "VALUES (?, ?, ?, ?, ?)", + (1, "Task 2", "Second task", "in_progress", 0), + ) + db.conn.commit() + + +@pytest.fixture +def workspace_with_tokens(tmp_path): + """Create a workspace with seeded token usage data.""" + codeframe_dir = tmp_path / ".codeframe" + codeframe_dir.mkdir() + db = Database(codeframe_dir / "state.db") + db.initialize() + + _seed_project_and_tasks(db) + + records = [ + TokenUsage( + task_id=1, + agent_id="react-agent", + project_id=1, + model_name="claude-sonnet-4-5", + input_tokens=1000, + output_tokens=500, + estimated_cost_usd=0.0105, + call_type=CallType.TASK_EXECUTION, + timestamp=datetime(2026, 3, 10, 10, 0, 0, tzinfo=timezone.utc), + ), + TokenUsage( + task_id=1, + agent_id="react-agent", + project_id=1, + model_name="claude-sonnet-4-5", + input_tokens=2000, + output_tokens=800, + estimated_cost_usd=0.018, + call_type=CallType.TASK_EXECUTION, + timestamp=datetime(2026, 3, 10, 11, 0, 0, tzinfo=timezone.utc), + ), + TokenUsage( + task_id=2, + agent_id="react-agent", + project_id=1, + model_name="claude-haiku-4", + input_tokens=500, + output_tokens=200, + estimated_cost_usd=0.0012, + call_type=CallType.CODE_REVIEW, + timestamp=datetime(2026, 3, 12, 10, 0, 0, tzinfo=timezone.utc), + ), + ] + + for record in records: + db.save_token_usage(record) + + db.close() + return tmp_path + + +@pytest.fixture +def empty_workspace(tmp_path): + """Create a workspace with initialized DB but no token data.""" + codeframe_dir = tmp_path / ".codeframe" + codeframe_dir.mkdir() + db = Database(codeframe_dir / "state.db") + db.initialize() + db.close() + return tmp_path + + +# ============================================================================= +# cf stats tokens +# ============================================================================= + + +class TestStatsTokens: + """Tests for 'cf stats tokens' command.""" + + def test_stats_tokens_no_workspace(self, tmp_path, monkeypatch): + """Should show error when no workspace exists.""" + monkeypatch.chdir(tmp_path) + result = runner.invoke(stats_app, ["tokens"]) + assert result.exit_code == 1 + assert "No workspace found" in result.output + + def test_stats_tokens_empty(self, empty_workspace, monkeypatch): + """Should show zeros when no token data exists.""" + monkeypatch.chdir(empty_workspace) + result = runner.invoke(stats_app, ["tokens"]) + assert result.exit_code == 0 + assert "0" in result.output + + def test_stats_tokens_with_data(self, workspace_with_tokens, monkeypatch): + """Should show correct summary with seeded data.""" + monkeypatch.chdir(workspace_with_tokens) + result = runner.invoke(stats_app, ["tokens"]) + assert result.exit_code == 0 + # Total tokens: 1000+500 + 2000+800 + 500+200 = 5000 + assert "5,000" in result.output or "5000" in result.output + # Should show input/output breakdown + assert "Input" in result.output + assert "Output" in result.output + + def test_stats_tokens_task_filter(self, workspace_with_tokens, monkeypatch): + """Should show per-task breakdown when --task is provided.""" + monkeypatch.chdir(workspace_with_tokens) + result = runner.invoke(stats_app, ["tokens", "--task", "1"]) + assert result.exit_code == 0 + # Task 1 tokens: 1000+500 + 2000+800 = 4300 + assert "4,300" in result.output or "4300" in result.output + + +# ============================================================================= +# cf stats costs +# ============================================================================= + + +class TestStatsCosts: + """Tests for 'cf stats costs' command.""" + + def test_stats_costs_no_workspace(self, tmp_path, monkeypatch): + """Should show error when no workspace exists.""" + monkeypatch.chdir(tmp_path) + result = runner.invoke(stats_app, ["costs"]) + assert result.exit_code == 1 + assert "No workspace found" in result.output + + def test_stats_costs_default(self, workspace_with_tokens, monkeypatch): + """Should show all-time costs.""" + monkeypatch.chdir(workspace_with_tokens) + result = runner.invoke(stats_app, ["costs"]) + assert result.exit_code == 0 + assert "$" in result.output + # Total cost: 0.0105 + 0.018 + 0.0012 = 0.0297 + assert "0.0297" in result.output + + def test_stats_costs_period_month(self, workspace_with_tokens, monkeypatch): + """Should respect period filter for 'month'.""" + monkeypatch.chdir(workspace_with_tokens) + result = runner.invoke(stats_app, ["costs", "--period", "month"]) + assert result.exit_code == 0 + assert "$" in result.output + + def test_stats_costs_period_week(self, workspace_with_tokens, monkeypatch): + """Should respect period filter for 'week'.""" + monkeypatch.chdir(workspace_with_tokens) + result = runner.invoke(stats_app, ["costs", "--period", "week"]) + assert result.exit_code == 0 + + def test_stats_costs_period_day(self, workspace_with_tokens, monkeypatch): + """Should respect period filter for 'day'.""" + monkeypatch.chdir(workspace_with_tokens) + result = runner.invoke(stats_app, ["costs", "--period", "day"]) + assert result.exit_code == 0 + + +# ============================================================================= +# cf stats export +# ============================================================================= + + +class TestStatsExport: + """Tests for 'cf stats export' command.""" + + def test_stats_export_no_workspace(self, tmp_path, monkeypatch): + """Should show error when no workspace exists.""" + monkeypatch.chdir(tmp_path) + output_file = str(tmp_path / "out.csv") + result = runner.invoke(stats_app, ["export", "--format", "csv", "--output", output_file]) + assert result.exit_code == 1 + + def test_stats_export_csv(self, workspace_with_tokens, monkeypatch): + """Should create a valid CSV file.""" + monkeypatch.chdir(workspace_with_tokens) + output_file = str(workspace_with_tokens / "tokens.csv") + result = runner.invoke(stats_app, ["export", "--format", "csv", "--output", output_file]) + assert result.exit_code == 0 + assert os.path.exists(output_file) + + with open(output_file) as f: + reader = csv.DictReader(f) + rows = list(reader) + assert len(rows) == 3 + assert "input_tokens" in rows[0] + + def test_stats_export_json(self, workspace_with_tokens, monkeypatch): + """Should create a valid JSON file.""" + monkeypatch.chdir(workspace_with_tokens) + output_file = str(workspace_with_tokens / "tokens.json") + result = runner.invoke(stats_app, ["export", "--format", "json", "--output", output_file]) + assert result.exit_code == 0 + assert os.path.exists(output_file) + + with open(output_file) as f: + data = json.load(f) + assert "records" in data + assert len(data["records"]) == 3 + + def test_stats_export_csv_task_filter(self, workspace_with_tokens, monkeypatch): + """Should export only records for a specific task.""" + monkeypatch.chdir(workspace_with_tokens) + output_file = str(workspace_with_tokens / "task1.csv") + result = runner.invoke( + stats_app, ["export", "--format", "csv", "--output", output_file, "--task", "1"] + ) + assert result.exit_code == 0 + assert os.path.exists(output_file) + + with open(output_file) as f: + reader = csv.DictReader(f) + rows = list(reader) + # Task 1 has 2 records + assert len(rows) == 2 diff --git a/tests/core/test_react_agent_tokens.py b/tests/core/test_react_agent_tokens.py new file mode 100644 index 00000000..f254d9a4 --- /dev/null +++ b/tests/core/test_react_agent_tokens.py @@ -0,0 +1,296 @@ +"""Tests for ReactAgent token usage tracking. + +Verifies that the ReactAgent accumulates token records during execution, +provides aggregation methods, and handles persistence failures gracefully. +""" + +from datetime import datetime, timezone +from unittest.mock import patch + +import pytest + +from codeframe.adapters.llm.base import ( + LLMResponse, + ToolCall, + ToolResult, +) +from codeframe.adapters.llm.mock import MockProvider +from codeframe.core.agent import AgentStatus +from codeframe.core.context import TaskContext +from codeframe.core.gates import GateResult, GateCheck, GateStatus +from codeframe.core.tasks import Task, TaskStatus +from codeframe.core.workspace import Workspace + +pytestmark = pytest.mark.v2 + + +# --------------------------------------------------------------------------- +# Fixtures +# --------------------------------------------------------------------------- + + +@pytest.fixture +def workspace(tmp_path): + """Create a minimal workspace for testing.""" + state_dir = tmp_path / ".codeframe" + state_dir.mkdir() + return Workspace( + id="ws-test", + repo_path=tmp_path, + state_dir=state_dir, + created_at=datetime(2026, 1, 1, tzinfo=timezone.utc), + tech_stack="Python with uv", + ) + + +@pytest.fixture +def mock_task(): + """Create a minimal task.""" + _ts = datetime(2026, 1, 1, tzinfo=timezone.utc) + return Task( + id="task-1", + workspace_id="ws-test", + prd_id=None, + title="Add hello function", + description="Create a hello() function that returns 'Hello, World!'", + status=TaskStatus.IN_PROGRESS, + priority=1, + created_at=_ts, + updated_at=_ts, + ) + + +@pytest.fixture +def mock_context(mock_task): + """Create a minimal TaskContext.""" + return TaskContext(task=mock_task) + + +@pytest.fixture +def provider(): + """Create a MockProvider.""" + return MockProvider() + + +def _gate_passed(): + """Return a GateResult that passed.""" + return GateResult( + passed=True, + checks=[GateCheck(name="ruff", status=GateStatus.PASSED)], + ) + + +# --------------------------------------------------------------------------- +# Tests +# --------------------------------------------------------------------------- + + +class TestReactAgentTokenAccumulation: + """Tests for token record accumulation during agent run.""" + + @patch("codeframe.core.react_agent.gates") + @patch("codeframe.core.react_agent.execute_tool") + @patch("codeframe.core.react_agent.TaskContextPackager") + def test_react_agent_accumulates_token_records( + self, mock_ctx_loader, mock_exec_tool, mock_gates, workspace, provider, mock_context + ): + """Verify tokens are collected during a run with tool calls.""" + from codeframe.core.react_agent import ReactAgent + + # First call: tool call with known token counts + provider.add_response( + LLMResponse( + content="", + tool_calls=[ToolCall(id="tc1", name="read_file", input={"path": "a.py"})], + stop_reason="tool_use", + model="claude-sonnet-4-20250514", + input_tokens=150, + output_tokens=50, + ) + ) + # Second call: text response (agent done) + provider.add_response( + LLMResponse( + content="I have completed the task.", + model="claude-sonnet-4-20250514", + input_tokens=200, + output_tokens=30, + ) + ) + + mock_ctx_loader.return_value.load_context.return_value = mock_context + mock_exec_tool.return_value = ToolResult(tool_call_id="tc1", content="file contents") + mock_gates.run.return_value = _gate_passed() + + agent = ReactAgent(workspace=workspace, llm_provider=provider) + status = agent.run("task-1") + + assert status == AgentStatus.COMPLETED + + records = agent.get_token_usage() + assert len(records) == 2 + + # First record: tool call iteration + assert records[0]["input_tokens"] == 150 + assert records[0]["output_tokens"] == 50 + assert records[0]["model"] == "claude-sonnet-4-20250514" + assert records[0]["call_type"] == "task_execution" + assert records[0]["iteration"] == 1 + + # Second record: text response iteration + assert records[1]["input_tokens"] == 200 + assert records[1]["output_tokens"] == 30 + assert records[1]["iteration"] == 2 + + @patch("codeframe.core.react_agent.gates") + @patch("codeframe.core.react_agent.execute_tool") + @patch("codeframe.core.react_agent.TaskContextPackager") + def test_react_agent_accumulates_verification_tokens( + self, mock_ctx_loader, mock_exec_tool, mock_gates, workspace, provider, mock_context + ): + """Verify tokens from verification fix loops are also recorded.""" + from codeframe.core.react_agent import ReactAgent + + # Main loop: text response (agent done quickly) + provider.add_response( + LLMResponse( + content="I have completed the task.", + model="claude-sonnet-4-20250514", + input_tokens=100, + output_tokens=20, + ) + ) + + # Verification fails first time, then LLM fixes, then passes + failed_gate = GateResult( + passed=False, + checks=[GateCheck(name="ruff", status=GateStatus.FAILED, output="error")], + ) + # First gate check fails, second passes + mock_gates.run.side_effect = [failed_gate, _gate_passed()] + + # Fix loop LLM response (text-only, no tools needed) + provider.add_response( + LLMResponse( + content="Fixed the issue.", + model="claude-sonnet-4-20250514", + input_tokens=300, + output_tokens=60, + ) + ) + + mock_ctx_loader.return_value.load_context.return_value = mock_context + mock_exec_tool.return_value = ToolResult(tool_call_id="tc1", content="ok") + + agent = ReactAgent(workspace=workspace, llm_provider=provider) + status = agent.run("task-1") + + assert status == AgentStatus.COMPLETED + + records = agent.get_token_usage() + # Should have at least the main loop call + the verification fix call + assert len(records) >= 2 + + # Check that verification fix tokens are recorded with correct call_type + verification_records = [r for r in records if r["call_type"] == "verification_fix"] + assert len(verification_records) >= 1 + assert verification_records[0]["input_tokens"] == 300 + assert verification_records[0]["output_tokens"] == 60 + + +class TestReactAgentGetTotalTokens: + """Tests for get_total_tokens() aggregation method.""" + + @patch("codeframe.core.react_agent.gates") + @patch("codeframe.core.react_agent.execute_tool") + @patch("codeframe.core.react_agent.TaskContextPackager") + def test_react_agent_get_total_tokens( + self, mock_ctx_loader, mock_exec_tool, mock_gates, workspace, provider, mock_context + ): + """Verify get_total_tokens() returns correct aggregation.""" + from codeframe.core.react_agent import ReactAgent + + # Two iterations with known token counts + provider.add_response( + LLMResponse( + content="", + tool_calls=[ToolCall(id="tc1", name="read_file", input={"path": "a.py"})], + stop_reason="tool_use", + model="claude-sonnet-4-20250514", + input_tokens=100, + output_tokens=40, + ) + ) + provider.add_response( + LLMResponse( + content="Done.", + model="claude-sonnet-4-20250514", + input_tokens=200, + output_tokens=60, + ) + ) + + mock_ctx_loader.return_value.load_context.return_value = mock_context + mock_exec_tool.return_value = ToolResult(tool_call_id="tc1", content="ok") + mock_gates.run.return_value = _gate_passed() + + agent = ReactAgent(workspace=workspace, llm_provider=provider) + agent.run("task-1") + + totals = agent.get_total_tokens() + assert totals["input_tokens"] == 300 + assert totals["output_tokens"] == 100 + assert totals["total_tokens"] == 400 + assert "estimated_cost_usd" in totals + assert isinstance(totals["estimated_cost_usd"], float) + assert totals["estimated_cost_usd"] >= 0.0 + + def test_get_total_tokens_empty(self, workspace, provider): + """get_total_tokens() returns zeros when no calls have been made.""" + from codeframe.core.react_agent import ReactAgent + + agent = ReactAgent(workspace=workspace, llm_provider=provider) + totals = agent.get_total_tokens() + assert totals["input_tokens"] == 0 + assert totals["output_tokens"] == 0 + assert totals["total_tokens"] == 0 + assert totals["estimated_cost_usd"] == 0.0 + + +class TestReactAgentTokenPersistenceFailure: + """Tests for graceful handling of token persistence failures.""" + + @patch("codeframe.core.react_agent.gates") + @patch("codeframe.core.react_agent.execute_tool") + @patch("codeframe.core.react_agent.TaskContextPackager") + def test_react_agent_token_persistence_failure_doesnt_crash( + self, mock_ctx_loader, mock_exec_tool, mock_gates, workspace, provider, mock_context + ): + """Verify that a failure in _persist_token_usage does not crash the agent.""" + from codeframe.core.react_agent import ReactAgent + + # Simple text response + provider.add_response( + LLMResponse( + content="I have completed the task.", + model="claude-sonnet-4-20250514", + input_tokens=100, + output_tokens=20, + ) + ) + + mock_ctx_loader.return_value.load_context.return_value = mock_context + mock_gates.run.return_value = _gate_passed() + + agent = ReactAgent(workspace=workspace, llm_provider=provider) + + # Make _persist_token_usage raise an exception + with patch.object(agent, "_persist_token_usage", side_effect=Exception("DB error")): + status = agent.run("task-1") + + # Agent should still complete successfully despite persistence failure + assert status == AgentStatus.COMPLETED + + # Token records should still be available in-memory + records = agent.get_token_usage() + assert len(records) == 1 diff --git a/tests/lib/test_metrics_tracker.py b/tests/lib/test_metrics_tracker.py index ce557f09..e1ca92da 100644 --- a/tests/lib/test_metrics_tracker.py +++ b/tests/lib/test_metrics_tracker.py @@ -175,10 +175,12 @@ def test_calculate_cost_haiku(): def test_calculate_cost_unknown_model(): - """Test that unknown model raises ValueError.""" - # When/Then: Unknown model raises ValueError - with pytest.raises(ValueError, match="Unknown model"): - MetricsTracker.calculate_cost("claude-unknown-99", 1000, 500) + """Test that unknown model returns $0 instead of crashing.""" + # When: We calculate cost for an unknown model + cost = MetricsTracker.calculate_cost("claude-unknown-99", 1000, 500) + + # Then: Returns 0.0 (graceful degradation) + assert cost == 0.0 # ============================================================================ @@ -430,3 +432,331 @@ async def test_get_token_usage_stats_with_end_date(tracker, db): # Then: No usages match assert result["total_cost_usd"] == 0.0 assert result["total_calls"] == 0 + + +# ============================================================================ +# Step 2: Sync recording, aggregation, export +# ============================================================================ + + +def test_record_token_usage_sync(tracker, db): + """Test synchronous recording of token usage.""" + # Given: A task exists + cursor = db.conn.cursor() + cursor.execute( + "INSERT INTO tasks (project_id, title, description, status) VALUES (?, ?, ?, ?)", + (1, "Test task", "Test description", "in_progress"), + ) + db.conn.commit() + task_id = cursor.lastrowid + + # When: We record token usage synchronously + usage_id = tracker.record_token_usage_sync( + task_id=task_id, + agent_id="backend-001", + project_id=1, + model_name="claude-sonnet-4-5", + input_tokens=1000, + output_tokens=500, + call_type=CallType.TASK_EXECUTION, + ) + + # Then: Token usage is saved to database + assert usage_id > 0 + + # And: We can retrieve it + cursor.execute("SELECT * FROM token_usage WHERE id = ?", (usage_id,)) + row = cursor.fetchone() + assert row is not None + assert row["task_id"] == task_id + assert row["agent_id"] == "backend-001" + assert row["model_name"] == "claude-sonnet-4-5" + assert row["input_tokens"] == 1000 + assert row["output_tokens"] == 500 + assert row["estimated_cost_usd"] > 0 + + +def test_record_token_usage_sync_negative_tokens(tracker): + """Test sync recording rejects negative token counts.""" + with pytest.raises(ValueError, match="Token counts cannot be negative"): + tracker.record_token_usage_sync( + task_id=1, + agent_id="agent-001", + project_id=1, + model_name="claude-sonnet-4-5", + input_tokens=-1, + output_tokens=500, + ) + + +def test_get_task_token_summary(tracker, db): + """Test getting aggregated token summary for a single task.""" + # Given: A task exists and has multiple token usages + cursor = db.conn.cursor() + cursor.execute( + "INSERT INTO tasks (project_id, title, description, status) VALUES (?, ?, ?, ?)", + (1, "Summary task", "Test", "in_progress"), + ) + db.conn.commit() + task_id = cursor.lastrowid + + tracker.record_token_usage_sync( + task_id=task_id, + agent_id="backend-001", + project_id=1, + model_name="claude-sonnet-4-5", + input_tokens=1000, + output_tokens=500, + call_type=CallType.TASK_EXECUTION, + ) + tracker.record_token_usage_sync( + task_id=task_id, + agent_id="backend-001", + project_id=1, + model_name="claude-haiku-4", + input_tokens=2000, + output_tokens=1000, + call_type=CallType.CODE_REVIEW, + ) + + # When: We get the task summary + summary = tracker.get_task_token_summary(task_id=task_id) + + # Then: Aggregated values are correct + assert summary["task_id"] == task_id + assert summary["total_input_tokens"] == 3000 + assert summary["total_output_tokens"] == 1500 + assert summary["total_tokens"] == 4500 + assert summary["call_count"] == 2 + assert summary["total_cost_usd"] > 0 + + +def test_get_task_token_summary_no_records(tracker): + """Test task summary with no records returns zeros.""" + summary = tracker.get_task_token_summary(task_id=999) + + assert summary["task_id"] == 999 + assert summary["total_tokens"] == 0 + assert summary["total_cost_usd"] == 0.0 + assert summary["call_count"] == 0 + + +def _create_task_helper(db): + """Helper to create a task and return its ID.""" + cursor = db.conn.cursor() + cursor.execute( + "INSERT INTO tasks (project_id, title, description, status) VALUES (?, ?, ?, ?)", + (1, "Test task", "Test", "in_progress"), + ) + db.conn.commit() + return cursor.lastrowid + + +def test_get_workspace_costs(tracker, db): + """Test getting aggregated costs across the workspace.""" + # Given: Token usages across different tasks/projects + tid1 = _create_task_helper(db) + tid2 = _create_task_helper(db) + tracker.record_token_usage_sync( + task_id=tid1, + agent_id="backend-001", + project_id=1, + model_name="claude-sonnet-4-5", + input_tokens=1_000_000, + output_tokens=500_000, + ) + tracker.record_token_usage_sync( + task_id=tid2, + agent_id="frontend-001", + project_id=1, + model_name="claude-haiku-4", + input_tokens=500_000, + output_tokens=250_000, + ) + + # When: We get workspace costs + result = tracker.get_workspace_costs() + + # Then: All records are aggregated + # Sonnet: $10.50, Haiku: $1.40 => Total: $11.90 + assert result["total_cost_usd"] == pytest.approx(11.90, abs=0.01) + assert result["total_tokens"] == 2_250_000 + assert result["total_calls"] == 2 + + +def test_get_workspace_costs_with_date_filter(tracker, db): + """Test workspace costs with date range filtering.""" + now = datetime.now(timezone.utc) + + # Recent usage + tid = _create_task_helper(db) + tracker.record_token_usage_sync( + task_id=tid, + agent_id="backend-001", + project_id=1, + model_name="claude-sonnet-4-5", + input_tokens=1000, + output_tokens=500, + ) + + # Old usage (manually backdated) + old_timestamp = now - timedelta(days=10) + cursor = db.conn.cursor() + cursor.execute( + """INSERT INTO token_usage + (agent_id, project_id, model_name, input_tokens, output_tokens, + estimated_cost_usd, call_type, timestamp) + VALUES (?, ?, ?, ?, ?, ?, ?, ?)""", + ("agent-old", 1, "claude-haiku-4", 500, 250, 0.001, "other", old_timestamp.isoformat()), + ) + db.conn.commit() + + # When: We filter to last 7 days + start = now - timedelta(days=7) + result = tracker.get_workspace_costs(start_date=start) + + # Then: Only recent usage is included + assert result["total_calls"] == 1 + + +def test_get_workspace_costs_empty(tracker): + """Test workspace costs with no records.""" + result = tracker.get_workspace_costs() + + assert result["total_cost_usd"] == 0.0 + assert result["total_tokens"] == 0 + assert result["total_calls"] == 0 + + +def test_export_to_csv(tracker, db, tmp_path): + """Test exporting token usage records to CSV.""" + # Given: Some token usage records + tid = _create_task_helper(db) + tracker.record_token_usage_sync( + task_id=tid, + agent_id="backend-001", + project_id=1, + model_name="claude-sonnet-4-5", + input_tokens=1000, + output_tokens=500, + ) + records = db.get_workspace_token_usage() + + # When: We export to CSV + output_path = tmp_path / "usage.csv" + tracker.export_to_csv(records, str(output_path)) + + # Then: CSV file is created with correct content + assert output_path.exists() + content = output_path.read_text() + lines = content.strip().split("\n") + assert len(lines) == 2 # header + 1 record + header = lines[0] + assert "task_id" in header + assert "model_name" in header + assert "input_tokens" in header + assert "estimated_cost_usd" in header + + +def test_export_to_csv_empty(tracker, tmp_path): + """Test exporting empty records to CSV.""" + output_path = tmp_path / "empty.csv" + tracker.export_to_csv([], str(output_path)) + + assert output_path.exists() + content = output_path.read_text() + lines = content.strip().split("\n") + assert len(lines) == 1 # header only + + +def test_export_to_json(tracker, db, tmp_path): + """Test exporting token usage records to JSON.""" + import json + + # Given: Some token usage records + tid = _create_task_helper(db) + tracker.record_token_usage_sync( + task_id=tid, + agent_id="backend-001", + project_id=1, + model_name="claude-sonnet-4-5", + input_tokens=1000, + output_tokens=500, + ) + records = db.get_workspace_token_usage() + + # When: We export to JSON + output_path = tmp_path / "usage.json" + tracker.export_to_json(records, str(output_path)) + + # Then: JSON file is created with correct structure + assert output_path.exists() + data = json.loads(output_path.read_text()) + assert "metadata" in data + assert "records" in data + assert data["metadata"]["record_count"] == 1 + assert "exported_at" in data["metadata"] + assert len(data["records"]) == 1 + assert data["records"][0]["model_name"] == "claude-sonnet-4-5" + + +def test_export_to_json_empty(tracker, tmp_path): + """Test exporting empty records to JSON.""" + import json + + output_path = tmp_path / "empty.json" + tracker.export_to_json([], str(output_path)) + + assert output_path.exists() + data = json.loads(output_path.read_text()) + assert data["metadata"]["record_count"] == 0 + assert len(data["records"]) == 0 + + +# ============================================================================ +# Step 5: Model name normalization +# ============================================================================ + + +def test_normalize_model_name_with_date_suffix(): + """Test that date suffixes are stripped from model names.""" + from codeframe.lib.metrics_tracker import normalize_model_name + + assert normalize_model_name("claude-sonnet-4-5-20250514") == "claude-sonnet-4-5" + assert normalize_model_name("claude-opus-4-20250514") == "claude-opus-4" + assert normalize_model_name("claude-haiku-4-20250514") == "claude-haiku-4" + + +def test_normalize_model_name_exact_match(): + """Test that exact model names pass through unchanged.""" + from codeframe.lib.metrics_tracker import normalize_model_name + + assert normalize_model_name("claude-sonnet-4-5") == "claude-sonnet-4-5" + assert normalize_model_name("claude-opus-4") == "claude-opus-4" + assert normalize_model_name("claude-haiku-4") == "claude-haiku-4" + + +def test_normalize_model_name_unknown_model(): + """Test that unknown models return as-is.""" + from codeframe.lib.metrics_tracker import normalize_model_name + + assert normalize_model_name("gpt-4-turbo") == "gpt-4-turbo" + assert normalize_model_name("some-unknown-model") == "some-unknown-model" + + +def test_calculate_cost_with_date_suffix(): + """Test that calculate_cost handles model names with date suffixes.""" + # Should work the same as without the suffix + cost_with_suffix = MetricsTracker.calculate_cost( + "claude-sonnet-4-5-20250514", 1000, 500 + ) + cost_without_suffix = MetricsTracker.calculate_cost( + "claude-sonnet-4-5", 1000, 500 + ) + assert cost_with_suffix == cost_without_suffix + + +def test_calculate_cost_unknown_model_returns_zero(): + """Test that unknown models return $0 cost instead of raising.""" + cost = MetricsTracker.calculate_cost("totally-unknown-model", 1000, 500) + assert cost == 0.0 diff --git a/tests/persistence/test_token_repository.py b/tests/persistence/test_token_repository.py new file mode 100644 index 00000000..65b50f7a --- /dev/null +++ b/tests/persistence/test_token_repository.py @@ -0,0 +1,214 @@ +"""Tests for TokenRepository query methods (Issue #314 Step 3). + +Tests for: +- get_task_token_summary: SQL aggregate for a single task +- get_batch_token_usage: Filter by list of task_ids +- get_workspace_token_usage: All records, no project filter +""" + +import pytest +from datetime import datetime, timedelta, timezone + +from codeframe.core.models import CallType, TokenUsage +from codeframe.persistence.database import Database + +pytestmark = pytest.mark.v2 + + +@pytest.fixture +def db(): + """Create in-memory database for testing.""" + database = Database(":memory:") + database.initialize() + + # Create test project + cursor = database.conn.cursor() + cursor.execute( + "INSERT INTO projects (name, description, workspace_path, status) VALUES (?, ?, ?, ?)", + ("test-project", "Test project", "/tmp/test", "active"), + ) + database.conn.commit() + + return database + + +def _create_task(db, project_id=1, task_id_hint=None): + """Helper to create a task and return its ID.""" + cursor = db.conn.cursor() + cursor.execute( + "INSERT INTO tasks (project_id, title, description, status) VALUES (?, ?, ?, ?)", + (project_id, f"Task {task_id_hint or 'x'}", "Test task", "in_progress"), + ) + db.conn.commit() + return cursor.lastrowid + + +def _save_usage(db, task_id=None, agent_id="agent-001", project_id=1, + model_name="claude-sonnet-4-5", input_tokens=1000, + output_tokens=500, cost=0.0105, call_type=CallType.TASK_EXECUTION, + timestamp=None): + """Helper to save a token usage record.""" + if timestamp is None: + timestamp = datetime.now(timezone.utc) + usage = TokenUsage( + task_id=task_id, + agent_id=agent_id, + project_id=project_id, + model_name=model_name, + input_tokens=input_tokens, + output_tokens=output_tokens, + estimated_cost_usd=cost, + actual_cost_usd=None, + call_type=call_type, + timestamp=timestamp, + ) + return db.save_token_usage(usage) + + +# ============================================================================ +# get_task_token_summary +# ============================================================================ + + +def test_get_task_token_summary_single_call(db): + """Test task summary with a single LLM call.""" + tid = _create_task(db) + _save_usage(db, task_id=tid, input_tokens=1000, output_tokens=500, cost=0.0105) + + summary = db.get_task_token_summary(task_id=tid) + + assert summary["task_id"] == tid + assert summary["total_input_tokens"] == 1000 + assert summary["total_output_tokens"] == 500 + assert summary["total_tokens"] == 1500 + assert summary["total_cost_usd"] == pytest.approx(0.0105, abs=1e-6) + assert summary["call_count"] == 1 + + +def test_get_task_token_summary_multiple_calls(db): + """Test task summary aggregates multiple LLM calls.""" + tid = _create_task(db) + _save_usage(db, task_id=tid, input_tokens=1000, output_tokens=500, cost=0.01) + _save_usage(db, task_id=tid, input_tokens=2000, output_tokens=1000, cost=0.02) + + summary = db.get_task_token_summary(task_id=tid) + + assert summary["total_input_tokens"] == 3000 + assert summary["total_output_tokens"] == 1500 + assert summary["total_tokens"] == 4500 + assert summary["total_cost_usd"] == pytest.approx(0.03, abs=1e-6) + assert summary["call_count"] == 2 + + +def test_get_task_token_summary_no_records(db): + """Test task summary returns zeros when no records exist.""" + summary = db.get_task_token_summary(task_id=999) + + assert summary["task_id"] == 999 + assert summary["total_input_tokens"] == 0 + assert summary["total_output_tokens"] == 0 + assert summary["total_tokens"] == 0 + assert summary["total_cost_usd"] == 0.0 + assert summary["call_count"] == 0 + + +def test_get_task_token_summary_excludes_other_tasks(db): + """Test task summary only includes records for the specified task.""" + tid1 = _create_task(db) + tid2 = _create_task(db) + _save_usage(db, task_id=tid1, input_tokens=1000, output_tokens=500, cost=0.01) + _save_usage(db, task_id=tid2, input_tokens=2000, output_tokens=1000, cost=0.02) + + summary = db.get_task_token_summary(task_id=tid1) + + assert summary["total_input_tokens"] == 1000 + assert summary["call_count"] == 1 + + +# ============================================================================ +# get_batch_token_usage +# ============================================================================ + + +def test_get_batch_token_usage(db): + """Test getting token usage for a batch of task IDs.""" + tid1 = _create_task(db) + tid2 = _create_task(db) + tid3 = _create_task(db) + _save_usage(db, task_id=tid1, input_tokens=100, output_tokens=50, cost=0.001) + _save_usage(db, task_id=tid2, input_tokens=200, output_tokens=100, cost=0.002) + _save_usage(db, task_id=tid3, input_tokens=300, output_tokens=150, cost=0.003) + + records = db.get_batch_token_usage(task_ids=[tid1, tid2]) + + assert len(records) == 2 + task_ids = {r["task_id"] for r in records} + assert task_ids == {tid1, tid2} + + +def test_get_batch_token_usage_with_date_filter(db): + """Test batch token usage with date filtering.""" + now = datetime.now(timezone.utc) + old = now - timedelta(days=10) + + tid1 = _create_task(db) + tid2 = _create_task(db) + _save_usage(db, task_id=tid1, input_tokens=100, output_tokens=50, cost=0.001, timestamp=now) + _save_usage(db, task_id=tid2, input_tokens=200, output_tokens=100, cost=0.002, timestamp=old) + + start = now - timedelta(days=1) + records = db.get_batch_token_usage(task_ids=[tid1, tid2], start_date=start) + + assert len(records) == 1 + assert records[0]["task_id"] == tid1 + + +def test_get_batch_token_usage_empty_list(db): + """Test batch token usage with empty task ID list.""" + tid = _create_task(db) + _save_usage(db, task_id=tid, input_tokens=100, output_tokens=50, cost=0.001) + + records = db.get_batch_token_usage(task_ids=[]) + + assert len(records) == 0 + + +# ============================================================================ +# get_workspace_token_usage +# ============================================================================ + + +def test_get_workspace_token_usage(db): + """Test getting all token usage across the workspace.""" + tid = _create_task(db) + _save_usage(db, task_id=tid, project_id=1, input_tokens=100, output_tokens=50, cost=0.001) + _save_usage(db, task_id=None, project_id=1, input_tokens=200, output_tokens=100, cost=0.002) + + records = db.get_workspace_token_usage() + + assert len(records) == 2 + + +def test_get_workspace_token_usage_with_date_filter(db): + """Test workspace token usage with date filtering.""" + now = datetime.now(timezone.utc) + old = now - timedelta(days=10) + + tid1 = _create_task(db) + tid2 = _create_task(db) + _save_usage(db, task_id=tid1, input_tokens=100, output_tokens=50, cost=0.001, timestamp=now) + _save_usage(db, task_id=tid2, input_tokens=200, output_tokens=100, cost=0.002, timestamp=old) + + start = now - timedelta(days=1) + end = now + timedelta(days=1) + records = db.get_workspace_token_usage(start_date=start, end_date=end) + + assert len(records) == 1 + assert records[0]["task_id"] == tid1 + + +def test_get_workspace_token_usage_empty(db): + """Test workspace token usage when no records exist.""" + records = db.get_workspace_token_usage() + + assert len(records) == 0