Merge 22d3a9bbaf into 3b2ed83152

litellm/proxy/management_endpoints/key_management_endpoints.py

@@ -161,6 +161,7 @@ from litellm.types.proxy.management_endpoints.key_management_endpoints import (
     CustomKeyPolicyRequest,
     FailedKeyUpdate,
     KeySearchWhere,
+    KeyTagWhere,
     SuccessfulKeyUpdate,
 )
 from litellm.types.router import Deployment
@@ -6545,6 +6546,7 @@ async def list_keys(
     project_id: str | None = Query(None, description="Filter keys by project ID"),
     access_group_id: str | None = Query(None, description="Filter keys by access group ID"),
     agent_id: str | None = Query(None, description="Filter keys by agent ID"),
+    tag: str | None = Query(None, description="Filter keys by tag name. Exact match on an entry of metadata.tags"),
     substring_matching: bool = Query(
         False,
         description="If true, match key_alias (any caller) and user_id (proxy admins only) as case-insensitive substrings instead of exact values. Defaults to false: /key/list matched these exactly before substring search was added, and an exact user_id filter must never return another user's keys.",
@@ -6675,6 +6677,7 @@ async def list_keys(
             project_id=project_id,
             access_group_id=access_group_id,
             agent_id=agent_id,
+            tag=tag,
             use_substring_matching=use_substring_matching,
             use_key_alias_substring_matching=use_key_alias_substring_matching,
             expires_filter=expires if isinstance(expires, str) else None,
@@ -6920,6 +6923,11 @@ def _build_key_search_where(search: str) -> KeySearchWhere:
     return search_where
 
 
+def _build_key_tag_where(tag: str) -> KeyTagWhere:
+    tag_where: Final[KeyTagWhere] = {"metadata": {"path": ("tags",), "array_contains": json.dumps((tag,))}}
+    return tag_where
+
+
 def _build_key_filter_conditions(
     user_id: str | None,
     team_id: str | None,
@@ -6933,6 +6941,7 @@ def _build_key_filter_conditions(
     project_id: str | None = None,
     access_group_id: str | None = None,
     agent_id: str | None = None,
+    tag: str | None = None,
     use_substring_matching: bool = False,
     use_key_alias_substring_matching: bool = False,
     expires_filter: str | None = None,
@@ -7025,7 +7034,7 @@ def _build_key_filter_conditions(
     elif len(or_conditions) == 1:
         where.update(or_conditions[0])
 
-    # Apply team_id, project_id and access_group_id as global AND filters so they
+    # Apply team_id, project_id, access_group_id and tag as global AND filters so they
     # narrow results across all visibility conditions (own keys, team keys, etc.)
     now: Final = datetime.now(timezone.utc)
     status_where: Final = _build_status_where_clause(status_filter, now)
@@ -7045,6 +7054,7 @@ def _build_key_filter_conditions(
         *(({"project_id": project_id},) if project_id else ()),
         *(({"access_group_ids": {"hasSome": [access_group_id]}},) if access_group_id else ()),
         *(({"agent_id": agent_id},) if agent_id and isinstance(agent_id, str) else ()),
+        *((_build_key_tag_where(tag),) if tag and isinstance(tag, str) else ()),
         *(
             (_build_expires_where_clause(expires_filter, now),)
             if expires_filter is not None and expires_filter in VALID_EXPIRES_FILTER_VALUES
@@ -7079,6 +7089,7 @@ async def _list_key_helper(
     project_id: str | None = None,
     access_group_id: str | None = None,
     agent_id: str | None = None,
+    tag: str | None = None,
     use_substring_matching: bool = False,
     use_key_alias_substring_matching: bool = False,
     expires_filter: str | None = None,
@@ -7119,6 +7130,7 @@ async def _list_key_helper(
         project_id=project_id,
         access_group_id=access_group_id,
         agent_id=agent_id,
+        tag=tag,
         use_substring_matching=use_substring_matching,
         use_key_alias_substring_matching=use_key_alias_substring_matching,
         expires_filter=expires_filter,

litellm/types/proxy/management_endpoints/key_management_endpoints.py

@@ -29,6 +29,17 @@ class KeySearchWhere(TypedDict):
     OR: ReadOnly[tuple[KeyTokenWhere, KeyAliasContainsWhere]]
 
 
+class JsonPathArrayContains(TypedDict):
+    path: ReadOnly[tuple[str, ...]]
+    array_contains: ReadOnly[str]
+
+
+class KeyTagWhere(TypedDict):
+    """Prisma filter behind `/key/list?tag=`: keys whose metadata.tags array holds the exact tag."""
+
+    metadata: ReadOnly[JsonPathArrayContains]
+
+
 class BulkUpdateKeyRequestItem(BaseModel):
     """One /key/bulk_update item; only the fields it carries are written."""
 

tests/unit/proxy/management_endpoints/test_key_management_endpoints.py

@@ -17182,6 +17182,12 @@ def _prisma_where_matches(row, where):
         elif field == "OR":
             if not any(_prisma_where_matches(row, child) for child in expected):
                 return False
+        elif isinstance(expected, dict) and "array_contains" in expected:
+            value = getattr(row, field)
+            for key in expected["path"]:
+                value = value.get(key) if isinstance(value, dict) else None
+            if not isinstance(value, list) or not all(item in value for item in json.loads(expected["array_contains"])):
+                return False
         elif isinstance(expected, dict):
             value = getattr(row, field)
             if "in" in expected and value not in expected["in"]:
@@ -17222,14 +17228,14 @@ _TEAM_A_KEYS = (
 )
 
 
-def _list_team_a_keys_as(user_role, members_with_roles, query):
+def _list_team_a_keys_as(user_role, members_with_roles, query, rows=_TEAM_A_KEYS):
     from fastapi import FastAPI
 
     from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
     from litellm.proxy.management_endpoints.key_management_endpoints import router
 
     mock_prisma_client = AsyncMock()
-    mock_prisma_client.db.litellm_verificationtoken = _InMemoryVerificationTokenTable(_TEAM_A_KEYS)
+    mock_prisma_client.db.litellm_verificationtoken = _InMemoryVerificationTokenTable(rows)
     mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock(
         return_value=LiteLLM_UserTable(user_id="alice", teams=["team-a"], organization_memberships=[])
     )
@@ -17276,6 +17282,53 @@ def test_list_keys_key_alias_stays_exact_without_substring_matching():
     ) == ["tok-alice-first"]
 
 
+def _tagged_team_key(token, key_alias, user_id, metadata):
+    return LiteLLM_VerificationToken(
+        token=token, key_alias=key_alias, user_id=user_id, team_id="team-a", metadata=metadata
+    )
+
+
+_TAGGED_TEAM_A_KEYS = (
+    _tagged_team_key("tok-alice-batch", "alice-batch", "alice", {"tags": ["prod-batch", "nightly"]}),
+    _tagged_team_key("tok-bob-batch", "bob-batch", "bob", {"tags": ["prod-batch"]}),
+    _tagged_team_key("tok-svc-batch", "svc-batch", None, {"tags": ["prod-batch"]}),
+    _tagged_team_key("tok-alice-similar", "alice-similar", "alice", {"tags": ["prod-batch-2"]}),
+    _tagged_team_key("tok-alice-scalar", "alice-scalar", "alice", {"tags": "prod-batch"}),
+    _tagged_team_key("tok-alice-untagged", "alice-untagged", "alice", {}),
+)
+
+
+def test_list_keys_tag_returns_only_keys_carrying_that_exact_tag():
+    keys = _list_team_a_keys_as(
+        LitellmUserRoles.INTERNAL_USER, _ALICE_TEAM_ADMIN, "tag=prod-batch", rows=_TAGGED_TEAM_A_KEYS
+    )
+    assert keys == ["tok-alice-batch", "tok-bob-batch", "tok-svc-batch"]
+
+
+def test_list_keys_tag_composes_with_other_filters():
+    keys = _list_team_a_keys_as(
+        LitellmUserRoles.INTERNAL_USER,
+        _ALICE_TEAM_ADMIN,
+        "tag=prod-batch&key_alias=alice&substring_matching=true",
+        rows=_TAGGED_TEAM_A_KEYS,
+    )
+    assert keys == ["tok-alice-batch"]
+
+
+def test_list_keys_tag_stays_within_caller_visibility():
+    keys = _list_team_a_keys_as(
+        LitellmUserRoles.INTERNAL_USER, _ALICE_TEAM_MEMBER, "tag=prod-batch", rows=_TAGGED_TEAM_A_KEYS
+    )
+    assert keys == ["tok-alice-batch", "tok-svc-batch"]
+
+
+def test_list_keys_tag_without_keys_returns_none():
+    keys = _list_team_a_keys_as(
+        LitellmUserRoles.INTERNAL_USER, _ALICE_TEAM_ADMIN, "tag=unused-tag", rows=_TAGGED_TEAM_A_KEYS
+    )
+    assert keys == []
+
+
 @pytest.mark.asyncio
 async def test_list_keys_search_is_honored_for_non_admin():
     """LIT-4741: unlike substring_matching, `search` is not admin-gated. A non-admin's

ui/litellm-dashboard/src/app/(dashboard)/hooks/keys/useKeys.test.ts

@@ -424,6 +424,22 @@ describe("useKeys", () => {
     expect(result.current.data?.keys[0].project_id).toBe("project-1");
   });
 
+  it("should pass the tag filter to the API", async () => {
+    mockFetch.mockResolvedValueOnce({
+      ok: true,
+      json: async () => mockKeysResponse,
+    });
+
+    const { result } = renderHook(() => useKeys(1, 10, { tag: "prod batch" }), { wrapper });
+
+    await waitFor(() => {
+      expect(result.current.isLoading).toBe(false);
+    });
+
+    const callUrl = new URL(mockFetch.mock.calls[0][0], "http://localhost");
+    expect(callUrl.searchParams.get("tag")).toBe("prod batch");
+  });
+
   it("should pass both projectID and teamID filters to the API", async () => {
     mockFetch.mockResolvedValueOnce({
       ok: true,

ui/litellm-dashboard/src/app/(dashboard)/hooks/keys/useKeys.ts

@@ -37,6 +37,7 @@ export interface KeyListCallOptions {
   teamID?: string | null;
   projectID?: string | null;
   agentID?: string | null;
+  tag?: string | null;
   selectedKeyAlias?: string | null;
   userID?: string | null;
   keyHash?: string | null;
@@ -59,6 +60,7 @@ const keyListCall = async (accessToken: string, page: number, pageSize: number,
         team_id: options.teamID,
         project_id: options.projectID,
         agent_id: options.agentID,
+        tag: options.tag,
         organization_id: options.organizationID,
         key_alias: options.selectedKeyAlias,
         key_hash: options.keyHash,

ui/litellm-dashboard/src/app/(dashboard)/tag-management/_components/TagKeysSection.test.tsx

@@ -0,0 +1,84 @@
+import { render, screen, within } from "@testing-library/react";
+import { beforeEach, describe, expect, it, vi } from "vitest";
+
+import type { KeyResponse } from "@/components/key_team_helpers/key_list";
+
+import TagKeysSection, { TAG_KEYS_PAGE_SIZE } from "./TagKeysSection";
+
+const mockUseKeys = vi.fn();
+vi.mock("@/app/(dashboard)/hooks/keys/useKeys", () => ({
+  useKeys: (...args: unknown[]) => mockUseKeys(...args),
+}));
+
+const makeKey = (overrides: Partial<KeyResponse>): KeyResponse =>
+  ({ token: "tok", key_alias: "", key_name: "sk-...0000", team_id: null, spend: 0, ...overrides }) as KeyResponse;
+
+const loaded = (keys: KeyResponse[], totalCount = keys.length) => ({
+  data: { keys, total_count: totalCount, current_page: 1, total_pages: 1 },
+  isLoading: false,
+  isError: false,
+});
+
+describe("TagKeysSection", () => {
+  beforeEach(() => {
+    mockUseKeys.mockReset();
+  });
+
+  it("should request the first page of keys filtered by the tag name", () => {
+    mockUseKeys.mockReturnValue(loaded([]));
+    render(<TagKeysSection tagName="prod-batch" />);
+    expect(mockUseKeys).toHaveBeenCalledWith(1, TAG_KEYS_PAGE_SIZE, { tag: "prod-batch" });
+  });
+
+  it("should list each key with its alias linking to the key's page, its team and its spend", () => {
+    const teamKey: Partial<KeyResponse> = { token: "tok-1", key_alias: "batch-key-1", team_id: "team-a", spend: 1.5 };
+    mockUseKeys.mockReturnValue(loaded([makeKey(teamKey), makeKey({ token: "tok-2", key_alias: "batch-key-2" })]));
+    render(<TagKeysSection tagName="prod-batch" />);
+
+    expect(screen.getByRole("link", { name: "batch-key-1" })).toHaveAttribute("href", "/ui/api-keys?key=tok-1");
+    expect(screen.getByRole("link", { name: "batch-key-2" })).toHaveAttribute("href", "/ui/api-keys?key=tok-2");
+    const firstRow = screen.getByRole("row", { name: /batch-key-1/ });
+    expect(within(firstRow).getByText("team-a")).toBeInTheDocument();
+    expect(within(firstRow).getByText("1.5000")).toBeInTheDocument();
+    expect(within(screen.getByRole("row", { name: /batch-key-2/ })).getByText("-")).toBeInTheDocument();
+  });
+
+  it("should fall back to the masked key name when a key has no alias", () => {
+    mockUseKeys.mockReturnValue(loaded([makeKey({ token: "tok-3", key_alias: "", key_name: "sk-...wxyz" })]));
+    render(<TagKeysSection tagName="prod-batch" />);
+    expect(screen.getByRole("link", { name: "sk-...wxyz" })).toHaveAttribute("href", "/ui/api-keys?key=tok-3");
+  });
+
+  it("should say no virtual keys use the tag when none carry it", () => {
+    mockUseKeys.mockReturnValue(loaded([]));
+    render(<TagKeysSection tagName="prod-batch" />);
+    expect(screen.getByText("No virtual keys use this tag")).toBeInTheDocument();
+    expect(screen.queryByRole("table")).not.toBeInTheDocument();
+  });
+
+  it("should show a loading message while the keys are loading", () => {
+    mockUseKeys.mockReturnValue({ data: undefined, isLoading: true, isError: false });
+    render(<TagKeysSection tagName="prod-batch" />);
+    expect(screen.getByText("Loading virtual keys...")).toBeInTheDocument();
+    expect(screen.queryByText("No virtual keys use this tag")).not.toBeInTheDocument();
+  });
+
+  it("should show an error message instead of the empty message when the keys fail to load", () => {
+    mockUseKeys.mockReturnValue({ data: undefined, isLoading: false, isError: true });
+    render(<TagKeysSection tagName="prod-batch" />);
+    expect(screen.getByText("Could not load the virtual keys for this tag")).toBeInTheDocument();
+    expect(screen.queryByText("No virtual keys use this tag")).not.toBeInTheDocument();
+  });
+
+  it("should say how many keys are shown when the tag has more keys than one page", () => {
+    mockUseKeys.mockReturnValue(loaded([makeKey({ token: "tok-1", key_alias: "batch-key-1" })], 250));
+    render(<TagKeysSection tagName="prod-batch" />);
+    expect(screen.getByText("Showing the 1 most recently created of 250 keys")).toBeInTheDocument();
+  });
+
+  it("should not show a count note when every key fits on the page", () => {
+    mockUseKeys.mockReturnValue(loaded([makeKey({ token: "tok-1", key_alias: "batch-key-1" })]));
+    render(<TagKeysSection tagName="prod-batch" />);
+    expect(screen.queryByText(/most recently created/)).not.toBeInTheDocument();
+  });
+});

ui/litellm-dashboard/src/app/(dashboard)/tag-management/_components/TagKeysSection.tsx

@@ -0,0 +1,69 @@
+"use client";
+
+import React from "react";
+
+import { useKeys } from "@/app/(dashboard)/hooks/keys/useKeys";
+import { Card, CardContent, CardTitle } from "@/components/ui/card";
+import { Table, TableBody, TableCell, TableHead, TableHeader, TableRow } from "@/components/ui/table";
+import { formatNumberWithCommas } from "@/utils/dataUtils";
+import { keyDetailHref } from "@/utils/entityLinks";
+
+export const TAG_KEYS_PAGE_SIZE = 100;
+
+interface TagKeysSectionProps {
+  tagName: string;
+}
+
+const TagKeysSection: React.FC<TagKeysSectionProps> = ({ tagName }) => {
+  const { data, isLoading, isError } = useKeys(1, TAG_KEYS_PAGE_SIZE, { tag: tagName });
+  const keys = data?.keys ?? [];
+  const totalCount = data?.total_count ?? 0;
+
+  const renderBody = () => {
+    if (isLoading) return <p className="mt-4 text-sm text-muted-foreground">Loading virtual keys...</p>;
+    if (isError) return <p className="mt-4 text-sm text-destructive">Could not load the virtual keys for this tag</p>;
+    if (keys.length === 0) return <p className="mt-4 text-sm text-muted-foreground">No virtual keys use this tag</p>;
+    return (
+      <>
+        {totalCount > keys.length && (
+          <p className="mt-4 text-sm text-muted-foreground">
+            Showing the {keys.length} most recently created of {totalCount} keys
+          </p>
+        )}
+        <Table className="mt-4">
+          <TableHeader>
+            <TableRow>
+              <TableHead>Key</TableHead>
+              <TableHead>Team</TableHead>
+              <TableHead className="text-right">Spend (USD)</TableHead>
+            </TableRow>
+          </TableHeader>
+          <TableBody>
+            {keys.map((key) => (
+              <TableRow key={key.token}>
+                <TableCell className="font-medium">
+                  <a href={keyDetailHref(key.token)} className="text-info hover:underline">
+                    {key.key_alias || key.key_name}
+                  </a>
+                </TableCell>
+                <TableCell>{key.team_id ?? "-"}</TableCell>
+                <TableCell className="text-right">{formatNumberWithCommas(key.spend, 4)}</TableCell>
+              </TableRow>
+            ))}
+          </TableBody>
+        </Table>
+      </>
+    );
+  };
+
+  return (
+    <Card>
+      <CardContent>
+        <CardTitle>Virtual Keys</CardTitle>
+        {renderBody()}
+      </CardContent>
+    </Card>
+  );
+};
+
+export default TagKeysSection;

ui/litellm-dashboard/src/app/(dashboard)/tag-management/_components/tag_info.integration.test.tsx

@@ -1,5 +1,6 @@
 import { fireEvent, render, screen } from "@testing-library/react";
 import userEvent from "@testing-library/user-event";
+import { QueryClient, QueryClientProvider } from "@tanstack/react-query";
 import { beforeEach, describe, expect, it, vi } from "vitest";
 
 import { tagInfoCall, tagUpdateCall } from "@/components/networking";
@@ -10,6 +11,14 @@ import TagInfoView from "./tag_info";
 vi.mock("@/components/networking", () => ({
   tagInfoCall: vi.fn(),
   tagUpdateCall: vi.fn(),
+  getProxyBaseUrl: () => "",
+  getGlobalLitellmHeaderName: () => "Authorization",
+  deriveErrorMessage: (errorData: unknown) => JSON.stringify(errorData),
+  handleError: vi.fn(),
+}));
+
+vi.mock("@/app/(dashboard)/hooks/useAuthorized", () => ({
+  default: () => ({ accessToken: "sk-test" }),
 }));
 
 vi.mock("@/components/organisms/create_key_button", () => ({
@@ -34,9 +43,20 @@ const tag: Tag = {
   litellm_budget_table: { max_budget: 10, budget_duration: "7d", tpm_limit: 1000, rpm_limit: 60 },
 };
 
+const keyListFetch = vi.fn();
+
+const renderTagInfo = (editTag: boolean) => {
+  const queryClient = new QueryClient({ defaultOptions: { queries: { retry: false } } });
+  render(
+    <QueryClientProvider client={queryClient}>
+      <TagInfoView tagId="prod-tag" onClose={vi.fn()} accessToken="sk-test" is_admin editTag={editTag} />
+    </QueryClientProvider>,
+  );
+};
+
 const renderEditor = async () => {
   const user = userEvent.setup();
-  render(<TagInfoView tagId="prod-tag" onClose={vi.fn()} accessToken="sk-test" is_admin editTag />);
+  renderTagInfo(true);
   const nameInput = await screen.findByLabelText("Tag Name");
   return { user, nameInput };
 };
@@ -46,6 +66,11 @@ describe("TagInfoView save payload", () => {
     vi.clearAllMocks();
     mockTagInfoCall.mockResolvedValue({ "prod-tag": tag });
     mockTagUpdateCall.mockResolvedValue(undefined);
+    keyListFetch.mockResolvedValue({
+      ok: true,
+      json: async () => ({ keys: [], total_count: 0, current_page: 1, total_pages: 0 }),
+    });
+    vi.stubGlobal("fetch", keyListFetch);
   });
 
   it("leaves the budget unchanged while the budget section is collapsed", async () => {
@@ -184,3 +209,46 @@ describe("TagInfoView save payload", () => {
     expect(mockTagUpdateCall).not.toHaveBeenCalled();
   });
 });
+
+describe("TagInfoView virtual keys", () => {
+  beforeEach(() => {
+    vi.clearAllMocks();
+    mockTagInfoCall.mockResolvedValue({ "prod-tag": tag });
+    vi.stubGlobal("fetch", keyListFetch);
+  });
+
+  it("should list the keys that /key/list returns for this tag, each linking to its key page", async () => {
+    keyListFetch.mockResolvedValue({
+      ok: true,
+      json: async () => ({
+        keys: [{ token: "tok-1", key_alias: "batch-key-1", key_name: "sk-...0001", team_id: "team-a", spend: 2 }],
+        total_count: 1,
+        current_page: 1,
+        total_pages: 1,
+      }),
+    });
+    renderTagInfo(false);
+
+    expect(await screen.findByRole("link", { name: "batch-key-1" })).toHaveAttribute("href", "/ui/api-keys?key=tok-1");
+    const requestUrl = new URL(keyListFetch.mock.calls[0][0], "http://localhost");
+    expect(requestUrl.pathname).toBe("/key/list");
+    expect(requestUrl.searchParams.get("tag")).toBe("prod-tag");
+  });
+
+  it("should say no virtual keys use the tag when /key/list returns none", async () => {
+    keyListFetch.mockResolvedValue({
+      ok: true,
+      json: async () => ({ keys: [], total_count: 0, current_page: 1, total_pages: 0 }),
+    });
+    renderTagInfo(false);
+
+    expect(await screen.findByText("No virtual keys use this tag")).toBeInTheDocument();
+  });
+
+  it("should show an error in the section when /key/list fails", async () => {
+    keyListFetch.mockResolvedValue({ ok: false, json: async () => ({ error: "boom" }) });
+    renderTagInfo(false);
+
+    expect(await screen.findByText("Could not load the virtual keys for this tag")).toBeInTheDocument();
+  });
+});

ui/litellm-dashboard/src/app/(dashboard)/tag-management/_components/tag_info.tsx

@@ -22,6 +22,7 @@ import { Textarea } from "@/components/ui/textarea";
 import { useZodForm } from "@/lib/forms/useZodForm";
 import { copyToClipboard as utilCopyToClipboard } from "@/utils/dataUtils";
 import { CheckIcon, ChevronRight, CopyIcon } from "lucide-react";
+import TagKeysSection from "./TagKeysSection";
 
 const tagEditShape = {
   name: z.string().min(1, "Please input a tag name"),
@@ -349,6 +350,8 @@ const TagInfoView: React.FC<TagInfoViewProps> = ({ tagId, onClose, accessToken,
               </CardContent>
             </Card>
           )}
+
+          <TagKeysSection tagName={tagDetails.name} />
         </div>
       )}
     </div>

ui/litellm-dashboard/src/lib/http/schema.d.ts

@@ -61442,6 +61442,8 @@ export interface operations {
                 access_group_id?: string | null;
                 /** @description Filter keys by agent ID */
                 agent_id?: string | null;
+                /** @description Filter keys by tag name. Exact match on an entry of metadata.tags */
+                tag?: string | null;
                 /** @description If true, match key_alias (any caller) and user_id (proxy admins only) as case-insensitive substrings instead of exact values. Defaults to false: /key/list matched these exactly before substring search was added, and an exact user_id filter must never return another user's keys. */
                 substring_matching?: boolean;
                 /** @description Filter keys by expiration. 'expired' returns keys whose expires is in the past; 'active' returns keys that never expire or expire in the future. Omit to return keys regardless of expiration. */