Address comments

Signed-off-by: Nana Nosirova <10577112+nananosirova@users.noreply.github.com>

Signed-off-by: Juntao Wang <juntwang@redhat.com>

Author: nananosirova

mlflow/server/js/src/mcp-registry/components/MCPServerCardGrid.test.tsx

@@ -1,12 +1,21 @@
-import { describe, it, expect } from '@jest/globals';
+import { describe, it, expect, jest } from '@jest/globals';
 import { render, screen } from '@testing-library/react';
 import { IntlProvider } from 'react-intl';
 import { DesignSystemProvider } from '@databricks/design-system';
 import { testRoute, TestRouter } from '../../common/utils/RoutingTestUtils';
 import { MCPServerCardGrid } from './MCPServerCardGrid';
 import { createMockMCPServer } from '../test-utils';
 
-const renderGrid = (props: Partial<React.ComponentProps<typeof MCPServerCardGrid>> = {}) =>
+const noop = () => {};
+
+const defaultPaginationProps = {
+  hasNextPage: false,
+  hasPreviousPage: false,
+  onNextPage: noop,
+  onPreviousPage: noop,
+};
+
+const renderGrid = (props: React.ComponentProps<typeof MCPServerCardGrid>) =>
   render(
     <IntlProvider locale="en">
       <TestRouter
@@ -24,17 +33,17 @@ const renderGrid = (props: Partial<React.ComponentProps<typeof MCPServerCardGrid
 
 describe('MCPServerCardGrid', () => {
   it('renders loading spinner when isLoading is true', () => {
-    renderGrid({ isLoading: true });
+    renderGrid({ ...defaultPaginationProps, isLoading: true });
     expect(screen.getByText('Loading servers...')).toBeInTheDocument();
   });
 
   it('renders "No servers found" when filtered and no results', () => {
-    renderGrid({ servers: [], isFiltered: true });
+    renderGrid({ ...defaultPaginationProps, servers: [], isFiltered: true });
     expect(screen.getByText('No servers found')).toBeInTheDocument();
   });
 
   it('renders nothing when no servers and not filtered', () => {
-    const { container } = renderGrid({ servers: [] });
+    const { container } = renderGrid({ ...defaultPaginationProps, servers: [] });
     expect(container.firstChild).toBeNull();
   });
 
@@ -44,14 +53,29 @@ describe('MCPServerCardGrid', () => {
       createMockMCPServer({ name: 'server-b', display_name: 'Server B' }),
       createMockMCPServer({ name: 'server-c', display_name: 'Server C' }),
     ];
-    renderGrid({ servers });
+    renderGrid({ ...defaultPaginationProps, servers });
     expect(screen.getByText('Server A')).toBeInTheDocument();
     expect(screen.getByText('Server B')).toBeInTheDocument();
     expect(screen.getByText('Server C')).toBeInTheDocument();
   });
 
   it('does not render loading spinner when servers are present', () => {
-    renderGrid({ servers: [createMockMCPServer()], isLoading: false });
+    renderGrid({ ...defaultPaginationProps, servers: [createMockMCPServer()], isLoading: false });
     expect(screen.queryByText('Loading servers...')).not.toBeInTheDocument();
   });
+
+  it('renders pagination controls when servers are present', () => {
+    const servers = [createMockMCPServer()];
+    renderGrid({ ...defaultPaginationProps, servers, hasNextPage: true });
+    expect(screen.getByText('Next')).toBeInTheDocument();
+    expect(screen.getByText('Previous')).toBeInTheDocument();
+  });
+
+  it('calls onNextPage when Next is clicked', () => {
+    const onNextPage = jest.fn();
+    const servers = [createMockMCPServer()];
+    renderGrid({ ...defaultPaginationProps, servers, hasNextPage: true, onNextPage });
+    screen.getByText('Next').click();
+    expect(onNextPage).toHaveBeenCalledTimes(1);
+  });
 });

mlflow/server/js/src/mcp-registry/components/MCPServerCardGrid.tsx

@@ -1,4 +1,5 @@
-import { Empty, NoIcon, Spinner, useDesignSystemTheme } from '@databricks/design-system';
+import type { CursorPaginationProps } from '@databricks/design-system';
+import { CursorPagination, Empty, NoIcon, Spinner, useDesignSystemTheme } from '@databricks/design-system';
 import { FormattedMessage } from 'react-intl';
 
 import type { MCPServer } from '../types';
@@ -8,10 +9,20 @@ export const MCPServerCardGrid = ({
   servers,
   isLoading,
   isFiltered,
+  hasNextPage,
+  hasPreviousPage,
+  onNextPage,
+  onPreviousPage,
+  pageSizeSelect,
 }: {
   servers?: MCPServer[];
   isLoading?: boolean;
   isFiltered?: boolean;
+  hasNextPage: boolean;
+  hasPreviousPage: boolean;
+  onNextPage: () => void;
+  onPreviousPage: () => void;
+  pageSizeSelect?: CursorPaginationProps['pageSizeSelect'];
 }) => {
   const { theme } = useDesignSystemTheme();
 
@@ -55,17 +66,40 @@ export const MCPServerCardGrid = ({
   }
 
   return (
-    <div
-      css={{
-        display: 'grid',
-        gridTemplateColumns: 'repeat(auto-fill, minmax(280px, 1fr))',
-        gap: theme.spacing.md,
-        paddingTop: theme.spacing.md,
-      }}
-    >
-      {servers.map((server) => (
-        <MCPServerCard key={server.name} server={server} />
-      ))}
+    <div css={{ flex: 1, display: 'flex', flexDirection: 'column', minHeight: 0, overflow: 'hidden' }}>
+      <div
+        css={{
+          flex: '0 1 auto',
+          overflow: 'auto',
+          minHeight: 0,
+          display: 'grid',
+          gridTemplateColumns: 'repeat(auto-fill, minmax(280px, 1fr))',
+          gap: theme.spacing.md,
+          paddingTop: theme.spacing.md,
+        }}
+      >
+        {servers.map((server) => (
+          <MCPServerCard key={server.name} server={server} />
+        ))}
+      </div>
+      <div
+        css={{
+          flexShrink: 0,
+          display: 'flex',
+          justifyContent: 'flex-end',
+          paddingTop: theme.spacing.sm,
+          paddingBottom: theme.spacing.sm,
+        }}
+      >
+        <CursorPagination
+          hasNextPage={hasNextPage}
+          hasPreviousPage={hasPreviousPage}
+          onNextPage={onNextPage}
+          onPreviousPage={onPreviousPage}
+          pageSizeSelect={pageSizeSelect}
+          componentId="mlflow.mcp_registry.grid.pagination"
+        />
+      </div>
     </div>
   );
 };

mlflow/server/js/src/mcp-registry/hooks/useMCPServersListQuery.ts

@@ -1,6 +1,6 @@
 import type { QueryFunctionContext } from '@mlflow/mlflow/src/common/utils/reactQueryHooks';
 import { useQuery } from '@mlflow/mlflow/src/common/utils/reactQueryHooks';
-import { useCallback, useEffect, useRef, useState } from 'react';
+import { useCallback, useEffect, useMemo, useRef, useState } from 'react';
 import { useLocalStorage } from '@databricks/web-shared/hooks';
 import type { CursorPaginationProps } from '@databricks/design-system';
 import { MCPRegistryApi } from '../api';
@@ -50,15 +50,18 @@ export const useMCPServersListQuery = ({ searchFilter }: { searchFilter?: string
     previousPageTokens.current = [];
   }, [searchFilter]);
 
-  const pageSizeSelect = useMemo<CursorPaginationProps['pageSizeSelect']>(() => ({
-    options: [10, 25, 50, 100],
-    default: pageSize,
-    onChange(newPageSize) {
-      setPageSize(newPageSize);
-      setCurrentPageToken(undefined);
-      previousPageTokens.current = [];
-    },
-  }), [pageSize]);
+  const pageSizeSelect = useMemo<CursorPaginationProps['pageSizeSelect']>(
+    () => ({
+      options: [10, 25, 50, 100],
+      default: pageSize,
+      onChange(newPageSize) {
+        setPageSize(newPageSize);
+        setCurrentPageToken(undefined);
+        previousPageTokens.current = [];
+      },
+    }),
+    [pageSize, setPageSize],
+  );
 
   const queryResult = useQuery<SearchMCPServersResponse, Error, SearchMCPServersResponse, MCPServersListQueryKey>(
     ['mcp_servers_list', { searchFilter, pageToken: currentPageToken, pageSize }],

mlflow/server/js/src/mcp-registry/pages/MCPRegistryPage.test.tsx

@@ -250,8 +250,7 @@ describe('MCPRegistryPage', () => {
       { timeout: 2000 },
     );
 
-    // Should only have made one additional call after debounce, not 6
-    expect(callCount).toBeLessThanOrEqual(initialCallCount + 2);
+    expect(callCount).toBeLessThan(initialCallCount + 6);
   });
 
   it('keeps previous data visible while loading new search results', async () => {

mlflow/server/js/src/mcp-registry/pages/MCPRegistryPage.tsx

@@ -2,7 +2,6 @@ import { useCallback, useState } from 'react';
 import {
   Alert,
   Button,
-  CursorPagination,
   Empty,
   GridIcon,
   Header,
@@ -190,29 +189,16 @@ const MCPRegistryPage = () => {
               isEmptyState ? (
                 serversEmptyState
               ) : (
-                <div css={{ flex: 1, display: 'flex', flexDirection: 'column', minHeight: 0, overflow: 'hidden' }}>
-                  <div css={{ flex: '0 1 auto', overflow: 'auto', minHeight: 0 }}>
-                    <MCPServerCardGrid servers={servers} isLoading={isLoading} isFiltered={Boolean(searchFilter)} />
-                  </div>
-                  <div
-                    css={{
-                      flexShrink: 0,
-                      display: 'flex',
-                      justifyContent: 'flex-end',
-                      paddingTop: theme.spacing.sm,
-                      paddingBottom: theme.spacing.sm,
-                    }}
-                  >
-                    <CursorPagination
-                      hasNextPage={hasNextPage}
-                      hasPreviousPage={hasPreviousPage}
-                      onNextPage={onNextPage}
-                      onPreviousPage={onPreviousPage}
-                      pageSizeSelect={pageSizeSelect}
-                      componentId="mlflow.mcp_registry.grid.pagination"
-                    />
-                  </div>
-                </div>
+                <MCPServerCardGrid
+                  servers={servers}
+                  isLoading={isLoading}
+                  isFiltered={Boolean(searchFilter)}
+                  hasNextPage={hasNextPage}
+                  hasPreviousPage={hasPreviousPage}
+                  onNextPage={onNextPage}
+                  onPreviousPage={onPreviousPage}
+                  pageSizeSelect={pageSizeSelect}
+                />
               )
             ) : (
               <MCPServerListTable