cloud: production control plane (API, auth/RBAC, queue, sandbox, memory/KB, fallback/cost, webhooks)

cloud/cost.py

@@ -0,0 +1,158 @@
+"""Cost estimation and cost recording.
+
+Zero third-party dependencies; standard library only.
+
+Actual token counts come from the model/provider response when
+available; when they are not, token counts are estimates (see
+``estimate_tokens``) and are stored separately as ``source="estimated"``
+so estimates are never mixed into actual totals.
+"""
+
+import os
+import sqlite3
+import time
+from datetime import datetime, timezone
+
+
+def estimate_tokens(prompt_text, output_text) -> int:
+    """Estimate total tokens as (chars of prompt + output) // 4."""
+    prompt_text = prompt_text or ""
+    output_text = output_text or ""
+    return (len(prompt_text) + len(output_text)) // 4
+
+
+def fetch_account_usage(usage_client) -> dict:
+    """Fetch account usage through an injected ``usage_client`` callable.
+
+    This is a call wrapper around the app-server account/usage/read
+    concept: the real ``usage_client`` is injected by the control plane
+    once the billing route is confirmed available. This function itself
+    never runs a model turn and never touches the network.
+
+    On success it returns ``{"available": True, "data": <raw data>}``.
+    If ``usage_client`` raises, or returns None / an empty dict,
+    it returns ``{"available": False, "data": None}`` and never raises
+    to the caller. On a relay, usage being unreadable is a normal
+    outcome (available=False), not a bug.
+    """
+    try:
+        data = usage_client()
+    except Exception:
+        return {"available": False, "data": None}
+    if data is None:
+        return {"available": False, "data": None}
+    if isinstance(data, dict) and not data:
+        return {"available": False, "data": None}
+    return {"available": True, "data": data}
+
+
+class CostStore:
+    """SQLite store for per-task / per-project / per-day cost records."""
+
+    _SOURCES = ("actual", "estimated")
+
+    def __init__(self, db_path):
+        self.db_path = db_path
+        if db_path != ":memory:":
+            parent = os.path.dirname(os.path.abspath(db_path))
+            os.makedirs(parent, exist_ok=True)
+        self._conn = sqlite3.connect(db_path)
+        self._conn.execute(
+            """CREATE TABLE IF NOT EXISTS cost_records (
+                id INTEGER PRIMARY KEY AUTOINCREMENT,
+                task_id TEXT,
+                project_id TEXT,
+                model TEXT,
+                input_tokens INTEGER,
+                output_tokens INTEGER,
+                total_tokens INTEGER,
+                usd_micros INTEGER,
+                source TEXT,
+                fallback_count INTEGER DEFAULT 0,
+                day TEXT,
+                created_at REAL
+            )"""
+        )
+        self._conn.commit()
+
+    def record(self, task_id, project_id, model, input_tokens=0,
+               output_tokens=0, usd_micros=None, source="actual",
+               fallback_count=0):
+        """Insert one cost record; returns its row id.
+
+        ``total_tokens`` is computed as input + output, and ``day`` is
+        the UTC date (YYYY-MM-DD) of ``created_at``.
+        """
+        if source not in self._SOURCES:
+            raise ValueError(
+                "source must be one of %r, got %r" % (self._SOURCES, source)
+            )
+        input_tokens = int(input_tokens or 0)
+        output_tokens = int(output_tokens or 0)
+        total_tokens = input_tokens + output_tokens
+        created_at = time.time()
+        day = datetime.fromtimestamp(created_at, tz=timezone.utc).strftime("%Y-%m-%d")
+        cursor = self._conn.execute(
+            """INSERT INTO cost_records
+               (task_id, project_id, model, input_tokens, output_tokens,
+                total_tokens, usd_micros, source, fallback_count, day,
+                created_at)
+               VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)""",
+            (task_id, project_id, model, input_tokens, output_tokens,
+             total_tokens, usd_micros, source, int(fallback_count or 0),
+             day, created_at),
+        )
+        self._conn.commit()
+        return cursor.lastrowid
+
+    def record_estimated(self, task_id, project_id, model, prompt_text,
+                         output_text):
+        """Record an estimated cost row; returns its row id.
+
+        The token split cannot be recovered from a character estimate,
+        so ``input_tokens`` holds the whole estimate and
+        ``output_tokens`` is 0; ``source`` is ``"estimated"`` and
+        ``usd_micros`` is None.
+        """
+        total = estimate_tokens(prompt_text, output_text)
+        return self.record(
+            task_id, project_id, model,
+            input_tokens=total, output_tokens=0,
+            usd_micros=None, source="estimated", fallback_count=0,
+        )
+
+    def _report(self, where, params):
+        row = self._conn.execute(
+            """SELECT
+                 COALESCE(SUM(total_tokens), 0),
+                 COALESCE(SUM(usd_micros), 0),
+                 COUNT(*),
+                 COALESCE(SUM(fallback_count), 0),
+                 COALESCE(SUM(CASE WHEN source = 'actual'
+                                   THEN total_tokens ELSE 0 END), 0),
+                 COALESCE(SUM(CASE WHEN source = 'estimated'
+                                   THEN total_tokens ELSE 0 END), 0)
+               FROM cost_records WHERE %s""" % where,
+            params,
+        ).fetchone()
+        total_tokens, usd_micros, turns, fallback_count, actual, estimated = row
+        return {
+            "total_tokens": int(total_tokens),
+            "usd_micros": int(usd_micros),
+            "turns": int(turns),
+            "fallback_count": int(fallback_count),
+            "actual_tokens": int(actual),
+            "estimated_tokens": int(estimated),
+        }
+
+    def report_by_task(self, task_id) -> dict:
+        """Aggregate cost records for one task."""
+        return self._report("task_id = ?", (task_id,))
+
+    def report_by_project(self, project_id) -> dict:
+        """Aggregate cost records for one project."""
+        return self._report("project_id = ?", (project_id,))
+
+    def report_by_day(self, day) -> dict:
+        """Aggregate cost records for one UTC day (YYYY-MM-DD)."""
+        return self._report("day = ?", (day,))