CWFHEALTH-5046: Add retry logic to MCP server registration

src/utils/common.py

@@ -7,9 +7,10 @@
 from typing import Any, cast
 
 from llama_stack.core.library_client import AsyncLlamaStackAsLibraryClient
-from llama_stack_client import AsyncLlamaStackClient
+from llama_stack_client import APIConnectionError, AsyncLlamaStackClient
 
 from client import AsyncLlamaStackClientHolder
+from constants import DEFAULT_MAX_RETRIES, DEFAULT_RETRY_DELAY
 from models.config import Configuration, ModelContextProtocolServer
 
 
@@ -48,17 +49,31 @@ async def register_mcp_servers_async(
             AsyncLlamaStackAsLibraryClient, AsyncLlamaStackClientHolder().get_client()
         )
         await client.initialize()
-        await _register_mcp_toolgroups_async(client, configuration.mcp_servers, logger)
+        await _register_mcp_toolgroups_async(
+            client,
+            configuration.mcp_servers,
+            logger,
+            max_retries=configuration.llama_stack.max_retries,
+            retry_delay=configuration.llama_stack.retry_delay,
+        )
     else:
         # Service client - also use async interface
         client = AsyncLlamaStackClientHolder().get_client()
-        await _register_mcp_toolgroups_async(client, configuration.mcp_servers, logger)
+        await _register_mcp_toolgroups_async(
+            client,
+            configuration.mcp_servers,
+            logger,
+            max_retries=configuration.llama_stack.max_retries,
+            retry_delay=configuration.llama_stack.retry_delay,
+        )
 
 
 async def _register_mcp_toolgroups_async(
     client: AsyncLlamaStackClient,
     mcp_servers: list[ModelContextProtocolServer],
     logger: Logger,
+    max_retries: int = DEFAULT_MAX_RETRIES,
+    retry_delay: int = DEFAULT_RETRY_DELAY,
 ) -> None:
     """
     Register MCP (Model Context Protocol) toolgroups with a LlamaStack async client.
@@ -69,9 +84,9 @@ async def _register_mcp_toolgroups_async(
     `toolgroup_id`=`mcp.name`, `provider_id`=`mcp.provider_id`, and
     `mcp_endpoint` containing the server `url`.
 
-    This function performs network calls against the provided async client and does not
-    catch exceptions raised by those calls — any exceptions from the client (e.g., RPC
-    or HTTP errors) will propagate to the caller.
+    This function performs network calls against the provided async client and retries
+    on ``APIConnectionError`` up to ``max_retries`` times with a fixed ``retry_delay``
+    between attempts. All other exceptions propagate immediately.
 
     Parameters:
     ----------
@@ -81,27 +96,45 @@ async def _register_mcp_toolgroups_async(
                                                         to ensure are registered.
         logger (Logger): Logger used for debug messages about registration
                          progress.
+        max_retries (int): Maximum number of connection attempts before giving up.
+        retry_delay (int): Delay in seconds between retry attempts.
+
+    Raises:
+    ------
+        APIConnectionError: If the client is unreachable after all retries.
     """
-    # Get registered tools
-    registered_toolgroups = await client.toolgroups.list()
-    registered_toolgroups_ids = [
-        tool_group.provider_resource_id for tool_group in registered_toolgroups
-    ]
-    logger.debug("Registered toolgroups: %s", registered_toolgroups_ids)
-
-    # Register toolgroups for MCP servers if not already registered
-    for mcp in mcp_servers:
-        if mcp.name not in registered_toolgroups_ids:
-            logger.debug("Registering MCP server: %s, %s", mcp.name, mcp.url)
-
-            registration_params = {
-                "toolgroup_id": mcp.name,
-                "provider_id": mcp.provider_id,
-                "mcp_endpoint": {"uri": mcp.url},
-            }
-
-            await client.toolgroups.register(**registration_params)
-            logger.debug("MCP server %s registered successfully", mcp.name)
+    if max_retries < 1:
+        raise ValueError("max_retries must be >= 1")
+
+    for attempt in range(max_retries):
+        try:
+            registered_toolgroups = await client.toolgroups.list()
+            registered_toolgroups_ids = [
+                tool_group.provider_resource_id for tool_group in registered_toolgroups
+            ]
+            logger.debug("Registered toolgroups: %s", registered_toolgroups_ids)
+
+            for mcp in mcp_servers:
+                if mcp.name not in registered_toolgroups_ids:
+                    logger.debug("Registering MCP server: %s, %s", mcp.name, mcp.url)
+                    registration_params = {
+                        "toolgroup_id": mcp.name,
+                        "provider_id": mcp.provider_id,
+                        "mcp_endpoint": {"uri": mcp.url},
+                    }
+                    await client.toolgroups.register(**registration_params)
+                    logger.debug("MCP server %s registered successfully", mcp.name)
+            return
+        except APIConnectionError:
+            if attempt == max_retries - 1:
+                raise
+            logger.warning(
+                "MCP server registration failed (attempt %d/%d), retrying in %ds...",
+                attempt + 1,
+                max_retries,
+                retry_delay,
+            )
+            await asyncio.sleep(retry_delay)
 
 
 def run_once_async(func: Callable[..., Any]) -> Callable[..., Any]:

tests/unit/utils/test_common.py

@@ -3,6 +3,7 @@
 from logging import Logger
 
 import pytest
+from llama_stack_client import APIConnectionError
 from pydantic import AnyHttpUrl
 from pytest_mock import MockerFixture
 
@@ -14,6 +15,7 @@
     UserDataCollection,
 )
 from utils.common import (
+    _register_mcp_toolgroups_async,
     register_mcp_servers_async,
 )
 
@@ -412,3 +414,199 @@ async def test_register_mcp_servers_async_with_library_client(
         provider_id="model-context-protocol",
         mcp_endpoint={"uri": "http://localhost:8080"},
     )
+
+
+@pytest.mark.asyncio
+async def test_register_mcp_toolgroups_retries_on_connection_error(
+    mocker: MockerFixture,
+) -> None:
+    """Test that _register_mcp_toolgroups_async retries on APIConnectionError."""
+    mock_client = mocker.AsyncMock()
+    mock_logger = mocker.Mock(spec=Logger)
+    mock_sleep = mocker.patch("utils.common.asyncio.sleep")
+
+    mock_tool = mocker.Mock()
+    mock_tool.provider_resource_id = "existing-server"
+
+    mock_client.toolgroups.list.side_effect = [
+        APIConnectionError(request=mocker.MagicMock()),
+        APIConnectionError(request=mocker.MagicMock()),
+        [mock_tool],
+    ]
+
+    mcp_servers = [
+        ModelContextProtocolServer(
+            name="new-server",
+            url="http://localhost:8080",
+        ),  # pyright: ignore[reportCallIssue]
+    ]
+
+    await _register_mcp_toolgroups_async(
+        mock_client, mcp_servers, mock_logger, max_retries=5, retry_delay=1
+    )
+
+    assert mock_client.toolgroups.list.call_count == 3
+    assert mock_sleep.call_count == 2
+    mock_sleep.assert_called_with(1)
+    assert mock_logger.warning.call_count == 2
+
+
+@pytest.mark.asyncio
+async def test_register_mcp_toolgroups_raises_after_max_retries(
+    mocker: MockerFixture,
+) -> None:
+    """Test that _register_mcp_toolgroups_async raises after all retries are exhausted."""
+    mock_client = mocker.AsyncMock()
+    mock_logger = mocker.Mock(spec=Logger)
+    mock_sleep = mocker.patch("utils.common.asyncio.sleep")
+
+    mock_client.toolgroups.list.side_effect = APIConnectionError(
+        request=mocker.MagicMock()
+    )
+
+    mcp_servers = [
+        ModelContextProtocolServer(
+            name="new-server",
+            url="http://localhost:8080",
+        ),  # pyright: ignore[reportCallIssue]
+    ]
+
+    with pytest.raises(APIConnectionError):
+        await _register_mcp_toolgroups_async(
+            mock_client, mcp_servers, mock_logger, max_retries=3, retry_delay=1
+        )
+
+    assert mock_client.toolgroups.list.call_count == 3
+    assert mock_sleep.call_count == 2
+
+
+@pytest.mark.asyncio
+async def test_register_mcp_toolgroups_retries_on_register_error(
+    mocker: MockerFixture,
+) -> None:
+    """Test retry when register() raises APIConnectionError."""
+    mock_client = mocker.AsyncMock()
+    mock_logger = mocker.Mock(spec=Logger)
+    mock_sleep = mocker.patch("utils.common.asyncio.sleep")
+
+    mock_client.toolgroups.list.return_value = []
+    mock_client.toolgroups.register.side_effect = [
+        APIConnectionError(request=mocker.MagicMock()),
+        None,
+    ]
+
+    mcp_servers = [
+        ModelContextProtocolServer(
+            name="new-server",
+            url="http://localhost:8080",
+        ),  # pyright: ignore[reportCallIssue]
+    ]
+
+    await _register_mcp_toolgroups_async(
+        mock_client, mcp_servers, mock_logger, max_retries=3, retry_delay=1
+    )
+
+    assert mock_client.toolgroups.list.call_count == 2
+    assert mock_client.toolgroups.register.call_count == 2
+    assert mock_sleep.call_count == 1
+    mock_sleep.assert_called_with(1)
+
+
+@pytest.mark.asyncio
+async def test_register_mcp_toolgroups_single_attempt_success(
+    mocker: MockerFixture,
+) -> None:
+    """Test max_retries=1 with successful first attempt."""
+    mock_client = mocker.AsyncMock()
+    mock_logger = mocker.Mock(spec=Logger)
+    mock_sleep = mocker.patch("utils.common.asyncio.sleep")
+
+    mock_client.toolgroups.list.return_value = []
+    mock_client.toolgroups.register.return_value = None
+
+    mcp_servers = [
+        ModelContextProtocolServer(
+            name="new-server",
+            url="http://localhost:8080",
+        ),  # pyright: ignore[reportCallIssue]
+    ]
+
+    await _register_mcp_toolgroups_async(
+        mock_client, mcp_servers, mock_logger, max_retries=1, retry_delay=1
+    )
+
+    mock_client.toolgroups.list.assert_called_once()
+    mock_client.toolgroups.register.assert_called_once()
+    mock_sleep.assert_not_called()
+    mock_logger.warning.assert_not_called()
+
+
+@pytest.mark.asyncio
+async def test_register_mcp_toolgroups_single_attempt_failure(
+    mocker: MockerFixture,
+) -> None:
+    """Test max_retries=1 raises immediately without sleeping."""
+    mock_client = mocker.AsyncMock()
+    mock_logger = mocker.Mock(spec=Logger)
+    mock_sleep = mocker.patch("utils.common.asyncio.sleep")
+
+    mock_client.toolgroups.list.side_effect = APIConnectionError(
+        request=mocker.MagicMock()
+    )
+
+    mcp_servers = [
+        ModelContextProtocolServer(
+            name="new-server",
+            url="http://localhost:8080",
+        ),  # pyright: ignore[reportCallIssue]
+    ]
+
+    with pytest.raises(APIConnectionError):
+        await _register_mcp_toolgroups_async(
+            mock_client, mcp_servers, mock_logger, max_retries=1, retry_delay=1
+        )
+
+    mock_client.toolgroups.list.assert_called_once()
+    mock_sleep.assert_not_called()
+    mock_logger.warning.assert_not_called()
+
+
+@pytest.mark.asyncio
+async def test_register_mcp_toolgroups_non_connection_error_propagates(
+    mocker: MockerFixture,
+) -> None:
+    """Test that non-APIConnectionError exceptions propagate without retry."""
+    mock_client = mocker.AsyncMock()
+    mock_logger = mocker.Mock(spec=Logger)
+    mock_sleep = mocker.patch("utils.common.asyncio.sleep")
+
+    mock_client.toolgroups.list.side_effect = RuntimeError("unexpected")
+
+    mcp_servers = [
+        ModelContextProtocolServer(
+            name="new-server",
+            url="http://localhost:8080",
+        ),  # pyright: ignore[reportCallIssue]
+    ]
+
+    with pytest.raises(RuntimeError, match="unexpected"):
+        await _register_mcp_toolgroups_async(
+            mock_client, mcp_servers, mock_logger, max_retries=5, retry_delay=1
+        )
+
+    mock_client.toolgroups.list.assert_called_once()
+    mock_sleep.assert_not_called()
+
+
+@pytest.mark.asyncio
+async def test_register_mcp_toolgroups_invalid_max_retries(
+    mocker: MockerFixture,
+) -> None:
+    """Test that max_retries < 1 raises ValueError."""
+    mock_client = mocker.AsyncMock()
+    mock_logger = mocker.Mock(spec=Logger)
+
+    with pytest.raises(ValueError, match="max_retries must be >= 1"):
+        await _register_mcp_toolgroups_async(
+            mock_client, [], mock_logger, max_retries=0
+        )