From 6a7eaf0ff8839c35c48913ec21571cb20bb89169 Mon Sep 17 00:00:00 2001
From: InferenceBot <inference-bot@inference.coop>
Date: Sun, 1 Nov 2026 18:00:25 +0000
Subject: [PATCH] broker: per-key spend (chat separate) in /broker/usage

Members need to see how much each API key has spent. LiteLLM's per-key
'spend' counter is cycle-only (resets at budget_reset_at), so derive it
from the LiteLLM_SpendLogs table instead: sum spend grouped by api_key,
matching rows to member_api_keys.hash by name and labelling unmatched
rows 'chat' (the portal-held chat key is never in member_api_keys).

broker_get_usage now returns key_spend: [{name, is_chat, spend, tokens,
requests}], API keys by spend desc with the chat row last.
---
 app/main.py | 81 +++++++++++++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 81 insertions(+)

diff --git a/app/main.py b/app/main.py
index e0658cf..110ec4f 100644
--- a/app/main.py
+++ b/app/main.py
@@ -3522,6 +3522,74 @@ async def broker_set_matrix_mxid(request: Request):
     return {"status": "saved", "matrix_mxid": mxid, "invite_sent": invite_sent}
 
 
+def member_per_key_spend(email: str, team_id: str) -> list[dict]:
+    """Per-key spend for a member, derived from LiteLLM spend LOGS.
+
+    The per-key `spend` counter LiteLLM exposes (`LiteLLM_VerificationToken.spend`,
+    `/spend/keys`) is CYCLE-ONLY: it resets at the team's `budget_reset_at`, so it
+    cannot answer "how much has this key spent?". We instead sum `spend` over the
+    raw rows in `LiteLLM_SpendLogs`, grouped by the key identifier in `api_key`
+    (the key's LiteLLM hash — the same value we store as `member_api_keys.hash`).
+
+    A row whose `api_key` matches one of the member's stored `member_api_keys`
+    rows is labelled with that key's name. Any row that matches no stored key
+    belongs to the portal-held CHAT key (alias `member:<email>`, minted at
+    provisioning, never in member_api_keys) and is grouped under the label
+    "chat" — so a member sees "chat: $X, <keyname>: $Y".
+
+    Returns [{"name", "is_chat", "spend", "tokens", "requests"}], API keys by
+    spend descending with the chat row last (it is secondary in the keys view).
+    Empty list when the DB is not configured — never raises for the caller.
+    """
+    if not litellm_db_ready():
+        return []
+
+    conn = get_db()
+    try:
+        rows = conn.execute(
+            "SELECT name, hash FROM member_api_keys WHERE email = ?", (email,)
+        ).fetchall()
+    finally:
+        conn.close()
+    by_hash = {r[1]: r[0] for r in rows if r[1]}
+
+    pg = _pg_connect()
+    try:
+        with pg.cursor() as cur:
+            cur.execute(
+                'SELECT api_key, COALESCE(SUM(spend), 0), '
+                'COALESCE(SUM(completion_tokens), 0), COUNT(*) '
+                'FROM "LiteLLM_SpendLogs" WHERE team_id = %s GROUP BY api_key',
+                (team_id,),
+            )
+            raw = cur.fetchall()
+    finally:
+        pg.close()
+
+    # Aggregate by a stable internal id. Rows matching no stored key are the chat
+    # key; the sentinel can't collide with a real name (key names are
+    # [A-Za-z0-9._-], so a member key literally named "chat" stays separate).
+    agg: dict = {}
+    for api_key, spend, tokens, reqs in raw:
+        name = by_hash.get(api_key)
+        cid = name if name else "__chat__"
+        a = agg.setdefault(
+            cid,
+            {"name": name or "chat", "is_chat": name is None,
+             "spend": 0.0, "tokens": 0, "requests": 0},
+        )
+        a["spend"] += float(spend or 0.0)
+        a["tokens"] += int(tokens or 0)
+        a["requests"] += int(reqs or 0)
+
+    out = list(agg.values())
+    # API keys by spend desc, then the chat key last.
+    out.sort(key=lambda x: (x["is_chat"], -x["spend"]))
+    for r in out:
+        r["spend"] = round(r["spend"], 2)
+    return out
+
+
 async def broker_get_usage(email: str, days: int = 0) -> dict:
     """Return the member's spend vs. balance, plus per-model and token detail.
 
@@ -3569,6 +3637,14 @@ async def broker_get_usage(email: str, days: int = 0) -> dict:
                  "calls": m["requests"]}
                 for m in usage["models"]
             ]
+            # Per-key spend (chat key shown separately), from the spend-log rows —
+            # never LiteLLM's cycle-only per-key counter. Best-effort: a failure
+            # must not blank the whole usage view.
+            try:
+                key_spend = member_per_key_spend(email, team_id)
+            except Exception as e:
+                logger.warning("broker_get_usage: per-key spend failed: %s", e)
+                key_spend = []
             return {
                 "email": email,
                 "balance": balance,
@@ -3581,6 +3657,7 @@ async def broker_get_usage(email: str, days: int = 0) -> dict:
                 "total_tokens": usage["totals"]["total_tokens"],
                 "reset_at": reset_at,
                 "models": model_list,
+                "key_spend": key_spend,
             }
         except Exception as e:
             logger.warning("broker_get_usage: DB path failed (%s); falling back to HTTP", e)
@@ -3654,6 +3731,10 @@ async def broker_get_usage(email: str, days: int = 0) -> dict:
         "total_tokens": total_tokens,
         "reset_at": reset_at,
         "models": model_list,
+        # Per-key spend needs the Postgres spend-log table (LiteLLM's exposed
+        # per-key counter is cycle-only); this HTTP fallback runs only when the
+        # DB is unconfigured, so there is nothing accurate to report here.
+        "key_spend": [],
     }