From 2ac821087ecdc086ec71142177d198a08e05d9c7 Mon Sep 17 00:00:00 2001 From: inference-bot Date: Mon, 14 Sep 2026 09:14:26 -0600 Subject: [PATCH] Teams migration: per-member LiteLLM team + fresh sk- key (fixes hash-key bug breaking chat for 3 members) --- app/main.py | 158 ++++++++++++++++++++++++++++++++++++++++++---------- 1 file changed, 130 insertions(+), 28 deletions(-) diff --git a/app/main.py b/app/main.py index 6c67924..cdb1a48 100644 --- a/app/main.py +++ b/app/main.py @@ -71,6 +71,9 @@ OC_PERSONAL_TOKEN = os.environ.get("OC_PERSONAL_TOKEN", "") # regardless of their $10/15/20 contribution. Governance decision (Loomio). MEMBER_BUDGET = float(os.environ.get("MEMBER_BUDGET", "15.0")) +# Models every member key/team may access. +MEMBER_MODELS = ["deepseek-v4-flash", "gpt-oss-120b", "glm-5-3-flash"] + # The "members" group in Cloudron (group-based access control). # Members are assigned to this group, which grants access to the chat app. MEMBERS_GROUP_ID = os.environ.get("MEMBERS_GROUP_ID", "") @@ -358,6 +361,21 @@ def store_member(email: str, key_token: str, cloudron_user_id: str, slug: str = conn.close() +def rekey_member(email: str, key_token: str) -> None: + """Update only the key token, preserving cloudron_user_id and slug. + + Used by the Teams migration to swap in a fresh sk- key without clobbering + the member's existing Cloudron identity. + """ + conn = get_db() + conn.execute( + "UPDATE members SET key_token = ?, active = 1 WHERE email = ?", + (key_token, email), + ) + conn.commit() + conn.close() + + def get_member_key(email: str) -> str | None: conn = get_db() row = conn.execute( @@ -507,7 +525,8 @@ 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) - key_token = await litellm_create_key(email, MEMBER_BUDGET) + team_id = await litellm_get_or_create_team(email, MEMBER_BUDGET) + key_token = await litellm_create_key(email, MEMBER_BUDGET, team_id) store_member(email, key_token, user_id, slug) # Send our own welcome email with the account-setup link (no OAuth step). @@ -520,50 +539,61 @@ async def provision_member(email: str, name: str, slug: str = "") -> str: return user_id -async def litellm_find_key_by_alias(alias: str) -> str | None: - """Find an existing key's token by its alias (key/list returns tokens).""" +async def litellm_get_or_create_team(email: str, budget: float) -> str: + """Return the team_id for a member's budget team, creating it if needed. + + Each member gets a LiteLLM "team" (a budget container) with a $15/30d + budget. The team holds the chat key + any member-managed API keys, so all + spend draws from the one pool. + + Idempotent: looks up the existing team by alias before creating. Note that + /team/new is NOT idempotent (it creates a new team each call), so we must + check /team/list first. + """ + alias = f"member:{email}" async with httpx.AsyncClient() as client: r = await client.get( - f"{LITELLM_BASE}/key/list", + f"{LITELLM_BASE}/team/list", headers={"Authorization": f"Bearer {LITELLM_MASTER_KEY}"}, ) r.raise_for_status() - tokens = r.json().get("keys", []) - for token in tokens: - info = await client.get( - f"{LITELLM_BASE}/key/info", - headers={"Authorization": f"Bearer {LITELLM_MASTER_KEY}"}, - params={"key": token}, - ) - if info.status_code == 200: - data = info.json().get("info", {}) - if data.get("key_alias") == alias: - return token - return None + for t in r.json(): + if t.get("team_alias") == alias: + return t["team_id"] + + r = await client.post( + f"{LITELLM_BASE}/team/new", + headers={"Authorization": f"Bearer {LITELLM_MASTER_KEY}"}, + json={ + "team_alias": alias, + "max_budget": budget, + "budget_duration": "30d", + "models": MEMBER_MODELS, + }, + ) + r.raise_for_status() + return r.json()["team_id"] -async def litellm_create_key(email: str, budget: float) -> str: - """Create a LiteLLM virtual key for a member with a budget cap. +async def litellm_create_key(email: str, budget: float, team_id: str) -> str: + """Create a fresh LiteLLM virtual key under the member's team. - Idempotent: if a key with this alias already exists, return its token. + Returns the ACTUAL sk- key (only revealed at generation time). We always + 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). """ alias = f"member:{email}" - - # If the key already exists (e.g. from a prior partial attempt), reuse it. - existing = await litellm_find_key_by_alias(alias) - if existing: - logger.info("Reusing existing LiteLLM key for %s", email) - return existing - async with httpx.AsyncClient() as client: r = await client.post( f"{LITELLM_BASE}/key/generate", headers={"Authorization": f"Bearer {LITELLM_MASTER_KEY}"}, json={ "key_alias": alias, + "team_id": team_id, "max_budget": budget, "budget_duration": "30d", - "models": ["deepseek-v4-flash", "gpt-oss-120b", "glm-5-3-flash"], + "models": MEMBER_MODELS, }, ) r.raise_for_status() @@ -571,7 +601,11 @@ async def litellm_create_key(email: str, budget: float) -> str: async def litellm_disable_key(key_token: str) -> None: - """Delete a member's LiteLLM key (on payment lapse).""" + """Delete a member's LiteLLM key (on payment lapse). + + /key/delete accepts either the sk- key or its SHA256 hash, so this works + for both current keys and legacy hash-keyed members. + """ if not key_token: return async with httpx.AsyncClient() as client: @@ -790,6 +824,74 @@ async def admin_provision(token: str, request: Request): return {"status": "provisioned", "email": email, "user_id": user_id} +@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 + sk- key. Fixes a bug where some members had SHA256 hashes (from /key/list) + stored as their key, which cannot authenticate. + + For each active member: delete any existing key(s) by alias, create/reuse a + team, generate a fresh sk- key under it, and store the sk- token. Idempotent + but re-issues keys, so call once. + """ + if not WEBHOOK_TOKEN or not secrets.compare_digest(token, WEBHOOK_TOKEN): + raise HTTPException(401, "Invalid token") + + conn = get_db() + rows = conn.execute("SELECT email, key_token FROM members WHERE active = 1").fetchall() + conn.close() + + results = [] + for email, old_token in rows: + 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. + team_id = await litellm_get_or_create_team(email, MEMBER_BUDGET) + new_key = await litellm_create_key(email, MEMBER_BUDGET, 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) + except Exception as e: + results.append({"email": email, "status": "error", "error": str(e)}) + logger.warning("Migrate failed for %s: %s", email, e) + + return {"status": "migrated", "results": results} + + +async def litellm_delete_keys_by_alias(email: str) -> None: + """Delete all LiteLLM keys whose alias matches member:. + + /key/list returns SHA256 hashes; /key/delete accepts hashes. This clears + both legacy hash-keyed and current sk- keyed entries for a member. + """ + alias = f"member:{email}" + async with httpx.AsyncClient() as client: + r = await client.get( + f"{LITELLM_BASE}/key/list", + headers={"Authorization": f"Bearer {LITELLM_MASTER_KEY}"}, + ) + r.raise_for_status() + hashes = r.json().get("keys", []) + to_delete = [] + for h in hashes: + info = await client.get( + f"{LITELLM_BASE}/key/info", + headers={"Authorization": f"Bearer {LITELLM_MASTER_KEY}"}, + params={"key": h}, + ) + if info.status_code == 200: + data = info.json().get("info", {}) + if data.get("key_alias") == alias: + to_delete.append(h) + if to_delete: + await client.post( + f"{LITELLM_BASE}/key/delete", + headers={"Authorization": f"Bearer {LITELLM_MASTER_KEY}"}, + json={"keys": to_delete}, + ) + + async def reconcile_memberships() -> dict: """Reconcile the portal against the live Open Collective member list.