multi-forge 0.2.0__py3-none-any.whl

forge/proxy/cost_logger.py

@@ -0,0 +1,132 @@
+"""PID-sharded JSONL cost log writer.
+
+Each proxy process writes to its own shard file to avoid interprocess
+locking. The CLI aggregates across shards at query time.
+
+Location: ~/.forge/costs/requests/YYYY-MM_<pid>.jsonl
+"""
+
+from __future__ import annotations
+
+import json
+import logging
+import os
+import threading
+from datetime import datetime, timezone
+from pathlib import Path
+from typing import Any
+
+from forge.core.paths import get_forge_home
+
+logger = logging.getLogger(__name__)
+
+_lock = threading.Lock()
+
+
+def _pid_suffix() -> str:
+    return str(os.getpid())
+
+
+def _costs_dir() -> Path:
+    return get_forge_home() / "costs" / "requests"
+
+
+def _current_log_path() -> Path:
+    month = datetime.now(timezone.utc).strftime("%Y-%m")
+    return _costs_dir() / f"{month}_{_pid_suffix()}.jsonl"
+
+
+def log_request_cost(
+    *,
+    proxy_id: str,
+    model: str,
+    tier: str,
+    input_tokens: int,
+    output_tokens: int,
+    cached_tokens: int,
+    cost_micros: int,
+    latency_ms: float,
+    failed: bool,
+    request_id: str,
+    pricing_source: str = "catalog",
+) -> None:
+    """Append a cost record to the PID-sharded JSONL log.
+
+    Best-effort: write failures are logged but never block the request.
+    """
+    record: dict[str, Any] = {
+        "ts": datetime.now(timezone.utc).strftime("%Y-%m-%dT%H:%M:%SZ"),
+        "proxy_id": proxy_id,
+        "model": model,
+        "tier": tier,
+        "input_tokens": input_tokens,
+        "output_tokens": output_tokens,
+        "cached_tokens": cached_tokens,
+        "cost_micros": cost_micros,
+        "estimated": True,
+        "pricing_source": pricing_source,
+        "latency_ms": round(latency_ms, 1),
+        "failed": failed,
+        "request_id": request_id,
+    }
+
+    try:
+        from forge.core.state import open_secure_append
+
+        log_path = _current_log_path()
+        log_path.parent.mkdir(parents=True, exist_ok=True)
+
+        with _lock:
+            with open_secure_append(log_path) as f:
+                f.write(json.dumps(record, separators=(",", ":")) + "\n")
+    except Exception as e:
+        logger.warning("Failed to write cost log: %s", e)
+
+
+def read_cost_logs(
+    period_start: datetime | None = None,
+    period_end: datetime | None = None,
+) -> list[dict[str, Any]]:
+    """Read and aggregate cost records from all PID shards.
+
+    Args:
+        period_start: Only include records at or after this time (UTC).
+        period_end: Only include records before this time (UTC).
+
+    Returns:
+        List of cost record dicts, sorted by timestamp.
+    """
+    costs_dir = _costs_dir()
+    if not costs_dir.is_dir():
+        return []
+
+    records: list[dict[str, Any]] = []
+    for path in sorted(costs_dir.glob("*.jsonl")):
+        try:
+            with open(path) as f:
+                for line in f:
+                    line = line.strip()
+                    if not line:
+                        continue
+                    try:
+                        record = json.loads(line)
+                    except json.JSONDecodeError:
+                        continue
+
+                    if period_start or period_end:
+                        ts_str = record.get("ts", "")
+                        try:
+                            ts = datetime.fromisoformat(ts_str.rstrip("Z").removesuffix("+00:00") + "+00:00")
+                        except (ValueError, TypeError):
+                            continue
+                        if period_start and ts < period_start:
+                            continue
+                        if period_end and ts >= period_end:
+                            continue
+
+                    records.append(record)
+        except OSError as e:
+            logger.warning("Failed to read cost log %s: %s", path, e)
+
+    records.sort(key=lambda r: r.get("ts", ""))
+    return records

forge/proxy/cost_tracker.py

@@ -0,0 +1,242 @@
+"""Spend cap enforcement with JSONL-bootstrapped tracking.
+
+On proxy startup, reads the current (and previous) month's cost JSONL
+logs to initialize in-memory spend counters. Caps are enforced per
+request via check_cap().
+
+Two enforcement modes:
+  post   -- block once accumulated spend already exceeds the cap
+  strict -- estimate incoming request cost and block if projected total exceeds
+"""
+
+from __future__ import annotations
+
+import logging
+import time
+from collections import deque
+from dataclasses import dataclass
+from datetime import datetime, timezone
+from pathlib import Path
+
+logger = logging.getLogger(__name__)
+
+_MICROS_PER_DOLLAR = 1_000_000
+_24H_SECONDS = 86400
+
+
+@dataclass
+class CapResult:
+    """Result of a spend cap check."""
+
+    exceeded: bool
+    cap_type: str | None = None  # "daily" or "monthly"
+    current_micros: int = 0
+    limit_micros: int = 0
+    projected: bool = False  # True if this is a pre-flight estimate
+
+
+class CostTracker:
+    """In-memory spend tracking with cap enforcement.
+
+    Thread-safe via the proxy's single-threaded async event loop
+    (all calls happen on the main thread in FastAPI/uvicorn).
+    """
+
+    def __init__(
+        self,
+        *,
+        daily_cap_usd: float | None = None,
+        monthly_cap_usd: float | None = None,
+        cap_mode: str = "post",
+        on_cap_hit: str = "reject",
+    ) -> None:
+        self.daily_cap_micros = int(daily_cap_usd * _MICROS_PER_DOLLAR) if daily_cap_usd is not None else None
+        self.monthly_cap_micros = int(monthly_cap_usd * _MICROS_PER_DOLLAR) if monthly_cap_usd is not None else None
+        self.cap_mode = cap_mode
+        self.on_cap_hit = on_cap_hit
+
+        self._daily_window: deque[tuple[float, int]] = deque()
+        self._monthly_total: int = 0
+        self._monthly_key: str = ""
+
+    @property
+    def has_caps(self) -> bool:
+        return self.daily_cap_micros is not None or self.monthly_cap_micros is not None
+
+    def bootstrap_from_logs(self, log_dir: Path, *, proxy_id: str | None = None) -> None:
+        """Read existing cost logs to initialize spend counters.
+
+        Reads current month + previous month (for rolling 24h window
+        at month boundaries). Scans all PID shards.
+
+        When proxy_id is set, only records matching that proxy are counted.
+        Records without a proxy_id field are skipped (pre-proxy-id logs).
+        When proxy_id is None, all records are counted (backward compat).
+        """
+        if not log_dir.is_dir():
+            return
+
+        now = datetime.now(timezone.utc)
+        current_month = now.strftime("%Y-%m")
+        self._monthly_key = current_month
+
+        if now.month == 1:
+            prev_month = f"{now.year - 1}-12"
+        else:
+            prev_month = f"{now.year}-{now.month - 1:02d}"
+
+        cutoff = time.time() - _24H_SECONDS
+
+        for path in sorted(log_dir.glob("*.jsonl")):
+            fname = path.stem  # e.g., "2026-05_12345"
+            file_month = fname.split("_")[0] if "_" in fname else fname
+
+            if file_month not in (current_month, prev_month):
+                continue
+
+            try:
+                with open(path) as f:
+                    for line in f:
+                        line = line.strip()
+                        if not line:
+                            continue
+                        try:
+                            record = self._parse_record(line)
+                        except Exception:
+                            continue
+                        if record is None:
+                            continue
+
+                        ts_unix, cost_micros, record_month, record_proxy_id = record
+
+                        if proxy_id is not None and record_proxy_id != proxy_id:
+                            continue
+
+                        if record_month == current_month:
+                            self._monthly_total += cost_micros
+
+                        if ts_unix >= cutoff:
+                            self._daily_window.append((ts_unix, cost_micros))
+            except OSError as e:
+                logger.warning("Failed to read cost log %s: %s", path, e)
+
+        daily_total = sum(c for _, c in self._daily_window)
+        logger.info(
+            "Cost tracker bootstrapped: daily=$%.2f, monthly=$%.2f (%d records in window)",
+            daily_total / _MICROS_PER_DOLLAR,
+            self._monthly_total / _MICROS_PER_DOLLAR,
+            len(self._daily_window),
+        )
+
+    @staticmethod
+    def _parse_record(line: str) -> tuple[float, int, str, str | None] | None:
+        """Parse a JSONL line into (unix_timestamp, cost_micros, month_key, record_proxy_id)."""
+        import json
+
+        data = json.loads(line)
+        ts_str = data.get("ts", "")
+        cost_micros = int(data.get("cost_micros", 0))
+        if cost_micros <= 0:
+            return None
+
+        try:
+            ts = datetime.fromisoformat(ts_str.rstrip("Z").removesuffix("+00:00") + "+00:00")
+        except (ValueError, TypeError):
+            return None
+
+        month_key = ts.strftime("%Y-%m")
+        record_proxy_id = data.get("proxy_id")
+        return ts.timestamp(), cost_micros, month_key, record_proxy_id
+
+    def record(self, cost_micros: int) -> None:
+        """Record a completed request's cost."""
+        if cost_micros <= 0:
+            return
+
+        now = time.time()
+        self._roll_month_if_needed()
+
+        self._monthly_total += cost_micros
+        self._daily_window.append((now, cost_micros))
+
+    def _roll_month_if_needed(self) -> None:
+        """Reset the calendar-month accumulator when UTC month changes."""
+        current_month = datetime.now(timezone.utc).strftime("%Y-%m")
+        if current_month != self._monthly_key:
+            self._monthly_total = 0
+            self._monthly_key = current_month
+
+    def _prune_daily_window(self) -> None:
+        """Remove entries older than 24 hours from the rolling window."""
+        cutoff = time.time() - _24H_SECONDS
+        while self._daily_window and self._daily_window[0][0] < cutoff:
+            self._daily_window.popleft()
+
+    def daily_spend_micros(self) -> int:
+        """Current rolling 24h spend in microdollars."""
+        self._prune_daily_window()
+        return sum(c for _, c in self._daily_window)
+
+    def monthly_spend_micros(self) -> int:
+        """Current calendar month spend in microdollars."""
+        self._roll_month_if_needed()
+        return self._monthly_total
+
+    def check_cap(self, projected_cost_micros: int = 0) -> CapResult:
+        """Check if spend would exceed any configured caps.
+
+        Args:
+            projected_cost_micros: Estimated cost of the pending request.
+                In strict mode, added to current spend for pre-flight check.
+                In post mode, ignored (only accumulated spend matters).
+
+        Returns:
+            CapResult indicating whether any cap is exceeded.
+        """
+        if not self.has_caps:
+            return CapResult(exceeded=False)
+
+        extra = projected_cost_micros if self.cap_mode == "strict" else 0
+
+        if self.daily_cap_micros is not None:
+            daily = self.daily_spend_micros() + extra
+            if daily >= self.daily_cap_micros:
+                return CapResult(
+                    exceeded=True,
+                    cap_type="daily",
+                    current_micros=daily,
+                    limit_micros=self.daily_cap_micros,
+                    projected=extra > 0,
+                )
+
+        if self.monthly_cap_micros is not None:
+            monthly = self.monthly_spend_micros() + extra
+            if monthly >= self.monthly_cap_micros:
+                return CapResult(
+                    exceeded=True,
+                    cap_type="monthly",
+                    current_micros=monthly,
+                    limit_micros=self.monthly_cap_micros,
+                    projected=extra > 0,
+                )
+
+        return CapResult(exceeded=False)
+
+    def cap_summary(self) -> dict[str, dict[str, float]]:
+        """Return current spend vs caps for CLI display."""
+        result: dict[str, dict[str, float]] = {}
+        if self.daily_cap_micros is not None:
+            daily = self.daily_spend_micros()
+            result["daily"] = {
+                "current_usd": daily / _MICROS_PER_DOLLAR,
+                "limit_usd": self.daily_cap_micros / _MICROS_PER_DOLLAR,
+                "percent": round(daily / self.daily_cap_micros * 100, 1) if self.daily_cap_micros > 0 else 0,
+            }
+        if self.monthly_cap_micros is not None:
+            monthly = self.monthly_spend_micros()
+            result["monthly"] = {
+                "current_usd": monthly / _MICROS_PER_DOLLAR,
+                "limit_usd": self.monthly_cap_micros / _MICROS_PER_DOLLAR,
+                "percent": round(monthly / self.monthly_cap_micros * 100, 1) if self.monthly_cap_micros > 0 else 0,
+            }
+        return result