fix(perf): push cf stats aggregation into SQL, stream exports (#752)

codeframe/cli/stats_commands.py

@@ -112,32 +112,14 @@ def tokens(
 
             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
-
+            # Workspace-wide summary — aggregated in SQL, then totals
+            # derived from the handful of per-model rows.
+            by_model = db.get_costs_by_model()
+
+            total_input = sum(m["input_tokens"] for m in by_model)
+            total_output = sum(m["output_tokens"] for m in by_model)
+            total_cost = sum(m["total_cost_usd"] for m in by_model)
+            total_calls = sum(m["call_count"] for m in by_model)
             total_tokens = total_input + total_output
 
             console.print("\n[bold]Workspace Token Usage Summary[/bold]\n")
@@ -150,24 +132,24 @@ def 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)))
+            summary_table.add_row("LLM Calls", str(total_calls))
 
             console.print(summary_table)
 
-            if model_stats:
+            if by_model:
                 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():
+                for m in by_model:
                     model_table.add_row(
-                        model_name,
-                        _format_number(stats["input_tokens"] + stats["output_tokens"]),
-                        f"${stats['cost_usd']:.4f}",
-                        str(stats["calls"]),
+                        m["model_name"],
+                        _format_number(m["input_tokens"] + m["output_tokens"]),
+                        f"${m['total_cost_usd']:.4f}",
+                        str(m["call_count"]),
                     )
 
                 console.print(model_table)
@@ -214,24 +196,13 @@ def costs(
             )
             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
+        # Per-model rollup aggregated in SQL; totals summed over the
+        # small per-model result set.
+        by_model = db.get_costs_by_model(start_date=start_date, end_date=end_date)
+
+        total_cost = sum(m["total_cost_usd"] for m in by_model)
+        total_tokens = sum(m["input_tokens"] + m["output_tokens"] for m in by_model)
+        total_calls = sum(m["call_count"] for m in by_model)
 
         period_label = f" ({period})" if period else " (all time)"
         console.print(f"\n[bold]Cost Report{period_label}[/bold]\n")
@@ -242,24 +213,24 @@ def costs(
 
         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)))
+        table.add_row("LLM Calls", str(total_calls))
 
         console.print(table)
 
-        if model_costs:
+        if by_model:
             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():
+            for m in by_model:
                 model_table.add_row(
-                    model_name,
-                    f"${stats['cost_usd']:.4f}",
-                    _format_number(stats["tokens"]),
-                    str(stats["calls"]),
+                    m["model_name"],
+                    f"${m['total_cost_usd']:.4f}",
+                    _format_number(m["input_tokens"] + m["output_tokens"]),
+                    str(m["call_count"]),
                 )
 
             console.print(model_table)
@@ -293,19 +264,21 @@ def export_data(
 
     db = _get_db()
     try:
+        if format not in ("csv", "json"):
+            console.print(f"[red]Error:[/red] Unknown format '{format}'. Use 'csv' or 'json'.")
+            raise typer.Exit(1)
+
         if task is not None:
             records = db.get_batch_token_usage(task_ids=[task])
         else:
-            records = db.get_workspace_token_usage()
+            # Stream the workspace table — never buffered into a list.
+            records = db.get_token_usage_iter()
 
         if format == "csv":
-            MetricsTracker.export_to_csv(records, output)
-        elif format == "json":
-            MetricsTracker.export_to_json(records, output)
+            n = MetricsTracker.export_to_csv(records, output)
         else:
-            console.print(f"[red]Error:[/red] Unknown format '{format}'. Use 'csv' or 'json'.")
-            raise typer.Exit(1)
+            n = MetricsTracker.export_to_json(records, output)
 
-        console.print(f"Exported {len(records)} records to {output}")
+        console.print(f"Exported {n} records to {output}")
     finally:
         db.close()

codeframe/lib/metrics_tracker.py

@@ -38,9 +38,11 @@
 import csv
 import json
 import logging
+import os
 import re
+import tempfile
 from datetime import datetime, timedelta, timezone
-from typing import Any, Dict, List, Optional, Union
+from typing import Any, Callable, Dict, Iterable, Optional, TextIO, Union
 from codeframe.core.models import CallType, TokenUsage
 from codeframe.platform_store.database import Database
 
@@ -334,71 +336,125 @@ def get_workspace_costs(
                 "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),
+        # Aggregation is pushed into SQL: the per-model rollup returns a
+        # handful of rows; summing those in Python is O(models), not O(records).
+        by_model = self.db.get_costs_by_model(start_date=start_date, end_date=end_date)
+
+        total_cost = sum(m["total_cost_usd"] for m in by_model)
+        total_tokens = sum(m["input_tokens"] + m["output_tokens"] for m in by_model)
+        total_calls = sum(m["call_count"] for m in by_model)
+
+        return {
+            "total_cost_usd": round(total_cost, 6),
+            "total_tokens": total_tokens,
+            "total_calls": total_calls,
         }
 
-        for record in records:
-            result["total_cost_usd"] += record["estimated_cost_usd"]
-            result["total_tokens"] += record["input_tokens"] + record["output_tokens"]
+    @staticmethod
+    def _atomic_stream_write(
+        output_path: str, write_fn: Callable[[TextIO], int]
+    ) -> int:
+        """Stream through a temp file in the same dir, then ``os.replace``.
 
-        result["total_cost_usd"] = round(result["total_cost_usd"], 6)
-        return result
+        The exporters write incrementally, so a mid-stream failure (source
+        iterator raises, disk fills) must not leave a truncated file at
+        ``output_path``. Writing to a sibling temp file and atomically renaming
+        on success means readers only ever see a complete export; the temp file
+        is unlinked on any failure.
+
+        Args:
+            output_path: Final destination path.
+            write_fn: Callback that writes to the open file and returns a count.
+
+        Returns:
+            Whatever ``write_fn`` returns (the record count).
+        """
+        directory = os.path.dirname(os.path.abspath(output_path))
+        # mkstemp creates the temp file 0600 and os.replace preserves that mode,
+        # so exports land owner-only (not umask-derived 0644). Intentional: token
+        # spend data is mildly sensitive and this file is the user's named output.
+        fd, tmp_path = tempfile.mkstemp(dir=directory, prefix=".export-", suffix=".tmp")
+        try:
+            with os.fdopen(fd, "w", newline="") as f:
+                count = write_fn(f)
+            os.replace(tmp_path, output_path)
+            return count
+        except BaseException:
+            try:
+                os.unlink(tmp_path)
+            except OSError:
+                pass
+            raise
 
     @staticmethod
-    def export_to_csv(records: List[Dict[str, Any]], output_path: str) -> None:
-        """Export token usage records to a CSV file.
+    def export_to_csv(records: Iterable[Dict[str, Any]], output_path: str) -> int:
+        """Stream token usage records to a CSV file.
+
+        Consumes ``records`` lazily (accepts the ``get_token_usage_iter``
+        generator) so a large table is never buffered into a list. Written
+        atomically — a partial write never lands at ``output_path``.
 
         Args:
-            records: List of token usage record dictionaries
-            output_path: Path to write the CSV file
+            records: Iterable of token usage record dictionaries.
+            output_path: Path to write the CSV file.
+
+        Returns:
+            Number of rows written.
         """
         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:
+        def _write(f: TextIO) -> int:
             writer = csv.DictWriter(f, fieldnames=fieldnames, extrasaction="ignore")
             writer.writeheader()
+            count = 0
             for record in records:
                 writer.writerow(record)
+                count += 1
+            return count
+
+        return MetricsTracker._atomic_stream_write(output_path, _write)
 
     @staticmethod
-    def export_to_json(records: List[Dict[str, Any]], output_path: str) -> None:
-        """Export token usage records to a JSON file with metadata.
+    def export_to_json(records: Iterable[Dict[str, Any]], output_path: str) -> int:
+        """Stream token usage records to a JSON file with metadata.
+
+        Writes the ``records`` array incrementally as it consumes the iterator,
+        so the whole table is never held in memory at once. ``metadata`` (with
+        ``record_count``) is written last; JSON object key order is not
+        semantically meaningful, so consumers using ``json.load`` are unaffected.
+        Written atomically — a partial write never lands at ``output_path``.
 
         Args:
-            records: List of token usage record dictionaries
-            output_path: Path to write the JSON file
+            records: Iterable of token usage record dictionaries.
+            output_path: Path to write the JSON file.
+
+        Returns:
+            Number of records written.
         """
-        # 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": {
+        def _write(f: TextIO) -> int:
+            count = 0
+            f.write('{\n  "records": [')
+            for record in records:
+                row = {
+                    k: (v.isoformat() if isinstance(v, datetime) else v)
+                    for k, v in dict(record).items()
+                }
+                f.write("\n    " if count == 0 else ",\n    ")
+                f.write(json.dumps(row, default=str))
+                count += 1
+            f.write("\n  ],\n")
+            metadata = {
                 "exported_at": datetime.now(timezone.utc).isoformat(),
-                "record_count": len(serializable_records),
-            },
-            "records": serializable_records,
-        }
+                "record_count": count,
+            }
+            f.write('  "metadata": ' + json.dumps(metadata) + "\n}\n")
+            return count
 
-        with open(output_path, "w") as f:
-            json.dump(data, f, indent=2, default=str)
+        return MetricsTracker._atomic_stream_write(output_path, _write)
 
     async def get_project_costs(
         self,

codeframe/platform_store/database.py

@@ -271,6 +271,14 @@ 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 get_token_usage_iter(self, *args, **kwargs):
+        """Delegate to token_usage.get_token_usage_iter()."""
+        return self.token_usage.get_token_usage_iter(*args, **kwargs)
+
+    def get_costs_by_model(self, *args, **kwargs):
+        """Delegate to token_usage.get_costs_by_model()."""
+        return self.token_usage.get_costs_by_model(*args, **kwargs)
+
     # ----- Audit log -----
     def create_audit_log(self, *args, **kwargs):
         """Delegate to audit_logs.create_audit_log()."""