Limit by tenant and org

src/app/api/tenants/[name]/budget/route.ts

@@ -0,0 +1,126 @@
+import { NextRequest, NextResponse } from "next/server";
+import { z } from "zod";
+import { getSessionUser, canMutate } from "@/lib/session";
+import { getTenant } from "@/lib/k8s";
+import { canUserSeeTenant } from "@/lib/visibility";
+import { findKeyByAlias, updateKeyBudget } from "@/lib/litellm";
+import { safeError } from "@/lib/errors";
+
+/**
+ * Update the per-tenant budget — operates on the LiteLLM virtual
+ * key, NOT on the team.
+ *
+ * Why per-key
+ * -----------
+ * Each tenant in an org has its own virtual key
+ * (`key_alias = tenant.metadata.name`); the team that owns those
+ * keys is org-scoped and shared across all the org's tenants. A
+ * budget on the team would cap the whole org; a budget on the key
+ * caps just this one tenant. Customers landing on the tenant detail
+ * page reasonably expect "edit budget" to mean "the budget of THIS
+ * tenant" — so we put it on the key.
+ *
+ * The team-level (org-wide) budget is a separate control that lives
+ * in /settings (not yet implemented) — the two coexist: LiteLLM
+ * applies whichever cap is hit first.
+ *
+ * Schema:
+ *   - maxBudget: number > 0 (set a cap), or null (remove the cap).
+ *   - budgetDuration: one of "30d", "1mo", "1y", or null (lifetime).
+ *
+ * Authorization: owners and platform admins.
+ */
+
+const patchSchema = z.object({
+  // > 0 because LiteLLM rejects 0 and a zero cap would lock the key
+  // out instantly. Upper bound 1M as a typo guard.
+  maxBudget: z.number().positive().max(1_000_000).nullable(),
+  budgetDuration: z.enum(["30d", "1mo", "1y"]).nullable(),
+});
+
+export async function PATCH(
+  req: NextRequest,
+  { params }: { params: Promise<{ name: string }> }
+) {
+  const user = await getSessionUser();
+  if (!user) {
+    return NextResponse.json({ error: "Unauthorized" }, { status: 401 });
+  }
+  if (!canMutate(user)) {
+    return NextResponse.json({ error: "Forbidden" }, { status: 403 });
+  }
+
+  const { name } = await params;
+  const tenant = await getTenant(name);
+  if (!tenant) {
+    return NextResponse.json({ error: "Not found" }, { status: 404 });
+  }
+  if (!(await canUserSeeTenant(user, tenant))) {
+    // Don't leak existence — same 404 a non-visible tenant gets.
+    return NextResponse.json({ error: "Not found" }, { status: 404 });
+  }
+
+  const teamId = tenant.status?.litellmTeamId;
+  if (!teamId) {
+    return NextResponse.json(
+      {
+        error:
+          "Tenant has no LiteLLM team yet. Please wait until provisioning completes.",
+      },
+      { status: 409 }
+    );
+  }
+
+  const body = await req.json().catch(() => null);
+  const parsed = patchSchema.safeParse(body);
+  if (!parsed.success) {
+    return NextResponse.json(
+      { error: "Invalid input", details: parsed.error.flatten() },
+      { status: 400 }
+    );
+  }
+
+  // Defensive: removing the cap should null out the duration too —
+  // a reset cadence on an unlimited budget is meaningless and would
+  // confuse LiteLLM's bookkeeping.
+  const maxBudget = parsed.data.maxBudget;
+  const budgetDuration =
+    maxBudget === null ? null : parsed.data.budgetDuration;
+
+  // Look up the key by alias (= tenant name). The token returned is
+  // what /key/update wants in the `key` field.
+  let keyInfo;
+  try {
+    keyInfo = await findKeyByAlias(teamId, name);
+  } catch (e: any) {
+    console.error("Failed to look up tenant key:", e);
+    return NextResponse.json(
+      { error: safeError(e, "Failed to look up tenant key") },
+      { status: 500 }
+    );
+  }
+  if (!keyInfo) {
+    return NextResponse.json(
+      {
+        error:
+          "Tenant has no virtual key yet. Please wait until provisioning completes.",
+      },
+      { status: 409 }
+    );
+  }
+
+  try {
+    await updateKeyBudget(keyInfo.token, { maxBudget, budgetDuration });
+    return NextResponse.json({
+      message: maxBudget === null ? "Budget removed." : "Budget updated.",
+      maxBudget,
+      budgetDuration,
+    });
+  } catch (e: any) {
+    console.error("Failed to update key budget:", e);
+    return NextResponse.json(
+      { error: safeError(e, "Failed to update budget") },
+      { status: 500 }
+    );
+  }
+}