Flexible balance: per-member balance (DB) synced to team budget; keys inherit team; add /admin/set-balance top-up endpoint

app/main.py

@@ -166,9 +166,40 @@ def get_db() -> sqlite3.Connection:
     cols = [r[1] for r in conn.execute("PRAGMA table_info(members)").fetchall()]
     if "slug" not in cols:
         conn.execute("ALTER TABLE members ADD COLUMN slug TEXT")
+    # Migration: add balance column (USD) — the member's current credit
+    # allowance. Defaults to MEMBER_BUDGET; can be topped up by payments.
+    # This is the flexible counter to the fixed $15/month allowance.
+    if "balance" not in cols:
+        conn.execute("ALTER TABLE members ADD COLUMN balance REAL")
     return conn
 
 
+def get_member_balance(email: str) -> float:
+    """Return the member's current balance (USD), defaulting to MEMBER_BUDGET.
+
+    The balance is the flexible allowance — it starts at MEMBER_BUDGET but can
+    be adjusted (topped up by payments) without changing code.
+    """
+    conn = get_db()
+    row = conn.execute(
+        "SELECT balance FROM members WHERE email = ?", (email,)
+    ).fetchone()
+    conn.close()
+    if row and row[0] is not None:
+        return float(row[0])
+    return MEMBER_BUDGET
+
+
+def set_member_balance(email: str, balance: float) -> None:
+    """Set the member's balance (USD). Used by top-up / billing adjustments."""
+    conn = get_db()
+    conn.execute(
+        "UPDATE members SET balance = ? WHERE email = ?", (balance, email)
+    )
+    conn.commit()
+    conn.close()
+
+
 async def fetch_members() -> list[dict]:
     """Return the full active-member list: email, name, slug, role.
 
@@ -535,9 +566,10 @@ async def cloudron_get_invite_link(user_id: str) -> str:
 async def provision_member(email: str, name: str, slug: str = "") -> str:
     """Provision a member end-to-end and send our own welcome email.
 
-    Creates the Cloudron user (members group), issues a LiteLLM key, stores the
+    Creates the Cloudron user (members group), issues a LiteLLM key under the
-    mapping, and sends our branded welcome email containing the Cloudron
+    member's team (budget = the member's current balance, not a hardcoded $15),
-    account-setup link (instead of Cloudron's default invite email).
+    stores the mapping, and sends our branded welcome email containing the
+    Cloudron account-setup link (instead of Cloudron's default invite email).
 
     Idempotent: if the member already exists, reuses the existing user/key.
     Returns the Cloudron user id.
@@ -545,8 +577,9 @@ async def provision_member(email: str, name: str, slug: str = "") -> str:
     user_id = await cloudron_create_user(email, name)
     await cloudron_set_group(user_id)
     await cloudron_set_active(user_id, True)
-    team_id = await litellm_get_or_create_team(email, MEMBER_BUDGET)
+    balance = get_member_balance(email)
-    key_token = await litellm_create_key(email, MEMBER_BUDGET, team_id)
+    team_id = await litellm_get_or_create_team(email, balance)
+    key_token = await litellm_create_key(email, balance, team_id)
     store_member(email, key_token, user_id, slug)
 
     # Send our own welcome email with the account-setup link (no OAuth step).
@@ -595,6 +628,36 @@ async def litellm_get_or_create_team(email: str, budget: float) -> str:
         return r.json()["team_id"]
 
 
+async def litellm_update_team_budget(team_id: str, budget: float) -> None:
+    """Set a team's max_budget (the enforced spend cap).
+
+    This is the primitive behind "payment → add to balance": a top-up adjusts
+    the team's max_budget, which every key under the team (chat + API) draws
+    from. Keys inherit the team budget (their own max_budget is ignored when a
+    team_id is set), so the team is the single source of truth for allowance.
+    """
+    async with httpx.AsyncClient() as client:
+        r = await client.post(
+            f"{LITELLM_BASE}/team/update",
+            headers={"Authorization": f"Bearer {LITELLM_MASTER_KEY}"},
+            json={"team_id": team_id, "max_budget": budget},
+        )
+        r.raise_for_status()
+
+
+async def sync_team_budget_from_balance(email: str) -> float:
+    """Apply the member's stored balance to their team's max_budget.
+
+    Reads the balance from the members table and pushes it to LiteLLM. Returns
+    the applied budget. This keeps the portal DB and LiteLLM in sync whenever
+    a balance changes (top-up, reset, pricing-model change).
+    """
+    balance = get_member_balance(email)
+    team_id = await litellm_get_or_create_team(email, balance)
+    await litellm_update_team_budget(team_id, balance)
+    return balance
+
+
 async def litellm_create_key(email: str, budget: float, team_id: str) -> str:
     """Create a fresh LiteLLM virtual key under the member's team.
 
@@ -602,6 +665,12 @@ async def litellm_create_key(email: str, budget: float, team_id: str) -> str:
     create a new key rather than reusing — the sk- value is unrecoverable from
     a stored hash, so "reuse" would store a useless hash (a bug that previously
     broke chat for three members).
+
+    NOTE: the key does NOT set its own max_budget — it inherits the team's
+    budget (LiteLLM ignores a key's max_budget when team_id is set anyway).
+    The team is the single source of truth for allowance, so a balance change
+    (top-up) updates the team once and every key follows. The `budget` arg is
+    retained only for the team-creation fallback in the migration path.
     """
     alias = f"member:{email}"
     async with httpx.AsyncClient() as client:
@@ -611,8 +680,6 @@ async def litellm_create_key(email: str, budget: float, team_id: str) -> str:
             json={
                 "key_alias": alias,
                 "team_id": team_id,
-                "max_budget": budget,
-                "budget_duration": "30d",
                 "models": MEMBER_MODELS,
             },
         )
@@ -920,6 +987,41 @@ async def admin_provision(token: str, request: Request):
     return {"status": "provisioned", "email": email, "user_id": user_id}
 
 
+@app.post("/admin/set-balance/{token}")
+async def admin_set_balance(token: str, request: Request):
+    """Set (or add to) a member's credit balance and sync it to LiteLLM.
+
+    The flexible alternative to a fixed $15/month allowance. A payment (from OC
+    or another system) is recorded by adjusting the balance here, which pushes
+    the new cap to the member's team (every chat/API key under it follows).
+
+    Body: {"email": "...", "balance": 25.0}            → set absolute balance
+          {"email": "...", "add": 10.0}                → add to current balance
+    """
+    if not WEBHOOK_TOKEN or not secrets.compare_digest(token, WEBHOOK_TOKEN):
+        raise HTTPException(401, "Invalid token")
+
+    body = await request.json()
+    email = (body.get("email") or "").strip().lower()
+    if not email:
+        raise HTTPException(400, "Missing email")
+
+    current = get_member_balance(email)
+    if "add" in body:
+        new_balance = current + float(body["add"])
+    elif "balance" in body:
+        new_balance = float(body["balance"])
+    else:
+        raise HTTPException(400, "Provide 'balance' or 'add'")
+
+    if new_balance < 0:
+        raise HTTPException(400, "Balance cannot be negative")
+
+    set_member_balance(email, new_balance)
+    applied = await sync_team_budget_from_balance(email)
+    return {"status": "updated", "email": email, "balance": new_balance, "team_budget": applied}
+
+
 @app.post("/admin/migrate-teams/{token}")
 async def admin_migrate_teams(token: str):
     """One-off migration: put every active member under a LiteLLM team + fresh
@@ -942,9 +1044,10 @@ async def admin_migrate_teams(token: str):
         try:
             # Delete any existing key(s) for this alias (by hash or sk-).
             await litellm_delete_keys_by_alias(email)
-            # (Re)create the team and a fresh sk- key.
+            # (Re)create the team and a fresh sk- key, at the member's balance.
-            team_id = await litellm_get_or_create_team(email, MEMBER_BUDGET)
+            balance = get_member_balance(email)
-            new_key = await litellm_create_key(email, MEMBER_BUDGET, team_id)
+            team_id = await litellm_get_or_create_team(email, balance)
+            new_key = await litellm_create_key(email, balance, team_id)
             rekey_member(email, new_key)  # preserve cloudron_user_id + slug
             results.append({"email": email, "status": "rekeyed"})
             logger.info("Migrated %s to team %s", email, team_id)