From 7be62f268c3eef8160b1779195129af67b7a3b9d Mon Sep 17 00:00:00 2001
From: Ogulcan <ogulcansarioglu1@gmail.com>
Date: Fri, 22 May 2026 12:27:24 +0100
Subject: [PATCH] Fix agents key list page size

---
 .../src/components/agents.test.tsx            | 62 ++++++++++++++++
 .../src/components/agents.tsx                 | 71 +++++++++++++------
 2 files changed, 112 insertions(+), 21 deletions(-)

diff --git a/ui/litellm-dashboard/src/components/agents.test.tsx b/ui/litellm-dashboard/src/components/agents.test.tsx
index 3b0f4341da8..0710fe343cd 100644
--- a/ui/litellm-dashboard/src/components/agents.test.tsx
+++ b/ui/litellm-dashboard/src/components/agents.test.tsx
@@ -89,6 +89,68 @@ describe("AgentsPanel", () => {
     });
   });
 
+  it("should fetch agent keys with the key list API max page size", async () => {
+    vi.mocked(networking.getAgentsList).mockResolvedValueOnce({
+      agents: [
+        {
+          agent_id: "agent-1",
+          agent_name: "Test Agent",
+          litellm_params: { model: "gpt-4o" },
+        },
+      ],
+    });
+
+    render(<AgentsPanel accessToken="test-token" userRole="Admin" />);
+
+    await waitFor(() => {
+      expect(networking.keyListCall).toHaveBeenCalledWith(
+        "test-token",
+        null,
+        null,
+        null,
+        null,
+        null,
+        1,
+        100,
+      );
+    });
+  });
+
+  it("should keep fetching key list pages until visible agent keys are found", async () => {
+    vi.mocked(networking.getAgentsList).mockResolvedValueOnce({
+      agents: [
+        {
+          agent_id: "agent-1",
+          agent_name: "Test Agent",
+          litellm_params: { model: "gpt-4o" },
+        },
+      ],
+    });
+    vi.mocked(networking.keyListCall)
+      .mockResolvedValueOnce({ keys: [], total_pages: 2 })
+      .mockResolvedValueOnce({
+        keys: [{ agent_id: "agent-1", key_alias: "agent-key", token: "1234567890" }],
+        total_pages: 2,
+      });
+
+    render(<AgentsPanel accessToken="test-token" userRole="Admin" />);
+
+    await waitFor(() => {
+      expect(networking.keyListCall).toHaveBeenCalledTimes(2);
+      expect(networking.keyListCall).toHaveBeenLastCalledWith(
+        "test-token",
+        null,
+        null,
+        null,
+        null,
+        null,
+        2,
+        100,
+      );
+      expect(screen.getByText("Active")).toBeInTheDocument();
+    });
+  });
+
   it("should call getAgentsList with health_check=true when toggle is enabled", async () => {
     render(<AgentsPanel accessToken="test-token" userRole="Admin" />);
     await waitFor(() => {
diff --git a/ui/litellm-dashboard/src/components/agents.tsx b/ui/litellm-dashboard/src/components/agents.tsx
index 542d890e12f..1a70c9879d6 100644
--- a/ui/litellm-dashboard/src/components/agents.tsx
+++ b/ui/litellm-dashboard/src/components/agents.tsx
@@ -33,6 +33,37 @@ interface AgentsResponse {
   agents: Agent[];
 }
 
+interface AgentKeyRecord {
+  agent_id?: string;
+  key_alias?: string;
+  token?: string;
+}
+
+interface KeyListResponse {
+  keys?: AgentKeyRecord[];
+  total_pages?: number;
+}
+
+const KEY_LIST_PAGE_SIZE = 100;
+
+const addAgentKeysToMap = (
+  keys: AgentKeyRecord[],
+  keyInfoMap: Record<string, AgentKeyInfo>,
+  missingAgentIds: Set<string>,
+) => {
+  for (const key of keys) {
+    const agentId = key.agent_id;
+    if (agentId && missingAgentIds.has(agentId)) {
+      keyInfoMap[agentId] = {
+        has_key: true,
+        key_alias: key.key_alias,
+        token_prefix: key.token ? `${key.token.slice(0, 8)}…` : undefined,
+      };
+      missingAgentIds.delete(agentId);
+    }
+  }
+};
+
 const AgentsPanel: React.FC<AgentsPanelProps> = ({ accessToken, userRole, teams }) => {
   const [agentsList, setAgentsList] = useState<Agent[]>([]);
   const [keyInfoMap, setKeyInfoMap] = useState<Record<string, AgentKeyInfo>>({});
@@ -64,29 +95,27 @@ const AgentsPanel: React.FC<AgentsPanelProps> = ({ accessToken, userRole, teams
   const fetchKeysForAgents = async () => {
     if (!accessToken) return;
     try {
-      const { keys = [] } = await keyListCall(
-        accessToken,
-        null,
-        null,
-        null,
-        null,
-        null,
-        1,
-        500
-      );
       const map: Record<string, AgentKeyInfo> = {};
-      for (const key of keys) {
-        const agentId = (key as { agent_id?: string }).agent_id;
-        if (agentId && !map[agentId]) {
-          map[agentId] = {
-            has_key: true,
-            key_alias: (key as { key_alias?: string }).key_alias,
-            token_prefix: (key as { token?: string }).token
-              ? `${(key as { token: string }).token.slice(0, 8)}…`
-              : undefined,
-          };
-        }
+      const missingAgentIds = new Set(agentsList.map((agent) => agent.agent_id));
+      let page = 1;
+      let totalPages = 1;
+
+      while (page <= totalPages && missingAgentIds.size > 0) {
+        const { keys = [], total_pages = 1 }: KeyListResponse = await keyListCall(
+          accessToken,
+          null,
+          null,
+          null,
+          null,
+          null,
+          page,
+          KEY_LIST_PAGE_SIZE
+        );
+        addAgentKeysToMap(keys, map, missingAgentIds);
+        totalPages = total_pages;
+        page += 1;
       }
+
       setKeyInfoMap(map);
     } catch (error) {
       console.error("Error fetching keys for agents:", error);