Limit by tenant and org

src/lib/litellm.ts

@@ -93,6 +93,94 @@ export async function listTeams(): Promise<any[]> {
   return Array.isArray(data) ? data : data?.data ?? data?.teams ?? [];
 }
 
+/**
+ * Find a virtual key on a team by its alias and return its current
+ * state (token, spend, budget cap, reset cadence). Returns null if
+ * the alias doesn't match any key on the team.
+ *
+ * Why we need this
+ * ----------------
+ * Per-tenant budgets live on the virtual key, not the team. The
+ * portal needs to:
+ *   1. Display the current key's `max_budget` / `budget_duration` /
+ *      `spend` on the tenant detail page.
+ *   2. Pass the key's `token` to `/key/update` when the customer
+ *      changes the budget.
+ *
+ * The token is opaque to the customer; the operator's
+ * `FindKeyByAlias` does the same lookup for stale-key cleanup. We
+ * mirror its API call here.
+ */
+export async function findKeyByAlias(
+  teamId: string,
+  keyAlias: string
+): Promise<{
+  token: string;
+  spend: number;
+  maxBudget: number | null;
+  budgetDuration: string | null;
+} | null> {
+  const data = await litellmFetch(
+    `/key/list?team_id=${encodeURIComponent(teamId)}&return_full_object=true&include_team_keys=true`
+  );
+  const keys: any[] = Array.isArray(data?.keys)
+    ? data.keys
+    : Array.isArray(data?.data)
+      ? data.data
+      : Array.isArray(data)
+        ? data
+        : [];
+  for (const k of keys) {
+    if (typeof k !== "object" || k === null) continue;
+    const alias = k.key_alias ?? k.keyAlias;
+    if (alias !== keyAlias) continue;
+    if (typeof k.token !== "string" || !k.token) continue;
+    return {
+      token: k.token,
+      spend: typeof k.spend === "number" ? k.spend : Number(k.spend) || 0,
+      maxBudget:
+        typeof k.max_budget === "number"
+          ? k.max_budget
+          : k.max_budget == null
+            ? null
+            : Number(k.max_budget) || null,
+      budgetDuration:
+        typeof k.budget_duration === "string" ? k.budget_duration : null,
+    };
+  }
+  return null;
+}
+
+/**
+ * Update a virtual key's budget cap and reset duration.
+ *
+ * Pass `maxBudget: null` to remove the cap. Pass `budgetDuration:
+ * null` to make the budget never reset (lifetime cap).
+ *
+ * Identified by `key` parameter — accepts either the raw `sk-...`
+ * token or its hash (LiteLLM accepts both shapes on /key/update).
+ * The portal flow uses the hash returned by `findKeyByAlias`.
+ */
+export async function updateKeyBudget(
+  key: string,
+  changes: {
+    maxBudget?: number | null;
+    budgetDuration?: string | null;
+  }
+): Promise<void> {
+  const body: Record<string, any> = { key };
+  if (changes.maxBudget !== undefined) {
+    body.max_budget = changes.maxBudget;
+  }
+  if (changes.budgetDuration !== undefined) {
+    body.budget_duration = changes.budgetDuration;
+  }
+  await litellmFetch("/key/update", {
+    method: "POST",
+    body: JSON.stringify(body),
+  });
+}
+
 /**
  * Get LiteLLM health status.
  */