127 lines
4.0 KiB
TypeScript
127 lines
4.0 KiB
TypeScript
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 }
|
|
);
|
|
}
|
|
}
|