From 1aa531767ca67b2b791d9640bd23e11da8cd783a Mon Sep 17 00:00:00 2001 From: moe-berri Date: Fri, 23 Oct 2026 15:45:17 -0700 Subject: [PATCH] feat(ui): sort virtual keys by expiry, activity, spend and limits (#48074) --- .../key_management_endpoints.py | 29 ++++--- tests/integration/management/test_key_list.py | 80 ++++++++++++++++++- .../VirtualKeysPage/VirtualKeysTable.test.tsx | 51 ++++++++++++ .../VirtualKeysPage/keyTableColumns.tsx | 35 +++++--- .../key_team_helpers/key_sorting.ts | 6 ++ .../team/TeamVirtualKeysTable.test.tsx | 49 +++++++++++- .../components/team/TeamVirtualKeysTable.tsx | 33 ++++++-- 7 files changed, 252 insertions(+), 31 deletions(-) create mode 100644 ui/litellm-dashboard/src/components/key_team_helpers/key_sorting.ts diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index 824b310a1dd..eebc537509b 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -7397,20 +7397,26 @@ async def key_aliases( ) -def _validate_sort_params(sort_by: str | None, sort_order: str) -> dict[str, str] | None: - order_by: Final[dict[str, str]] = {} - +def _validate_sort_params( + sort_by: str | None, sort_order: str +) -> Mapping[str, object] | Sequence[Mapping[str, object]] | None: if sort_by is None: return None # Validate sort_by is a valid column - valid_columns: Final = [ + valid_columns: Final = ( "spend", "max_budget", "created_at", "updated_at", "token", "key_alias", - ] + "expires", + "last_active", + "budget_reset_at", + "total_spend", + "tpm_limit", + "rpm_limit", + ) if sort_by not in valid_columns: raise HTTPException( status_code=400, @@ -7424,9 +7430,14 @@ def _validate_sort_params(sort_by: str | None, sort_order: str) -> dict[str, str detail={"error": "Invalid sort order. Must be 'asc' or 'desc'"}, ) - order_by[sort_by] = sort_order.lower() - - return order_by + direction: Final = sort_order.lower() + nullable_sort_columns: Final = ("expires", "last_active", "budget_reset_at", "tpm_limit", "rpm_limit") + if sort_by not in (*nullable_sort_columns, "total_spend"): + return {sort_by: direction} + return [ + {sort_by: {"sort": direction, "nulls": "last"}} if sort_by in nullable_sort_columns else {sort_by: direction}, + {"token": "asc"}, + ] def _build_expires_where_clause(expires_filter: str, now: datetime) -> dict[str, object]: @@ -7668,7 +7679,7 @@ async def list_key_helper( verbose_proxy_logger.debug("Pagination: skip=%s, take=%s", skip, size) - order_by: Final[dict[str, str] | None] = ( + order_by: Final = ( _validate_sort_params(sort_by, sort_order) if sort_by is not None and isinstance(sort_by, str) else None ) diff --git a/tests/integration/management/test_key_list.py b/tests/integration/management/test_key_list.py index b70a3782eda..4634314817b 100644 --- a/tests/integration/management/test_key_list.py +++ b/tests/integration/management/test_key_list.py @@ -1,11 +1,16 @@ -from typing import FrozenSet +from datetime import datetime, timedelta, timezone +from itertools import chain +from typing import Final, FrozenSet import pytest from litellm.proxy.utils import hash_token -from tests.integration._support.actors import TEAM_ALPHA, Actor +from tests.integration._support.actors import TEAM_ALPHA, Actor, ManagementDatabase, ScratchResources from tests.integration._support.actors import create_scratch_key +from tests.integration._support.client import Gateway, Scenario +from litellm.repositories.table_repositories import DeletedVerificationTokenRepository +from litellm.repositories.verification_token_repository import VerificationTokenRepository pytestmark = pytest.mark.asyncio(loop_scope="session") @@ -202,3 +207,74 @@ async def test_key_list_non_admin_cannot_filter_other_team(gateway_client, actor }, ) assert resp.status_code == 403, resp.text + + +@pytest.mark.parametrize( + "field", ("expires", "last_active", "budget_reset_at", "total_spend", "tpm_limit", "rpm_limit") +) +@pytest.mark.parametrize("direction", ("asc", "desc")) +@pytest.mark.parametrize("deleted", (False, True), ids=("live", "deleted")) +async def test_key_list_sorts_across_pages_with_nulls_last( + field: str, + direction: str, + deleted: bool, + gateway: Gateway, + scenario: Scenario, + db: ManagementDatabase, + resources: ScratchResources, +) -> None: + team_admin: Final = scenario.user(user_role="internal_user") + owner: Final = scenario.user(user_role="internal_user") + team_id: Final = scenario.team( + members_with_roles=[{"user_id": team_admin, "role": "admin"}, {"user_id": owner, "role": "user"}] + ) + team_admin_key: Final = scenario.key(user_id=team_admin) + table: Final = ( + DeletedVerificationTokenRepository(db.prisma).table if deleted else VerificationTokenRepository(db.prisma).table + ) + values: Final = (20, 3, 3, 100, 0 if field == "total_spend" else None) + tokens: Final = tuple(hash_token(f"{resources.prefix}-{index}") for index in range(len(values))) + rows: Final = [ + { + "token": token, + "key_alias": resources.prefix, + "team_id": team_id, + "user_id": owner, + field: ( + datetime(2030, 1, 1, tzinfo=timezone.utc) + timedelta(days=value) + if field in ("expires", "last_active", "budget_reset_at") and value is not None + else value + ), + } + for token, value in zip(tokens, values, strict=True) + ] + await table.create_many(rows) + try: + valued: Final = tuple((value, token) for value, token in zip(values, tokens, strict=True) if value is not None) + expected: Final = [ + token + for _, token in sorted(valued, key=lambda item: (item[0] if direction == "asc" else -item[0], item[1])) + ] + [token for value, token in zip(values, tokens, strict=True) if value is None] + for caller in (gateway.key, team_admin_key): + pages: Final = [ + gateway.client.get( + "/key/list", + headers={"Authorization": f"Bearer {caller}"}, + params={ + "key_alias": resources.prefix, + "team_id": team_id, + "include_team_keys": "true", + "sort_by": field, + "sort_order": direction, + "size": "2", + "page": str(page), + **({"status": "deleted"} if deleted else {}), + }, + ) + for page in (1, 2, 3) + ] + assert all(page.status_code == 200 for page in pages), [page.text for page in pages] + assert all(page.json()["total_count"] == len(tokens) for page in pages) + assert list(chain.from_iterable(page.json()["keys"] for page in pages)) == expected + finally: + await table.delete_many(where={"token": {"in": tokens}}) diff --git a/ui/litellm-dashboard/src/components/VirtualKeysPage/VirtualKeysTable.test.tsx b/ui/litellm-dashboard/src/components/VirtualKeysPage/VirtualKeysTable.test.tsx index ccf7f453132..47fd554d196 100644 --- a/ui/litellm-dashboard/src/components/VirtualKeysPage/VirtualKeysTable.test.tsx +++ b/ui/litellm-dashboard/src/components/VirtualKeysPage/VirtualKeysTable.test.tsx @@ -427,6 +427,57 @@ it("sorts by spend ascending when 'Spend ascending' is chosen from the Spend / B }); }); +it.each([ + ["expires", "Expires", "asc"], + ["last_active", "Last Active", "asc"], + ["budget_reset_at", "Budget Reset", "asc"], + ["total_spend", "Lifetime Spend", "desc"], +])("sorts %s in both directions and resets pagination", async (field, label, direction) => { + const user = userEvent.setup(); + renderWithProviders(, { searchParams: { page: "2" } }); + if (field === "expires") { + await user.click(screen.getByRole("button", { name: "Columns" })); + await user.click(await screen.findByText("Expires")); + await user.keyboard("{Escape}"); + } + await user.click(screen.getByRole("button", { name: label, exact: true })); + await waitFor(() => { + expect(mockUseKeys).toHaveBeenLastCalledWith( + 1, + 50, + expect.objectContaining({ sortBy: field, sortOrder: direction }), + ); + }); + await user.click(screen.getByRole("button", { name: label, exact: true })); + await waitFor(() => { + expect(mockUseKeys).toHaveBeenLastCalledWith( + 1, + 50, + expect.objectContaining({ sortBy: field, sortOrder: direction === "asc" ? "desc" : "asc" }), + ); + }); +}); + +it.each([ + ["TPM", "tpm_limit"], + ["RPM", "rpm_limit"], +])("sorts by %s through the rate-limit menu", async (label, field) => { + const user = userEvent.setup(); + renderWithProviders(); + await user.click(screen.getByRole("button", { name: "Columns" })); + await user.click(await screen.findByText("Rate Limits")); + await user.keyboard("{Escape}"); + await chooseSelectOption( + user, + screen.getByRole("button", { name: "Sort options for TPM or RPM" }), + `${label} ascending`, + "menuitem", + ); + await waitFor(() => { + expect(mockUseKeys).toHaveBeenLastCalledWith(1, 50, expect.objectContaining({ sortBy: field, sortOrder: "asc" })); + }); +}); + it("clicking the key cell deep-links via ?key=", async () => { const onUrlUpdate = vi.fn(); renderWithProviders(, { onUrlUpdate }); diff --git a/ui/litellm-dashboard/src/components/VirtualKeysPage/keyTableColumns.tsx b/ui/litellm-dashboard/src/components/VirtualKeysPage/keyTableColumns.tsx index 9751b9ebf2c..547298fa8cc 100644 --- a/ui/litellm-dashboard/src/components/VirtualKeysPage/keyTableColumns.tsx +++ b/ui/litellm-dashboard/src/components/VirtualKeysPage/keyTableColumns.tsx @@ -1,5 +1,6 @@ "use client"; +import type { ReactNode } from "react"; import { Info } from "lucide-react"; import { ColumnDef } from "@tanstack/react-table"; @@ -28,6 +29,7 @@ import { orgDetailHref, teamDetailHref } from "@/utils/entityLinks"; import { KeyResponse, Team } from "../key_team_helpers/key_list"; import { Organization } from "../networking"; +import { KEY_RATE_LIMIT_SORT_FIELDS } from "../key_team_helpers/key_sorting"; interface KeyStatus { tone: StatusTone; @@ -46,6 +48,11 @@ export const KEY_TABLE_SORT_FIELDS: readonly string[] = [ "created_at", "updated_at", ...SPEND_BUDGET_SORT_FIELDS.map((field) => field.id), + "expires", + "last_active", + "budget_reset_at", + "total_spend", + ...KEY_RATE_LIMIT_SORT_FIELDS.map((field) => field.id), ]; const getKeyStatus = (key: KeyResponse): KeyStatus => { @@ -77,7 +84,7 @@ const getKeyStatus = (key: KeyResponse): KeyStatus => { }; }; -const InfoHeader = ({ label, tooltip }: { label: string; tooltip: string }) => ( +const InfoHeader = ({ label, tooltip }: { label: ReactNode; tooltip: string }) => ( {label} @@ -252,23 +259,25 @@ export const getKeyTableColumns = ({ id: "last_active", accessorKey: "last_active", meta: { title: "Last Active" }, - header: () => ( + header: ({ column }) => ( } tooltip="This is a new field and is not backfilled. Only new key usage will update this value." /> ), size: 130, - enableSorting: false, + enableSorting: true, + sortDescFirst: false, cell: (info) => , }, { id: "expires", accessorKey: "expires", meta: { title: "Expires" }, - header: "Expires", + header: ({ column }) => , size: 120, - enableSorting: false, + enableSorting: true, + sortDescFirst: false, cell: (info) => , }, { @@ -313,23 +322,25 @@ export const getKeyTableColumns = ({ id: "total_spend", accessorKey: "total_spend", meta: { title: "Lifetime Spend" }, - header: () => ( + header: ({ column }) => ( } tooltip="Cumulative spend across every budget period. Budget resets do not touch this value. Lifetime tracking started with LiteLLM v1.103.0 on September 19, 2026, so keys created earlier only count spend since that upgrade." /> ), size: 130, - enableSorting: false, + enableSorting: true, + sortDescFirst: true, cell: (info) => , }, { id: "budget_reset_at", accessorKey: "budget_reset_at", meta: { title: "Budget Reset" }, - header: "Budget Reset", + header: ({ column }) => , size: 130, - enableSorting: false, + enableSorting: true, + sortDescFirst: false, cell: (info) => , }, { @@ -350,7 +361,7 @@ export const getKeyTableColumns = ({ { id: "rate_limits", meta: { title: "Rate Limits" }, - header: "Rate Limits", + header: ({ table }) => , size: 140, enableSorting: false, cell: ({ row }) => { diff --git a/ui/litellm-dashboard/src/components/key_team_helpers/key_sorting.ts b/ui/litellm-dashboard/src/components/key_team_helpers/key_sorting.ts new file mode 100644 index 00000000000..b8a780efb21 --- /dev/null +++ b/ui/litellm-dashboard/src/components/key_team_helpers/key_sorting.ts @@ -0,0 +1,6 @@ +import type { DataTableSortField } from "@/components/shared/DataTable"; + +export const KEY_RATE_LIMIT_SORT_FIELDS: DataTableSortField[] = [ + { id: "tpm_limit", label: "TPM" }, + { id: "rpm_limit", label: "RPM" }, +]; diff --git a/ui/litellm-dashboard/src/components/team/TeamVirtualKeysTable.test.tsx b/ui/litellm-dashboard/src/components/team/TeamVirtualKeysTable.test.tsx index 5c3af3f6c3f..bb6204b5219 100644 --- a/ui/litellm-dashboard/src/components/team/TeamVirtualKeysTable.test.tsx +++ b/ui/litellm-dashboard/src/components/team/TeamVirtualKeysTable.test.tsx @@ -1,7 +1,7 @@ import userEvent from "@testing-library/user-event"; import { beforeEach, describe, expect, it, vi, MockedFunction } from "vitest"; import { QueryClient, QueryObserver } from "@tanstack/react-query"; -import { fireEvent, renderWithProviders, screen, waitFor, within } from "../../../tests/test-utils"; +import { chooseSelectOption, fireEvent, renderWithProviders, screen, waitFor, within } from "../../../tests/test-utils"; import { TeamVirtualKeysTable } from "./TeamVirtualKeysTable"; import { KeysResponse, useKeys } from "@/app/(dashboard)/hooks/keys/useKeys"; import { useTeamMemberBudgets } from "@/app/(dashboard)/hooks/teams/useTeamMemberBudgets"; @@ -283,6 +283,53 @@ describe("TeamVirtualKeysTable", () => { ); }); + it.each([ + ["expires", "Expires", "asc"], + ["last_active", "Last Active", "asc"], + ["budget_reset_at", "Budget Reset", "asc"], + ["total_spend", "Lifetime Spend", "desc"], + ])("sorts %s in both directions while preserving the team scope", async (field, label, direction) => { + const user = userEvent.setup(); + renderWithProviders(); + await user.click(screen.getByRole("button", { name: label, exact: true })); + await waitFor(() => { + expect(mockUseKeys).toHaveBeenLastCalledWith( + 1, + 50, + expect.objectContaining({ teamID: "team-1", sortBy: field, sortOrder: direction }), + ); + }); + await user.click(screen.getByRole("button", { name: label, exact: true })); + await waitFor(() => { + expect(mockUseKeys).toHaveBeenLastCalledWith( + 1, + 50, + expect.objectContaining({ teamID: "team-1", sortBy: field, sortOrder: direction === "asc" ? "desc" : "asc" }), + ); + }); + }); + + it.each([ + ["TPM", "tpm_limit"], + ["RPM", "rpm_limit"], + ])("sorts by %s through the rate-limit menu", async (label, field) => { + const user = userEvent.setup(); + renderWithProviders(); + await chooseSelectOption( + user, + screen.getByRole("button", { name: "Sort options for TPM or RPM" }), + `${label} ascending`, + "menuitem", + ); + await waitFor(() => { + expect(mockUseKeys).toHaveBeenLastCalledWith( + 1, + 50, + expect.objectContaining({ teamID: "team-1", sortBy: field, sortOrder: "asc" }), + ); + }); + }); + it("resets to the first page when the sort changes", async () => { const user = userEvent.setup(); mockUseKeys.mockImplementation( diff --git a/ui/litellm-dashboard/src/components/team/TeamVirtualKeysTable.tsx b/ui/litellm-dashboard/src/components/team/TeamVirtualKeysTable.tsx index 8ea559990c0..29a1cc232b6 100644 --- a/ui/litellm-dashboard/src/components/team/TeamVirtualKeysTable.tsx +++ b/ui/litellm-dashboard/src/components/team/TeamVirtualKeysTable.tsx @@ -16,6 +16,7 @@ import { DataTableFilterDrawer, DataTableFilterField, DataTableSortHeader, + DataTableMultiSortHeader, DataTableToolbar, } from "@/components/shared/DataTable"; import { Badge } from "@/components/ui/badge"; @@ -31,6 +32,7 @@ import DefaultProxyAdminTag from "../common_components/DefaultProxyAdminTag"; import React, { useCallback, useEffect, useMemo, useState } from "react"; import { getModelDisplayName } from "../key_team_helpers/fetch_available_models_team_key"; import { deriveKeyModelScope } from "../key_scope"; +import { KEY_RATE_LIMIT_SORT_FIELDS } from "../key_team_helpers/key_sorting"; import { KeyResponse, Team } from "../key_team_helpers/key_list"; import { Organization } from "../networking"; import KeyInfoView from "../templates/key_info_view"; @@ -278,17 +280,21 @@ export function TeamVirtualKeysTable({ teamId, teamAlias, organization }: TeamVi { id: "last_active", accessorKey: "last_active", - header: "Last Active", + meta: { title: "Last Active" }, + header: ({ column }) => , size: 130, - enableSorting: false, + enableSorting: true, + sortDescFirst: false, cell: (info) => , }, { id: "expires", accessorKey: "expires", - header: "Expires", + meta: { title: "Expires" }, + header: ({ column }) => , size: 120, - enableSorting: false, + enableSorting: true, + sortDescFirst: false, cell: (info) => , }, { @@ -300,6 +306,16 @@ export function TeamVirtualKeysTable({ teamId, teamAlias, organization }: TeamVi enableSorting: true, cell: (info) => , }, + { + id: "total_spend", + accessorKey: "total_spend", + meta: { title: "Lifetime Spend", numeric: true }, + header: ({ column }) => , + size: 130, + enableSorting: true, + sortDescFirst: true, + cell: (info) => , + }, { id: "max_budget", accessorKey: "max_budget", @@ -327,9 +343,11 @@ export function TeamVirtualKeysTable({ teamId, teamAlias, organization }: TeamVi { id: "budget_reset_at", accessorKey: "budget_reset_at", - header: "Budget Reset", + meta: { title: "Budget Reset" }, + header: ({ column }) => , size: 130, - enableSorting: false, + enableSorting: true, + sortDescFirst: false, cell: (info) => , }, { @@ -428,7 +446,8 @@ export function TeamVirtualKeysTable({ teamId, teamAlias, organization }: TeamVi }, { id: "rate_limits", - header: "Rate Limits", + meta: { title: "Rate Limits" }, + header: ({ table }) => , size: 140, enableSorting: false, cell: ({ row }) => {