traceforge-pcb 1.0.0__py3-none-any.whl

traceforge/remote/orgs.py

@@ -0,0 +1,510 @@
+"""Organisations for the hosted server: seats, roles, pooled quotas, spend caps, audit log.
+
+Everything lives in the server's SQLite database next to the API keys (``KeyStore``). A
+*member* belongs to one organisation with a role (``admin`` or ``member``); API keys can be
+tied to a member (``keys.member_id``), and then every turn is checked against
+
+* the member's monthly spend cap (USD) and token cap,
+* the organisation's pooled monthly token and spend quotas.
+
+The audit log records who ran what (member, key, mode, model, tool names, token counts,
+cost). It stores no prompt text, no tool arguments and no model output.
+
+Cost comes from :mod:`traceforge.prices`; a model without a price has unknown cost. A cost
+limit applies fail-closed: a member or organisation with a USD limit cannot use an unpriced
+model (error code ``unpriced_model``), because its spend could not be counted.
+"""
+
+from __future__ import annotations
+
+import hashlib
+import json
+import secrets
+import sqlite3
+import time
+import uuid
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Any
+
+from traceforge.config import data_home
+
+ROLES = ("admin", "member")
+INVITE_TTL = 7 * 24 * 3600
+
+
+class OrgError(Exception):
+    """A request the store refuses; ``code`` is stable for clients."""
+
+    def __init__(self, code: str, message: str):
+        super().__init__(message)
+        self.code = code
+        self.message = message
+
+
+def _hash(secret: str) -> str:
+    return hashlib.sha256(secret.encode()).hexdigest()
+
+
+def month_start(now: float | None = None) -> float:
+    t = time.gmtime(now if now is not None else time.time())
+    import calendar
+
+    return float(calendar.timegm((t.tm_year, t.tm_mon, 1, 0, 0, 0)))
+
+
+@dataclass
+class Check:
+    """Result of a pre-flight or running limit check."""
+
+    ok: bool
+    code: str = ""
+    message: str = ""
+
+
+_SCHEMA = """
+CREATE TABLE IF NOT EXISTS orgs (
+  id TEXT PRIMARY KEY, name TEXT UNIQUE, created REAL,
+  quota_tokens INTEGER, quota_usd REAL);
+CREATE TABLE IF NOT EXISTS members (
+  id TEXT PRIMARY KEY, org_id TEXT, email TEXT, name TEXT, role TEXT, created REAL,
+  removed INTEGER DEFAULT 0, cap_usd REAL, cap_tokens INTEGER, sso_subject TEXT,
+  UNIQUE (org_id, email));
+CREATE TABLE IF NOT EXISTS invites (
+  hash TEXT PRIMARY KEY, org_id TEXT, email TEXT, role TEXT, created REAL, expires REAL,
+  used INTEGER DEFAULT 0);
+CREATE TABLE IF NOT EXISTS audit (
+  id INTEGER PRIMARY KEY AUTOINCREMENT, ts REAL, org_id TEXT, member_id TEXT, key_id TEXT,
+  event TEXT, mode TEXT, model TEXT, tools TEXT, input_tokens INTEGER, output_tokens INTEGER,
+  cache_read_tokens INTEGER, cache_write_tokens INTEGER, requests INTEGER, cost_usd REAL,
+  cost_known INTEGER, status TEXT, detail TEXT, by_model TEXT);
+CREATE INDEX IF NOT EXISTS audit_member ON audit (member_id, ts);
+CREATE INDEX IF NOT EXISTS audit_org ON audit (org_id, ts);
+"""
+
+
+class OrgStore:
+    def __init__(self, path: Path | None = None):
+        from traceforge.remote.store import KeyStore
+
+        self.path = path or data_home() / "server.db"
+        KeyStore(self.path)  # creates the keys table this store extends
+        with self._db() as db:
+            db.executescript(_SCHEMA)
+            columns = {r[1] for r in db.execute("PRAGMA table_info(keys)")}
+            if columns and "member_id" not in columns:
+                db.execute("ALTER TABLE keys ADD COLUMN member_id TEXT")
+
+    def _db(self) -> sqlite3.Connection:
+        db = sqlite3.connect(self.path, timeout=10)
+        db.row_factory = sqlite3.Row
+        return db
+
+    # ------------------------------------------------------------ orgs / members
+
+    def create_org(
+        self, name: str, quota_tokens: int | None = None, quota_usd: float | None = None
+    ) -> dict:
+        org_id = "org_" + uuid.uuid4().hex[:10]
+        try:
+            with self._db() as db:
+                db.execute(
+                    "INSERT INTO orgs VALUES (?,?,?,?,?)",
+                    (org_id, name, time.time(), quota_tokens, quota_usd),
+                )
+        except sqlite3.IntegrityError as exc:
+            raise OrgError("exists", f"organisation {name!r} already exists") from exc
+        return self.org(org_id)  # type: ignore[return-value]
+
+    def org(self, name_or_id: str) -> dict | None:
+        with self._db() as db:
+            row = db.execute(
+                "SELECT * FROM orgs WHERE id = ? OR name = ?", (name_or_id, name_or_id)
+            ).fetchone()
+        return dict(row) if row else None
+
+    def need_org(self, name_or_id: str) -> dict:
+        org = self.org(name_or_id)
+        if org is None:
+            raise OrgError("no_such_org", f"no organisation {name_or_id!r}")
+        return org
+
+    def orgs(self) -> list[dict]:
+        with self._db() as db:
+            rows = db.execute("SELECT * FROM orgs ORDER BY created").fetchall()
+        return [{**dict(r), **self.usage(org_id=r["id"])} for r in rows]
+
+    def set_quota(
+        self, org: str, tokens: int | None = None, usd: float | None = None, clear: bool = False
+    ) -> None:
+        org_row = self.need_org(org)
+        with self._db() as db:
+            db.execute(
+                "UPDATE orgs SET quota_tokens = ?, quota_usd = ? WHERE id = ?",
+                (
+                    None if clear else tokens if tokens is not None else org_row["quota_tokens"],
+                    None if clear else usd if usd is not None else org_row["quota_usd"],
+                    org_row["id"],
+                ),
+            )
+
+    def add_member(
+        self,
+        org: str,
+        email: str,
+        role: str = "member",
+        name: str = "",
+        sso_subject: str | None = None,
+    ) -> dict:
+        if role not in ROLES:
+            raise OrgError("bad_role", f"role must be one of {ROLES}")
+        org_row = self.need_org(org)
+        email = email.strip().lower()
+        existing = self.member(org_row["id"], email, include_removed=True)
+        if existing and not existing["removed"]:
+            raise OrgError("exists", f"{email} is already a member of {org_row['name']}")
+        with self._db() as db:
+            if existing:  # re-adding a removed seat
+                db.execute(
+                    "UPDATE members SET removed = 0, role = ?, name = ?, sso_subject = ? WHERE id = ?",
+                    (role, name or existing["name"], sso_subject, existing["id"]),
+                )
+                member_id = existing["id"]
+            else:
+                member_id = "usr_" + uuid.uuid4().hex[:10]
+                db.execute(
+                    "INSERT INTO members (id, org_id, email, name, role, created, sso_subject)"
+                    " VALUES (?,?,?,?,?,?,?)",
+                    (
+                        member_id,
+                        org_row["id"],
+                        email,
+                        name or email,
+                        role,
+                        time.time(),
+                        sso_subject,
+                    ),
+                )
+        return self.member_by_id(member_id)  # type: ignore[return-value]
+
+    def member(self, org_id: str, email: str, include_removed: bool = False) -> dict | None:
+        with self._db() as db:
+            row = db.execute(
+                "SELECT * FROM members WHERE org_id = ? AND email = ?", (org_id, email.lower())
+            ).fetchone()
+        if row is None or (row["removed"] and not include_removed):
+            return None
+        return dict(row)
+
+    def member_by_id(self, member_id: str) -> dict | None:
+        with self._db() as db:
+            row = db.execute("SELECT * FROM members WHERE id = ?", (member_id,)).fetchone()
+        return dict(row) if row else None
+
+    def member_by_subject(self, org_id: str, subject: str) -> dict | None:
+        with self._db() as db:
+            row = db.execute(
+                "SELECT * FROM members WHERE org_id = ? AND sso_subject = ? AND removed = 0",
+                (org_id, subject),
+            ).fetchone()
+        return dict(row) if row else None
+
+    def members(self, org: str, include_removed: bool = False) -> list[dict]:
+        org_row = self.need_org(org)
+        with self._db() as db:
+            rows = db.execute(
+                "SELECT * FROM members WHERE org_id = ? ORDER BY created", (org_row["id"],)
+            ).fetchall()
+        out = []
+        for row in rows:
+            if row["removed"] and not include_removed:
+                continue
+            out.append({**dict(row), **self.usage(member_id=row["id"])})
+        return out
+
+    def _need_member(self, org: str, email: str) -> dict:
+        org_row = self.need_org(org)
+        member = self.member(org_row["id"], email)
+        if member is None:
+            raise OrgError("no_such_member", f"{email} is not a member of {org_row['name']}")
+        return member
+
+    def remove_member(self, org: str, email: str) -> None:
+        """Remove a seat and revoke every API key of that member."""
+        member = self._need_member(org, email)
+        if member["role"] == "admin" and self._admins(member["org_id"]) <= 1:
+            raise OrgError("last_admin", "cannot remove the only admin of an organisation")
+        with self._db() as db:
+            db.execute("UPDATE members SET removed = 1 WHERE id = ?", (member["id"],))
+            db.execute("UPDATE keys SET revoked = 1 WHERE member_id = ?", (member["id"],))
+
+    def _admins(self, org_id: str) -> int:
+        with self._db() as db:
+            return db.execute(
+                "SELECT COUNT(*) FROM members WHERE org_id = ? AND role = 'admin' AND removed = 0",
+                (org_id,),
+            ).fetchone()[0]
+
+    def set_role(self, org: str, email: str, role: str) -> None:
+        if role not in ROLES:
+            raise OrgError("bad_role", f"role must be one of {ROLES}")
+        member = self._need_member(org, email)
+        if member["role"] == "admin" and role != "admin" and self._admins(member["org_id"]) <= 1:
+            raise OrgError("last_admin", "cannot demote the only admin of an organisation")
+        with self._db() as db:
+            db.execute("UPDATE members SET role = ? WHERE id = ?", (role, member["id"]))
+
+    def set_cap(
+        self,
+        org: str,
+        email: str,
+        usd: float | None = None,
+        tokens: int | None = None,
+        clear: bool = False,
+    ) -> None:
+        member = self._need_member(org, email)
+        with self._db() as db:
+            db.execute(
+                "UPDATE members SET cap_usd = ?, cap_tokens = ? WHERE id = ?",
+                (
+                    None if clear else usd if usd is not None else member["cap_usd"],
+                    None if clear else tokens if tokens is not None else member["cap_tokens"],
+                    member["id"],
+                ),
+            )
+
+    # ------------------------------------------------------------------ invites
+
+    def invite(self, org: str, email: str, role: str = "member") -> str:
+        """A one-time invite token (shown once, stored hashed, valid for seven days)."""
+        if role not in ROLES:
+            raise OrgError("bad_role", f"role must be one of {ROLES}")
+        org_row = self.need_org(org)
+        token = "tfi_" + secrets.token_urlsafe(24)
+        with self._db() as db:
+            db.execute(
+                "INSERT INTO invites VALUES (?,?,?,?,?,?,0)",
+                (
+                    _hash(token),
+                    org_row["id"],
+                    email.strip().lower(),
+                    role,
+                    time.time(),
+                    time.time() + INVITE_TTL,
+                ),
+            )
+        return token
+
+    def accept_invite(self, token: str, name: str = "") -> tuple[str, dict]:
+        """Turn an invite into a seat and its first API key: ``(key, member)``."""
+        with self._db() as db:
+            row = db.execute("SELECT * FROM invites WHERE hash = ?", (_hash(token),)).fetchone()
+            if row is None or row["used"] or row["expires"] < time.time():
+                raise OrgError("bad_invite", "this invite is invalid, used or expired")
+            db.execute("UPDATE invites SET used = 1 WHERE hash = ?", (row["hash"],))
+        member = self.add_member(row["org_id"], row["email"], row["role"], name)
+        key, _ = self.issue_key(member, "invite")
+        return key, member
+
+    def issue_key(self, member: dict, label: str = "") -> tuple[str, dict]:
+        key = "tf_" + secrets.token_urlsafe(32)
+        key_id = uuid.uuid4().hex[:10]
+        with self._db() as db:
+            db.execute(
+                "INSERT INTO keys (id, name, hash, created, revoked, monthly_limit, member_id)"
+                " VALUES (?,?,?,?,0,NULL,?)",
+                (
+                    key_id,
+                    f"{member['email']} ({label})" if label else member["email"],
+                    _hash(key),
+                    time.time(),
+                    member["id"],
+                ),
+            )
+        return key, {"id": key_id, "member": member["id"]}
+
+    def member_of_key(self, key_id: str) -> dict | None:
+        with self._db() as db:
+            row = db.execute(
+                "SELECT m.*, o.name AS org_name FROM keys k JOIN members m ON m.id = k.member_id"
+                " JOIN orgs o ON o.id = m.org_id WHERE k.id = ? AND m.removed = 0",
+                (key_id,),
+            ).fetchone()
+        return dict(row) if row else None
+
+    # ------------------------------------------------------------------- usage
+
+    def usage(
+        self, member_id: str | None = None, org_id: str | None = None, since: float | None = None
+    ) -> dict[str, Any]:
+        """Totals for the current month: tokens, calls and USD (known part) per member or org."""
+        where, params = ["ts >= ?"], [since if since is not None else month_start()]
+        if member_id:
+            where.append("member_id = ?")
+            params.append(member_id)
+        if org_id:
+            where.append("org_id = ?")
+            params.append(org_id)
+        with self._db() as db:
+            row = db.execute(
+                "SELECT COALESCE(SUM(input_tokens + output_tokens + cache_read_tokens"
+                " + cache_write_tokens), 0), COALESCE(SUM(requests), 0),"
+                " COALESCE(SUM(cost_usd), 0), COALESCE(SUM(1 - cost_known), 0)"
+                f" FROM audit WHERE {' AND '.join(where)} AND event = 'turn'",
+                params,
+            ).fetchone()
+        return {
+            "tokens": int(row[0]),
+            "requests": int(row[1]),
+            "spent_usd": round(float(row[2]), 6),
+            "unpriced_turns": int(row[3]),
+        }
+
+    def usage_by_model(self, member_id: str | None = None, org_id: str | None = None) -> dict:
+        where, params = ["ts >= ?", "event = 'turn'"], [month_start()]
+        if member_id:
+            where.append("member_id = ?")
+            params.append(member_id)
+        if org_id:
+            where.append("org_id = ?")
+            params.append(org_id)
+        merged: dict[str, dict[str, int]] = {}
+        with self._db() as db:
+            for (blob,) in db.execute(
+                f"SELECT by_model FROM audit WHERE {' AND '.join(where)}", params
+            ):
+                for model, counts in json.loads(blob or "{}").items():
+                    slot = merged.setdefault(model, {})
+                    for k, v in counts.items():
+                        slot[k] = slot.get(k, 0) + int(v)
+        return merged
+
+    def check(
+        self, member: dict, model_priced: bool, inflight: dict[str, tuple[int, float]]
+    ) -> Check:
+        """Pre-flight: may this member start a turn now? ``inflight`` is spend of running turns."""
+        org = self.need_org(member["org_id"])
+        mine = self.usage(member_id=member["id"])
+        pool = self.usage(org_id=org["id"])
+        my_tokens, my_usd = inflight.get(member["id"], (0, 0.0))
+        org_tokens, org_usd = inflight.get(org["id"], (0, 0.0))
+        if member["cap_usd"] is not None and not model_priced:
+            return Check(
+                False,
+                "unpriced_model",
+                "this model has no price, so your spend cap cannot be enforced; pick another model",
+            )
+        if org["quota_usd"] is not None and not model_priced:
+            return Check(
+                False,
+                "unpriced_model",
+                "this model has no price, so the organisation's spend quota cannot be enforced; pick another model",
+            )
+        return self.verdict(member, org, mine, pool, (my_tokens, my_usd), (org_tokens, org_usd))
+
+    @staticmethod
+    def verdict(
+        member: dict, org: dict, mine: dict, pool: dict, mine_in: tuple, pool_in: tuple
+    ) -> Check:
+        spent = mine["spent_usd"] + mine_in[1]
+        tokens = mine["tokens"] + mine_in[0]
+        if member["cap_usd"] is not None and spent >= member["cap_usd"]:
+            return Check(
+                False,
+                "spend_cap",
+                f"your monthly spend cap of ${member['cap_usd']:.2f} is reached "
+                f"(${spent:.2f} spent); ask an admin to raise it",
+            )
+        if member["cap_tokens"] is not None and tokens >= member["cap_tokens"]:
+            return Check(
+                False, "spend_cap", f"your monthly token cap of {member['cap_tokens']:,} is reached"
+            )
+        if org["quota_usd"] is not None and pool["spent_usd"] + pool_in[1] >= org["quota_usd"]:
+            return Check(
+                False, "org_quota", f"{org['name']}'s pooled monthly spend quota is used up"
+            )
+        if org["quota_tokens"] is not None and pool["tokens"] + pool_in[0] >= org["quota_tokens"]:
+            return Check(
+                False, "org_quota", f"{org['name']}'s pooled monthly token quota is used up"
+            )
+        return Check(True)
+
+    # ------------------------------------------------------------------- audit
+
+    def record(
+        self,
+        *,
+        org_id: str | None,
+        member_id: str | None,
+        key_id: str | None,
+        event: str = "turn",
+        mode: str = "",
+        model: str = "",
+        tools: list[str] | None = None,
+        usage: dict[str, Any] | None = None,
+        cost_usd: float | None = None,
+        status: str = "ok",
+        detail: str = "",
+    ) -> None:
+        usage = usage or {}
+        with self._db() as db:
+            db.execute(
+                "INSERT INTO audit (ts, org_id, member_id, key_id, event, mode, model, tools,"
+                " input_tokens, output_tokens, cache_read_tokens, cache_write_tokens, requests,"
+                " cost_usd, cost_known, status, detail, by_model)"
+                " VALUES (?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?)",
+                (
+                    time.time(),
+                    org_id,
+                    member_id,
+                    key_id,
+                    event,
+                    mode,
+                    model,
+                    json.dumps(sorted(set(tools or []))),
+                    int(usage.get("input_tokens", 0)),
+                    int(usage.get("output_tokens", 0)),
+                    int(usage.get("cache_read_input_tokens", 0)),
+                    int(usage.get("cache_creation_input_tokens", 0)),
+                    int(usage.get("requests", 0)),
+                    cost_usd if cost_usd is not None else 0.0,
+                    1 if cost_usd is not None else 0,
+                    status,
+                    detail[:500],
+                    json.dumps(usage.get("by_model") or {}),
+                ),
+            )
+
+    def audit(self, org: str, limit: int = 50, member_id: str | None = None) -> list[dict]:
+        org_row = self.need_org(org)
+        query = (
+            "SELECT a.ts, a.event, a.mode, a.model, a.tools, a.input_tokens, a.output_tokens,"
+            " a.cache_read_tokens, a.requests, a.cost_usd, a.cost_known, a.status, a.detail,"
+            " m.email FROM audit a LEFT JOIN members m ON m.id = a.member_id WHERE a.org_id = ?"
+        )
+        params: list[Any] = [org_row["id"]]
+        if member_id:
+            query += " AND a.member_id = ?"
+            params.append(member_id)
+        query += " ORDER BY a.id DESC LIMIT ?"
+        params.append(limit)
+        with self._db() as db:
+            rows = db.execute(query, params).fetchall()
+        return [
+            {
+                "time": time.strftime("%Y-%m-%d %H:%M:%S", time.gmtime(r["ts"])) + "Z",
+                "user": r["email"] or "",
+                "event": r["event"],
+                "mode": r["mode"],
+                "model": r["model"],
+                "tools": json.loads(r["tools"] or "[]"),
+                "input_tokens": r["input_tokens"],
+                "output_tokens": r["output_tokens"],
+                "cached_tokens": r["cache_read_tokens"],
+                "requests": r["requests"],
+                "cost_usd": r["cost_usd"] if r["cost_known"] else None,
+                "status": r["status"],
+                "detail": r["detail"],
+            }
+            for r in rows
+        ]