feat: Add withCounts option to MCP server queries (#10092)

* Add withCounts option to MCP server queries

Introduces a 'withCounts' parameter to the useGetMCPServers hook, allowing conditional fetching of server counts. Updates MCPServersPage and McpComponent to utilize this option for more efficient data loading.

* test fix

---------

Co-authored-by: Mike Fortman <michael.fortman@datastax.com>

src/frontend/src/components/core/parameterRenderComponent/components/mcpComponent/index.tsx

@@ -1,4 +1,4 @@
-import { useEffect, useMemo, useRef, useState } from "react";
+import { useEffect, useMemo, useState } from "react";
 import { useAddMCPServer } from "@/controllers/API/queries/mcp/use-add-mcp-server";
 import { useGetMCPServers } from "@/controllers/API/queries/mcp/use-get-mcp-servers";
 import AddMcpServerModal from "@/modals/addMcpServerModal";
@@ -16,7 +16,8 @@ export default function McpComponent({
   editNode = false,
   id = "",
 }: InputProps<string, any>): JSX.Element {
-  const { data: mcpServers } = useGetMCPServers();
+  const [open, setOpen] = useState(false);
+  const { data: mcpServers } = useGetMCPServers({ withCounts: open });
   const { mutate: addMcpServer } = useAddMCPServer();
   const setErrorData = useAlertStore((state) => state.setErrorData);
   const options = useMemo(
@@ -36,7 +37,6 @@ export default function McpComponent({
       })),
     [mcpServers],
   );
-  const [open, setOpen] = useState(false);
   const [addOpen, setAddOpen] = useState(false);
   const [selectedItem, setSelectedItem] = useState<any[]>([]);
   const { name, config } = useMemo(

src/frontend/src/controllers/API/queries/mcp/use-get-mcp-servers.ts

@@ -12,9 +12,11 @@ export type getMCPServersResponse = Array<MCPServerInfoType>;
 
 export const useGetMCPServers: useQueryFunctionType<
   undefined,
-  getMCPServersResponse
+  getMCPServersResponse,
+  { withCounts?: boolean }
 > = (options) => {
   const { query, queryClient } = UseRequestProcessor();
+  const { withCounts, ...queryOptions } = options ?? {};
 
   // First fetch: action_count=false (fast)
   const responseFn = async () => {
@@ -64,11 +66,11 @@ export const useGetMCPServers: useQueryFunctionType<
   };
 
   const queryResult = query(["useGetMCPServers"], responseFn, {
-    ...options,
+    ...queryOptions,
   });
 
   useEffect(() => {
-    if (queryResult.data && queryResult.data.length > 0) {
+    if (withCounts && queryResult.data && queryResult.data.length > 0) {
       fetchWithCounts().then((countsData) => {
         if (!countsData || countsData.length === 0) return;
         // Merge by name
@@ -83,7 +85,7 @@ export const useGetMCPServers: useQueryFunctionType<
         );
       });
     }
-  }, [queryResult.data]);
+  }, [withCounts, queryResult.data]);
 
   return queryResult;
 };

src/frontend/src/pages/SettingsPage/pages/MCPServersPage/index.tsx

@@ -19,7 +19,7 @@ import type { MCPServerInfoType } from "@/types/mcp";
 import { cn } from "@/utils/utils";
 
 export default function MCPServersPage() {
-  const { data: servers } = useGetMCPServers();
+  const { data: servers } = useGetMCPServers({ withCounts: true });
   const { mutate: deleteServer } = useDeleteMCPServer();
   const setErrorData = useAlertStore((state) => state.setErrorData);
   const [addOpen, setAddOpen] = useState(false);

src/frontend/tests/core/unit/dropdownComponent.spec.ts

@@ -14,6 +14,10 @@ test(
 
     await page.getByTestId("blank-flow").click();
 
+    // Allow for legacy components
+    await page.getByTestId("sidebar-options-trigger").click();
+    await page.getByTestId("sidebar-legacy-switch").click();
+
     await page.getByTestId("sidebar-search-input").click();
     await page.getByTestId("sidebar-search-input").fill("amazon");
 

src/frontend/tests/core/unit/keyPairListComponent.spec.ts

@@ -12,6 +12,11 @@ test(
       timeout: 30000,
     });
     await page.getByTestId("blank-flow").click();
+
+    // Allow for legacy components
+    await page.getByTestId("sidebar-options-trigger").click();
+    await page.getByTestId("sidebar-legacy-switch").click();
+
     await page.getByTestId("sidebar-search-input").click();
     await page.getByTestId("sidebar-search-input").fill("amazon bedrock");