Implement requests retries (#457)

kevalmahajan · web-flow · commit 9800e3f33fae · 2025-07-15T15:54:40.000+01:00
* added retry mechanism for httpx requests

Signed-off-by: Keval Mahajan &lt;mahajankeval23@gmail.com&gt;

* linting

Signed-off-by: Keval Mahajan &lt;mahajankeval23@gmail.com&gt;

* configurable values for retries from settings

Signed-off-by: Keval Mahajan &lt;mahajankeval23@gmail.com&gt;

* added test cases for ResilientHttpClient

Signed-off-by: Keval Mahajan &lt;mahajankeval23@gmail.com&gt;

* minor change

Signed-off-by: Keval Mahajan &lt;mahajankeval23@gmail.com&gt;

---------

Signed-off-by: Keval Mahajan &lt;mahajankeval23@gmail.com&gt;
diff --git a/.env.example b/.env.example
@@ -126,6 +126,15 @@ ALLOWED_ORIGINS='["http://localhost", "http://localhost:4444"]'
 # Enable CORS handling in the gateway
 CORS_ENABLED=true
 
+#####################################
+# Retry Config for HTTP Requests
+#####################################
+
+RETRY_MAX_ATTEMPTS=3
+RETRY_BASE_DELAY=1.0  # seconds
+RETRY_MAX_DELAY=60.0  # seconds
+RETRY_JITTER_MAX=0.5  # fraction of delay
+
 #####################################
 # Logging
 #####################################
diff --git a/mcpgateway/admin.py b/mcpgateway/admin.py
@@ -68,6 +68,7 @@
     ToolService,
 )
 from mcpgateway.utils.create_jwt_token import get_jwt_token
+from mcpgateway.utils.retry_manager import ResilientHttpClient
 from mcpgateway.utils.verify_credentials import require_auth, require_basic_auth
 from mcpgateway.utils.error_formatter import ErrorFormatter
 
@@ -1383,7 +1384,7 @@ async def admin_test_gateway(request: GatewayTestRequest, user: str = Depends(re
     logger.debug(f"User {user} testing server at {request.base_url}.")
     try:
         start_time = time.monotonic()
-        async with httpx.AsyncClient(timeout=settings.federation_timeout, verify=not settings.skip_ssl_verify) as client:
+        async with ResilientHttpClient(client_args={"timeout": settings.federation_timeout, "verify": not settings.skip_ssl_verify}) as client:
             response = await client.request(method=request.method.upper(), url=full_url, headers=request.headers, json=request.body)
         latency_ms = int((time.monotonic() - start_time) * 1000)
         try:
diff --git a/mcpgateway/alembic/env.py b/mcpgateway/alembic/env.py
@@ -4,14 +4,13 @@
 from logging.config import fileConfig
 
 # Third-Party
-from alembic import context
-
 # this is the Alembic Config object, which provides
 # access to the values within the .ini file in use.
 from alembic.config import Config
 from sqlalchemy import engine_from_config, pool
 
 # First-Party
+from alembic import context
 from mcpgateway.config import settings
 from mcpgateway.db import Base
 
diff --git a/mcpgateway/alembic/versions/b77ca9d2de7e_uuid_pk_and_slug_refactor.py b/mcpgateway/alembic/versions/b77ca9d2de7e_uuid_pk_and_slug_refactor.py
@@ -12,11 +12,11 @@
 import uuid
 
 # Third-Party
-from alembic import op
 import sqlalchemy as sa
 from sqlalchemy.orm import Session
 
 # First-Party
+from alembic import op
 from mcpgateway.config import settings
 from mcpgateway.utils.create_slug import slugify
 
diff --git a/mcpgateway/alembic/versions/e4fc04d1a442_add_annotations_to_tables.py b/mcpgateway/alembic/versions/e4fc04d1a442_add_annotations_to_tables.py
@@ -11,9 +11,11 @@
 from typing import Sequence, Union
 
 # Third-Party
-from alembic import op
 import sqlalchemy as sa
 
+# First-Party
+from alembic import op
+
 # revision identifiers, used by Alembic.
 revision: str = "e4fc04d1a442"
 down_revision: Union[str, Sequence[str], None] = "b77ca9d2de7e"
diff --git a/mcpgateway/alembic/versions/e75490e949b1_add_improved_status_to_tables.py b/mcpgateway/alembic/versions/e75490e949b1_add_improved_status_to_tables.py
@@ -10,9 +10,11 @@
 from typing import Sequence, Union
 
 # Third-Party
-from alembic import op
 import sqlalchemy as sa
 
+# First-Party
+from alembic import op
+
 # Revision identifiers.
 revision: str = "e75490e949b1"
 down_revision: Union[str, Sequence[str], None] = "e4fc04d1a442"
diff --git a/mcpgateway/bootstrap_db.py b/mcpgateway/bootstrap_db.py
@@ -25,11 +25,11 @@
 import logging
 
 # Third-Party
-from alembic import command
 from alembic.config import Config
 from sqlalchemy import create_engine, inspect
 
 # First-Party
+from alembic import command
 from mcpgateway.config import settings
 from mcpgateway.db import Base
 
diff --git a/mcpgateway/cache/session_registry.py b/mcpgateway/cache/session_registry.py
@@ -33,7 +33,6 @@
 
 # Third-Party
 from fastapi import HTTPException, status
-import httpx
 
 # First-Party
 from mcpgateway import __version__
@@ -42,6 +41,7 @@
 from mcpgateway.models import Implementation, InitializeResult, ServerCapabilities
 from mcpgateway.services import PromptService, ResourceService, ToolService
 from mcpgateway.transports import SSETransport
+from mcpgateway.utils.retry_manager import ResilientHttpClient
 
 logger = logging.getLogger(__name__)
 
@@ -780,7 +780,7 @@ async def generate_response(self, message: json, transport: SSETransport, server
                 }
                 headers = {"Authorization": f"Bearer {user['token']}", "Content-Type": "application/json"}
                 rpc_url = base_url + "/rpc"
-                async with httpx.AsyncClient(timeout=settings.federation_timeout, verify=not settings.skip_ssl_verify) as client:
+                async with ResilientHttpClient(client_args={"timeout": settings.federation_timeout, "verify": not settings.skip_ssl_verify}) as client:
                     rpc_response = await client.post(
                         url=rpc_url,
                         json=rpc_input,
diff --git a/mcpgateway/config.py b/mcpgateway/config.py
@@ -134,6 +134,12 @@ class Settings(BaseSettings):
         "http://localhost:4444",
     }
 
+    # Max retries for HTTP requests
+    retry_max_attempts: int = 3
+    retry_base_delay: float = 1.0  # seconds
+    retry_max_delay: int = 60  # seconds
+    retry_jitter_max: float = 0.5  # fraction of base delay
+
     @field_validator("allowed_origins", mode="before")
     @classmethod
     def _parse_allowed_origins(cls, v):
diff --git a/mcpgateway/main.py b/mcpgateway/main.py
@@ -49,7 +49,6 @@
 from fastapi.responses import JSONResponse, RedirectResponse, StreamingResponse
 from fastapi.staticfiles import StaticFiles
 from fastapi.templating import Jinja2Templates
-import httpx
 from sqlalchemy import text
 from sqlalchemy.orm import Session
 from starlette.middleware.base import BaseHTTPMiddleware
@@ -125,6 +124,7 @@
 )
 from mcpgateway.utils.db_isready import wait_for_db_ready
 from mcpgateway.utils.redis_isready import wait_for_redis_ready
+from mcpgateway.utils.retry_manager import ResilientHttpClient
 from mcpgateway.utils.verify_credentials import require_auth, require_auth_override
 from mcpgateway.validation.jsonrpc import (
     JSONRPCError,
@@ -1848,7 +1848,8 @@ async def websocket_endpoint(websocket: WebSocket):
         while True:
             try:
                 data = await websocket.receive_text()
-                async with httpx.AsyncClient(timeout=settings.federation_timeout, verify=not settings.skip_ssl_verify) as client:
+                client_args = {"timeout": settings.federation_timeout, "verify": not settings.skip_ssl_verify}
+                async with ResilientHttpClient(client_args=client_args) as client:
                     response = await client.post(
                         f"http://localhost:{settings.port}/rpc",
                         json=json.loads(data),
diff --git a/mcpgateway/services/gateway_service.py b/mcpgateway/services/gateway_service.py
@@ -50,6 +50,7 @@
 from mcpgateway.schemas import GatewayCreate, GatewayRead, GatewayUpdate, ToolCreate
 from mcpgateway.services.tool_service import ToolService
 from mcpgateway.utils.create_slug import slugify
+from mcpgateway.utils.retry_manager import ResilientHttpClient
 from mcpgateway.utils.services_auth import decode_auth
 
 # logging.getLogger("httpx").setLevel(logging.WARNING)  # Disables httpx logs for regular health checks
@@ -106,7 +107,7 @@ class GatewayService:
     def __init__(self) -> None:
         """Initialize the gateway service."""
         self._event_subscribers: List[asyncio.Queue] = []
-        self._http_client = httpx.AsyncClient(timeout=settings.federation_timeout, verify=not settings.skip_ssl_verify)
+        self._http_client = ResilientHttpClient(client_args={"timeout": settings.federation_timeout, "verify": not settings.skip_ssl_verify})
         self._health_check_interval = GW_HEALTH_CHECK_INTERVAL
         self._health_check_task: Optional[asyncio.Task] = None
         self._active_gateways: Set[str] = set()  # Track active gateway URLs
diff --git a/mcpgateway/services/tool_service.py b/mcpgateway/services/tool_service.py
@@ -25,7 +25,6 @@
 from typing import Any, AsyncGenerator, Dict, List, Optional
 
 # Third-Party
-import httpx
 from mcp import ClientSession
 from mcp.client.sse import sse_client
 from mcp.client.streamable_http import streamablehttp_client
@@ -46,6 +45,7 @@
     ToolUpdate,
 )
 from mcpgateway.utils.create_slug import slugify
+from mcpgateway.utils.retry_manager import ResilientHttpClient
 from mcpgateway.utils.services_auth import decode_auth
 
 # Local
@@ -104,7 +104,7 @@ class ToolService:
     def __init__(self):
         """Initialize the tool service."""
         self._event_subscribers: List[asyncio.Queue] = []
-        self._http_client = httpx.AsyncClient(timeout=settings.federation_timeout, verify=not settings.skip_ssl_verify)
+        self._http_client = ResilientHttpClient(client_args={"timeout": settings.federation_timeout, "verify": not settings.skip_ssl_verify})
 
     async def initialize(self) -> None:
         """Initialize the service."""
diff --git a/mcpgateway/utils/retry_manager.py b/mcpgateway/utils/retry_manager.py
diff --git a/mcpgateway/wrapper.py b/mcpgateway/wrapper.py
diff --git a/tests/unit/mcpgateway/cache/test_session_registry.py b/tests/unit/mcpgateway/cache/test_session_registry.py
diff --git a/tests/unit/mcpgateway/test_main.py b/tests/unit/mcpgateway/test_main.py
diff --git a/tests/unit/mcpgateway/test_wrapper.py b/tests/unit/mcpgateway/test_wrapper.py
diff --git a/tests/unit/mcpgateway/utils/test_retry_manager.py b/tests/unit/mcpgateway/utils/test_retry_manager.py
