Hotfix/AI post limits deploy fixes

Dockerfile

@@ -1,7 +1,6 @@
 FROM uselagoon/python-3.12:latest@sha256:5ab457220705f7b4c072ee746b5920779a385a70175e0471b9a263c840ff1070
 
-RUN apk add bash --no-cache
-RUN apk add curl --no-cache
+RUN apk add bash curl postgresql-client --no-cache
 
 WORKDIR /app
 

app/api/private_ai_keys.py

@@ -23,9 +23,9 @@
     DBPoolPurchase,
     DBPrivateAIKey,
     DBRegion,
+    DBSpendCap,
     DBUser,
     DBTeam,
-    DBSpendCap,
 )
 from app.services.litellm import LiteLLMService
 from app.core.security import (
@@ -917,7 +917,17 @@ async def delete_private_ai_key(
     return {"message": "Private AI Key deleted successfully"}
 
 
-@router.get("/{key_id}/spend", response_model=PrivateAIKeySpendBasic)
+@router.get(
+    "/{key_id}/spend",
+    response_model=PrivateAIKeySpendBasic,
+    deprecated=True,
+    summary="Legacy: use GET /spend/{region_id}/key/{key_id} instead",
+    description=(
+        "Returns spend and budget metadata for a specific key. "
+        "This endpoint is legacy. Prefer /spend/{{region_id}}/key/{{key_id}} "
+        "which provides richer budget metadata and team-level context."
+    ),
+)
 async def get_private_ai_key_spend(
     key_id: int,
     current_user=Depends(get_current_user_from_auth),
@@ -942,6 +952,24 @@ async def get_private_ai_key_spend(
         data = await litellm_service.get_key_info(private_ai_key.litellm_token)
         info = data.get("info", {})
 
+        # Override max_budget with the value from spend_caps DB table if present.
+        # This ensures the configured cap (which may differ from LiteLLM's value
+        # for purchase-gated teams) is returned to the caller.
+        # The unique index for key-scope caps is on (region_id, key_id); team_id
+        # and user_id are metadata only and must NOT be used as lookup filters.
+        configured_cap = (
+            db.query(DBSpendCap.max_budget)
+            .filter(
+                DBSpendCap.scope == "key",
+                DBSpendCap.region_id == private_ai_key.region_id,
+                DBSpendCap.key_id == private_ai_key.id,
+            )
+            .first()
+        )
+        if configured_cap is not None and configured_cap[0] is not None:
+            info = dict(info)
+            info["max_budget"] = round(float(configured_cap[0]), 4)
+
         # Only set default for spend field
         spend_info = {"spend": info.get("spend", 0.0), **info}
 

app/api/spend.py

@@ -254,17 +254,16 @@ def _get_spend_cap_max_budget(
     user_id: int | None = None,
     key_id: int | None = None,
 ) -> float | None:
-    cap = (
-        db.query(DBSpendCap.max_budget)
-        .filter(
-            DBSpendCap.scope == scope,
-            DBSpendCap.region_id == region_id,
-            DBSpendCap.team_id == team_id,
-            DBSpendCap.user_id == user_id,
-            DBSpendCap.key_id == key_id,
-        )
-        .first()
-    )
+    # Match on unique-index columns only (see _upsert_spend_cap for rationale).
+    filters = [DBSpendCap.scope == scope, DBSpendCap.region_id == region_id]
+    if scope == "team":
+        filters.append(DBSpendCap.team_id == team_id)
+    elif scope == "team_member":
+        filters.extend([DBSpendCap.team_id == team_id, DBSpendCap.user_id == user_id])
+    elif scope == "key":
+        filters.append(DBSpendCap.key_id == key_id)
+
+    cap = db.query(DBSpendCap.max_budget).filter(*filters).first()
     if cap is None or cap[0] is None:
         return None
     return float(cap[0])
@@ -384,17 +383,20 @@ def _upsert_spend_cap(
     month_anchor: date | None = None,
     month_start_spend: float | None = None,
 ) -> None:
-    cap = (
-        db.query(DBSpendCap)
-        .filter(
-            DBSpendCap.scope == scope,
-            DBSpendCap.region_id == region_id,
-            DBSpendCap.team_id == team_id,
-            DBSpendCap.user_id == user_id,
-            DBSpendCap.key_id == key_id,
-        )
-        .first()
-    )
+    # Look up the existing row using the same columns as the partial unique
+    # index for this scope.  A previous implementation filtered on ALL four
+    # columns (team_id, user_id, key_id) which could miss a row whose
+    # team_id had been repaired from NULL to a real value, causing a
+    # UniqueViolation on INSERT.
+    filters = [DBSpendCap.scope == scope, DBSpendCap.region_id == region_id]
+    if scope == "team":
+        filters.append(DBSpendCap.team_id == team_id)
+    elif scope == "team_member":
+        filters.extend([DBSpendCap.team_id == team_id, DBSpendCap.user_id == user_id])
+    elif scope == "key":
+        filters.append(DBSpendCap.key_id == key_id)
+
+    cap = db.query(DBSpendCap).filter(*filters).first()
     if cap is None:
         cap = DBSpendCap(
             scope=scope,
@@ -403,6 +405,17 @@ def _upsert_spend_cap(
             user_id=user_id,
             key_id=key_id,
         )
+    else:
+        # Repair stale columns so the row stays consistent with the
+        # current key/team/user relationship. Normalize all relationship
+        # columns to the requested values, including clearing stale values
+        # to None when they are not part of the current scope.
+        if cap.team_id != team_id:
+            cap.team_id = team_id
+        if cap.user_id != user_id:
+            cap.user_id = user_id
+        if cap.key_id != key_id:
+            cap.key_id = key_id
     cap.max_budget = max_budget
     cap.budget_duration = budget_duration
     cap.month_anchor = month_anchor
@@ -421,17 +434,16 @@ def _delete_spend_cap(
     user_id: int | None = None,
     key_id: int | None = None,
 ) -> None:
-    (
-        db.query(DBSpendCap)
-        .filter(
-            DBSpendCap.scope == scope,
-            DBSpendCap.region_id == region_id,
-            DBSpendCap.team_id == team_id,
-            DBSpendCap.user_id == user_id,
-            DBSpendCap.key_id == key_id,
-        )
-        .delete()
-    )
+    # Match on unique-index columns only (see _upsert_spend_cap for rationale).
+    filters = [DBSpendCap.scope == scope, DBSpendCap.region_id == region_id]
+    if scope == "team":
+        filters.append(DBSpendCap.team_id == team_id)
+    elif scope == "team_member":
+        filters.extend([DBSpendCap.team_id == team_id, DBSpendCap.user_id == user_id])
+    elif scope == "key":
+        filters.append(DBSpendCap.key_id == key_id)
+
+    db.query(DBSpendCap).filter(*filters).delete()
     # Defer commit to the endpoint so DB changes and remote sync share one boundary.
     db.flush()