Merge pull request #19534 from BerriAI/litellm_ui_sorting_keys_fix_2

[Fix] UI - Virtual Keys Table: Sorting Shows Incorrect Entries
2026-08-05 14:23:44 +00:00 · 2026-01-21 15:09:43 -08:00
parent 3a49e4362f a1636640ce
commit 665d3f2f25
4 changed files with 350 additions and 21 deletions
@@ -2,7 +2,7 @@ import { describe, it, expect, vi, beforeEach } from "vitest";
 import { renderHook, waitFor } from "@testing-library/react";
 import { QueryClient, QueryClientProvider } from "@tanstack/react-query";
 import React, { ReactNode } from "react";
-import { useKeys } from "./useKeys";
+import { useKeys, useDeletedKeys } from "./useKeys";
 import type { KeyResponse } from "@/components/key_team_helpers/key_list";

 // Mock the networking utilities
@@ -397,3 +397,293 @@ describe("useKeys", () => {
    );
  });
 });
+
+describe("useDeletedKeys", () => {
+  let queryClient: QueryClient;
+
+  beforeEach(() => {
+    queryClient = new QueryClient({
+      defaultOptions: {
+        queries: {
+          retry: false,
+        },
+      },
+    });
+
+    // Reset all mocks
+    vi.clearAllMocks();
+
+    // Set default mock for useAuthorized (enabled state)
+    mockUseAuthorized.mockReturnValue({
+      accessToken: "test-access-token",
+      userRole: "Admin",
+      userId: "test-user-id",
+      token: "test-token",
+      userEmail: "test@example.com",
+      premiumUser: false,
+      disabledPersonalKeyCreation: null,
+      showSSOBanner: false,
+    });
+
+    // Reset fetch mock
+    mockFetch.mockClear();
+  });
+
+  const wrapper = ({ children }: { children: ReactNode }) =>
+    React.createElement(QueryClientProvider, { client: queryClient }, children);
+
+  it("should return deleted keys data when query is successful", async () => {
+    // Mock successful API call
+    mockFetch.mockResolvedValueOnce({
+      ok: true,
+      json: async () => mockKeysResponse,
+    });
+
+    const { result } = renderHook(() => useDeletedKeys(1, 10), { wrapper });
+
+    // Initially loading
+    expect(result.current.isLoading).toBe(true);
+    expect(result.current.data).toBeUndefined();
+
+    // Wait for success
+    await waitFor(() => {
+      expect(result.current.isLoading).toBe(false);
+      expect(result.current.isSuccess).toBe(true);
+    });
+
+    expect(result.current.data).toEqual(mockKeysResponse);
+    expect(result.current.error).toBeNull();
+    expect(mockFetch).toHaveBeenCalledTimes(1);
+    expect(mockFetch).toHaveBeenCalledWith(
+      "/key/list?page=1&size=10&status=deleted&return_full_object=true&include_team_keys=true&include_created_by_keys=true",
+      {
+        method: "GET",
+        headers: {
+          Authorization: "Bearer test-access-token",
+          "Content-Type": "application/json",
+        },
+      },
+    );
+  });
+
+  it("should pass status=deleted parameter to the API", async () => {
+    // Mock successful API call
+    mockFetch.mockResolvedValueOnce({
+      ok: true,
+      json: async () => mockKeysResponse,
+    });
+
+    const { result } = renderHook(() => useDeletedKeys(1, 10), { wrapper });
+
+    // Wait for success
+    await waitFor(() => {
+      expect(result.current.isLoading).toBe(false);
+    });
+
+    // Verify that status=deleted is included in the URL
+    const callUrl = mockFetch.mock.calls[0][0];
+    expect(callUrl).toContain("status=deleted");
+    expect(result.current.data).toEqual(mockKeysResponse);
+  });
+
+  it("should handle error when deleted keys API call fails", async () => {
+    const errorMessage = "Failed to fetch deleted keys";
+    const errorResponse = { error: errorMessage };
+
+    // Mock failed API call
+    mockFetch.mockResolvedValueOnce({
+      ok: false,
+      json: async () => errorResponse,
+    });
+
+    const { result } = renderHook(() => useDeletedKeys(1, 10), { wrapper });
+
+    // Initially loading
+    expect(result.current.isLoading).toBe(true);
+
+    // Wait for error
+    await waitFor(() => {
+      expect(result.current.isLoading).toBe(false);
+      expect(result.current.isError).toBe(true);
+    });
+
+    expect(result.current.error).toBeDefined();
+    expect(result.current.error?.message).toBe(errorMessage);
+    expect(result.current.data).toBeUndefined();
+    expect(mockFetch).toHaveBeenCalledTimes(1);
+    expect(mockFetch).toHaveBeenCalledWith(
+      "/key/list?page=1&size=10&status=deleted&return_full_object=true&include_team_keys=true&include_created_by_keys=true",
+      {
+        method: "GET",
+        headers: {
+          Authorization: "Bearer test-access-token",
+          "Content-Type": "application/json",
+        },
+      },
+    );
+  });
+
+  it("should not execute query when accessToken is missing", async () => {
+    // Mock missing accessToken
+    mockUseAuthorized.mockReturnValue({
+      accessToken: null,
+      userRole: "Admin",
+      userId: "test-user-id",
+      token: null,
+      userEmail: "test@example.com",
+      premiumUser: false,
+      disabledPersonalKeyCreation: null,
+      showSSOBanner: false,
+    });
+
+    const { result } = renderHook(() => useDeletedKeys(1, 10), { wrapper });
+
+    // Query should not execute
+    expect(result.current.isLoading).toBe(false);
+    expect(result.current.data).toBeUndefined();
+    expect(result.current.isFetched).toBe(false);
+
+    // API should not be called
+    expect(mockFetch).not.toHaveBeenCalled();
+  });
+
+  it("should pass correct page and pageSize parameters to the API", async () => {
+    // Mock successful API call
+    mockFetch.mockResolvedValueOnce({
+      ok: true,
+      json: async () => mockKeysResponse,
+    });
+
+    const page = 2;
+    const pageSize = 20;
+
+    const { result } = renderHook(() => useDeletedKeys(page, pageSize), { wrapper });
+
+    // Wait for success
+    await waitFor(() => {
+      expect(result.current.isLoading).toBe(false);
+    });
+
+    expect(mockFetch).toHaveBeenCalledWith(
+      `/key/list?page=${page}&size=${pageSize}&status=deleted&return_full_object=true&include_team_keys=true&include_created_by_keys=true`,
+      {
+        method: "GET",
+        headers: {
+          Authorization: "Bearer test-access-token",
+          "Content-Type": "application/json",
+        },
+      },
+    );
+  });
+
+  it("should return empty deleted keys array when API returns empty data", async () => {
+    // Mock API returning empty keys array
+    const emptyResponse = {
+      keys: [],
+      total_count: 0,
+      current_page: 1,
+      total_pages: 0,
+    };
+    mockFetch.mockResolvedValueOnce({
+      ok: true,
+      json: async () => emptyResponse,
+    });
+
+    const { result } = renderHook(() => useDeletedKeys(1, 10), { wrapper });
+
+    // Wait for success
+    await waitFor(() => {
+      expect(result.current.isLoading).toBe(false);
+      expect(result.current.isSuccess).toBe(true);
+    });
+
+    expect(result.current.data).toEqual(emptyResponse);
+    expect(mockFetch).toHaveBeenCalledWith(
+      "/key/list?page=1&size=10&status=deleted&return_full_object=true&include_team_keys=true&include_created_by_keys=true",
+      {
+        method: "GET",
+        headers: {
+          Authorization: "Bearer test-access-token",
+          "Content-Type": "application/json",
+        },
+      },
+    );
+  });
+
+  it("should handle network timeout error", async () => {
+    const timeoutError = new Error("Network timeout");
+
+    // Mock network timeout
+    mockFetch.mockRejectedValueOnce(timeoutError);
+
+    const { result } = renderHook(() => useDeletedKeys(1, 10), { wrapper });
+
+    // Wait for error
+    await waitFor(() => {
+      expect(result.current.isError).toBe(true);
+    });
+
+    expect(result.current.error).toEqual(timeoutError);
+    expect(result.current.data).toBeUndefined();
+  });
+
+  it("should handle pagination correctly", async () => {
+    const paginatedResponse = {
+      keys: [mockKeys[0]], // Only first key
+      total_count: 15,
+      current_page: 2,
+      total_pages: 2,
+    };
+    mockFetch.mockResolvedValueOnce({
+      ok: true,
+      json: async () => paginatedResponse,
+    });
+
+    const { result } = renderHook(() => useDeletedKeys(2, 10), { wrapper });
+
+    // Wait for success
+    await waitFor(() => {
+      expect(result.current.isLoading).toBe(false);
+    });
+
+    expect(result.current.data).toEqual(paginatedResponse);
+    expect(mockFetch).toHaveBeenCalledWith(
+      "/key/list?page=2&size=10&status=deleted&return_full_object=true&include_team_keys=true&include_created_by_keys=true",
+      {
+        method: "GET",
+        headers: {
+          Authorization: "Bearer test-access-token",
+          "Content-Type": "application/json",
+        },
+      },
+    );
+  });
+
+  it("should pass additional options along with status=deleted", async () => {
+    // Mock successful API call
+    mockFetch.mockResolvedValueOnce({
+      ok: true,
+      json: async () => mockKeysResponse,
+    });
+
+    const options = {
+      organizationID: "org-1",
+      teamID: "team-1",
+      selectedKeyAlias: "test-alias",
+    };
+
+    const { result } = renderHook(() => useDeletedKeys(1, 10, options), { wrapper });
+
+    // Wait for success
+    await waitFor(() => {
+      expect(result.current.isLoading).toBe(false);
+    });
+
+    const callUrl = mockFetch.mock.calls[0][0];
+    expect(callUrl).toContain("status=deleted");
+    expect(callUrl).toContain("organization_id=org-1");
+    expect(callUrl).toContain("team_id=team-1");
+    expect(callUrl).toContain("key_alias=test-alias");
+    expect(result.current.data).toEqual(mockKeysResponse);
+  });
+});
@@ -101,12 +101,16 @@ const keyListCall = async (
  }
 };

-export const useKeys = (page: number, pageSize: number): UseQueryResult<KeysResponse> => {
+export const useKeys = (
+  page: number,
+  pageSize: number,
+  options: KeyListCallOptions = {},
+): UseQueryResult<KeysResponse> => {
  const { accessToken } = useAuthorized();

  return useQuery<KeysResponse>({
-    queryKey: keyKeys.list({ page, limit: pageSize }),
-    queryFn: async () => await keyListCall(accessToken!, page, pageSize),
+    queryKey: keyKeys.list({ page, limit: pageSize, ...options }),
+    queryFn: async () => await keyListCall(accessToken!, page, pageSize, options),
    enabled: Boolean(accessToken),
    staleTime: 30000, // 30 seconds
    placeholderData: keepPreviousData,
@@ -71,12 +71,19 @@ export function VirtualKeysTable({ teams, organizations, onSortChange, currentSo
    pageSize: 50,
  });

+  // Extract sort parameters from sorting state
+  const sortBy = sorting.length > 0 ? sorting[0].id : null;
+  const sortOrder = sorting.length > 0 ? (sorting[0].desc ? "desc" : "asc") : null;
+
  const {
    data: keys,
    isPending: isLoading,
    isFetching,
    refetch,
-  } = useKeys(tablePagination.pageIndex + 1, tablePagination.pageSize);
+  } = useKeys(tablePagination.pageIndex + 1, tablePagination.pageSize, {
+    sortBy: sortBy || undefined,
+    sortOrder: sortOrder || undefined,
+  });
  const totalCount = keys?.total_count || 0;
  const [expandedAccordions, setExpandedAccordions] = useState<Record<string, boolean>>({});

@@ -110,6 +117,7 @@ export function VirtualKeysTable({ teams, organizations, onSortChange, currentSo
      id: "expander",
      header: () => null,
      size: 40,
+      enableSorting: false,
      cell: ({ row }) =>
        row.getCanExpand() ? (
          <button onClick={row.getToggleExpandedHandler()} style={{ cursor: "pointer" }}>
@@ -122,6 +130,7 @@ export function VirtualKeysTable({ teams, organizations, onSortChange, currentSo
      accessorKey: "token",
      header: "Key ID",
      size: 150,
+      enableSorting: true,
      cell: (info) => (
        <div className="overflow-hidden">
          <Tooltip title={info.getValue() as string}>
@@ -142,6 +151,7 @@ export function VirtualKeysTable({ teams, organizations, onSortChange, currentSo
      accessorKey: "key_alias",
      header: "Key Alias",
      size: 150,
+      enableSorting: true,
      cell: (info) => {
        const value = info.getValue() as string;
        const width = info.cell.column.getSize();
@@ -159,6 +169,7 @@ export function VirtualKeysTable({ teams, organizations, onSortChange, currentSo
      accessorKey: "key_name",
      header: "Secret Key",
      size: 120,
+      enableSorting: false,
      cell: (info) => <span className="font-mono text-xs">{info.getValue() as string}</span>,
    },
    {
@@ -166,6 +177,7 @@ export function VirtualKeysTable({ teams, organizations, onSortChange, currentSo
      accessorKey: "team_id",
      header: "Team Alias",
      size: 120,
+      enableSorting: false,
      cell: ({ row, getValue }) => {
        const teamId = getValue() as string;
        const team = teams?.find((t) => t.team_id === teamId);
@@ -177,6 +189,7 @@ export function VirtualKeysTable({ teams, organizations, onSortChange, currentSo
      accessorKey: "team_id",
      header: "Team ID",
      size: 120,
+      enableSorting: false,
      cell: (info) => (
        <Tooltip title={info.getValue() as string}>
          {info.getValue() ? `${(info.getValue() as string).slice(0, 7)}...` : "-"}
@@ -188,6 +201,7 @@ export function VirtualKeysTable({ teams, organizations, onSortChange, currentSo
      accessorKey: "organization_id",
      header: "Organization ID",
      size: 140,
+      enableSorting: false,
      cell: (info) => (info.getValue() ? info.renderValue() : "-"),
    },
    {
@@ -195,6 +209,7 @@ export function VirtualKeysTable({ teams, organizations, onSortChange, currentSo
      accessorKey: "user",
      header: "User Email",
      size: 160,
+      enableSorting: false,
      cell: (info) => {
        const user = info.getValue() as any;
        const value = user?.user_email;
@@ -213,6 +228,7 @@ export function VirtualKeysTable({ teams, organizations, onSortChange, currentSo
      accessorKey: "user_id",
      header: "User ID",
      size: 120,
+      enableSorting: false,
      cell: (info) => {
        const userId = info.getValue() as string | null;
        if (userId && userId.length > 15) {
@@ -230,6 +246,7 @@ export function VirtualKeysTable({ teams, organizations, onSortChange, currentSo
      accessorKey: "created_at",
      header: "Created At",
      size: 120,
+      enableSorting: true,
      cell: (info) => {
        const value = info.getValue();
        return value ? new Date(value as string).toLocaleDateString() : "-";
@@ -240,6 +257,7 @@ export function VirtualKeysTable({ teams, organizations, onSortChange, currentSo
      accessorKey: "created_by",
      header: "Created By",
      size: 120,
+      enableSorting: false,
      cell: (info) => {
        const value = info.getValue() as string | null;
        if (value && value.length > 15) {
@@ -257,6 +275,7 @@ export function VirtualKeysTable({ teams, organizations, onSortChange, currentSo
      accessorKey: "updated_at",
      header: "Updated At",
      size: 120,
+      enableSorting: true,
      cell: (info) => {
        const value = info.getValue();
        return value ? new Date(value as string).toLocaleDateString() : "Never";
@@ -267,6 +286,7 @@ export function VirtualKeysTable({ teams, organizations, onSortChange, currentSo
      accessorKey: "expires",
      header: "Expires",
      size: 120,
+      enableSorting: false,
      cell: (info) => {
        const value = info.getValue();
        return value ? new Date(value as string).toLocaleDateString() : "Never";
@@ -277,6 +297,7 @@ export function VirtualKeysTable({ teams, organizations, onSortChange, currentSo
      accessorKey: "spend",
      header: "Spend (USD)",
      size: 100,
+      enableSorting: true,
      cell: (info) => formatNumberWithCommas(info.getValue() as number, 4),
    },
    {
@@ -284,6 +305,7 @@ export function VirtualKeysTable({ teams, organizations, onSortChange, currentSo
      accessorKey: "max_budget",
      header: "Budget (USD)",
      size: 110,
+      enableSorting: true,
      cell: (info) => {
        const maxBudget = info.getValue() as number | null;
        if (maxBudget === null) {
@@ -297,6 +319,7 @@ export function VirtualKeysTable({ teams, organizations, onSortChange, currentSo
      accessorKey: "budget_reset_at",
      header: "Budget Reset",
      size: 130,
+      enableSorting: false,
      cell: (info) => {
        const value = info.getValue();
        return value ? new Date(value as string).toLocaleString() : "Never";
@@ -307,6 +330,7 @@ export function VirtualKeysTable({ teams, organizations, onSortChange, currentSo
      accessorKey: "models",
      header: "Models",
      size: 200,
+      enableSorting: false,
      cell: (info) => {
        const models = info.getValue() as string[];
        return (
@@ -391,6 +415,7 @@ export function VirtualKeysTable({ teams, organizations, onSortChange, currentSo
      id: "rate_limits",
      header: "Rate Limits",
      size: 140,
+      enableSorting: false,
      cell: ({ row }) => {
        const key = row.original;
        return (
@@ -491,11 +516,16 @@ export function VirtualKeysTable({ teams, organizations, onSortChange, currentSo
        const sortBy = sortState.id;
        const sortOrder = sortState.desc ? "desc" : "asc";
        console.log(`sortBy: ${sortBy}, sortOrder: ${sortOrder}`);
-        handleFilterChange({
-          ...filters,
-          "Sort By": sortBy,
-          "Sort Order": sortOrder,
-        });
+        // Update filters state without triggering debouncedSearch
+        // The useKeys hook will automatically refetch with the new sort parameters
+        handleFilterChange(
+          {
+            ...filters,
+            "Sort By": sortBy,
+            "Sort Order": sortOrder,
+          },
+          true, // skipDebounce - let useKeys handle the API call with correct page size
+        );
        onSortChange?.(sortBy, sortOrder);
      }
    },
@@ -601,12 +631,13 @@ export function VirtualKeysTable({ teams, organizations, onSortChange, currentSo
                            key={header.id}
                            data-header-id={header.id}
                            className={`py-1 h-8 relative hover:bg-gray-50 ${header.id === "actions"
-                                ? "sticky right-0 bg-white shadow-[-4px_0_8px_-6px_rgba(0,0,0,0.1)]"
-                                : ""
+                              ? "sticky right-0 bg-white shadow-[-4px_0_8px_-6px_rgba(0,0,0,0.1)]"
+                              : ""
                              }`}
                            style={{
                              width: header.getSize(),
                              position: "relative",
+                              cursor: header.column.getCanSort() ? "pointer" : "default",
                            }}
                            onMouseEnter={() => {
                              const resizer = document.querySelector(`[data-header-id="${header.id}"] .resizer`);
@@ -620,7 +651,7 @@ export function VirtualKeysTable({ teams, organizations, onSortChange, currentSo
                                (resizer as HTMLElement).style.opacity = "0";
                              }
                            }}
-                            onClick={header.column.getToggleSortingHandler()}
+                            onClick={header.column.getCanSort() ? header.column.getToggleSortingHandler() : undefined}
                          >
                            <div className="flex items-center justify-between gap-2">
                              <div className="flex items-center">
@@ -628,7 +659,7 @@ export function VirtualKeysTable({ teams, organizations, onSortChange, currentSo
                                  ? null
                                  : flexRender(header.column.columnDef.header, header.getContext())}
                              </div>
-                              {header.id !== "actions" && (
+                              {header.id !== "actions" && header.column.getCanSort() && (
                                <div className="w-4">
                                  {header.column.getIsSorted() ? (
                                    {
@@ -151,7 +151,7 @@ export function useFilterLogic({
    }
  }, [organizations]);

-  const handleFilterChange = (newFilters: Record<string, string>) => {
+  const handleFilterChange = (newFilters: Record<string, string>, skipDebounce: boolean = false) => {
    // Update filters state
    setFilters({
      "Team ID": newFilters["Team ID"] || "",
@@ -162,12 +162,16 @@ export function useFilterLogic({
      "Sort Order": newFilters["Sort Order"] || "desc",
    });

-    // Fetch keys based on new filters
-    const updatedFilters = {
-      ...filters,
-      ...newFilters,
-    };
-    debouncedSearch(updatedFilters);
+    // Only trigger debouncedSearch if skipDebounce is false
+    // This allows sorting to be handled by the parent component's useKeys hook
+    if (!skipDebounce) {
+      // Fetch keys based on new filters
+      const updatedFilters = {
+        ...filters,
+        ...newFilters,
+      };
+      debouncedSearch(updatedFilters);
+    }
  };

  const handleFilterReset = () => {