From ca03c889c9b72ec47a5d4f604626ac0b7cfdfadd Mon Sep 17 00:00:00 2001
From: Joshua Valluru <326636767+joshua-berri@users.noreply.github.com>
Date: Thu, 10 Sep 2026 20:26:53 -0700
Subject: [PATCH] fix(mcp): avoid caching cancelled OpenAPI health probes

---
 .../mcp_server/mcp_server_manager.py          | 11 ++--
 .../mcp_server/test_mcp_server_manager.py     | 51 ++++++++++++++++---
 2 files changed, 53 insertions(+), 9 deletions(-)

diff --git a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py
index d0b0c83a23c..17d3098bf4c 100644
--- a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py
+++ b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py
@@ -897,8 +897,6 @@ async def _openapi_spec_health(
         await asyncio.wait_for(load_openapi_spec_async(spec_path, max_bytes=10 * 1024 * 1024), timeout=timeout)
     except asyncio.TimeoutError:
         return "unhealthy", f"OpenAPI specification check timed out after {timeout} seconds"
-    except asyncio.CancelledError:
-        return "unknown", "OpenAPI specification check was cancelled"
     except HTTPStatusError as exc:
         return "unhealthy", f"OpenAPI specification request failed (HTTP {exc.response.status_code})"
     except HTTPResponseLimitError as exc:
@@ -920,7 +918,14 @@ class _OpenAPIHealthProbe:
         async with self.lock:
             if self.result is not None and self.clock() - self.checked_at < 30.0:
                 return self.result
-            status, error = await _openapi_spec_health(self.spec_path, timeout=MCP_HEALTH_CHECK_TIMEOUT)
+            try:
+                status, error = await _openapi_spec_health(self.spec_path, timeout=MCP_HEALTH_CHECK_TIMEOUT)
+            except asyncio.CancelledError:
+                return (
+                    "unknown",
+                    "OpenAPI specification check was cancelled",
+                    datetime.datetime.now(datetime.timezone.utc),
+                )
             self.result = (status, error, datetime.datetime.now(datetime.timezone.utc))
             self.checked_at = self.clock()
             return self.result
diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py
index e3dd92485fc..a6e43686e2c 100644
--- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py
+++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py
@@ -4575,12 +4575,12 @@ class TestMCPServerManager:
         await asyncio.wait_for(started.wait(), timeout=1)
         if cancel:
             task.cancel()
-        status, error = await task
-        assert status == ("unknown" if cancel else "unhealthy")
-        assert error == (
-            "OpenAPI specification check was cancelled"
-            if cancel else "OpenAPI specification check timed out after 0.1 seconds"
-        )
+            with pytest.raises(asyncio.CancelledError):
+                await task
+        else:
+            status, error = await task
+            assert status == "unhealthy"
+            assert error == "OpenAPI specification check timed out after 0.1 seconds"
         assert cancelled.is_set()
 
     @pytest.mark.asyncio
@@ -12862,3 +12862,42 @@ async def test_openapi_health_reports_size_limit_as_unknown_and_caches_failure(r
     assert cached.health_check_error == result.health_check_error
     assert cached.last_health_check == result.last_health_check
     assert route.call_count == 1
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("already_waiting", [False, True])
+async def test_openapi_health_cancellation_does_not_poison_cache(respx_mock, monkeypatch, already_waiting):
+    monkeypatch.setenv("DISABLE_AIOHTTP_TRANSPORT", "True")
+    manager = MCPServerManager()
+    server = MCPServer(
+        server_id="cancelled-cache", name="cancelled-cache", transport=MCPTransport.http,
+        spec_path="https://93.184.216.34/cancelled-cache.json", auth_type=MCPAuth.none,
+    )
+    manager.registry = {server.server_id: server}
+    started = asyncio.Event()
+    attempts = []
+
+    async def serve(request):
+        attempts.append(request.url)
+        if not started.is_set():
+            started.set()
+            await asyncio.Event().wait()
+        return httpx.Response(200, json={"paths": {}})
+
+    route = respx_mock.get(server.spec_path).mock(side_effect=serve)
+    leader = asyncio.create_task(manager.health_check_server(server.server_id))
+    await asyncio.wait_for(started.wait(), timeout=1)
+    follower = asyncio.create_task(manager.health_check_server(server.server_id)) if already_waiting else None
+    await asyncio.sleep(0)
+    leader.cancel()
+    cancelled = await leader
+    assert cancelled.status == "unknown"
+    assert cancelled.health_check_error == "OpenAPI specification check was cancelled"
+    recovered = await follower if follower is not None else await manager.health_check_server(server.server_id)
+    assert recovered.status == "healthy"
+    assert recovered.health_check_error is None
+    cached = await manager.health_check_server(server.server_id)
+    assert cached.last_health_check == recovered.last_health_check
+    assert cached.status == "healthy"
+    assert len(attempts) == 2
+    assert route.call_count == 1