diff --git a/tests/e2e/ui/helpers/navigation.ts b/tests/e2e/ui/helpers/navigation.ts index 4a7c4e7baa9..e0e7b4da396 100644 --- a/tests/e2e/ui/helpers/navigation.ts +++ b/tests/e2e/ui/helpers/navigation.ts @@ -73,3 +73,13 @@ export async function clickTeamId(page: PlaywrightPage, teamId: string): Promise await cell.click(); await expect(page.getByText("Back to Teams")).toBeVisible({ timeout: 10_000 }); } + +export async function openKeyDetail(page: PlaywrightPage, alias: string): Promise { + await page.getByPlaceholder("Search by key alias or ID").fill(alias); + const row = page.getByRole("row").filter({ hasText: alias }); + await expect(row, `key row "${alias}" never appeared on the Virtual Keys page`).toBeVisible({ timeout: 15_000 }); + await row.getByRole("button", { name: alias }).click(); + await expect(page.getByText("Back to Keys"), `key detail for "${alias}" never opened`).toBeVisible({ + timeout: 15_000, + }); +} diff --git a/tests/e2e/ui/helpers/traffic.ts b/tests/e2e/ui/helpers/traffic.ts index 7f8417cdffb..cb68747b364 100644 --- a/tests/e2e/ui/helpers/traffic.ts +++ b/tests/e2e/ui/helpers/traffic.ts @@ -1,4 +1,4 @@ -import { APIRequestContext, expect } from "@playwright/test"; +import { APIRequestContext, APIResponse, expect } from "@playwright/test"; /** Model names served by fixtures/config.yml, both backed by the mock LLM server. */ export const CHAT_MODEL_A = "fake-openai-gpt-4"; @@ -15,6 +15,9 @@ export const masterKey = (): string => process.env.LITELLM_MASTER_KEY || "sk-123 export const rootPath = (): string => process.env.SERVER_ROOT_PATH ?? ""; +/** Date.now() alone collides: `--repeat-each` starts its copies inside the same millisecond. */ +export const uniqueSuffix = (): string => `${Date.now()}-${Math.random().toString(36).slice(2, 8)}`; + interface ChatOptions { model: string; prompt: string; @@ -25,9 +28,8 @@ interface ChatOptions { traceId?: string; } -/** POST /v1/chat/completions and return the completion id (the Logs Request ID). */ -export async function sendChatCompletion(request: APIRequestContext, opts: ChatOptions): Promise { - const res = await request.post(`${rootPath()}/v1/chat/completions`, { +const postChatCompletion = (request: APIRequestContext, opts: ChatOptions): Promise => + request.post(`${rootPath()}/v1/chat/completions`, { headers: { Authorization: `Bearer ${opts.apiKey ?? masterKey()}`, "Content-Type": "application/json", @@ -39,12 +41,26 @@ export async function sendChatCompletion(request: APIRequestContext, opts: ChatO ...(opts.traceId ? { litellm_trace_id: opts.traceId } : {}), }, }); + +/** POST /v1/chat/completions and return the completion id (the Logs Request ID). */ +export async function sendChatCompletion(request: APIRequestContext, opts: ChatOptions): Promise { + const res = await postChatCompletion(request, opts); expect(res.ok(), `chat completion for ${opts.model} failed (${res.status()}): ${await res.text()}`).toBe(true); const body = await res.json(); expect(body.choices?.[0]?.message?.content).toContain(MOCK_RESPONSE_TEXT); return body.id as string; } +export interface ChatAttempt { + status: number; + body: string; +} + +export async function attemptChatCompletion(request: APIRequestContext, opts: ChatOptions): Promise { + const res = await postChatCompletion(request, opts); + return { status: res.status(), body: await res.text() }; +} + /** `key` is the sk- value to authenticate with; `token` is its hash, which spend aggregates are keyed by. */ export async function createVirtualKey( request: APIRequestContext, @@ -66,6 +82,33 @@ export async function createVirtualKey( }; } +export interface KeyInfo { + key_alias: string | null; + max_budget: number | null; + budget_duration: string | null; + budget_reset_at: string | null; + blocked: boolean | null; + models: string[]; + team_id: string | null; +} + +export async function readKeyInfo(request: APIRequestContext, token: string): Promise { + const res = await request.get(`${rootPath()}/key/info?key=${encodeURIComponent(token)}`, { + headers: { Authorization: `Bearer ${masterKey()}` }, + }); + expect(res.ok(), `GET /key/info for ${token} failed (${res.status()}): ${await res.text()}`).toBe(true); + const body = await res.json(); + return body.info as KeyInfo; +} + +export async function deleteVirtualKey(request: APIRequestContext, token: string): Promise { + const res = await request.post(`${rootPath()}/key/delete`, { + headers: { Authorization: `Bearer ${masterKey()}`, "Content-Type": "application/json" }, + data: { keys: [token] }, + }); + expect(res.ok(), `key delete for ${token} failed (${res.status()}): ${await res.text()}`).toBe(true); +} + /** Spend logs are flushed on a timer, so an assertion straight after a completion races the writer. */ export async function waitForSpendLog( request: APIRequestContext, diff --git a/tests/e2e/ui/tests/internal-user/internalUserKeyScope.spec.ts b/tests/e2e/ui/tests/internal-user/internalUserKeyScope.spec.ts new file mode 100644 index 00000000000..f923841257a --- /dev/null +++ b/tests/e2e/ui/tests/internal-user/internalUserKeyScope.spec.ts @@ -0,0 +1,208 @@ +import { test, expect, type APIRequestContext } from "@playwright/test"; +import { Page } from "../../fixtures/pages"; +import { + dismissFeedbackPopup, + navigateToPage, + openKeyDetail, +} from "../../helpers/navigation"; +import { + CHAT_MODEL_A, + CHAT_MODEL_B, + MOCK_RESPONSE_TEXT, + attemptChatCompletion, + createVirtualKey, + deleteVirtualKey, + masterKey, + readKeyInfo, + rootPath, + uniqueSuffix, +} from "../../helpers/traffic"; + +const MEMBER_PASSWORD = "E2e-Team-Member-Pass-1!"; + +interface CreatedTeam { + readonly team_id: string; +} + +function assertCreatedTeam(body: unknown): asserts body is CreatedTeam { + expect(body, "/team/new returned no team_id").toMatchObject({ + team_id: expect.any(String), + }); +} + +async function postAsMaster( + request: APIRequestContext, + path: string, + data: Record, +): Promise { + const res = await request.post(`${rootPath()}${path}`, { + headers: { + Authorization: `Bearer ${masterKey()}`, + "Content-Type": "application/json", + }, + data, + }); + expect( + res.ok(), + `POST ${path} failed (${res.status()}): ${await res.text()}`, + ).toBe(true); + return res.json(); +} + +test.describe("Internal User - own team key model scope", () => { + test.use({ storageState: { cookies: [], origins: [] } }); + + test("a team member narrows their own key's models and the proxy enforces it", async ({ + page, + request, + }) => { + const suffix = uniqueSuffix(); + const email = `team-member-${suffix}@test.local`; + const userId = `e2e-key-scope-user-${suffix}`; + const alias = `e2e-key-scope-${suffix}`; + + const team = await postAsMaster(request, "/team/new", { + team_alias: `E2E Key Scope ${suffix}`, + models: [CHAT_MODEL_A, CHAT_MODEL_B], + team_member_permissions: ["/key/generate", "/key/update", "/key/info"], + }); + assertCreatedTeam(team); + const teamId = team.team_id; + + try { + await postAsMaster(request, "/user/new", { + user_id: userId, + user_email: email, + user_role: "internal_user", + auto_create_key: false, + }); + await postAsMaster(request, "/user/update", { + user_id: userId, + password: MEMBER_PASSWORD, + }); + await postAsMaster(request, "/team/member_add", { + team_id: teamId, + member: { role: "user", user_id: userId }, + }); + + const created = await createVirtualKey(request, { + key_alias: alias, + team_id: teamId, + user_id: userId, + models: [], + }); + + try { + await page.goto("/ui/login"); + await page.getByPlaceholder("Enter your username").fill(email); + await page + .getByPlaceholder("Enter your password") + .fill(MEMBER_PASSWORD); + await page.getByRole("button", { name: "Login", exact: true }).click(); + await expect( + page.locator("a", { hasText: "Virtual Keys" }), + `${email} never reached the dashboard`, + ).toBeVisible({ timeout: 30_000 }); + await dismissFeedbackPopup(page); + + await navigateToPage(page, Page.ApiKeys); + await openKeyDetail(page, alias); + + await page.getByRole("tab", { name: "Settings" }).click(); + await page.getByRole("button", { name: "Edit Settings" }).click(); + + await page.getByRole("combobox", { name: "Select models" }).click(); + await expect( + page.getByRole("option", { name: CHAT_MODEL_A, exact: true }), + `the Models dropdown does not offer ${CHAT_MODEL_A} to a team member`, + ).toBeVisible({ timeout: 15_000 }); + await expect( + page.getByRole("option", { name: CHAT_MODEL_B, exact: true }), + `the Models dropdown does not offer ${CHAT_MODEL_B} to a team member`, + ).toBeVisible(); + + await page + .getByRole("option", { name: CHAT_MODEL_A, exact: true }) + .click(); + await page.keyboard.press("Escape"); + + const updated = page.waitForResponse( + (res) => + res.url().includes("/key/update") && + res.request().method() === "POST", + ); + await page.getByRole("button", { name: "Save Changes" }).click(); + const updateStatus = (await updated).status(); + expect( + updateStatus, + "a team member's own-key edit was refused", + ).toBeGreaterThanOrEqual(200); + expect( + updateStatus, + "a team member's own-key edit was refused", + ).toBeLessThan(300); + await expect( + page.getByText("Key updated successfully").first(), + ).toBeVisible({ timeout: 15_000 }); + + await expect + .poll( + async () => (await readKeyInfo(request, created.token)).models, + { + message: `the narrowed model scope never reached /key/info for ${alias}`, + timeout: 20_000, + }, + ) + .toEqual([CHAT_MODEL_A]); + + await expect + .poll( + async () => + await attemptChatCompletion(request, { + model: CHAT_MODEL_B, + prompt: `out of scope ${suffix}`, + apiKey: created.key, + }), + { + message: `${CHAT_MODEL_B} was still served after the key was narrowed to ${CHAT_MODEL_A}`, + timeout: 30_000, + }, + ) + .toMatchObject({ + status: 403, + body: expect.stringContaining(CHAT_MODEL_B), + }); + + const inScope = await attemptChatCompletion(request, { + model: CHAT_MODEL_A, + prompt: `in scope ${suffix}`, + apiKey: created.key, + }); + expect( + inScope, + `${CHAT_MODEL_A} is no longer served by the narrowed key`, + ).toMatchObject({ + status: 200, + body: expect.stringContaining(MOCK_RESPONSE_TEXT), + }); + } finally { + await deleteVirtualKey(request, created.token); + } + } finally { + await request.post(`${rootPath()}/user/delete`, { + headers: { + Authorization: `Bearer ${masterKey()}`, + "Content-Type": "application/json", + }, + data: { user_ids: [userId] }, + }); + await request.post(`${rootPath()}/team/delete`, { + headers: { + Authorization: `Bearer ${masterKey()}`, + "Content-Type": "application/json", + }, + data: { team_ids: [teamId] }, + }); + } + }); +}); diff --git a/tests/e2e/ui/tests/proxy-admin/keyBlocking.spec.ts b/tests/e2e/ui/tests/proxy-admin/keyBlocking.spec.ts new file mode 100644 index 00000000000..99a8065a797 --- /dev/null +++ b/tests/e2e/ui/tests/proxy-admin/keyBlocking.spec.ts @@ -0,0 +1,112 @@ +import { test as base, expect } from "@playwright/test"; +import { ADMIN_STORAGE_PATH } from "../../constants"; +import { Page } from "../../fixtures/pages"; +import { dismissFeedbackPopup, navigateToPage, openKeyDetail } from "../../helpers/navigation"; +import { + CHAT_MODEL_A, + MOCK_RESPONSE_TEXT, + attemptChatCompletion, + createVirtualKey, + deleteVirtualKey, + readKeyInfo, + sendChatCompletion, + uniqueSuffix, +} from "../../helpers/traffic"; + +interface ScopedKey { + alias: string; + token: string; + apiKey: string; +} + +const test = base.extend<{ scopedKey: ScopedKey }>({ + scopedKey: async ({ page }, use) => { + const alias = `e2e-block-key-${uniqueSuffix()}`; + const created = await createVirtualKey(page.request, { + key_alias: alias, + models: [CHAT_MODEL_A], + }); + await use({ alias, token: created.token, apiKey: created.key }); + await deleteVirtualKey(page.request, created.token); + }, +}); + +test.describe("Proxy Admin - Key blocking", () => { + test.use({ storageState: ADMIN_STORAGE_PATH }); + + test("blocking a key stops it serving and unblocking restores it", async ({ page, scopedKey }) => { + const { alias, token, apiKey } = scopedKey; + + await sendChatCompletion(page.request, { + model: CHAT_MODEL_A, + prompt: `pre-block ${alias}`, + apiKey, + }); + + await navigateToPage(page, Page.ApiKeys); + await dismissFeedbackPopup(page); + await openKeyDetail(page, alias); + + await page.getByRole("button", { name: "More key actions" }).click(); + await page.getByRole("menuitem", { name: "Block Key" }).click(); + const blockDialog = page.getByRole("dialog", { name: "Block Key" }); + await expect(blockDialog, "the Block Key confirmation never opened").toBeVisible({ timeout: 10_000 }); + await blockDialog.getByRole("button", { name: "Block", exact: true }).click(); + + await expect + .poll(async () => (await readKeyInfo(page.request, token)).blocked, { + message: "the key never came back blocked from /key/info", + timeout: 20_000, + }) + .toBe(true); + + await expect + .poll( + async () => + await attemptChatCompletion(page.request, { + model: CHAT_MODEL_A, + prompt: "blocked", + apiKey, + }), + { + message: "a blocked key was still served by /v1/chat/completions", + timeout: 30_000, + }, + ) + .toMatchObject({ status: 401, body: expect.stringContaining("blocked") }); + + await page.reload(); + await expect( + page.getByText("Blocked", { exact: true }), + "the reloaded key detail does not show the key as blocked", + ).toBeVisible({ timeout: 15_000 }); + + await page.getByRole("button", { name: "More key actions" }).click(); + await page.getByRole("menuitem", { name: "Unblock Key" }).click(); + const unblockDialog = page.getByRole("dialog", { name: "Unblock Key" }); + await expect(unblockDialog, "the Unblock Key confirmation never opened").toBeVisible({ timeout: 10_000 }); + await unblockDialog.getByRole("button", { name: "Unblock", exact: true }).click(); + + await expect + .poll(async () => (await readKeyInfo(page.request, token)).blocked, { + message: "the key never came back unblocked from /key/info", + timeout: 20_000, + }) + .toBe(false); + + await expect + .poll( + async () => + await attemptChatCompletion(page.request, { + model: CHAT_MODEL_A, + prompt: "unblocked", + apiKey, + }), + { + message: "an unblocked key is still refused by /v1/chat/completions", + timeout: 30_000, + }, + ) + .toMatchObject({ status: 200, body: expect.stringContaining(MOCK_RESPONSE_TEXT) }); + }); +}); diff --git a/tests/e2e/ui/tests/proxy-admin/keyBudgetWindow.spec.ts b/tests/e2e/ui/tests/proxy-admin/keyBudgetWindow.spec.ts new file mode 100644 index 00000000000..4e4d0a395c3 --- /dev/null +++ b/tests/e2e/ui/tests/proxy-admin/keyBudgetWindow.spec.ts @@ -0,0 +1,101 @@ +import { test as base, expect } from "@playwright/test"; +import { ADMIN_STORAGE_PATH, E2E_TEAM_CRUD_ID } from "../../constants"; +import { Page } from "../../fixtures/pages"; +import { dismissFeedbackPopup, navigateToPage, openKeyDetail } from "../../helpers/navigation"; +import { captureRequestBody } from "../../helpers/roundTrip"; +import { CHAT_MODEL_A, createVirtualKey, deleteVirtualKey, readKeyInfo, uniqueSuffix } from "../../helpers/traffic"; + +interface ScopedKey { + alias: string; + token: string; +} + +const test = base.extend<{ scopedKey: ScopedKey }>({ + scopedKey: async ({ page }, use) => { + const alias = `e2e-budget-window-${uniqueSuffix()}`; + const created = await createVirtualKey(page.request, { + key_alias: alias, + team_id: E2E_TEAM_CRUD_ID, + models: [CHAT_MODEL_A], + }); + await use({ alias, token: created.token }); + await deleteVirtualKey(page.request, created.token); + }, +}); + +test.describe("Proxy Admin - Key budget window", () => { + test.use({ storageState: ADMIN_STORAGE_PATH }); + + test("a monthly spend cap survives a reload, and clearing the window keeps the cap", async ({ page, scopedKey }) => { + const { alias, token } = scopedKey; + + const before = await readKeyInfo(page.request, token); + expect(before.max_budget, "a freshly generated key starts with no budget").toBeNull(); + + await navigateToPage(page, Page.ApiKeys); + await dismissFeedbackPopup(page); + await openKeyDetail(page, alias); + + await page.getByRole("tab", { name: "Settings" }).click(); + await page.getByRole("button", { name: "Edit Settings" }).click(); + + await page.getByRole("spinbutton", { name: "Max Budget (USD)" }).fill("12.5"); + await page.getByLabel("Reset Budget", { exact: true }).click(); + await page.getByRole("option", { name: "monthly", exact: true }).click(); + await page.getByRole("button", { name: "Save Changes" }).click(); + + await expect + .poll(async () => (await readKeyInfo(page.request, token)).max_budget, { + message: "the $12.50 cap never reached /key/info", + timeout: 20_000, + }) + .toBe(12.5); + await expect + .poll(async () => (await readKeyInfo(page.request, token)).budget_duration, { + message: "the monthly reset window never reached /key/info", + timeout: 20_000, + }) + .toBe("30d"); + + const capped = await readKeyInfo(page.request, token); + const resetAt = new Date(capped.budget_reset_at ?? ""); + expect(Number.isNaN(resetAt.getTime()), "a monthly window left the key with no budget_reset_at").toBe(false); + expect(resetAt.getTime(), "budget_reset_at was set in the past").toBeGreaterThan(Date.now()); + expect(resetAt.getUTCDate(), "a monthly window resets on the 1st, a daily one would not").toBe(1); + + await page.reload(); + await expect( + page.getByRole("paragraph").filter({ hasText: "of $12.50" }), + "the reloaded key detail does not render the $12.50 cap", + ).toBeVisible({ timeout: 15_000 }); + + await page.getByRole("tab", { name: "Settings" }).click(); + await expect( + page.getByTestId("budget-reset-value"), + "the reloaded key detail does not name the 30d reset window", + ).toHaveText(/Every 30d/, { timeout: 15_000 }); + + await page.getByRole("button", { name: "Edit Settings" }).click(); + await page.getByLabel("Reset Budget", { exact: true }).click(); + await page.getByRole("option", { name: "Never resets", exact: true }).click(); + + const cleared = await captureRequestBody(page, { method: "POST", urlIncludes: "/key/update" }, async () => { + await page.getByRole("button", { name: "Save Changes" }).click(); + }); + expect(cleared).toHaveProperty("budget_duration"); + expect(cleared.budget_duration, "clearing the window must send budget_duration: null explicitly").toBeNull(); + + await expect + .poll(async () => (await readKeyInfo(page.request, token)).budget_duration, { + message: "the reset window was never cleared on /key/info", + timeout: 20_000, + }) + .toBeNull(); + + const after = await readKeyInfo(page.request, token); + expect(after.budget_reset_at, "clearing the reset window left a stale next-reset timestamp").toBeNull(); + expect(after.max_budget, "clearing the reset window also wiped the spend cap").toBe(12.5); + expect(after.models, "editing the budget left the key's models untouched").toEqual(before.models); + expect(after.team_id, "editing the budget left the key's team untouched").toEqual(before.team_id); + }); +});