fix(mcp): avoid caching cancelled OpenAPI health probes

This commit is contained in:
Joshua Valluru 2026-09-10 20:26:53 -07:00
parent da1dfcdb24
commit ca03c889c9
2 changed files with 53 additions and 9 deletions

View file

@ -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

View file

@ -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