Merge pull request #18416 from BerriAI/litellm_ui_refactor_2

[Refactor] useQuery Hooks Now Depend on useAuthorized

ui/litellm-dashboard/src/app/(dashboard)/hooks/agents/useAgents.ts

@@ -3,10 +3,12 @@ import { AgentsResponse } from "@/components/agents/types";
 import { useQuery } from "@tanstack/react-query";
 import { createQueryKeys } from "../common/queryKeysFactory";
 import { all_admin_roles } from "@/utils/roles";
+import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized";
 
 const agentsKeys = createQueryKeys("agents");
 
-export const useAgents = (accessToken: string | null, userRole: string | null) => {
+export const useAgents = () => {
+  const { accessToken, userRole } = useAuthorized();
   return useQuery<AgentsResponse>({
     queryKey: agentsKeys.list({}),
     queryFn: async () => await getAgentsList(accessToken!),

ui/litellm-dashboard/src/app/(dashboard)/hooks/credentials/useCredentials.ts

@@ -1,10 +1,12 @@
 import { credentialListCall, CredentialsResponse } from "@/components/networking";
 import { useQuery } from "@tanstack/react-query";
 import { createQueryKeys } from "../common/queryKeysFactory";
+import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized";
 
 const credentialsKeys = createQueryKeys("credentials");
 
-export const useCredentials = (accessToken: string | null) => {
+export const useCredentials = () => {
+  const { accessToken } = useAuthorized();
   return useQuery<CredentialsResponse>({
     queryKey: credentialsKeys.list({}),
     queryFn: async () => await credentialListCall(accessToken!),

ui/litellm-dashboard/src/app/(dashboard)/hooks/customers/useCustomers.ts

@@ -2,7 +2,7 @@ import { allEndUsersCall } from "@/components/networking";
 import { useQuery } from "@tanstack/react-query";
 import { createQueryKeys } from "../common/queryKeysFactory";
 import { all_admin_roles } from "@/utils/roles";
-
+import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized";
 const customersKeys = createQueryKeys("customers");
 
 export interface Customer {
@@ -32,10 +32,11 @@ export interface Customer {
 
 export type CustomersResponse = Customer[];
 
-export const useCustomers = (accessToken: string | null, userRole: string | null) => {
+export const useCustomers = () => {
+  const { accessToken, userRole } = useAuthorized();
   return useQuery<CustomersResponse>({
     queryKey: customersKeys.list({}),
     queryFn: async () => await allEndUsersCall(accessToken!),
-    enabled: Boolean(accessToken) && all_admin_roles.includes(userRole || ""),
+    enabled: Boolean(accessToken) && all_admin_roles.includes(userRole!),
   });
 };

ui/litellm-dashboard/src/app/(dashboard)/hooks/mcpServers/useMCPAccessGroups.ts

@@ -1,13 +1,14 @@
 import { useQuery } from "@tanstack/react-query";
 import { createQueryKeys } from "../common/queryKeysFactory";
 import { fetchMCPAccessGroups } from "@/components/networking";
-
+import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized";
 const mcpAccessGroupsKeys = createQueryKeys("mcpAccessGroups");
 
-export const useMCPAccessGroups = (accessToken: string | null) => {
+export const useMCPAccessGroups = () => {
+  const { accessToken } = useAuthorized();
   return useQuery<string[]>({
     queryKey: mcpAccessGroupsKeys.list({}),
     queryFn: async () => await fetchMCPAccessGroups(accessToken!),
-    enabled: !!accessToken,
+    enabled: Boolean(accessToken),
   });
 };

ui/litellm-dashboard/src/app/(dashboard)/hooks/mcpServers/useMCPServers.ts

@@ -2,10 +2,12 @@ import { useQuery } from "@tanstack/react-query";
 import { createQueryKeys } from "../common/queryKeysFactory";
 import { fetchMCPServers } from "@/components/networking";
 import { MCPServer } from "@/components/mcp_tools/types";
+import useAuthorized from "../useAuthorized";
 
 const mcpServersKeys = createQueryKeys("mcpServers");
 
-export const useMCPServers = (accessToken: string | null) => {
+export const useMCPServers = () => {
+  const { accessToken } = useAuthorized();
   return useQuery<MCPServer[]>({
     queryKey: mcpServersKeys.list({}),
     queryFn: async () => await fetchMCPServers(accessToken!),

ui/litellm-dashboard/src/app/(dashboard)/hooks/models/useModels.ts

@@ -1,24 +1,26 @@
 import { useQuery } from "@tanstack/react-query";
 import { createQueryKeys } from "../common/queryKeysFactory";
 import { modelInfoCall, modelHubCall } from "@/components/networking";
-
+import useAuthorized from "../useAuthorized";
 const modelKeys = createQueryKeys("models");
 const modelHubKeys = createQueryKeys("modelHub");
 
-export const useModelsInfo = (accessToken: string | null, userID: string | null, userRole: string | null) => {
+export const useModelsInfo = () => {
+  const { accessToken, userId, userRole } = useAuthorized();
   return useQuery({
     queryKey: modelKeys.list({
       filters: {
-        ...(userID && { userID }),
+        ...(userId && { userId }),
         ...(userRole && { userRole }),
       },
     }),
-    queryFn: async () => await modelInfoCall(accessToken!, userID!, userRole!),
-    enabled: Boolean(accessToken && userID && userRole),
+    queryFn: async () => await modelInfoCall(accessToken!, userId!, userRole!),
+    enabled: Boolean(accessToken && userId && userRole),
   });
 };
 
-export const useModelHub = (accessToken: string | null) => {
+export const useModelHub = () => {
+  const { accessToken } = useAuthorized();
   return useQuery({
     queryKey: modelHubKeys.list({}),
     queryFn: async () => await modelHubCall(accessToken!),

ui/litellm-dashboard/src/app/(dashboard)/hooks/organizations/useOrganizations.ts

@@ -1,13 +1,16 @@
 import { useQuery, UseQueryResult } from "@tanstack/react-query";
 import { createQueryKeys } from "../common/queryKeysFactory";
 import { organizationListCall, Organization } from "@/components/networking";
+import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized";
 
 const organizationKeys = createQueryKeys("organizations");
 
-export const useOrganizations = (accessToken: string | null): UseQueryResult<Organization[]> => {
+export const useOrganizations = (): UseQueryResult<Organization[]> => {
+  const { accessToken } = useAuthorized();
+  const { userId, userRole } = useAuthorized();
   return useQuery<Organization[]>({
     queryKey: organizationKeys.list({}),
     queryFn: async () => await organizationListCall(accessToken!),
-    enabled: Boolean(accessToken),
+    enabled: Boolean(accessToken && userId && userRole),
   });
 };

ui/litellm-dashboard/src/app/(dashboard)/hooks/teams/useTeams.ts

@@ -7,11 +7,11 @@ import { createQueryKeys } from "@/app/(dashboard)/hooks/common/queryKeysFactory
 const teamKeys = createQueryKeys("teams");
 
 export const useTeams = (): UseQueryResult<Team[]> => {
-  const { accessToken, userId: userID, userRole } = useAuthorized();
+  const { accessToken, userId, userRole } = useAuthorized();
 
   return useQuery<Team[]>({
     queryKey: teamKeys.list({}),
-    queryFn: async () => await fetchTeams(accessToken!, userID, userRole, null),
+    queryFn: async () => await fetchTeams(accessToken!, userId, userRole, null),
     enabled: Boolean(accessToken),
   });
 };

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

@@ -5,6 +5,9 @@ import { afterEach, describe, expect, it, vi } from "vitest";
 import { QueryClient, QueryClientProvider } from "@tanstack/react-query";
 import useAuthorized from "./useAuthorized";
 
+// Unmock useAuthorized to test the actual implementation
+vi.unmock("@/app/(dashboard)/hooks/useAuthorized");
+
 const { replaceMock, clearTokenCookiesMock, getProxyBaseUrlMock, getUiConfigMock } = vi.hoisted(() => ({
   replaceMock: vi.fn(),
   clearTokenCookiesMock: vi.fn(),

ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/ModelsAndEndpointsView.test.tsx

@@ -37,19 +37,6 @@ vi.mock("@/app/(dashboard)/models-and-endpoints/components/ModelAnalyticsTab/Mod
   default: () => null,
 }));
 
-vi.mock("@/app/(dashboard)/hooks/useAuthorized", () => ({
-  default: () => ({
-    token: "123",
-    accessToken: "123",
-    userId: "user-1",
-    userEmail: "user@example.com",
-    userRole: "Admin",
-    premiumUser: false,
-    disabledPersonalKeyCreation: null,
-    showSSOBanner: false,
-  }),
-}));
-
 vi.mock("@/app/(dashboard)/hooks/useTeams", () => ({
   default: () => ({
     teams: [],

ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/ModelsAndEndpointsView.tsx

@@ -152,12 +152,8 @@ const ModelsAndEndpointsView: React.FC<ModelDashboardProps> = ({
   const [selectedTabIndex, setSelectedTabIndex] = useState(0);
 
   const queryClient = useQueryClient();
-  const {
-    data: modelDataResponse,
-    isLoading: isLoadingModels,
-    refetch: refetchModels,
-  } = useModelsInfo(accessToken, userID, userRole);
-  const { data: credentialsResponse } = useCredentials(accessToken);
+  const { data: modelDataResponse, isLoading: isLoadingModels, refetch: refetchModels } = useModelsInfo();
+  const { data: credentialsResponse } = useCredentials();
   const credentialsList = credentialsResponse?.credentials || [];
   const { data: uiSettings } = useUISettings(accessToken || "");
 

ui/litellm-dashboard/src/components/OldTeams.tsx

@@ -174,7 +174,7 @@ const Teams: React.FC<TeamProps> = ({
   premiumUser = false,
 }) => {
   console.log(`organizations: ${JSON.stringify(organizations)}`);
-  const { data: organizationsData } = useOrganizations(accessToken);
+  const { data: organizationsData } = useOrganizations();
   const [lastRefreshed, setLastRefreshed] = useState("");
   const [currentOrg, setCurrentOrg] = useState<Organization | null>(null);
   const [currentOrgForCreateTeam, setCurrentOrgForCreateTeam] = useState<Organization | null>(null);

ui/litellm-dashboard/src/components/UsagePage/components/UsagePageView.tsx

@@ -80,8 +80,8 @@ const UsagePage: React.FC<UsagePageProps> = ({ teams, organizations }) => {
   });
 
   const [allTags, setAllTags] = useState<EntityList[]>([]);
-  const { data: customers = [] } = useCustomers(accessToken, userRole);
-  const { data: agentsResponse } = useAgents(accessToken, userRole);
+  const { data: customers = [] } = useCustomers();
+  const { data: agentsResponse } = useAgents();
   const [modelViewType, setModelViewType] = useState<"groups" | "individual">("groups");
   const [isCloudZeroModalOpen, setIsCloudZeroModalOpen] = useState(false);
   const [isGlobalExportModalOpen, setIsGlobalExportModalOpen] = useState(false);

ui/litellm-dashboard/src/components/leftnav.tsx

@@ -57,7 +57,7 @@ interface MenuGroup {
 
 const Sidebar: React.FC<SidebarProps> = ({ setPage, defaultSelectedKey, collapsed = false }) => {
   const { userId, accessToken, userRole } = useAuthorized();
-  const { data: organizations } = useOrganizations(accessToken);
+  const { data: organizations } = useOrganizations();
 
   // Check if user is an org_admin
   const isOrgAdmin = useMemo(() => {

ui/litellm-dashboard/src/components/mcp_server_management/MCPServerSelector.tsx

@@ -23,8 +23,8 @@ const MCPServerSelector: React.FC<MCPServerSelectorProps> = ({
   placeholder = "Select MCP servers",
   disabled = false,
 }) => {
-  const { data: mcpServers = [], isLoading: serversLoading } = useMCPServers(accessToken);
-  const { data: accessGroups = [], isLoading: groupsLoading } = useMCPAccessGroups(accessToken);
+  const { data: mcpServers = [], isLoading: serversLoading } = useMCPServers();
+  const { data: accessGroups = [], isLoading: groupsLoading } = useMCPAccessGroups();
 
   const loading = serversLoading || groupsLoading;
 

ui/litellm-dashboard/src/components/mcp_server_management/MCPToolPermissions.test.tsx

@@ -71,7 +71,9 @@ describe("MCPToolPermissions", () => {
     });
 
     // Verify API calls
-    expect(networking.fetchMCPServers).toHaveBeenCalledWith(mockAccessToken);
+    // Note: useMCPServers uses useAuthorized() internally, which returns "123" from global mock
+    expect(networking.fetchMCPServers).toHaveBeenCalledWith("123");
+    // listMCPTools uses the accessToken prop directly
     expect(networking.listMCPTools).toHaveBeenCalledWith(mockAccessToken, mockServerId);
   });
 

ui/litellm-dashboard/src/components/mcp_server_management/MCPToolPermissions.tsx

@@ -21,7 +21,7 @@ const MCPToolPermissions: React.FC<MCPToolPermissionsProps> = ({
   onChange,
   disabled = false,
 }) => {
-  const { data: allServers = [] } = useMCPServers(accessToken);
+  const { data: allServers = [] } = useMCPServers();
   const [serverTools, setServerTools] = useState<Record<string, MCPTool[]>>({});
   const [loadingTools, setLoadingTools] = useState<Record<string, boolean>>({});
   const [toolErrors, setToolErrors] = useState<Record<string, string>>({});

ui/litellm-dashboard/src/components/mcp_tools/mcp_servers.test.tsx

@@ -32,7 +32,7 @@ const createQueryClient = () =>
 
 describe("MCPServers", () => {
   const defaultProps = {
-    accessToken: "test-token",
+    accessToken: "123",
     userRole: "Admin",
     userID: "admin-user-id",
   };
@@ -120,6 +120,7 @@ describe("MCPServers", () => {
     expect(getByText("test-server-2")).toBeInTheDocument();
 
     // Verify the API was called
-    expect(networking.fetchMCPServers).toHaveBeenCalledWith("test-token");
+    // Note: useMCPServers uses useAuthorized() internally, which returns "123" from global mock
+    expect(networking.fetchMCPServers).toHaveBeenCalledWith("123");
   });
 });

ui/litellm-dashboard/src/components/mcp_tools/mcp_servers.tsx

@@ -19,7 +19,7 @@ const EDIT_OAUTH_UI_STATE_KEY = "litellm-mcp-oauth-edit-state";
 const { Option } = Select;
 
 const MCPServers: React.FC<MCPServerProps> = ({ accessToken, userRole, userID }) => {
-  const { data: mcpServers, isLoading: isLoadingServers, refetch, dataUpdatedAt } = useMCPServers(accessToken);
+  const { data: mcpServers, isLoading: isLoadingServers, refetch, dataUpdatedAt } = useMCPServers();
 
   // Log allowed_tools from fetched servers
   React.useEffect(() => {

ui/litellm-dashboard/src/components/model_add/credentials.tsx

@@ -32,7 +32,7 @@ interface CredentialsPanelProps {
 
 const CredentialsPanel: React.FC<CredentialsPanelProps> = ({ uploadProps }) => {
   const { accessToken } = useAuthorized();
-  const { data: credentialsResponse, refetch: refetchCredentials } = useCredentials(accessToken);
+  const { data: credentialsResponse, refetch: refetchCredentials } = useCredentials();
   const credentialList = credentialsResponse?.credentials || [];
 
   const [isAddModalOpen, setIsAddModalOpen] = useState(false);

ui/litellm-dashboard/src/components/model_info_view.tsx

@@ -86,7 +86,7 @@ export default function ModelInfoView({
   const isAdmin = userRole === "Admin";
   const isAutoRouter = modelData?.litellm_params?.auto_router_config != null;
 
-  const { data: modelsInfoData } = useModelsInfo(accessToken, userID, userRole);
+  const { data: modelsInfoData } = useModelsInfo();
   console.log("modelsInfoData, ", modelsInfoData);
   const usingExistingCredential =
     modelData?.litellm_params?.litellm_credential_name != null &&

ui/litellm-dashboard/tests/setupTests.ts

@@ -33,6 +33,20 @@ vi.mock("@tremor/react", async (importOriginal) => {
   };
 });
 
+// Global mock for useAuthorized hook to avoid repeating the same mock in every test file
+vi.mock("@/app/(dashboard)/hooks/useAuthorized", () => ({
+  default: () => ({
+    token: "123",
+    accessToken: "123",
+    userId: "user-1",
+    userEmail: "user@example.com",
+    userRole: "Admin",
+    premiumUser: false,
+    disabledPersonalKeyCreation: null,
+    showSSOBanner: false,
+  }),
+}));
+
 afterEach(() => {
   cleanup();
 });